torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

torchrl/modules/models/recipes/impala.py

@@ -0,0 +1,185 @@
+# 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 torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tensordict import TensorDictBase
+
+
+# TODO: code small architecture ref in Impala paper
+
+
+class _ResNetBlock(nn.Module):
+    def __init__(
+        self,
+        num_ch,
+    ):
+        super().__init__()
+        resnet_block = []
+        resnet_block.append(nn.ReLU(inplace=True))
+        resnet_block.append(
+            nn.LazyConv2d(
+                out_channels=num_ch,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+        )
+        resnet_block.append(nn.ReLU(inplace=True))
+        resnet_block.append(
+            nn.Conv2d(
+                in_channels=num_ch,
+                out_channels=num_ch,
+                kernel_size=3,
+                stride=1,
+                padding=1,
+            )
+        )
+        self.seq = nn.Sequential(*resnet_block)
+
+    def forward(self, x):
+        x += self.seq(x)
+        return x
+
+
+class _ConvNetBlock(nn.Module):
+    def __init__(self, num_ch):
+        super().__init__()
+
+        conv = nn.LazyConv2d(
+            out_channels=num_ch,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+        mp = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.feats_conv = nn.Sequential(conv, mp)
+        self.resnet1 = _ResNetBlock(num_ch=num_ch)
+        self.resnet2 = _ResNetBlock(num_ch=num_ch)
+
+    def forward(self, x):
+        x = self.feats_conv(x)
+        x = self.resnet1(x)
+        x = self.resnet1(x)
+        return x
+
+
+class ImpalaNet(nn.Module):  # noqa: D101
+    def __init__(
+        self,
+        num_actions,
+        channels=(16, 32, 32),
+        out_features=256,
+        use_lstm=False,
+        batch_first=True,
+        clamp_reward=True,
+        one_hot=False,
+    ):
+        super().__init__()
+        self.batch_first = batch_first
+        self.use_lstm = use_lstm
+        self.clamp_reward = clamp_reward
+        self.one_hot = one_hot
+        self.num_actions = num_actions
+
+        layers = [_ConvNetBlock(num_ch) for num_ch in channels]
+        layers += [nn.ReLU(inplace=True)]
+        self.convs = nn.Sequential(*layers)
+        self.fc = nn.Sequential(nn.LazyLinear(out_features), nn.ReLU(inplace=True))
+
+        # FC output size + last reward.
+        core_output_size = out_features + 1
+
+        if use_lstm:
+            self.core = nn.LSTM(
+                core_output_size,
+                out_features,
+                num_layers=1,
+                batch_first=batch_first,
+            )
+            core_output_size = out_features
+
+        self.policy = nn.Linear(core_output_size, self.num_actions)
+        self.baseline = nn.Linear(core_output_size, 1)
+
+    def forward(self, x, reward, done, core_state=None, mask=None):  # noqa: D102
+        if self.batch_first:
+            B, T, *x_shape = x.shape
+            batch_shape = torch.Size([B, T])
+        else:
+            T, B, *x_shape = x.shape
+            batch_shape = torch.Size([T, B])
+        if mask is None:
+            x = x.view(-1, *x.shape[-3:])
+        else:
+            x = x[mask]
+            if x.ndimension() != 4:
+                raise RuntimeError(
+                    f"masked input should have 4 dimensions but got {x.ndimension()} instead"
+                )
+        x = self.convs(x)
+        x = x.view(B * T, -1)
+        x = self.fc(x)
+
+        if mask is None:
+            if self.batch_first:
+                x = x.view(B, T, -1)
+            else:
+                x = x.view(T, B, -1)
+        else:
+            x = self._allocate_masked_x(x, mask)
+
+        if self.clamp_reward:
+            reward = torch.clamp(reward, -1, 1)
+        reward = reward.unsqueeze(-1)
+
+        core_input = torch.cat([x, reward], dim=-1)
+
+        if self.use_lstm:
+            core_output, _ = self.core(core_input, core_state)
+        else:
+            core_output = core_input
+
+        policy_logits = self.policy(core_output)
+        baseline = self.baseline(core_output)
+
+        softmax_vals = F.softmax(policy_logits, dim=-1)
+        action = torch.multinomial(
+            softmax_vals.view(-1, softmax_vals.shape[-1]), num_samples=1
+        ).view(softmax_vals.shape[:-1])
+        if self.one_hot:
+            action = F.one_hot(action, policy_logits.shape[-1])
+
+        if policy_logits.shape[:2] != batch_shape:
+            raise RuntimeError("policy_logits and batch-shape mismatch")
+        if baseline.shape[:2] != batch_shape:
+            raise RuntimeError("baseline and batch-shape mismatch")
+        if action.shape[:2] != batch_shape:
+            raise RuntimeError("action and batch-shape mismatch")
+
+        return (action, policy_logits, baseline), core_state
+
+    def _allocate_masked_x(self, x, mask):
+        x_empty = torch.zeros(
+            *mask.shape[:2], x.shape[-1], device=x.device, dtype=x.dtype
+        )
+        x_empty[mask] = x
+        return x_empty
+
+
+class ImpalaNetTensorDict(ImpalaNet):  # noqa: D101
+    observation_key = "pixels"
+
+    def forward(self, tensordict: TensorDictBase):  # noqa: D102
+        x = tensordict.get(self.observation_key)
+        done = tensordict.get("done").squeeze(-1)
+        reward = tensordict.get("reward").squeeze(-1)
+        mask = tensordict.get(("collector", "mask"))
+        core_state = (
+            tensordict.get("core_state") if "core_state" in tensordict.keys() else None
+        )
+        return super().forward(x, reward, done, core_state=core_state, mask=mask)

torchrl/modules/models/utils.py

@@ -0,0 +1,162 @@
+# 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 inspect
+import warnings
+from collections.abc import Callable, Sequence
+
+import torch
+from torch import nn
+from torchrl.data.utils import DEVICE_TYPING
+from torchrl.modules.models.exploration import NoisyLazyLinear, NoisyLinear
+
+LazyMapping = {
+    nn.Linear: nn.LazyLinear,
+    NoisyLinear: NoisyLazyLinear,
+}
+
+
+class SqueezeLayer(nn.Module):
+    """Squeezing layer.
+
+    Squeezes some given singleton dimensions of an input tensor.
+
+    Args:
+         dims (iterable): dimensions to be squeezed
+            default: (-1,)
+
+    """
+
+    def __init__(self, dims: Sequence[int] = (-1,)):
+        super().__init__()
+        for dim in dims:
+            if dim >= 0:
+                raise RuntimeError("dims must all be < 0")
+        self.dims = dims
+
+    def forward(self, input: torch.Tensor) -> torch.Tensor:  # noqa: D102
+        for dim in self.dims:
+            if input.shape[dim] != 1:
+                raise RuntimeError(
+                    f"Tried to squeeze an input over dims {self.dims} with shape {input.shape}"
+                )
+            input = input.squeeze(dim)
+        return input
+
+
+class Squeeze2dLayer(SqueezeLayer):
+    """Squeezing layer for convolutional neural networks.
+
+    Squeezes the last two singleton dimensions of an input tensor.
+
+    """
+
+    def __init__(self):
+        super().__init__((-2, -1))
+
+
+class SquashDims(nn.Module):
+    """A squashing layer.
+
+    Flattens the N last dimensions of an input tensor.
+
+    Args:
+        ndims_in (int): number of dimensions to be flattened.
+            default = 3
+
+    Examples:
+        >>> from torchrl.modules.models.utils import SquashDims
+        >>> import torch
+        >>> x = torch.randn(1, 2, 3, 4)
+        >>> print(SquashDims()(x).shape)
+        torch.Size([1, 24])
+
+    """
+
+    def __init__(self, ndims_in: int = 3):
+        super().__init__()
+        self.ndims_in = ndims_in
+
+    def forward(self, value: torch.Tensor) -> torch.Tensor:
+        value = value.flatten(-self.ndims_in, -1)
+        return value
+
+
+def _find_depth(depth: int | None, *list_or_ints: Sequence):
+    """Find depth based on a sequence of inputs and a depth indicator.
+
+    If the depth is None, it is inferred by the length of one (or more) matching
+    lists of integers.
+    Raises an exception if depth does not match the list lengths or if lists lengths
+    do not match.
+
+    Args:
+        depth (int, optional): depth of the network
+        *list_or_ints (lists of int or int): if depth is None, at least one of
+            these inputs must be a list of ints of the length of the desired
+            network.
+    """
+    if depth is None:
+        for item in list_or_ints:
+            if isinstance(item, (list, tuple)):
+                depth = len(item)
+    if depth is None:
+        raise ValueError(
+            f"depth=None requires one of the input args (kernel_sizes, strides, "
+            f"num_cells) to be a a list or tuple. Got {tuple(type(item) for item in list_or_ints)}"
+        )
+    return depth
+
+
+def create_on_device(
+    module_class: type[nn.Module] | Callable,
+    device: DEVICE_TYPING | None,
+    *args,
+    **kwargs,
+) -> nn.Module:
+    """Create a new instance of :obj:`module_class` on :obj:`device`.
+
+    The new instance is created directly on the device if its constructor supports this.
+
+    Args:
+        module_class (Type[nn.Module]): the class of module to be created.
+        device (DEVICE_TYPING): device to create the module on.
+        *args: positional arguments to be passed to the module constructor.
+        **kwargs: keyword arguments to be passed to the module constructor.
+
+    """
+    fullargspec = inspect.getfullargspec(module_class.__init__)
+    if "device" in fullargspec.args or "device" in fullargspec.kwonlyargs:
+        return module_class(*args, device=device, **kwargs)
+    else:
+        result = module_class(*args, **kwargs)
+        if hasattr(result, "to"):
+            result = result.to(device)
+        return result
+
+
+def _reset_parameters_recursive(module, warn_if_no_op: bool = True) -> bool:
+    """Recursively resets the parameters of a :class:`~torch.nn.Module` in-place.
+
+    Args:
+        module (torch.nn.Module): the module to reset.
+        warn_if_no_op (bool, optional): whether to raise a warning in case this is a no-op.
+            Defaults to ``True``.
+
+    Returns: whether any parameter has been reset.
+
+    """
+    any_reset = False
+    for layer in module.children():
+        if hasattr(layer, "reset_parameters"):
+            layer.reset_parameters()
+            any_reset |= True
+        any_reset |= _reset_parameters_recursive(layer, warn_if_no_op=False)
+    if warn_if_no_op and not any_reset:
+        warnings.warn(
+            "_reset_parameters_recursive was called without the parameters argument and did not find any parameters to reset"
+        )
+    return any_reset

torchrl/modules/planners/__init__.py

@@ -0,0 +1,10 @@
+# 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 .cem import CEMPlanner
+from .common import MPCPlannerBase
+from .mppi import MPPIPlanner
+
+__all__ = ["CEMPlanner", "MPCPlannerBase", "MPPIPlanner"]

torchrl/modules/planners/cem.py

@@ -0,0 +1,228 @@
+# 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 typing import TYPE_CHECKING
+
+import torch
+from tensordict import TensorDict, TensorDictBase
+from torchrl.modules.planners.common import MPCPlannerBase
+
+if TYPE_CHECKING:
+    from torchrl.envs.common import EnvBase
+
+
+class CEMPlanner(MPCPlannerBase):
+    """CEMPlanner Module.
+
+    Reference: The cross-entropy method for optimization, Botev et al. 2013
+
+    This module will perform a CEM planning step when given a TensorDict
+    containing initial states.
+    The CEM planning step is performed by sampling actions from a Gaussian
+    distribution with zero mean and unit variance.
+    The sampled actions are then used to perform a rollout in the environment.
+    The cumulative rewards obtained with the rollout is then
+    ranked. We select the top-k episodes and use their actions to update the
+    mean and standard deviation of the actions distribution.
+    The CEM planning step is repeated for a specified number of steps.
+
+    A call to the module returns the actions that empirically maximised the
+    returns given a planning horizon
+
+    Args:
+        env (EnvBase): The environment to perform the planning step on (can be
+            `ModelBasedEnv` or :obj:`EnvBase`).
+        planning_horizon (int): The length of the simulated trajectories
+        optim_steps (int): The number of optimization steps used by the MPC
+            planner
+        num_candidates (int): The number of candidates to sample from the
+            Gaussian distributions.
+        top_k (int): The number of top candidates to use to
+            update the mean and standard deviation of the Gaussian distribution.
+        reward_key (str, optional): The key in the TensorDict to use to
+            retrieve the reward. Defaults to "reward".
+        action_key (str, optional): The key in the TensorDict to use to store
+            the action. Defaults to "action"
+
+    Examples:
+        >>> from tensordict import TensorDict
+        >>> from torchrl.data import Composite, Unbounded
+        >>> from torchrl.envs.model_based import ModelBasedEnvBase
+        >>> from torchrl.modules import SafeModule
+        >>> class MyMBEnv(ModelBasedEnvBase):
+        ...     def __init__(self, world_model, device="cpu", dtype=None, batch_size=None):
+        ...         super().__init__(world_model, device=device, dtype=dtype, batch_size=batch_size)
+        ...         self.state_spec = Composite(
+        ...             hidden_observation=Unbounded((4,))
+        ...         )
+        ...         self.observation_spec = Composite(
+        ...             hidden_observation=Unbounded((4,))
+        ...         )
+        ...         self.action_spec = Unbounded((1,))
+        ...         self.reward_spec = Unbounded((1,))
+        ...
+        ...     def _reset(self, tensordict: TensorDict) -> TensorDict:
+        ...         tensordict = TensorDict(
+        ...             {},
+        ...             batch_size=self.batch_size,
+        ...             device=self.device,
+        ...         )
+        ...         tensordict = tensordict.update(
+        ...             self.full_state_spec.rand())
+        ...         tensordict = tensordict.update(
+        ...             self.full_action_spec.rand())
+        ...         tensordict = tensordict.update(
+        ...             self.full_observation_spec.rand())
+        ...         return tensordict
+        ...
+        >>> from torchrl.modules import MLP, WorldModelWrapper
+        >>> import torch.nn as nn
+        >>> world_model = WorldModelWrapper(
+        ...     SafeModule(
+        ...         MLP(out_features=4, activation_class=nn.ReLU, activate_last_layer=True, depth=0),
+        ...         in_keys=["hidden_observation", "action"],
+        ...         out_keys=["hidden_observation"],
+        ...     ),
+        ...     SafeModule(
+        ...         nn.Linear(4, 1),
+        ...         in_keys=["hidden_observation"],
+        ...         out_keys=["reward"],
+        ...     ),
+        ... )
+        >>> env = MyMBEnv(world_model)
+        >>> # Build a planner and use it as actor
+        >>> planner = CEMPlanner(env, 10, 11, 7, 3)
+        >>> env.rollout(5, planner)
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                hidden_observation: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        hidden_observation: Tensor(shape=torch.Size([5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                        reward: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([5]),
+                    device=cpu,
+                    is_shared=False),
+                terminated: Tensor(shape=torch.Size([5, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([5]),
+            device=cpu,
+            is_shared=False)
+    """
+
+    def __init__(
+        self,
+        env: EnvBase,
+        planning_horizon: int,
+        optim_steps: int,
+        num_candidates: int,
+        top_k: int,
+        reward_key: str = ("next", "reward"),
+        action_key: str = "action",
+    ):
+        super().__init__(env=env, action_key=action_key)
+        self.planning_horizon = planning_horizon
+        self.optim_steps = optim_steps
+        self.num_candidates = num_candidates
+        self.top_k = top_k
+        self.reward_key = reward_key
+
+    def planning(self, tensordict: TensorDictBase) -> torch.Tensor:
+        batch_size = tensordict.batch_size
+        action_shape = (
+            *batch_size,
+            self.num_candidates,
+            self.planning_horizon,
+            *self.action_spec.shape,
+        )
+        action_stats_shape = (
+            *batch_size,
+            1,
+            self.planning_horizon,
+            *self.action_spec.shape,
+        )
+        action_topk_shape = (
+            *batch_size,
+            self.top_k,
+            self.planning_horizon,
+            *self.action_spec.shape,
+        )
+        TIME_DIM = len(self.action_spec.shape) - 3
+        K_DIM = len(self.action_spec.shape) - 4
+        expanded_original_tensordict = (
+            tensordict.unsqueeze(-1)
+            .expand(*batch_size, self.num_candidates)
+            .to_tensordict()
+        )
+        _action_means = torch.zeros(
+            *action_stats_shape,
+            device=tensordict.device,
+            dtype=self.env.action_spec.dtype,
+        )
+        _action_stds = torch.ones_like(_action_means)
+        container = TensorDict(
+            {
+                "tensordict": expanded_original_tensordict,
+                "stats": TensorDict(
+                    {
+                        "_action_means": _action_means,
+                        "_action_stds": _action_stds,
+                    },
+                    [*batch_size, 1, self.planning_horizon],
+                ),
+            },
+            batch_size,
+        )
+
+        for _ in range(self.optim_steps):
+            actions_means = container.get(("stats", "_action_means"))
+            actions_stds = container.get(("stats", "_action_stds"))
+            actions = actions_means + actions_stds * torch.randn(
+                *action_shape,
+                device=actions_means.device,
+                dtype=actions_means.dtype,
+            )
+            actions = self.env.action_spec.project(actions)
+            optim_tensordict = container.get("tensordict").clone()
+            policy = _PrecomputedActionsSequentialSetter(actions)
+            optim_tensordict = self.env.rollout(
+                max_steps=self.planning_horizon,
+                policy=policy,
+                auto_reset=False,
+                tensordict=optim_tensordict,
+            )
+
+            sum_rewards = optim_tensordict.get(self.reward_key).sum(
+                dim=TIME_DIM, keepdim=True
+            )
+            _, top_k = sum_rewards.topk(self.top_k, dim=K_DIM)
+            top_k = top_k.expand(action_topk_shape)
+            best_actions = actions.gather(K_DIM, top_k)
+            container.set_(
+                ("stats", "_action_means"), best_actions.mean(dim=K_DIM, keepdim=True)
+            )
+            container.set_(
+                ("stats", "_action_stds"), best_actions.std(dim=K_DIM, keepdim=True)
+            )
+        action_means = container.get(("stats", "_action_means"))
+        return action_means[..., 0, 0, :]
+
+
+class _PrecomputedActionsSequentialSetter:
+    def __init__(self, actions):
+        self.actions = actions
+        self.cmpt = 0
+
+    def __call__(self, tensordict):
+        # checks that the step count is lower or equal to the horizon
+        if self.cmpt >= self.actions.shape[-2]:
+            raise ValueError("Precomputed actions sequence is too short")
+        tensordict = tensordict.set("action", self.actions[..., self.cmpt, :])
+        self.cmpt += 1
+        return tensordict

torchrl/modules/planners/common.py

@@ -0,0 +1,73 @@
+# 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 abc
+from typing import TYPE_CHECKING
+
+import torch
+from tensordict import TensorDictBase
+
+from torchrl.modules import SafeModule
+
+if TYPE_CHECKING:
+    from torchrl.envs.common import EnvBase
+
+
+class MPCPlannerBase(SafeModule, metaclass=abc.ABCMeta):
+    """MPCPlannerBase abstract Module.
+
+    This class inherits from :obj:`SafeModule`. Provided a :obj:`TensorDict`, this module will perform a Model Predictive Control (MPC) planning step.
+    At the end of the planning step, the :obj:`MPCPlanner` will return a proposed action.
+
+    Args:
+        env (EnvBase): The environment to perform the planning step on (Can be :obj:`ModelBasedEnvBase` or :obj:`EnvBase`).
+        action_key (str, optional): The key that will point to the computed action.
+    """
+
+    def __init__(
+        self,
+        env: EnvBase,
+        action_key: str = "action",
+    ):
+        # Check if env is stateless
+        if env.batch_locked:
+            raise ValueError(
+                "Environment is batch_locked. MPCPlanners need an environment that accepts batched inputs with any batch size"
+            )
+        out_keys = [action_key]
+        in_keys = list(env.observation_spec.keys(True, True))
+        super().__init__(env, in_keys=in_keys, out_keys=out_keys)
+        self.env = env
+        self.action_spec = env.action_spec
+        self.to(env.device)
+
+    @abc.abstractmethod
+    def planning(self, td: TensorDictBase) -> torch.Tensor:
+        """Performs the MPC planning step.
+
+        Args:
+            td (TensorDict): The TensorDict to perform the planning step on.
+        """
+        raise NotImplementedError()
+
+    def forward(
+        self,
+        tensordict: TensorDictBase,
+        tensordict_out: TensorDictBase | None = None,
+        **kwargs,
+    ) -> TensorDictBase:
+        if "params" in kwargs or "vmap" in kwargs:
+            raise ValueError(
+                "MPCPlannerBase does not currently support functional programming."
+            )
+        action = self.planning(tensordict)
+        action = self.action_spec.project(action)
+        tensordict_out = self._write_to_tensordict(
+            tensordict,
+            (action,),
+            tensordict_out,
+        )
+        return tensordict_out