amend

Author: vmoens

sota-implementations/grpo/grpo_utils.py

@@ -548,7 +548,7 @@ def make_env(cfg: DictConfig, devices: list[int] | None = None):
             AddThinkingPrompt(
                 cond=lambda td: td["reward"] <= reward_threshold
                 and td["step_count"] < max_steps,
-                role="user",
+                role="assistant",
                 edit_last_turn=False,
                 zero_reward=False,
                 undo_done=True,

torchrl/envs/llm/datasets/ifeval.py

@@ -8,7 +8,6 @@
 
 import torch
 from tensordict import NonTensorData, NonTensorStack, TensorClass, TensorDict
-from torchrl._utils import logger as torchrl_logger
 from torchrl.data import Composite, NonTensor, Unbounded
 from torchrl.envs import StepCounter
 from torchrl.envs.llm.chat import DatasetChatEnv