torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

torchrl/envs/custom/llm.py

@@ -0,0 +1,214 @@
+# 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
+
+from collections.abc import Callable
+
+import torch
+
+from tensordict import NestedKey, set_list_to_stack, TensorDict, TensorDictBase
+from tensordict.tensorclass import NonTensorData, NonTensorStack
+
+from torchrl.data.map.hash import SipHash
+from torchrl.data.tensor_specs import (
+    Categorical as CategoricalSpec,
+    Composite,
+    NonTensor,
+    Unbounded,
+)
+from torchrl.envs import EnvBase
+from torchrl.envs.utils import _StepMDP
+
+
+class LLMHashingEnv(EnvBase):
+    """A text generation environment that uses a hashing module to identify unique observations.
+
+    The primary goal of this environment is to identify token chains using a hashing function.
+    This allows the data to be stored in a :class:`~torchrl.data.MCTSForest` using nothing but hashes as node
+    identifiers, or easily prune repeated token chains in a data structure.
+
+    .. The following figure gives an overview of this workflow:
+    .. .. figure:: /_static/img/rollout-llm.png
+    ..     :alt: Data collection loop with our LLM environment.
+
+    Args:
+        vocab_size (int): The size of the vocabulary. Can be omitted if the tokenizer is passed.
+
+    Keyword Args:
+        hashing_module (Callable[[torch.Tensor], torch.Tensor], optional):
+            A hashing function that takes a tensor as input and returns a hashed tensor.
+            Defaults to :class:`~torchrl.data.SipHash` if not provided.
+        observation_key (NestedKey, optional): The key for the observation in the TensorDict.
+            Defaults to "observation".
+        text_output (bool, optional): Whether to include the text output in the observation.
+            Defaults to `True`.
+        tokenizer (transformers.Tokenizer | None, optional):
+            A tokenizer function that converts text to tensors.
+            Only used when `text_output` is `True`.
+            Must implement the following methods: `decode` and `batch_decode`.
+            Defaults to ``None``.
+        text_key (NestedKey | None, optional): The key for the text output in the TensorDict.
+            Defaults to "text".
+
+    Examples:
+        >>> from tensordict import TensorDict
+        >>> from torchrl.envs import LLMHashingEnv
+        >>> from transformers import GPT2Tokenizer
+        >>> tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
+        >>> x = tokenizer(["Check out TorchRL!"])["input_ids"]
+        >>> env = LLMHashingEnv(tokenizer=tokenizer)
+        >>> td = TensorDict(observation=x, batch_size=[1])
+        >>> td = env.reset(td)
+        >>> print(td)
+        TensorDict(
+            fields={
+                done: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                hash: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                observation: Tensor(shape=torch.Size([1, 5]), device=cpu, dtype=torch.int64, is_shared=False),
+                terminated: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                text: NonTensorStack(
+                    ['Check out TorchRL!'],
+                    batch_size=torch.Size([1]),
+                    device=None)},
+            batch_size=torch.Size([1]),
+            device=None,
+            is_shared=False)
+
+    """
+
+    def __init__(
+        self,
+        vocab_size: int | None = None,
+        *,
+        hashing_module: Callable[[torch.Tensor], torch.Tensor] = None,
+        observation_key: NestedKey = "observation",
+        text_output: bool = True,
+        tokenizer: Callable[[str | list[str]], torch.Tensor] | None = None,
+        text_key: NestedKey | None = "text",
+    ):
+        super().__init__()
+        if vocab_size is None:
+            if tokenizer is None:
+                raise TypeError(
+                    "You must provide a vocab_size integer if tokenizer is `None`."
+                )
+            vocab_size = tokenizer.vocab_size
+        self._batch_locked = False
+        if hashing_module is None:
+            hashing_module = SipHash()
+
+        self._hashing_module = hashing_module
+        self._tokenizer = tokenizer
+        self.observation_key = observation_key
+        observation_spec = {
+            observation_key: CategoricalSpec(n=vocab_size, shape=(-1,)),
+            "hashing": Unbounded(shape=(1,), dtype=torch.int64),
+        }
+        self.text_output = text_output
+        if not text_output:
+            text_key = None
+        elif text_key is None:
+            text_key = "text"
+        if text_key is not None:
+            observation_spec[text_key] = NonTensor(shape=())
+            self.text_key = text_key
+        self.observation_spec = Composite(observation_spec)
+        self.action_spec = Composite(action=CategoricalSpec(vocab_size, shape=(1,)))
+        _StepMDP(self)
+
+    @set_list_to_stack(True)
+    def make_tensordict(self, input: str | list[str]) -> TensorDict:
+        """Converts a string or list of strings in a TensorDict with appropriate shape and device."""
+        list_len = len(input) if isinstance(input, list) else 0
+        tensordict = TensorDict(
+            {self.observation_key: self._tokenizer(input)}, device=self.device
+        )
+        if list_len:
+            tensordict.batch_size = [list_len]
+        return self.reset(tensordict)
+
+    def _reset(self, tensordict: TensorDictBase):
+        """Initializes the environment with a given observation.
+
+        Args:
+            tensordict (TensorDictBase): A TensorDict containing the initial observation.
+
+        Returns:
+            A TensorDict containing the initial observation, its hash, and other relevant information.
+
+        """
+        out = tensordict.empty()
+        obs = tensordict.get(self.observation_key, None)
+        if obs is None:
+            raise RuntimeError(
+                f"Resetting the {type(self).__name__} environment requires a prompt."
+            )
+        if self.text_output:
+            if obs.ndim > 1:
+                text = self._tokenizer.batch_decode(obs)
+                text = NonTensorStack.from_list(text)
+            else:
+                text = self._tokenizer.decode(obs)
+                text = NonTensorData(text)
+            out.set(self.text_key, text)
+
+        if obs.ndim > 1:
+            out.set("hashing", self._hashing_module(obs).unsqueeze(-1))
+        else:
+            out.set("hashing", self._hashing_module(obs.unsqueeze(0)).transpose(0, -1))
+
+        if not self.full_done_spec.is_empty():
+            out.update(self.full_done_spec.zero(tensordict.shape))
+        else:
+            out.set("done", torch.zeros((*tensordict.batch_size, 1), dtype=torch.bool))
+            out.set(
+                "terminated", torch.zeros((*tensordict.batch_size, 1), dtype=torch.bool)
+            )
+        return out
+
+    def _step(self, tensordict):
+        """Takes an action (i.e., the next token to generate) and returns the next observation and reward.
+
+        Args:
+            tensordict: A TensorDict containing the current observation and action.
+
+        Returns:
+            A TensorDict containing the next observation, its hash, and other relevant information.
+        """
+        out = tensordict.empty()
+        action = tensordict.get("action")
+        obs = torch.cat([tensordict.get(self.observation_key), action], -1)
+        kwargs = {self.observation_key: obs}
+
+        catval = torch.cat([tensordict.get("hashing"), action], -1)
+        if obs.ndim > 1:
+            new_hash = self._hashing_module(catval).unsqueeze(-1)
+        else:
+            new_hash = self._hashing_module(catval.unsqueeze(0)).transpose(0, -1)
+
+        if self.text_output:
+            if obs.ndim > 1:
+                text = self._tokenizer.batch_decode(obs)
+                text = NonTensorStack.from_list(text)
+            else:
+                text = self._tokenizer.decode(obs)
+                text = NonTensorData(text)
+            kwargs[self.text_key] = text
+        kwargs.update(
+            {
+                "hashing": new_hash,
+                "done": torch.zeros((*tensordict.batch_size, 1), dtype=torch.bool),
+                "terminated": torch.zeros(
+                    (*tensordict.batch_size, 1), dtype=torch.bool
+                ),
+            }
+        )
+        return out.update(kwargs)
+
+    def _set_seed(self, *args) -> None:
+        """Sets the seed for the environment's randomness.
+
+        .. note:: This environment has no randomness, so this method does nothing.
+        """

torchrl/envs/custom/pendulum.py

@@ -0,0 +1,401 @@
+# 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 numpy as np
+
+import torch
+from tensordict import TensorDict, TensorDictBase
+from torchrl.data.tensor_specs import Bounded, Composite, Unbounded
+from torchrl.envs.common import EnvBase
+from torchrl.envs.utils import make_composite_from_td
+
+
+class PendulumEnv(EnvBase):
+    """A stateless Pendulum environment.
+
+    See the Pendulum tutorial for more details: :ref:`tutorial <pendulum_tuto>`.
+
+    Specs:
+        >>> env = PendulumEnv()
+        >>> env.specs
+        Composite(
+            output_spec: Composite(
+                full_observation_spec: Composite(
+                    th: BoundedContinuous(
+                        shape=torch.Size([]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    thdot: BoundedContinuous(
+                        shape=torch.Size([]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    params: Composite(
+                        max_speed: UnboundedDiscrete(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.int64,
+                            domain=discrete),
+                        max_torque: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        dt: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        g: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        m: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        l: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        device=None,
+                        shape=torch.Size([])),
+                    device=None,
+                    shape=torch.Size([])),
+                full_reward_spec: Composite(
+                    reward: UnboundedContinuous(
+                        shape=torch.Size([1]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    device=None,
+                    shape=torch.Size([])),
+                full_done_spec: Composite(
+                    done: Categorical(
+                        shape=torch.Size([1]),
+                        space=CategoricalBox(n=2),
+                        device=cpu,
+                        dtype=torch.bool,
+                        domain=discrete),
+                    terminated: Categorical(
+                        shape=torch.Size([1]),
+                        space=CategoricalBox(n=2),
+                        device=cpu,
+                        dtype=torch.bool,
+                        domain=discrete),
+                    device=None,
+                    shape=torch.Size([])),
+                device=None,
+                shape=torch.Size([])),
+            input_spec: Composite(
+                full_state_spec: Composite(
+                    th: BoundedContinuous(
+                        shape=torch.Size([]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    thdot: BoundedContinuous(
+                        shape=torch.Size([]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    params: Composite(
+                        max_speed: UnboundedDiscrete(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.int64,
+                            domain=discrete),
+                        max_torque: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        dt: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        g: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        m: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        l: UnboundedContinuous(
+                            shape=torch.Size([]),
+                            space=ContinuousBox(
+                                low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                                high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                        device=None,
+                        shape=torch.Size([])),
+                    device=None,
+                    shape=torch.Size([])),
+                full_action_spec: Composite(
+                    action: BoundedContinuous(
+                        shape=torch.Size([1]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    device=None,
+                    shape=torch.Size([])),
+                device=None,
+                shape=torch.Size([])),
+            device=None,
+            shape=torch.Size([]))
+
+    """
+
+    DEFAULT_X = np.pi
+    DEFAULT_Y = 1.0
+
+    metadata = {
+        "render_modes": ["human", "rgb_array"],
+        "render_fps": 30,
+    }
+    batch_locked = False
+    rng = None
+
+    def __init__(self, td_params=None, seed=None, device=None):
+        if td_params is None:
+            td_params = self.gen_params(device=self.device)
+
+        super().__init__(device=device)
+        self._make_spec(td_params)
+        if seed is None:
+            seed = torch.empty((), dtype=torch.int64).random_(generator=self.rng).item()
+        self.set_seed(seed)
+
+    @classmethod
+    def _step(cls, tensordict):
+        th, thdot = tensordict["th"], tensordict["thdot"]  # th := theta
+
+        g_force = tensordict["params", "g"]
+        mass = tensordict["params", "m"]
+        length = tensordict["params", "l"]
+        dt = tensordict["params", "dt"]
+        u = tensordict["action"].squeeze(-1)
+        u = u.clamp(
+            -tensordict["params", "max_torque"], tensordict["params", "max_torque"]
+        )
+        costs = cls.angle_normalize(th) ** 2 + 0.1 * thdot**2 + 0.001 * (u**2)
+
+        new_thdot = (
+            thdot
+            + (3 * g_force / (2 * length) * th.sin() + 3.0 / (mass * length**2) * u)
+            * dt
+        )
+        new_thdot = new_thdot.clamp(
+            -tensordict["params", "max_speed"], tensordict["params", "max_speed"]
+        )
+        new_th = th + new_thdot * dt
+        reward = -costs.view(*tensordict.shape, 1)
+        done = torch.zeros_like(reward, dtype=torch.bool)
+        out = TensorDict(
+            {
+                "th": new_th,
+                "thdot": new_thdot,
+                "params": tensordict["params"],
+                "reward": reward,
+                "done": done,
+            },
+            tensordict.shape,
+        )
+        return out
+
+    def _reset(self, tensordict):
+        batch_size = (
+            tensordict.batch_size if tensordict is not None else self.batch_size
+        )
+        if tensordict is None or "params" not in tensordict:
+            # if no ``tensordict`` is passed, we generate a single set of hyperparameters
+            # Otherwise, we assume that the input ``tensordict`` contains all the relevant
+            # parameters to get started.
+            tensordict = self.gen_params(batch_size=batch_size, device=self.device)
+        elif "th" in tensordict and "thdot" in tensordict:
+            # we can hard-reset the env too
+            return tensordict
+        out = self._reset_random_data(
+            tensordict.shape, batch_size, tensordict["params"]
+        )
+        return out
+
+    def _reset_random_data(self, shape, batch_size, params):
+
+        high_th = torch.tensor(self.DEFAULT_X, device=self.device)
+        high_thdot = torch.tensor(self.DEFAULT_Y, device=self.device)
+        low_th = -high_th
+        low_thdot = -high_thdot
+
+        # for non batch-locked environments, the input ``tensordict`` shape dictates the number
+        # of simulators run simultaneously. In other contexts, the initial
+        # random state's shape will depend upon the environment batch-size instead.
+        th = (
+            torch.rand(shape, generator=self.rng, device=self.device)
+            * (high_th - low_th)
+            + low_th
+        )
+        thdot = (
+            torch.rand(shape, generator=self.rng, device=self.device)
+            * (high_thdot - low_thdot)
+            + low_thdot
+        )
+        out = TensorDict(
+            {
+                "th": th,
+                "thdot": thdot,
+                "params": params,
+            },
+            batch_size=batch_size,
+        )
+        return out
+
+    def _make_spec(self, td_params):
+        # Under the hood, this will populate self.output_spec["observation"]
+        self.observation_spec = Composite(
+            th=Bounded(
+                low=-torch.pi,
+                high=torch.pi,
+                shape=(),
+                dtype=torch.float32,
+            ),
+            thdot=Bounded(
+                low=-td_params["params", "max_speed"],
+                high=td_params["params", "max_speed"],
+                shape=(),
+                dtype=torch.float32,
+            ),
+            # we need to add the ``params`` to the observation specs, as we want
+            # to pass it at each step during a rollout
+            params=make_composite_from_td(
+                td_params["params"], unsqueeze_null_shapes=False
+            ),
+            shape=(),
+        )
+        # since the environment is stateless, we expect the previous output as input.
+        # For this, ``EnvBase`` expects some state_spec to be available
+        self.state_spec = self.observation_spec.clone()
+        # action-spec will be automatically wrapped in input_spec when
+        # `self.action_spec = spec` will be called supported
+        self.action_spec = Bounded(
+            low=-td_params["params", "max_torque"],
+            high=td_params["params", "max_torque"],
+            shape=(1,),
+            dtype=torch.float32,
+        )
+        self.reward_spec = Unbounded(shape=(*td_params.shape, 1))
+
+    def make_composite_from_td(td):
+        # custom function to convert a ``tensordict`` in a similar spec structure
+        # of unbounded values.
+        composite = Composite(
+            {
+                key: make_composite_from_td(tensor)
+                if isinstance(tensor, TensorDictBase)
+                else Unbounded(
+                    dtype=tensor.dtype, device=tensor.device, shape=tensor.shape
+                )
+                for key, tensor in td.items()
+            },
+            shape=td.shape,
+        )
+        return composite
+
+    def _set_seed(self, seed: int) -> None:
+        rng = torch.Generator(device=self.device)
+        rng.manual_seed(seed)
+        self.rng = rng
+
+    @staticmethod
+    def gen_params(g=10.0, batch_size=None, device=None) -> TensorDictBase:
+        """Returns a ``tensordict`` containing the physical parameters such as gravitational force and torque or speed limits."""
+        if batch_size is None:
+            batch_size = []
+        td = TensorDict(
+            {
+                "params": TensorDict(
+                    {
+                        "max_speed": 8,
+                        "max_torque": 2.0,
+                        "dt": 0.05,
+                        "g": g,
+                        "m": 1.0,
+                        "l": 1.0,
+                    },
+                    [],
+                )
+            },
+            [],
+            device=device,
+        )
+        if batch_size:
+            td = td.expand(batch_size).contiguous()
+        return td
+
+    @staticmethod
+    def angle_normalize(x):
+        return ((x + torch.pi) % (2 * torch.pi)) - torch.pi