torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/envs/transforms/gym_transforms.py

@@ -0,0 +1,203 @@
+# 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.
+
+"""Gym-specific transforms."""
+
+from __future__ import annotations
+
+import warnings
+
+import torch
+from tensordict import TensorDictBase
+from tensordict.utils import expand_as_right, NestedKey
+from torchrl.data.tensor_specs import Unbounded
+
+from torchrl.envs.transforms.transforms import FORWARD_NOT_IMPLEMENTED, Transform
+
+
+class EndOfLifeTransform(Transform):
+    """Registers the end-of-life signal from a Gym env with a `lives` method.
+
+    Proposed by DeepMind for the DQN and co. It helps value estimation.
+
+    Args:
+        eol_key (NestedKey, optional): the key where the end-of-life signal should
+            be written. Defaults to ``"end-of-life"``.
+        done_key (NestedKey, optional): a "done" key in the parent env done_spec,
+            where the done value can be retrieved. This key must be unique and its
+            shape must match the shape of the end-of-life entry. Defaults to ``"done"``.
+        eol_attribute (str, optional): the location of the "lives" in the gym env.
+            Defaults to ``"unwrapped.ale.lives"``. Supported attribute types are
+            integer/array-like objects or callables that return these values.
+
+    .. note::
+        This transform should be used with gym envs that have a ``env.unwrapped.ale.lives``.
+
+    Examples:
+        >>> from torchrl.envs.libs.gym import GymEnv
+        >>> from torchrl.envs.transforms.transforms import TransformedEnv
+        >>> env = GymEnv("ALE/Breakout-v5")
+        >>> env.rollout(100)
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([100, 4]), device=cpu, dtype=torch.int64, is_shared=False),
+                done: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        pixels: Tensor(shape=torch.Size([100, 210, 160, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
+                        reward: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([100]),
+                    device=cpu,
+                    is_shared=False),
+                pixels: Tensor(shape=torch.Size([100, 210, 160, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
+                terminated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([100]),
+            device=cpu,
+            is_shared=False)
+        >>> eol_transform = EndOfLifeTransform()
+        >>> env = TransformedEnv(env, eol_transform)
+        >>> env.rollout(100)
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([100, 4]), device=cpu, dtype=torch.int64, is_shared=False),
+                done: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                eol: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                lives: Tensor(shape=torch.Size([100]), device=cpu, dtype=torch.int64, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        end-of-life: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        lives: Tensor(shape=torch.Size([100]), device=cpu, dtype=torch.int64, is_shared=False),
+                        pixels: Tensor(shape=torch.Size([100, 210, 160, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
+                        reward: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([100]),
+                    device=cpu,
+                    is_shared=False),
+                pixels: Tensor(shape=torch.Size([100, 210, 160, 3]), device=cpu, dtype=torch.uint8, is_shared=False),
+                terminated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([100, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([100]),
+            device=cpu,
+            is_shared=False)
+
+    The typical usage of this transform is to replace the "done" state by "end-of-life"
+    within the loss module. The end-of-life signal isn't registered within the ``done_spec``
+    because it should not instruct the env to reset.
+
+    Examples:
+        >>> from torchrl.objectives import DQNLoss
+        >>> module = torch.nn.Identity() # used as a placeholder
+        >>> loss = DQNLoss(module, action_space="categorical")
+        >>> loss.set_keys(done="end-of-life", terminated="end-of-life")
+        >>> # equivalently
+        >>> eol_transform.register_keys(loss)
+    """
+
+    NO_PARENT_ERR = "The {} transform is being executed without a parent env. This is currently not supported."
+
+    def __init__(
+        self,
+        eol_key: NestedKey = "end-of-life",
+        lives_key: NestedKey = "lives",
+        done_key: NestedKey = "done",
+        eol_attribute="unwrapped.ale.lives",
+    ):
+        super().__init__(in_keys=[done_key], out_keys=[eol_key, lives_key])
+        self.eol_key = eol_key
+        self.lives_key = lives_key
+        self.done_key = done_key
+        self.eol_attribute = eol_attribute.split(".")
+
+    def _get_lives(self):
+        from torchrl.envs.libs.gym import GymWrapper
+
+        base_env = self.parent.base_env
+        if not isinstance(base_env, GymWrapper):
+            warnings.warn(
+                f"The base_env is not a gym env. Compatibility of {type(self)} is not guaranteed with "
+                f"environment types that do not inherit from GymWrapper.",
+                category=UserWarning,
+            )
+        # getattr falls back on _env by default
+        lives = getattr(base_env, self.eol_attribute[0])
+        for att in self.eol_attribute[1:]:
+            if isinstance(lives, list):
+                # For SerialEnv (and who knows Parallel one day)
+                lives = [getattr(_lives, att) for _lives in lives]
+            else:
+                lives = getattr(lives, att)
+        if callable(lives):
+            lives = lives()
+        elif isinstance(lives, list) and all(callable(_lives) for _lives in lives):
+            lives = torch.as_tensor([_lives() for _lives in lives])
+        return lives
+
+    def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+        return next_tensordict
+
+    def _step(self, tensordict, next_tensordict):
+        parent = self.parent
+        if parent is None:
+            raise RuntimeError(self.NO_PARENT_ERR.format(type(self)))
+
+        lives = self._get_lives()
+        end_of_life = torch.as_tensor(
+            tensordict.get(self.lives_key) > lives, device=self.parent.device
+        )
+        done = next_tensordict.get(self.done_key, None)  # TODO: None soon to be removed
+        if done is None:
+            raise KeyError(
+                f"The done value pointed by {self.done_key} cannot be found in tensordict with keys {tensordict.keys(True, True)}. "
+                f"Make sure to pass the appropriate done_key to the {type(self)} transform."
+            )
+        end_of_life = expand_as_right(end_of_life, done) | done
+        next_tensordict.set(self.eol_key, end_of_life)
+        next_tensordict.set(self.lives_key, lives)
+        return next_tensordict
+
+    def _reset(self, tensordict, tensordict_reset):
+        parent = self.parent
+        if parent is None:
+            raise RuntimeError(self.NO_PARENT_ERR.format(type(self)))
+        lives = self._get_lives()
+        end_of_life = False
+        tensordict_reset.set(
+            self.eol_key,
+            torch.as_tensor(end_of_life).expand(
+                parent.full_done_spec[self.done_key].shape
+            ),
+        )
+        tensordict_reset.set(self.lives_key, lives)
+        return tensordict_reset
+
+    def transform_observation_spec(self, observation_spec):
+        full_done_spec = self.parent.output_spec["full_done_spec"]
+        observation_spec[self.eol_key] = full_done_spec[self.done_key].clone()
+        observation_spec[self.lives_key] = Unbounded(
+            self.parent.batch_size,
+            device=self.parent.device,
+            dtype=torch.int64,
+        )
+        return observation_spec
+
+    def register_keys(
+        self, loss_or_advantage: torchrl.objectives.common.LossModule  # noqa
+    ):
+        """Registers the end-of-life key at appropriate places within the loss.
+
+        Args:
+            loss_or_advantage (torchrl.objectives.LossModule or torchrl.objectives.value.ValueEstimatorBase): a module to instruct what the end-of-life key is.
+
+        """
+        loss_or_advantage.set_keys(done=self.eol_key, terminated=self.eol_key)
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        raise RuntimeError(FORWARD_NOT_IMPLEMENTED.format(type(self)))

torchrl/envs/transforms/module.py

@@ -0,0 +1,341 @@
+# 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
+from contextlib import nullcontext
+from typing import overload, TYPE_CHECKING
+
+import torch
+from tensordict import TensorDictBase
+from tensordict.nn import TensorDictModuleBase
+from torchrl._utils import logger as torchrl_logger
+
+from torchrl.data.tensor_specs import TensorSpec
+from torchrl.envs.transforms.ray_service import _RayServiceMetaClass, RayTransform
+from torchrl.envs.transforms.transforms import Transform
+
+if TYPE_CHECKING:
+    from torchrl.weight_update import WeightSyncScheme
+
+__all__ = ["ModuleTransform", "RayModuleTransform"]
+
+
+class RayModuleTransform(RayTransform):
+    """Ray-based ModuleTransform for distributed processing.
+
+    This transform creates a Ray actor that wraps a ModuleTransform,
+    allowing module execution in a separate Ray worker process.
+
+    Args:
+        weight_sync_scheme: Optional weight synchronization scheme for updating
+            the module's weights from a parent collector. When provided, the scheme
+            is initialized on the receiver side (the Ray actor) and can receive
+            weight updates via torch.distributed.
+        **kwargs: Additional arguments passed to RayTransform and ModuleTransform.
+
+    Example:
+        >>> from torchrl.weight_update import RayModuleTransformScheme
+        >>> scheme = RayModuleTransformScheme()
+        >>> transform = RayModuleTransform(module=my_module, weight_sync_scheme=scheme)
+        >>> # The scheme can then be registered with a collector for weight updates
+    """
+
+    def __init__(self, *, weight_sync_scheme=None, **kwargs):
+        self._weight_sync_scheme = weight_sync_scheme
+        super().__init__(**kwargs)
+
+        # After actor is created, initialize the scheme on the receiver side
+        if weight_sync_scheme is not None:
+            # Store transform reference in the scheme for sender initialization
+            weight_sync_scheme._set_transform(self)
+
+            weight_sync_scheme.init_on_sender()
+
+            # Initialize receiver in the actor
+            torchrl_logger.debug(
+                "Setting up weight sync scheme on sender -- sender will do the remote call"
+            )
+            weight_sync_scheme.connect()
+
+    @property
+    def in_keys(self):
+        return self._ray.get(self._actor._getattr.remote("in_keys"))
+
+    @property
+    def out_keys(self):
+        return self._ray.get(self._actor._getattr.remote("out_keys"))
+
+    def _create_actor(self, **kwargs):
+        import ray
+
+        remote = self._ray.remote(ModuleTransform)
+        ray_kwargs = {}
+        num_gpus = self._num_gpus
+        if num_gpus is not None:
+            ray_kwargs["num_gpus"] = num_gpus
+        num_cpus = self._num_cpus
+        if num_cpus is not None:
+            ray_kwargs["num_cpus"] = num_cpus
+        actor_name = self._actor_name
+        if actor_name is not None:
+            ray_kwargs["name"] = actor_name
+        if ray_kwargs:
+            remote = remote.options(**ray_kwargs)
+        actor = remote.remote(**kwargs)
+        # wait till the actor is ready
+        ray.get(actor._ready.remote())
+        return actor
+
+    @overload
+    def update_weights(self, state_dict: dict[str, torch.Tensor]) -> None:
+        ...
+
+    @overload
+    def update_weights(self, params: TensorDictBase) -> None:
+        ...
+
+    def update_weights(self, *args, **kwargs) -> None:
+        import ray
+
+        if self._update_weights_method == "tensordict":
+            try:
+                td = kwargs.get("params", args[0])
+            except IndexError:
+                raise ValueError("params must be provided")
+            return ray.get(self._actor._update_weights_tensordict.remote(params=td))
+        elif self._update_weights_method == "state_dict":
+            try:
+                state_dict = kwargs.get("state_dict", args[0])
+            except IndexError:
+                raise ValueError("state_dict must be provided")
+            return ray.get(
+                self._actor._update_weights_state_dict.remote(state_dict=state_dict)
+            )
+        else:
+            raise ValueError(
+                f"Invalid update_weights_method: {self._update_weights_method}"
+            )
+
+
+class ModuleTransform(Transform, metaclass=_RayServiceMetaClass):
+    """A transform that wraps a module.
+
+    Keyword Args:
+        module (TensorDictModuleBase): The module to wrap. Exclusive with `module_factory`. At least one of `module` or `module_factory` must be provided.
+        module_factory (Callable[[], TensorDictModuleBase]): The factory to create the module. Exclusive with `module`. At least one of `module` or `module_factory` must be provided.
+        no_grad (bool, optional): Whether to use gradient computation. Default is `False`.
+        inverse (bool, optional): Whether to use the inverse of the module. Default is `False`.
+        device (torch.device, optional): The device to use. Default is `None`.
+        use_ray_service (bool, optional): Whether to use Ray service. Default is `False`.
+        num_gpus (int, optional): The number of GPUs to use if using Ray. Default is `None`.
+        num_cpus (int, optional): The number of CPUs to use if using Ray. Default is `None`.
+        actor_name (str, optional): The name of the actor to use. Default is `None`. If an actor name is provided and
+            an actor with this name already exists, the existing actor will be used.
+        observation_spec_transform (TensorSpec or Callable[[TensorSpec], TensorSpec]): either a new spec for the observation
+            after it has been transformed by the module, or a function that modifies the existing spec.
+            Defaults to `None` (observation specs remain unchanged).
+        done_spec_transform (TensorSpec or Callable[[TensorSpec], TensorSpec]): either a new spec for the done
+            after it has been transformed by the module, or a function that modifies the existing spec.
+            Defaults to `None` (done specs remain unchanged).
+        reward_spec_transform (TensorSpec or Callable[[TensorSpec], TensorSpec]): either a new spec for the reward
+            after it has been transformed by the module, or a function that modifies the existing spec.
+            Defaults to `None` (reward specs remain unchanged).
+        state_spec_transform (TensorSpec or Callable[[TensorSpec], TensorSpec]): either a new spec for the state
+            after it has been transformed by the module, or a function that modifies the existing spec.
+            Defaults to `None` (state specs remain unchanged).
+        action_spec_transform (TensorSpec or Callable[[TensorSpec], TensorSpec]): either a new spec for the action
+            after it has been transformed by the module, or a function that modifies the existing spec.
+            Defaults to `None` (action specs remain unchanged).
+    """
+
+    _RayServiceClass = RayModuleTransform
+
+    def __init__(
+        self,
+        *,
+        module: TensorDictModuleBase | None = None,
+        module_factory: Callable[[], TensorDictModuleBase] | None = None,
+        no_grad: bool = False,
+        inverse: bool = False,
+        device: torch.device | None = None,
+        use_ray_service: bool = False,  # noqa
+        actor_name: str | None = None,  # noqa
+        num_gpus: int | None = None,
+        num_cpus: int | None = None,
+        observation_spec_transform: TensorSpec
+        | Callable[[TensorSpec], TensorSpec]
+        | None = None,
+        action_spec_transform: TensorSpec
+        | Callable[[TensorSpec], TensorSpec]
+        | None = None,
+        reward_spec_transform: TensorSpec
+        | Callable[[TensorSpec], TensorSpec]
+        | None = None,
+        done_spec_transform: TensorSpec
+        | Callable[[TensorSpec], TensorSpec]
+        | None = None,
+        state_spec_transform: TensorSpec
+        | Callable[[TensorSpec], TensorSpec]
+        | None = None,
+    ):
+        super().__init__()
+        if module is None and module_factory is None:
+            raise ValueError(
+                "At least one of `module` or `module_factory` must be provided."
+            )
+        if module is not None and module_factory is not None:
+            raise ValueError(
+                "Only one of `module` or `module_factory` must be provided."
+            )
+        self.module = module if module is not None else module_factory()
+        self.no_grad = no_grad
+        self.inverse = inverse
+        self.device = device
+        self.observation_spec_transform = observation_spec_transform
+        self.action_spec_transform = action_spec_transform
+        self.reward_spec_transform = reward_spec_transform
+        self.done_spec_transform = done_spec_transform
+        self.state_spec_transform = state_spec_transform
+
+    @property
+    def in_keys(self) -> list[str]:
+        return self._in_keys()
+
+    def _in_keys(self):
+        return self.module.in_keys if not self.inverse else []
+
+    @in_keys.setter
+    def in_keys(self, value: list[str] | None):
+        if value is not None:
+            raise RuntimeError(f"in_keys {value} cannot be set for ModuleTransform")
+
+    @property
+    def out_keys(self) -> list[str]:
+        return self._out_keys()
+
+    def _out_keys(self):
+        return self.module.out_keys if not self.inverse else []
+
+    @property
+    def in_keys_inv(self) -> list[str]:
+        return self._in_keys_inv()
+
+    def _in_keys_inv(self):
+        return self.module.out_keys if self.inverse else []
+
+    @in_keys_inv.setter
+    def in_keys_inv(self, value: list[str]):
+        if value is not None:
+            raise RuntimeError(f"in_keys_inv {value} cannot be set for ModuleTransform")
+
+    @property
+    def out_keys_inv(self) -> list[str]:
+        return self._out_keys_inv()
+
+    def _out_keys_inv(self):
+        return self.module.in_keys if self.inverse else []
+
+    @out_keys_inv.setter
+    def out_keys_inv(self, value: list[str] | None):
+        if value is not None:
+            raise RuntimeError(
+                f"out_keys_inv {value} cannot be set for ModuleTransform"
+            )
+
+    @out_keys.setter
+    def out_keys(self, value: list[str] | None):
+        if value is not None:
+            raise RuntimeError(f"out_keys {value} cannot be set for ModuleTransform")
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        return self._call(tensordict)
+
+    def _call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if self.inverse:
+            return tensordict
+        with torch.no_grad() if self.no_grad else nullcontext():
+            with (
+                tensordict.to(self.device)
+                if self.device is not None
+                else nullcontext(tensordict)
+            ) as td:
+                return self.module(td)
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        if not self.inverse:
+            return tensordict
+        with torch.no_grad() if self.no_grad else nullcontext():
+            with (
+                tensordict.to(self.device)
+                if self.device is not None
+                else nullcontext(tensordict)
+            ) as td:
+                return self.module(td)
+
+    def _update_weights_tensordict(self, params: TensorDictBase) -> None:
+        params.to_module(self.module)
+
+    def _update_weights_state_dict(self, state_dict: dict[str, torch.Tensor]) -> None:
+        self.module.load_state_dict(state_dict)
+
+    def _init_weight_sync_scheme(self, scheme: WeightSyncScheme, model_id: str) -> None:
+        """Initialize weight sync scheme on the receiver side (called in Ray actor).
+
+        This method is called by RayModuleTransform after the actor is created
+        to set up the receiver side of the weight synchronization scheme.
+
+        Args:
+            scheme: The weight sync scheme instance (e.g., RayModuleTransformScheme).
+            model_id: Identifier for the model being synchronized.
+        """
+        torchrl_logger.debug(f"Initializing weight sync scheme for {model_id=}")
+        scheme.init_on_receiver(model_id=model_id, context=self)
+        torchrl_logger.debug(f"Setup weight sync scheme for {model_id=}")
+        scheme.connect()
+        self._weight_sync_scheme = scheme
+
+    def _receive_weights_scheme(self):
+        self._weight_sync_scheme.receive()
+
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        if self.observation_spec_transform is not None:
+            if isinstance(self.observation_spec_transform, TensorSpec):
+                return self.observation_spec_transform
+            else:
+                return self.observation_spec_transform(observation_spec)
+        return observation_spec
+
+    def transform_action_spec(self, action_spec: TensorSpec) -> TensorSpec:
+        if self.action_spec_transform is not None:
+            if isinstance(self.action_spec_transform, TensorSpec):
+                return self.action_spec_transform
+            else:
+                return self.action_spec_transform(action_spec)
+        return action_spec
+
+    def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec:
+        if self.reward_spec_transform is not None:
+            if isinstance(self.reward_spec_transform, TensorSpec):
+                return self.reward_spec_transform
+            else:
+                return self.reward_spec_transform(reward_spec)
+        return reward_spec
+
+    def transform_done_spec(self, done_spec: TensorSpec) -> TensorSpec:
+        if self.done_spec_transform is not None:
+            if isinstance(self.done_spec_transform, TensorSpec):
+                return self.done_spec_transform
+            else:
+                return self.done_spec_transform(done_spec)
+        return done_spec
+
+    def transform_state_spec(self, state_spec: TensorSpec) -> TensorSpec:
+        if self.state_spec_transform is not None:
+            if isinstance(self.state_spec_transform, TensorSpec):
+                return self.state_spec_transform
+            else:
+                return self.state_spec_transform(state_spec)
+        return state_spec