torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/envs/model_based/__init__.py

@@ -0,0 +1,9 @@
+# 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 .common import ModelBasedEnvBase
+from .dreamer import DreamerDecoder, DreamerEnv
+
+__all__ = ["ModelBasedEnvBase", "DreamerDecoder", "DreamerEnv"]

torchrl/envs/model_based/common.py

@@ -0,0 +1,180 @@
+# 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
+import warnings
+
+import torch
+from tensordict import TensorDict
+from tensordict.nn import TensorDictModule
+from torchrl.data.utils import DEVICE_TYPING
+from torchrl.envs.common import EnvBase
+
+
+class ModelBasedEnvBase(EnvBase):
+    """Basic environment for Model Based RL sota-implementations.
+
+    Wrapper around the model of the MBRL algorithm.
+    It is meant to give an env framework to a world model (including but not limited to observations, reward, done state and safety constraints models).
+    and to behave as a classical environment.
+
+    This is a base class for other environments and it should not be used directly.
+
+    Example:
+        >>> import torch
+        >>> from tensordict import TensorDict
+        >>> from torchrl.data import Composite, Unbounded
+        >>> 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.observation_spec = Composite(
+        ...             hidden_observation=Unbounded((4,))
+        ...         )
+        ...         self.state_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.state_spec.rand())
+        ...         tensordict = tensordict.update(self.observation_spec.rand())
+        ...         return tensordict
+        >>> # This environment is used as follows:
+        >>> import torch.nn as nn
+        >>> from torchrl.modules import MLP, WorldModelWrapper
+        >>> world_model = WorldModelWrapper(
+        ...     TensorDictModule(
+        ...         MLP(out_features=4, activation_class=nn.ReLU, activate_last_layer=True, depth=0),
+        ...         in_keys=["hidden_observation", "action"],
+        ...         out_keys=["hidden_observation"],
+        ...     ),
+        ...     TensorDictModule(
+        ...         nn.Linear(4, 1),
+        ...         in_keys=["hidden_observation"],
+        ...         out_keys=["reward"],
+        ...     ),
+        ... )
+        >>> env = MyMBEnv(world_model)
+        >>> tensordict = env.rollout(max_steps=10)
+        >>> print(tensordict)
+        TensorDict(
+            fields={
+                action: Tensor(torch.Size([10, 1]), dtype=torch.float32),
+                done: Tensor(torch.Size([10, 1]), dtype=torch.bool),
+                hidden_observation: Tensor(torch.Size([10, 4]), dtype=torch.float32),
+                next: LazyStackedTensorDict(
+                    fields={
+                        hidden_observation: Tensor(torch.Size([10, 4]), dtype=torch.float32)},
+                    batch_size=torch.Size([10]),
+                    device=cpu,
+                    is_shared=False),
+                reward: Tensor(torch.Size([10, 1]), dtype=torch.float32)},
+            batch_size=torch.Size([10]),
+            device=cpu,
+            is_shared=False)
+
+
+    Properties:
+        observation_spec (Composite): sampling spec of the observations;
+        action_spec (TensorSpec): sampling spec of the actions;
+        reward_spec (TensorSpec): sampling spec of the rewards;
+        input_spec (Composite): sampling spec of the inputs;
+        batch_size (torch.Size): batch_size to be used by the env. If not set, the env accept tensordicts of all batch sizes.
+        device (torch.device): device where the env input and output are expected to live
+
+    Args:
+        world_model (nn.Module): model that generates world states and its corresponding rewards;
+        params (List[torch.Tensor], optional): list of parameters of the world model;
+        buffers (List[torch.Tensor], optional): list of buffers of the world model;
+        device (torch.device, optional): device where the env input and output are expected to live
+        dtype (torch.dtype, optional): dtype of the env input and output
+        batch_size (torch.Size, optional): number of environments contained in the instance
+        run_type_check (bool, optional): whether to run type checks on the step of the env
+
+    Methods:
+        step (TensorDict -> TensorDict): step in the environment
+        reset (TensorDict, optional -> TensorDict): reset the environment
+        set_seed (int -> int): sets the seed of the environment
+        rand_step (TensorDict, optional -> TensorDict): random step given the action spec
+        rollout (Callable, ... -> TensorDict): executes a rollout in the environment with the given policy (or random
+            steps if no policy is provided)
+
+    """
+
+    def __init__(
+        self,
+        world_model: TensorDictModule,
+        params: list[torch.Tensor] | None = None,
+        buffers: list[torch.Tensor] | None = None,
+        device: DEVICE_TYPING = "cpu",
+        batch_size: torch.Size | None = None,
+        run_type_checks: bool = False,
+        allow_done_after_reset: bool = False,
+    ):
+        super().__init__(
+            device=device,
+            batch_size=batch_size,
+            run_type_checks=run_type_checks,
+            allow_done_after_reset=allow_done_after_reset,
+        )
+        self.world_model = world_model.to(self.device)
+        self.world_model_params = params
+        self.world_model_buffers = buffers
+
+    @classmethod
+    def __new__(cls, *args, **kwargs):
+        return super().__new__(
+            cls, *args, _inplace_update=False, _batch_locked=False, **kwargs
+        )
+
+    def set_specs_from_env(self, env: EnvBase):
+        """Sets the specs of the environment from the specs of the given environment."""
+        device = self.device
+        output_spec = env.output_spec.clone()
+        input_spec = env.input_spec.clone()
+        if device is not None:
+            output_spec = output_spec.to(device)
+            input_spec = input_spec.to(device)
+        self.__dict__["_output_spec"] = output_spec
+        self.__dict__["_input_spec"] = input_spec
+        self.empty_cache()
+
+    def _step(
+        self,
+        tensordict: TensorDict,
+    ) -> TensorDict:
+        # step method requires to be immutable
+        tensordict_out = tensordict.clone(recurse=False)
+        # Compute world state
+        if self.world_model_params is not None:
+            tensordict_out = self.world_model(
+                tensordict_out,
+                params=self.world_model_params,
+                buffers=self.world_model_buffers,
+            )
+        else:
+            tensordict_out = self.world_model(tensordict_out)
+        # done can be missing, it will be filled by `step`
+        # Convert to list for torch.compile compatibility (dynamo can't unpack _CompositeSpecKeysView)
+        keys_to_select = (
+            list(self.observation_spec.keys())
+            + list(self.full_done_spec.keys())
+            + list(self.full_reward_spec.keys())
+        )
+        tensordict_out = tensordict_out.select(*keys_to_select, strict=False)
+        return tensordict_out
+
+    @abc.abstractmethod
+    def _reset(self, tensordict: TensorDict, **kwargs) -> TensorDict:
+        raise NotImplementedError
+
+    def _set_seed(self, seed: int | None) -> None:
+        warnings.warn("Set seed isn't needed for model based environments")

torchrl/envs/model_based/dreamer.py

@@ -0,0 +1,112 @@
+# 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
+from tensordict import TensorDict
+from tensordict.nn import TensorDictModule
+from torchrl.data.tensor_specs import Composite
+from torchrl.data.utils import DEVICE_TYPING
+from torchrl.envs.common import EnvBase
+from torchrl.envs.model_based import ModelBasedEnvBase
+from torchrl.envs.transforms.transforms import Transform
+
+
+class DreamerEnv(ModelBasedEnvBase):
+    """Dreamer simulation environment.
+
+    This environment is used for imagination rollouts in Dreamer training.
+    It never terminates (done is always False) since imagination runs for a
+    fixed horizon. The done-checking methods are overridden to avoid CUDA
+    synchronization overhead from Python control flow on CUDA tensors.
+    """
+
+    def __init__(
+        self,
+        world_model: TensorDictModule,
+        prior_shape: tuple[int, ...],
+        belief_shape: tuple[int, ...],
+        obs_decoder: TensorDictModule = None,
+        device: DEVICE_TYPING = "cpu",
+        batch_size: torch.Size | None = None,
+    ):
+        super().__init__(
+            world_model,
+            device=device,
+            batch_size=batch_size,
+            # Skip done validation in reset() — imagination never terminates.
+            allow_done_after_reset=True,
+        )
+        self.obs_decoder = obs_decoder
+        self.prior_shape = prior_shape
+        self.belief_shape = belief_shape
+
+    def any_done(self, tensordict) -> bool:
+        """Returns False — imagination rollouts never terminate.
+
+        Overridden to avoid CUDA sync from `done.any()` in parent class.
+        """
+        return False
+
+    def maybe_reset(self, tensordict):
+        """No-op — imagination rollouts don't need partial resets.
+
+        Overridden to avoid CUDA sync from done checks in parent class.
+        """
+        return tensordict
+
+    def set_specs_from_env(self, env: EnvBase):
+        """Sets the specs of the environment from the specs of the given environment."""
+        super().set_specs_from_env(env)
+        self.action_spec = self.action_spec.to(self.device)
+        self.state_spec = Composite(
+            state=self.observation_spec["state"],
+            belief=self.observation_spec["belief"],
+            shape=env.batch_size,
+        )
+
+    def _reset(self, tensordict=None, **kwargs) -> TensorDict:
+        batch_size = tensordict.batch_size if tensordict is not None else []
+        device = tensordict.device if tensordict is not None else self.device
+        if tensordict is None:
+            td = self.state_spec.rand(shape=batch_size)
+            # why don't we reuse actions taken at those steps?
+            td.set("action", self.action_spec.rand(shape=batch_size))
+            td[("next", "reward")] = self.reward_spec.rand(shape=batch_size)
+            td.update(self.observation_spec.rand(shape=batch_size))
+            if device is not None:
+                td = td.to(device, non_blocking=True)
+                if torch.cuda.is_available() and device.type == "cpu":
+                    torch.cuda.synchronize()
+                elif torch.backends.mps.is_available():
+                    torch.mps.synchronize()
+        else:
+            td = tensordict.clone()
+        return td
+
+    def decode_obs(self, tensordict: TensorDict, compute_latents=False) -> TensorDict:
+        if self.obs_decoder is None:
+            raise ValueError("No observation decoder provided")
+        if compute_latents:
+            tensordict = self.world_model(tensordict)
+        return self.obs_decoder(tensordict)
+
+
+class DreamerDecoder(Transform):
+    """A transform to record the decoded observations in Dreamer.
+
+    Examples:
+        >>> model_based_env = DreamerEnv(...)
+        >>> model_based_env_eval = model_based_env.append_transform(DreamerDecoder())
+    """
+
+    def _call(self, next_tensordict):
+        return self.parent.base_env.obs_decoder(next_tensordict)
+
+    def _reset(self, tensordict, tensordict_reset):
+        return self._call(tensordict_reset)
+
+    def transform_observation_spec(self, observation_spec):
+        return observation_spec

torchrl/envs/transforms/__init__.py

@@ -0,0 +1,147 @@
+# 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 .gym_transforms import EndOfLifeTransform
+from .module import ModuleTransform
+from .r3m import R3MTransform
+from .ray_service import RayTransform
+from .rb_transforms import MultiStepTransform
+from .transforms import (
+    ActionDiscretizer,
+    ActionMask,
+    AutoResetEnv,
+    AutoResetTransform,
+    BatchSizeTransform,
+    BinarizeReward,
+    BurnInTransform,
+    CatFrames,
+    CatTensors,
+    CenterCrop,
+    ClipTransform,
+    Compose,
+    ConditionalPolicySwitch,
+    ConditionalSkip,
+    Crop,
+    DeviceCastTransform,
+    DiscreteActionProjection,
+    DoubleToFloat,
+    DTypeCastTransform,
+    ExcludeTransform,
+    FiniteTensorDictCheck,
+    FlattenObservation,
+    FrameSkipTransform,
+    GrayScale,
+    gSDENoise,
+    Hash,
+    InitTracker,
+    LineariseRewards,
+    MultiAction,
+    NoopResetEnv,
+    ObservationNorm,
+    ObservationTransform,
+    PermuteTransform,
+    PinMemoryTransform,
+    RandomCropTensorDict,
+    RemoveEmptySpecs,
+    RenameTransform,
+    Resize,
+    Reward2GoTransform,
+    RewardClipping,
+    RewardScaling,
+    RewardSum,
+    SelectTransform,
+    SignTransform,
+    SqueezeTransform,
+    Stack,
+    StepCounter,
+    TargetReturn,
+    TensorDictPrimer,
+    TimeMaxPool,
+    Timer,
+    Tokenizer,
+    ToTensorImage,
+    TrajCounter,
+    Transform,
+    TransformedEnv,
+    UnaryTransform,
+    UnsqueezeTransform,
+    VecGymEnvTransform,
+    VecNorm,
+)
+from .vc1 import VC1Transform
+from .vecnorm import VecNormV2
+from .vip import VIPRewardTransform, VIPTransform
+
+__all__ = [
+    "ActionDiscretizer",
+    "ActionMask",
+    "AutoResetEnv",
+    "AutoResetTransform",
+    "BatchSizeTransform",
+    "BinarizeReward",
+    "BurnInTransform",
+    "CatFrames",
+    "CatTensors",
+    "CenterCrop",
+    "ClipTransform",
+    "Compose",
+    "ConditionalPolicySwitch",
+    "ConditionalSkip",
+    "Crop",
+    "DTypeCastTransform",
+    "DeviceCastTransform",
+    "DiscreteActionProjection",
+    "DoubleToFloat",
+    "EndOfLifeTransform",
+    "ExcludeTransform",
+    "FiniteTensorDictCheck",
+    "FlattenObservation",
+    "FrameSkipTransform",
+    "GrayScale",
+    "Hash",
+    "InitTracker",
+    "LineariseRewards",
+    "ModuleTransform",
+    "MultiAction",
+    "MultiStepTransform",
+    "NoopResetEnv",
+    "ObservationNorm",
+    "ObservationTransform",
+    "PermuteTransform",
+    "PinMemoryTransform",
+    "R3MTransform",
+    "RandomCropTensorDict",
+    "RayTransform",
+    "RemoveEmptySpecs",
+    "RenameTransform",
+    "Resize",
+    "Reward2GoTransform",
+    "RewardClipping",
+    "RewardScaling",
+    "RewardSum",
+    "SelectTransform",
+    "SignTransform",
+    "SqueezeTransform",
+    "Stack",
+    "StepCounter",
+    "TargetReturn",
+    "TensorDictPrimer",
+    "TimeMaxPool",
+    "Timer",
+    "ToTensorImage",
+    "Tokenizer",
+    "TrajCounter",
+    "Transform",
+    "TransformedEnv",
+    "UnaryTransform",
+    "UnsqueezeTransform",
+    "VC1Transform",
+    "VIPRewardTransform",
+    "VIPTransform",
+    "VecGymEnvTransform",
+    "VecNorm",
+    "VecNormV2",
+    "gSDENoise",
+]

torchrl/envs/transforms/functional.py

@@ -0,0 +1,48 @@
+# 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 torch import Tensor
+
+
+# copied from torchvision
+def _get_image_num_channels(img: Tensor) -> int:
+    if img.ndim == 2:
+        return 1
+    elif img.ndim > 2:
+        return img.shape[-3]
+
+    raise TypeError(f"Input ndim should be 2 or more. Got {img.ndim}")
+
+
+def _assert_channels(img: Tensor, permitted: list[int]) -> None:
+    c = _get_image_num_channels(img)
+    if c not in permitted:
+        raise TypeError(
+            f"Input image tensor permitted channel values are {permitted}, but found "
+            f"{c} (full shape: {img.shape})"
+        )
+
+
+def rgb_to_grayscale(img: Tensor, num_output_channels: int = 1) -> Tensor:
+    """Turns an RGB image into grayscale."""
+    if img.ndim < 3:
+        raise TypeError(
+            "Input image tensor should have at least 3 dimensions, but found"
+            "{}".format(img.ndim)
+        )
+    _assert_channels(img, [3])
+
+    if num_output_channels not in (1, 3):
+        raise ValueError("num_output_channels should be either 1 or 3")
+
+    r, g, b = img.unbind(dim=-3)
+    l_img = (0.2989 * r + 0.587 * g + 0.114 * b).to(img.dtype)
+    l_img = l_img.unsqueeze(dim=-3)
+
+    if num_output_channels == 3:
+        return l_img.expand(img.shape)
+
+    return l_img