[Feature] AddThinkingPrompt transform (#3027)

vmoens · web-flow · commit ed051bc3e5b3 · 2025-06-23T10:23:40.000+01:00
diff --git a/test/llm/test_envs.py b/test/llm/test_envs.py
@@ -1130,6 +1130,128 @@ def test_async_mcp_tools(self):
             env_pool.close()
 
 
+class TestThinkingPrompt:
+    @pytest.fixture(autouse=True, scope="class")
+    def base_env(self):
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
+        env = GSM8KEnv(shuffle=False, tokenizer=tokenizer, max_steps=10)
+        return env
+
+    @pytest.mark.skipif(not _has_transformers, reason="requires transformers")
+    @pytest.mark.skipif(not _has_datasets, reason="requires gsm8k")
+    @pytest.mark.parametrize(
+        "role,edit_last_turn",
+        [("assistant", True), ("assistant", False), ("user", False)],
+    )
+    @pytest.mark.parametrize("zero_reward", [True, False])
+    @pytest.mark.parametrize("undo_done", [True, False])
+    @pytest.mark.parametrize("random_prompt", [True, False])
+    def test_thinking_prompt_wrong_answer(
+        self,
+        role,
+        edit_last_turn,
+        zero_reward,
+        undo_done,
+        random_prompt,
+        tmp_path,
+        base_env,
+    ):
+        from torchrl.envs.llm.transforms import AddThinkingPrompt
+
+        if isinstance(base_env.transform[-1], AddThinkingPrompt):
+            base_env.transform.pop()
+        env = base_env.reset_dataloader()
+        env = base_env.append_transform(
+            AddThinkingPrompt(
+                cond=lambda td: td["reward"] < 50,
+                role=role,
+                edit_last_turn=edit_last_turn,
+                zero_reward=zero_reward,
+                undo_done=undo_done,
+                random_prompt=random_prompt,
+            )
+        )
+        reset = env.reset()
+        assert reset[0]["history"][-1].content.startswith(
+            "Natalia sold clips to 48 of her friends in April"
+        )
+        policy_anser = (
+            "<think>Let me solve this step by step. Natalia sold clips to 48 friends in April. Then she sold half as many in May. Half of 48 is 24. So in May she sold 24 clips. "
+            "To find the total, I need to add April and May: 48 + 24 = 72. Therefore, Natalia sold 72 clips altogether in April and May.</think>\n<answer>322 clips</answer><|im_end|>"
+        )
+        reset["text_response"] = [policy_anser]
+        s = env.step(reset)
+        if zero_reward:
+            assert (s["next", "reward"] == 0).all()
+        else:
+            assert (s["next", "reward"] != 0).all()
+        if undo_done:
+            assert (s["next", "done"] == 0).all()
+        else:
+            assert (s["next", "done"] != 0).all()
+        if edit_last_turn:
+            assert s["next", "history"].shape == (1, 3)
+        else:
+            assert s["next", "history"].shape == (1, 4)
+        if role == "assistant":
+            assert s[0]["next", "history", "role"][-1] == "assistant"
+        else:
+            assert s[0]["next", "history", "role"][-1] == "user"
+
+    @pytest.mark.skipif(not _has_transformers, reason="requires transformers")
+    @pytest.mark.skipif(not _has_datasets, reason="requires gsm8k")
+    @pytest.mark.parametrize(
+        "role,edit_last_turn",
+        [("assistant", True), ("assistant", False), ("user", False)],
+    )
+    @pytest.mark.parametrize("zero_reward", [True, False])
+    @pytest.mark.parametrize("undo_done", [True, False])
+    @pytest.mark.parametrize("random_prompt", [True, False])
+    def test_thinking_prompt_correct_answer(
+        self,
+        role,
+        edit_last_turn,
+        zero_reward,
+        undo_done,
+        random_prompt,
+        tmp_path,
+        base_env,
+    ):
+        # checks that if cond returns False, nothing is changed
+        from torchrl.envs.llm.transforms import AddThinkingPrompt
+
+        if isinstance(base_env.transform[-1], AddThinkingPrompt):
+            base_env.transform.pop()
+        env = base_env
+        env = env.reset_dataloader()
+        env = env.append_transform(
+            AddThinkingPrompt(
+                cond=lambda td: td["reward"] < 50,
+                role=role,
+                edit_last_turn=edit_last_turn,
+                zero_reward=zero_reward,
+                undo_done=undo_done,
+                random_prompt=random_prompt,
+            )
+        )
+        reset = env.reset()
+        assert reset[0]["history"][-1].content.startswith(
+            "Natalia sold clips to 48 of her friends in April"
+        )
+        policy_anser = (
+            "<think>Let me solve this step by step. Natalia sold clips to 48 friends in April. Then she sold half as many in May. Half of 48 is 24. So in May she sold 24 clips. "
+            "To find the total, I need to add April and May: 48 + 24 = 72. Therefore, Natalia sold 72 clips altogether in April and May.</think>\n<answer>72</answer><|im_end|>"
+        )
+        reset["text_response"] = [policy_anser]
+        s = env.step(reset)
+        assert (s["next", "reward"] != 0).all(), s["next", "reward"]
+        assert s[0]["next", "history", "role"][-1] == "assistant"
+        assert s["next", "done"].all()
+        assert len(s[0]["next", "history", "content"]) == 3
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
diff --git a/torchrl/envs/llm/__init__.py b/torchrl/envs/llm/__init__.py
@@ -15,6 +15,7 @@
 from .libs import make_mlgym, MLGymWrapper
 from .reward import GSM8KRewardParser, IFEvalScoreData, IfEvalScorer
 from .transforms import (
+    AddThinkingPrompt,
     as_nested_tensor,
     as_padded_tensor,
     BrowserTransform,
@@ -33,6 +34,7 @@
     "ChatEnv",
     "DataLoadingPrimer",
     "DatasetChatEnv",
+    "AddThinkingPrompt",
     "GSM8KEnv",
     "GSM8KPrepareQuestion",
     "GSM8KRewardParser",
diff --git a/torchrl/envs/llm/chat.py b/torchrl/envs/llm/chat.py
@@ -284,6 +284,7 @@ class DatasetChatEnv(TransformedEnv):
 
     Keyword Args:
         dataset (str): The name of the dataset.
+        shuffle (bool, optional): Whether to shuffle the dataset. Defaults to `True`.
         name (str, optional): name of the dataset configuration.
         split (str, optional): the split to use (usually from `"train"`, `"val"` or `"test"`). Defaults to `None` (no split).
         num_envs (int, optional): The number of environments to create. Defaults to `1`.
@@ -317,6 +318,7 @@ def __init__(
         self,
         *,
         dataset: str,
+        shuffle: bool = True,
         name: str | None = None,
         split: Literal["train", "val", "test"] | None = None,
         num_envs: int = 1,
@@ -355,7 +357,7 @@ def __init__(
         dataloader = DataLoader(  # noqa: TOR401
             dataset,
             batch_size=batch_size_dl,
-            shuffle=True,
+            shuffle=shuffle,
             collate_fn=collate_fn,
             generator=generator,
         )
@@ -375,3 +377,14 @@ def __init__(
             apply_template=apply_template,
         )
         return super().__init__(env_base, primer)
+
+    def reset_dataloader(self):
+        """Reset the dataloader.
+
+        This is useful when the dataloader is not infinite and we want to reset it.
+
+        Returns:
+            self: The environment itself.
+        """
+        self.transform[0].reset_dataloader()
+        return self
diff --git a/torchrl/envs/llm/datasets/gsm8k.py b/torchrl/envs/llm/datasets/gsm8k.py
@@ -135,6 +135,7 @@ class GSM8KEnv(DatasetChatEnv):
 
     Keyword Args:
         dataset (str, optional): The name of the dataset. Defaults to `"gsm8k"`.
+        shuffle (bool, optional): Whether to shuffle the dataset. Defaults to `True`.
         num_envs (int, optional): The number of environments to create. Defaults to `1`.
         repeats (int | None, optional): The number of times to repeat each sample from the dataset (mainly for Monte-Carlo
             based value estimation). If `None`, the dataset is not repeated. Defaults to `None`.
@@ -284,12 +285,13 @@ class GSM8KEnv(DatasetChatEnv):
     SYSTEM_PROMPT = """A conversation between User and Assistant. The user asks a question, and the Assistant solves it.
 The assistant first thinks about the reasoning process in the mind and then provides the user with the answer.
 The reasoning process and answer are enclosed within <think></think> and <answer></answer> tags, respectively,
-i.e., <think>reasoning process here</think> <answer>answer here</answer>."""
+i.e., <think>reasoning process here</think> <answer>answer here</answer>. The answer should be a number."""
 
     def __init__(
         self,
         *,
         dataset: str = "gsm8k",
+        shuffle: bool = True,
         num_envs: int = 1,
         repeats: int | None = None,
         batch_size_dl: int = 1,
@@ -307,6 +309,7 @@ def __init__(
             collate_fn = _collate_fn
         super().__init__(
             dataset=dataset,
+            shuffle=shuffle,
             name="main",
             num_envs=num_envs,
             repeats=repeats,
diff --git a/torchrl/envs/llm/datasets/ifeval.py b/torchrl/envs/llm/datasets/ifeval.py
@@ -41,6 +41,7 @@ class IFEvalEnv(DatasetChatEnv):
 
     Keyword Args:
         dataset (str, optional): The name of the dataset. Defaults to `"google/IFeval"`.
+        shuffle (bool, optional): Whether to shuffle the dataset. Defaults to `True`.
         num_envs (int, optional): The number of environments to create. Defaults to `1`.
         repeats (int | None, optional): The number of times to repeat each sample from the dataset (mainly for Monte-Carlo
             based value estimation). If `None`, the dataset is not repeated. Defaults to `None`.
@@ -146,6 +147,7 @@ def __init__(
         self,
         *,
         dataset: str = "google/IFeval",
+        shuffle: bool = True,
         num_envs: int = 1,
         repeats: int | None = None,
         batch_size_dl: int = 1,
@@ -163,6 +165,7 @@ def __init__(
             collate_fn = _collate_fn
         super().__init__(
             dataset=dataset,
+            shuffle=shuffle,
             num_envs=num_envs,
             repeats=repeats,
             batch_size_dl=batch_size_dl,
diff --git a/torchrl/envs/llm/reward/gsm8k.py b/torchrl/envs/llm/reward/gsm8k.py
@@ -20,6 +20,7 @@ class GSM8KRewardParser(Transform):
         in_keys (list of NestedKey): the input keys. Defaults to `["text_response", "answer"]`.
         out_keys (list of NestedKey): the output keys. Defaults to `[ "reward_answer", "reward_think", "reward_right", "reward_contained", "reward", "success"]`.
         eos_token (str): the end of sentence token. Defaults to `tokenizer.eos_token` if not provided.
+        set_done_if_answer (bool): whether to set the done flag to `True` when an answer is present. Defaults to `True`.
 
     """
 
@@ -29,10 +30,18 @@ def __init__(
         in_keys: list[NestedKey] | None = None,
         out_keys: list[NestedKey] | None = None,
         eos_token: str | None = None,
+        set_done_if_answer: bool = True,
     ):
         super().__init__()
         self.tokenizer = tokenizer
-        self.eos_token = eos_token if eos_token is not None else tokenizer.eos_token
+        self.eos_token = (
+            eos_token
+            if eos_token is not None
+            else tokenizer.eos_token
+            if tokenizer is not None
+            else None
+        )
+        self.set_done_if_answer = set_done_if_answer
         if in_keys is None:
             in_keys = ["text_response", "answer"]
         if not isinstance(in_keys, list) or len(in_keys) != 2:
@@ -118,7 +127,20 @@ def _step(
             tds = tds.add(
                 next_td_exist, default=torch.zeros((), device=next_tensordict.device)
             )
-        return next_tensordict.update(tds)
+        next_tensordict = next_tensordict.update(tds)
+        if (
+            self.set_done_if_answer
+            and (reward_answer := (next_tensordict["reward_answer"] > 0)).any()
+        ):
+            done = next_tensordict.get("done")
+            if done is not None:
+                next_tensordict.set("done", reward_answer.view_as(done) | done)
+            terminated = next_tensordict.get("terminated")
+            if terminated is not None:
+                next_tensordict.set(
+                    "terminated", reward_answer.view_as(terminated) | terminated
+                )
+        return next_tensordict
 
     def transform_reward_spec(self, reward_spec: Composite) -> Composite:
         shape = reward_spec.shape + (1, 1)
diff --git a/torchrl/envs/llm/transforms/__init__.py b/torchrl/envs/llm/transforms/__init__.py
@@ -8,6 +8,7 @@
 from .format import TemplateTransform
 from .kl import KLRewardTransform, RetrieveLogProb
 from .policy_version import PolicyVersion
+from .reason import AddThinkingPrompt
 from .tokenizer import Tokenizer
 from .tools import MCPToolTransform, PythonInterpreter
 
@@ -19,6 +20,7 @@
     "MCPToolTransform",
     "PolicyVersion",
     "PythonInterpreter",
+    "AddThinkingPrompt",
     "TemplateTransform",
     "Tokenizer",
     "as_nested_tensor",
diff --git a/torchrl/envs/llm/transforms/dataloading.py b/torchrl/envs/llm/transforms/dataloading.py
@@ -447,6 +447,18 @@ def __init__(
         )
         self._reset_key = "_reset"
 
+    def reset_dataloader(self):
+        """Reset the dataloader.
+
+        This is useful when the dataloader is not infinite and we want to reset it.
+
+        Returns:
+            self: The transform itself.
+        """
+        self._queue.clear()
+        self.endless_dataloader = self._endless_iter(self.dataloader)
+        return self
+
     @classmethod
     def _endless_iter(self, obj):
         while True:
diff --git a/torchrl/envs/llm/transforms/reason.py b/torchrl/envs/llm/transforms/reason.py
