torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/envs/llm/transforms/policy_version.py

@@ -0,0 +1,189 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import uuid
+from dataclasses import dataclass
+from datetime import datetime
+from typing import cast, Literal
+
+import torch
+from tensordict import NonTensorData, TensorDictBase
+from torchrl.data.tensor_specs import Composite, NonTensor, Unbounded
+from torchrl.envs.transforms.transforms import Transform
+
+
+@dataclass
+class VersionChange:
+    """Records a single version change event."""
+
+    timestamp: datetime
+    old_version: str | int | None
+    new_version: str | int
+
+
+class PolicyVersion(Transform):
+    """A transform that keeps track of the version of the policy.
+
+    This transform is used to track policy versions during training, particularly in asynchronous
+    settings where policy weights are updated periodically. It is designed to work seamlessly with
+    :class:`~torchrl.collectors.llm.LLMCollector` to ensure data collection and training remain in sync.
+
+    The version can be either a UUID (string) or an integer counter. When used with :class:`~torchrl.collectors.llm.LLMCollector`,
+    the version is automatically incremented each time the policy weights are updated.
+
+    Example usage with :class:`~torchrl.collectors.llm.LLMCollector`:
+
+    .. code-block:: python
+
+        >>> # Create a policy version tracker
+        >>> policy_version = PolicyVersion(version_type="int")  # or "uuid" for UUID-based versioning
+        >>> # Create collector with version tracking
+        >>> collector = LLMCollector(
+        ...     env=env,
+        ...     policy=policy,
+        ...     track_policy_version=policy_version,  # Pass the version tracker
+        ...     # ... other arguments
+        ... )
+        >>> # The version will be automatically incremented when weights are updated
+        >>> collector.update_policy_weights_(new_weights)
+        >>> # The version is stored in the collected data
+        >>> for batch in collector:
+        ...     current_version = batch["policy_version"]
+
+    Args:
+        version_type: The type of versioning to use. Can be either:
+            - str or "uuid": Uses UUID4 for versions (good for distributed systems)
+            - int or "int": Uses incrementing integers (good for debugging)
+    """
+
+    def __init__(self, version_type: type | Literal["uuid", "int"] = int):
+        super().__init__()
+        self.version_type = version_type
+        self.version_history: list[VersionChange] = []  # Track version changes
+        self._current_version: str | int | None = None
+        self._increment_version(init=True)
+        self.cal_on_reset = True
+
+    @property
+    def version(self) -> str | int:
+        """The current version of the policy."""
+        if self._current_version is None:
+            raise RuntimeError("Version not initialized")
+        return self._current_version
+
+    @version.setter
+    def version(self, value: str | int) -> None:
+        self._current_version = value
+
+    def increment_version(self) -> None:
+        """Increment the version number.
+
+        This is called automatically by LLMCollector when policy weights are updated.
+        Can also be called manually if needed.
+        """
+        self._increment_version()
+
+    def _increment_version(self, init: bool = False) -> str | int:
+        """Internal method to handle version incrementing with history tracking."""
+        old_version = self._current_version
+        if self.version_type in (str, "uuid"):
+            self._increment_version_uuid(init)
+        elif self.version_type in (int, "int"):
+            self._increment_version_int(init)
+        else:
+            raise ValueError(f"Invalid version type: {self.version_type}")
+
+        # Record the version change
+        self.version_history.append(
+            VersionChange(
+                timestamp=datetime.now(),
+                old_version=old_version,
+                new_version=self.version,
+            )
+        )
+        return self.version
+
+    def _increment_version_uuid(self, init: bool = False) -> None:
+        """Generate a new UUID version.
+
+        Args:
+            init: If True, this is the initial version creation.
+        """
+        self.version = str(uuid.uuid4())
+
+    def _increment_version_int(self, init: bool = False) -> None:
+        """Increment the integer version counter.
+
+        Args:
+            init: If True, initialize counter to 0, otherwise increment by 1.
+        """
+        if init:
+            self.version = 0
+        else:
+            # Cast to int to ensure type safety
+            current = cast(int, self.version)
+            self.version = current + 1
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        """Reset the environment and update version in the new tensordict.
+
+        Args:
+            tensordict: The current tensordict
+            tensordict_reset: The tensordict to reset to
+
+        Returns:
+            The reset tensordict with updated version
+        """
+        tensordict_reset = self._step(None, tensordict_reset)
+        return tensordict_reset
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        """Add the current version to the tensordict.
+
+        This method is called on each environment step to ensure the collected
+        data is tagged with the correct policy version.
+
+        Args:
+            tensordict: The tensordict to update with version info
+
+        Returns:
+            The tensordict with added version information
+        """
+        if self.version_type in (str, "uuid"):
+            version = NonTensorData(self.version).expand(next_tensordict.shape)
+        elif self.version_type in (int, "int"):
+            # Cast to float for torch.full
+            version = torch.full(next_tensordict.shape, float(cast(int, self.version)))
+        else:
+            raise ValueError(f"Invalid version type: {self.version_type}")
+
+        next_tensordict.set("policy_version", version)
+        return next_tensordict
+
+    def transform_observation_spec(self, spec: Composite) -> Composite:
+        """Update the environment spec to include the version field.
+
+        Args:
+            spec: The environment spec to update
+
+        Returns:
+            Updated spec including the version field
+        """
+        if self.version_type in (str, "uuid"):
+            spec["policy_version"] = NonTensor(
+                example_data=uuid.uuid4(), shape=spec.shape, device=spec.device
+            )
+        elif self.version_type in (int, "int"):
+            spec["policy_version"] = Unbounded(
+                shape=spec.shape, dtype=torch.int64, device=spec.device
+            )
+        else:
+            raise ValueError(f"Invalid version type: {self.version_type}")
+        return spec

torchrl/envs/llm/transforms/reason.py

@@ -0,0 +1,323 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import re
+from collections.abc import Callable
+from typing import Literal
+
+from tensordict import lazy_stack, TensorDictBase
+from torchrl._utils import logger as torchrl_logger
+
+from torchrl.data.llm.history import History
+from torchrl.envs import Transform
+from torchrl.envs.common import EnvBase
+
+
+class AddThinkingPrompt(Transform):
+    """A transform that adds thinking prompts to encourage the LLM to reconsider its response.
+
+    This transform can either add a new thinking prompt as a separate message or edit the last
+    assistant response to include a thinking prompt before the final answer. This is useful for
+    training LLMs to self-correct and think more carefully when their initial responses are
+    incorrect or incomplete.
+
+    Args:
+        cond (Callable[[TensorDictBase], bool], optional): Condition function that determines
+            when to add the thinking prompt. Takes a tensordict and returns `True` if the prompt
+            should be added.
+        prompt (str, optional): The thinking prompt to add. If None, a default prompt is used.
+            Defaults to `"But wait, let me think about this more carefully..."`.
+        random_prompt (bool, optional): Whether to randomly select from predefined prompts.
+            Defaults to `False`.
+        role (Literal["user", "assistant"], optional): The role for the thinking prompt.
+            If `"assistant"`, the prompt is added to the assistant's response. If `"user"`, it's
+            added as a separate user message. Defaults to `"assistant"`.
+        edit_last_turn (bool, optional): Whether to edit the last assistant response instead
+            of adding a new message. Only works with `role="assistant"`. Defaults to `True`.
+        zero_reward (bool, optional): Whether to zero out the reward when the thinking prompt
+            is added. If `None`, defaults to the value of `edit_last_turn`. Defaults to the same value as `edit_last_turn`.
+        undo_done (bool, optional): Whether to undo the done flag when the thinking prompt
+            is added. Defaults to `True`.
+        egocentric (bool, optional): Whether the thinking prompt is written from the perspective of the assistant.
+            Defaults to `None`, which means that the prompt is written from the perspective of the user if `role="user"`
+            and from the perspective of the assistant if `role="assistant"`.
+
+    Examples:
+        >>> from torchrl.envs.llm.transforms import AddThinkingPrompt
+        >>> from torchrl.envs.llm import GSM8KEnv
+        >>> from transformers import AutoTokenizer
+        >>> import torch
+        >>>
+        >>> # Create environment with thinking prompt transform
+        >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
+        >>> env = GSM8KEnv(tokenizer=tokenizer, max_steps=10)
+        >>> env = env.append_transform(
+        ...     AddThinkingPrompt(
+        ...         cond=lambda td: td["reward"] < 50,
+        ...         role="assistant",
+        ...         edit_last_turn=True,
+        ...         zero_reward=True,
+        ...         undo_done=True
+        ...     )
+        ... )
+        >>>
+        >>> # Test with wrong answer (low reward)
+        >>> reset = env.reset()
+        >>> wrong_answer = (
+        ...     "<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>"
+        ...     "<answer>322 clips</answer><|im_end|>"
+        ... )
+        >>> reset["text_response"] = [wrong_answer]
+        >>> s = env.step(reset)
+        >>> assert (s["next", "reward"] == 0).all()  # Reward zeroed
+        >>> assert (s["next", "done"] == 0).all()    # Done undone
+        >>> assert s["next", "history"].shape == (1, 3)  # History modified
+        >>>
+        >>> # Test with correct answer (high reward)
+        >>> reset = env.reset()
+        >>> correct_answer = (
+        ...     "<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>"
+        ...     "<answer>72</answer><|im_end|>"
+        ... )
+        >>> reset["text_response"] = [correct_answer]
+        >>> s = env.step(reset)
+        >>> assert (s["next", "reward"] != 0).all()  # Reward not zeroed
+        >>> assert s["next", "done"].all()           # Done remains True
+        >>> assert s["next", "history"].shape == (1, 3)  # History unchanged
+    """
+
+    # Predefined thinking prompts
+    DEFAULT_PROMPTS_EG = [
+        "But wait, let me think about this more carefully...",
+        "Actually, let me reconsider this...",
+        "But we can do better. Let me think about it step by step...",
+        "Wait, I need to double-check my reasoning...",
+        "Actually, let me think about it more carefully...",
+        "It looks like I made a mistake. Let me think about it step by step...",
+    ]
+    DEFAULT_PROMPTS_COG = [
+        "But wait, think about this more carefully...",
+        "Actually, reconsider this...",
+        "But we can do better. Let's think about it step by step...",
+        "Wait, you need to double-check your reasoning...",
+        "Actually, think about it more carefully...",
+        "It looks like you made a mistake. Can you see what went wrong? Let's think about it step by step...",
+    ]
+
+    def __init__(
+        self,
+        cond: Callable[[TensorDictBase], bool],
+        prompt: str | None = None,
+        random_prompt: bool = False,
+        role: Literal["user", "assistant"] = "assistant",
+        edit_last_turn: bool = True,
+        zero_reward: bool | None = None,
+        undo_done: bool = True,
+        egocentric: bool | None = None,
+    ) -> None:
+        super().__init__()
+
+        # Set condition and role
+        self.cond = cond
+        self.role = role
+        if egocentric is None:
+            egocentric = role == "assistant"
+        self.egocentric = egocentric
+
+        # Set the prompt
+        if prompt is None:
+            prompt = (
+                self.DEFAULT_PROMPTS_EG[0]
+                if egocentric
+                else self.DEFAULT_PROMPTS_COG[0]
+            )
+        self._prompt = prompt
+        self.random_prompt = random_prompt
+
+        # Validate edit_last_turn constraint
+        if edit_last_turn and role != "assistant":
+            raise ValueError("edit_last_turn can only be used with role='assistant'")
+        self.edit_last_turn = edit_last_turn
+
+        # Set zero_reward behavior
+        if zero_reward is None:
+            zero_reward = edit_last_turn
+        self.zero_reward = zero_reward
+        self.undo_done = undo_done
+
+    @property
+    def prompt(self) -> str:
+        if self.random_prompt:
+            import random
+
+            return random.choice(
+                self.DEFAULT_PROMPTS_EG if self.egocentric else self.DEFAULT_PROMPTS_COG
+            )
+        return self._prompt
+
+    def _step(
+        self, tensordict: TensorDictBase, next_tensordict: TensorDictBase
+    ) -> TensorDictBase:
+        """Process the tensordict and add thinking prompts based on the condition.
+
+        Args:
+            tensordict: The current tensordict
+            next_tensordict: The next tensordict containing the most recent history and reward
+
+        Returns:
+            The modified next_tensordict
+        """
+        # Handle batch dimensions
+        if next_tensordict.batch_dims >= 1:
+            ntds = []
+            for td, next_td in zip(tensordict.unbind(0), next_tensordict.unbind(0)):
+                ntds.append(self._step(td, next_td))
+            next_tensordict.update(lazy_stack(ntds))
+            return next_tensordict
+
+        # Check that base_env is on history mode
+        parent = self.parent
+        if parent is None:
+            raise RuntimeError("AddThinkingPrompt must be used with a ChatEnv")
+        base_env = parent.base_env
+        if base_env.input_mode != "history":
+            raise RuntimeError(
+                "AddThinkingPrompt must be used with a ChatEnv in history mode"
+            )
+
+        # Check if we should add the thinking prompt
+        if self.cond(next_tensordict):
+            torchrl_logger.info("Adding thinking prompt.")
+            history: History = next_tensordict["history"].prompt
+            last_turn = history[..., -1]
+
+            if self.edit_last_turn:
+
+                # Edit the last assistant response
+                content = last_turn.content
+                modified_content = self._replace_answer_with_prompt(content)
+
+                # Create new history entry with modified content
+                new_turn = History(
+                    role="assistant",
+                    content=modified_content,
+                    batch_size=last_turn.batch_size,
+                    device=last_turn.device,
+                )
+
+                # Replace the last turn in history
+                history = history[..., :-1].append(new_turn)
+                next_tensordict["history"].prompt = history
+
+            else:
+                # Add a new message
+                prompt = self.prompt
+
+                history = history.append(History(role=self.role, content=prompt))
+                next_tensordict["history"].prompt = history
+
+            if self.undo_done:
+                parent: EnvBase = self.parent
+                if parent is not None:
+                    done_keys = parent.done_keys
+                    for key in done_keys:
+                        done = next_tensordict.get(key)
+                        if done is not None:
+                            next_tensordict.set(key, done.zero_())
+
+            # Zero out reward if requested
+            if self.zero_reward:
+                parent: EnvBase = self.parent
+                if parent is not None:
+                    reward_keys = parent.reward_keys
+                    for key in reward_keys:
+                        reward = next_tensordict.get(key)
+                        if reward is not None:
+                            next_tensordict.set(key, reward.zero_())
+        else:
+            torchrl_logger.info("Not adding thinking prompt.")
+        return next_tensordict
+
+    def _replace_answer_with_prompt(self, content: str) -> str:
+        """Replace the last answer section with a thinking prompt.
+
+        This method uses regex to find and replace the last <answer>...</answer> section
+        with the thinking prompt, preserving any content before the answer tag.
+        Only the last answer block is replaced to avoid interfering with earlier
+        examples or instructions that might contain answer tags.
+
+        Args:
+            content: The original content string
+
+        Returns:
+            The modified content with the last answer replaced by the thinking prompt
+        """
+        # Pattern to match <answer>...</answer> with optional EOS token
+        # Use non-greedy matching and be more specific about the end
+        answer_pattern = r"<answer>.*?</answer>(?:\s*<\|im_end\|>)?"
+
+        # Check if there's an answer tag
+        if "<answer>" in content:
+            # Find all matches to get the last one
+            matches = list(re.finditer(answer_pattern, content, flags=re.DOTALL))
+
+            if matches:
+                # Get the last match
+                last_match = matches[-1]
+                start, end = last_match.span()
+
+                # Replace only the last answer section with the thinking prompt
+                prompt = self.prompt
+                modified_content = content[:start] + prompt + content[end:]
+
+                # Clean up any trailing whitespace
+                modified_content = modified_content.rstrip()
+
+                # Ensure we end with the EOS token if the original content had it
+                if content.endswith("<|im_end|>"):
+                    modified_content = modified_content.rstrip() + "<|im_end|>"
+
+                # Ensure proper spacing around the prompt
+                if not modified_content.endswith(prompt):
+                    # If the prompt wasn't properly inserted, append it
+                    modified_content = content.rstrip()
+                    if modified_content.endswith("<|im_end|>"):
+                        modified_content = modified_content[
+                            : -len("<|im_end|>")
+                        ].rstrip()
+                    modified_content = modified_content + "\n\n" + prompt + "<|im_end|>"
+            else:
+                # No matches found, just append the prompt
+                prompt = self.prompt
+                modified_content = content.rstrip() + "\n\n" + prompt
+
+        else:
+            # No answer tag found, just append the prompt
+            prompt = self.prompt
+            modified_content = content.rstrip() + "\n\n" + prompt
+
+        return modified_content
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        """Reset the transform state.
+
+        Args:
+            tensordict: The current tensordict
+            tensordict_reset: The reset tensordict
+
+        Returns:
+            The reset tensordict
+        """
+        return tensordict_reset