amend

Author: vmoens

sota-implementations/grpo/grpo_utils.py

@@ -12,9 +12,8 @@
 
 from torchrl._utils import logger as torchrl_logger
 from torchrl.collectors.llm.weight_update.vllm import vLLMUpdater
-from torchrl.envs.llm import GSM8KEnv, KLRewardTransform, RetrieveKL
+from torchrl.envs.llm import GSM8KEnv, KLRewardTransform, RetrieveKL, AddThinkingPrompt
 from torchrl.envs.llm.datasets.ifeval import IFEvalEnv
-from torchrl.envs.llm.transforms.enhanced_reasoning import EnhancedReasoningTransform
 from torchrl.modules.llm import TransformersWrapper, vLLMWrapper
 from transformers.models.auto.modeling_auto import AutoModelForCausalLM
 from transformers.tokenization_utils import PreTrainedTokenizer
@@ -524,6 +523,7 @@ def make_env(cfg: DictConfig, devices: list[int] | None = None):
     max_steps = cfg.env.max_steps if cfg.env.reasoning else 1
     if cfg.env.dataset == "gsm8k":
         # Reward scale is 0.0 to 100
+        reward_threshold=20
         env = GSM8KEnv(
             repeats=cfg.env.repeats,
             tokenizer=train_tokenizer,
@@ -533,6 +533,7 @@ def make_env(cfg: DictConfig, devices: list[int] | None = None):
         )
     elif cfg.env.dataset == "ifeval":  # ifeval
         # Reward scale is 0.0 to 2.2
+        reward_threshold=1.0
         env = IFEvalEnv(
             repeats=cfg.env.repeats,
             tokenizer=train_tokenizer,
@@ -544,24 +545,14 @@ def make_env(cfg: DictConfig, devices: list[int] | None = None):
         raise NotImplementedError(f"Dataset {cfg.env.dataset} not implemented")
     if cfg.env.reasoning:
         env = env.append_transform(
-            # AddThinkingPrompt(
-            #     cond=lambda td: td["reward"] <= reward_threshold
-            #     and td["step_count"] < max_steps,
-            #     role="assistant",
-            #     edit_last_turn=True,
-            #     zero_reward=True,
-            #     undo_done=True,
-            #     random_prompt=True,
-            # ),
-            EnhancedReasoningTransform(
-                cond=lambda td: td["reward"] <= 1.0 and td["step_count"] < 3,
-                strategy="user_guidance",  # User tells assistant to reconsider
-                reward_threshold=1.0,
-                max_steps=3,
+            AddThinkingPrompt(
+                cond=lambda td, reward_threshol=reward_threshold, max_steps=max_steps: td["reward"] <= reward_threshold and td["step_count"] < max_steps,
+                role="assistant",
+                edit_last_turn=True,
                 zero_reward=True,
                 undo_done=True,
                 random_prompt=True,
-            )
+            ),
         )
         env = env.append_transform(
             # RetrieveKL will be lazily initialized in the collector.

test/llm/test_wrapper.py

@@ -1366,8 +1366,8 @@ def test_kl_computation_transform(
 
         # Create KLComputation transform
         kl_transform = KLComputation(
-            gen_log_probs_key=("gen_log_probs", "full"),
-            ref_log_probs_key=("ref_log_probs", "full"),
+            gen_log_probs_full_key=("gen_log_probs", "full"),
+            ref_log_probs_full_key=("ref_log_probs", "full"),
             kl_key="kl",
             add_to_reward=True,
             coeff=1.0,

torchrl/envs/llm/transforms/enhanced_reasoning.py

@@ -859,7 +859,7 @@ def __init__(
             )
         t1 = RetrieveLogProb(
             gen_model,
-            log_probs_key=gen_log_probs_full_key,
+            log_probs_full_key=gen_log_probs_full_key,
             assistant_only=assistant_only,
             tokenizer_kwargs=tokenizer_kwargs,
             detach=detach,
@@ -870,7 +870,7 @@ def __init__(
         )
         t2 = RetrieveLogProb(
             ref_model,
-            log_probs_key=ref_log_probs_full_key,
+            log_probs_full_key=ref_log_probs_full_key,
             assistant_only=assistant_only,
             tokenizer_kwargs=tokenizer_kwargs,
             detach=detach,

torchrl/envs/llm/transforms/kl.py

@@ -261,6 +261,10 @@ def forward(self, tensordict: TensorDictBase) -> GRPOLossOutput:
             raise ValueError(
                 f"advantage and log_weight must have the same number of dimensions, got {advantage.ndim=} and {log_weight.ndim=}"
             )
+        print(f"log_weight: {log_weight.shape}")
+        print(f"advantage: {advantage.shape}")
+        print(f"mask: {mask.shape}")
+        print(f"data: {tensordict}")
         gain1 = log_weight.exp() * advantage
 
         log_weight_clip = log_weight.clamp(*self._clip_bounds)