torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/record/recorder.py

@@ -0,0 +1,554 @@
+# 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 importlib.util
+import math
+from collections.abc import Callable, Sequence
+from copy import copy
+
+import numpy as np
+import torch
+from tensordict import NonTensorData, TensorDictBase
+from tensordict.utils import NestedKey
+from torchrl._utils import _can_be_pickled
+from torchrl.data.tensor_specs import NonTensor, TensorSpec, Unbounded
+from torchrl.data.utils import CloudpickleWrapper
+from torchrl.envs import EnvBase
+from torchrl.envs.transforms import ObservationTransform, Transform
+from torchrl.record.loggers import Logger
+
+_has_tv = importlib.util.find_spec("torchvision", None) is not None
+
+
+class VideoRecorder(ObservationTransform):
+    """Video Recorder transform.
+
+    Will record a series of observations from an environment and write them
+    to a Logger object when needed.
+
+    Args:
+        logger (Logger): a Logger instance where the video
+            should be written. To save the video under a memmap tensor or an mp4 file, use
+            the :class:`~torchrl.record.loggers.CSVLogger` class.
+        tag (str): the video tag in the logger.
+        in_keys (Sequence of NestedKey, optional): keys to be read to produce the video.
+            Default is :obj:`"pixels"`.
+        skip (int): frame interval in the output video.
+            Default is ``2`` if the transform has a parent environment, and ``1`` if not.
+        center_crop (int, optional): value of square center crop.
+        make_grid (bool, optional): if ``True``, a grid is created assuming that a
+            tensor of shape [B x W x H x 3] is provided, with B being the batch
+            size. Default is ``True`` if the transform has a parent environment, and ``False``
+            if not.
+        out_keys (sequence of NestedKey, optional): destination keys. Defaults
+            to ``in_keys`` if not provided.
+        fps (int, optional): Frames per second of the output video. Defaults to the logger predefined ``fps``,
+            and overrides it if provided.
+        **kwargs (Dict[str, Any], optional): additional keyword arguments for
+            :meth:`~torchrl.record.loggers.Logger.log_video`.
+
+    Examples:
+        The following example shows how to save a rollout under a video. First a few imports:
+
+        >>> from torchrl.record import VideoRecorder
+        >>> from torchrl.record.loggers.csv import CSVLogger
+        >>> from torchrl.envs import TransformedEnv, DMControlEnv
+
+        The video format is chosen in the logger. Wandb and tensorboard will take care of that
+        on their own, CSV accepts various video formats.
+
+        >>> logger = CSVLogger(exp_name="cheetah", log_dir="cheetah_videos", video_format="mp4")
+
+        Some envs (eg, Atari games) natively return images, some require the user to ask for them.
+        Check :class:`~torchrl.envs.GymEnv` or :class:`~torchrl.envs.DMControlEnv` to see how to render images
+        in these contexts.
+
+        >>> base_env = DMControlEnv("cheetah", "run", from_pixels=True)
+        >>> env = TransformedEnv(base_env, VideoRecorder(logger=logger, tag="run_video"))
+        >>> env.rollout(100)
+
+        All transforms have a dump function, mostly a no-op except for ``VideoRecorder``, and :class:`~torchrl.envs.transforms.Compose`
+        which will dispatch the `dumps` to all its members.
+
+        >>> env.transform.dump()
+
+        The transform can also be used within a dataset to save the video collected. Unlike in the environment case,
+        images will come in a batch. The ``skip`` argument will enable to save the images only at specific intervals.
+
+            >>> from torchrl.data.datasets import OpenXExperienceReplay
+            >>> from torchrl.envs import Compose
+            >>> from torchrl.record import VideoRecorder, CSVLogger
+            >>> # Create a logger that saves videos as mp4 using 24 frames per sec
+            >>> logger = CSVLogger("./dump", video_format="mp4", video_fps=24)
+            >>> # We use the VideoRecorder transform to save register the images coming from the batch.
+            >>> #  Setting the fps to 12 overrides the one set in the logger, not doing so keeps it unchanged.
+            >>> t = VideoRecorder(logger=logger, tag="pixels", in_keys=[("next", "observation", "image")], fps=12)
+            >>> # Each batch of data will have 10 consecutive videos of 200 frames each (maximum, since strict_length=False)
+            >>> dataset = OpenXExperienceReplay("cmu_stretch", batch_size=2000, slice_len=200,
+            ...             download=True, strict_length=False,
+            ...             transform=t)
+            >>> # Get a batch of data and visualize it
+            >>> for data in dataset:
+            ...     t.dump()
+            ...     break
+
+
+    Our video is available under ``./cheetah_videos/cheetah/videos/run_video_0.mp4``!
+
+    """
+
+    def __init__(
+        self,
+        logger: Logger,
+        tag: str,
+        in_keys: Sequence[NestedKey] | None = None,
+        skip: int | None = None,
+        center_crop: int | None = None,
+        make_grid: bool | None = None,
+        out_keys: Sequence[NestedKey] | None = None,
+        fps: int | None = None,
+        **kwargs,
+    ) -> None:
+        if in_keys is None:
+            in_keys = ["pixels"]
+        if out_keys is None:
+            out_keys = copy(in_keys)
+        super().__init__(in_keys=in_keys, out_keys=out_keys)
+        video_kwargs = {}
+        video_kwargs.update(kwargs)
+        if fps is not None:
+            video_kwargs["fps"] = fps
+        self.video_kwargs = video_kwargs
+        self.iter = 0
+        self.skip = skip
+        self.logger = logger
+        self.tag = tag
+        self.count = 0
+        self.center_crop = center_crop
+        self.make_grid = make_grid
+        if center_crop and not _has_tv:
+            raise ImportError(
+                "Could not load center_crop from torchvision. Make sure torchvision is installed."
+            )
+        self.obs = []
+
+    @property
+    def make_grid(self):
+        make_grid = self._make_grid
+        if make_grid is None:
+            if self.parent is not None:
+                self._make_grid = True
+                return True
+            self._make_grid = False
+            return False
+        return make_grid
+
+    @make_grid.setter
+    def make_grid(self, value):
+        self._make_grid = value
+
+    @property
+    def skip(self):
+        skip = self._skip
+        if skip is None:
+            if self.parent is not None:
+                self._skip = 2
+                return 2
+            self._skip = 1
+            return 1
+        return skip
+
+    @skip.setter
+    def skip(self, value):
+        self._skip = value
+
+    def _apply_transform(self, observation: torch.Tensor) -> torch.Tensor:
+        if isinstance(observation, NonTensorData):
+            observation_trsf = torch.tensor(observation.data)
+        else:
+            observation_trsf = observation
+        self.count += 1
+        if self.count % self.skip == 0:
+            if (
+                observation_trsf.ndim >= 3
+                and observation_trsf.shape[-3] == 3
+                and observation_trsf.shape[-2] > 3
+                and observation_trsf.shape[-1] > 3
+            ):
+                # permute the channels to the last dim
+                observation_trsf = observation_trsf.permute(
+                    *range(observation_trsf.ndim - 3), -2, -1, -3
+                )
+            if not (
+                observation_trsf.shape[-1] == 3 or observation_trsf.ndimension() == 2
+            ):
+                raise RuntimeError(
+                    f"Invalid observation shape, got: {observation.shape}"
+                )
+            observation_trsf = observation_trsf.clone()
+
+            if observation.ndimension() == 2:
+                observation_trsf = observation.unsqueeze(-3)
+            else:
+                if observation_trsf.shape[-1] != 3:
+                    raise RuntimeError(
+                        "observation_trsf is expected to have 3 dimensions, "
+                        f"got {observation_trsf.ndimension()} instead"
+                    )
+                trailing_dim = range(observation_trsf.ndimension() - 3)
+                observation_trsf = observation_trsf.permute(*trailing_dim, -1, -3, -2)
+            if self.center_crop:
+                if not _has_tv:
+                    raise ImportError(
+                        "Could not import torchvision, `center_crop` not available. "
+                        "Make sure torchvision is installed in your environment."
+                    )
+                from torchvision.transforms.functional import (
+                    center_crop as center_crop_fn,
+                )
+
+                observation_trsf = center_crop_fn(
+                    observation_trsf, [self.center_crop, self.center_crop]
+                )
+            if self.make_grid and observation_trsf.ndimension() >= 4:
+                if not _has_tv:
+                    raise ImportError(
+                        "Could not import torchvision, `make_grid` not available. "
+                        "Make sure torchvision is installed in your environment."
+                    )
+                from torchvision.utils import make_grid
+
+                obs_flat = observation_trsf.flatten(0, -4)
+                observation_trsf = make_grid(
+                    obs_flat, nrow=int(math.ceil(math.sqrt(obs_flat.shape[0])))
+                )
+                self.obs.append(observation_trsf.to("cpu", torch.uint8))
+            elif observation_trsf.ndimension() >= 4:
+                self.obs.extend(observation_trsf.to("cpu", torch.uint8).flatten(0, -4))
+            else:
+                self.obs.append(observation_trsf.to("cpu", torch.uint8))
+        return observation
+
+    def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
+        return self._call(tensordict)
+
+    def dump(self, suffix: str | None = None, step: int | None = None) -> None:
+        """Writes the video to the ``self.logger`` attribute.
+
+        Calling ``dump`` when no image has been stored in a no-op.
+
+        Args:
+            suffix (str, optional): a suffix for the video to be recorded.
+            step (int, optional): the step to log the video at. If not provided,
+                uses an internal counter that increments with each dump call.
+        """
+        if self.obs:
+            obs = torch.stack(self.obs, 0).unsqueeze(0).cpu()
+        else:
+            obs = None
+        self.obs = []
+        if obs is not None:
+            if suffix is None:
+                tag = self.tag
+            else:
+                tag = "_".join([self.tag, suffix])
+            if self.logger is not None:
+                self.logger.log_video(
+                    name=tag,
+                    video=obs,
+                    step=step if step is not None else self.iter,
+                    **self.video_kwargs,
+                )
+        self.iter += 1
+        self.count = 0
+        self.obs = []
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        self._call(tensordict_reset)
+        return tensordict_reset
+
+
+class TensorDictRecorder(Transform):
+    """TensorDict recorder.
+
+    When the 'dump' method is called, this class will save a stack of the tensordict resulting from :obj:`env.step(td)` in a
+    file with a prefix defined by the out_file_base argument.
+
+    Args:
+        out_file_base (str): a string defining the prefix of the file where the tensordict will be written.
+        skip_reset (bool): if ``True``, the first TensorDict of the list will be discarded (usually the tensordict
+            resulting from the call to :obj:`env.reset()`)
+            default: True
+        skip (int): frame interval for the saved tensordict.
+            default: 4
+
+    """
+
+    def __init__(
+        self,
+        out_file_base: str,
+        skip_reset: bool = True,
+        skip: int = 4,
+        in_keys: Sequence[str] | None = None,
+    ) -> None:
+        if in_keys is None:
+            in_keys = []
+
+        super().__init__(in_keys=in_keys)
+        self.iter = 0
+        self.out_file_base = out_file_base
+        self.td = []
+        self.skip_reset = skip_reset
+        self.skip = skip
+        self.count = 0
+
+    def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+        self.count += 1
+        if self.count % self.skip == 0:
+            _td = next_tensordict
+            if self.in_keys:
+                _td = next_tensordict.select(*self.in_keys).to_tensordict()
+            self.td.append(_td)
+        return next_tensordict
+
+    def dump(self, suffix: str | None = None) -> None:
+        if suffix is None:
+            tag = self.tag
+        else:
+            tag = "_".join([self.tag, suffix])
+
+        td = self.td
+        if self.skip_reset:
+            td = td[1:]
+        torch.save(
+            torch.stack(td, 0).contiguous(),
+            f"{tag}_tensordict.t",
+        )
+        self.iter += 1
+        self.count = 0
+        del self.td
+        self.td = []
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        self._call(tensordict_reset)
+        return tensordict_reset
+
+
+class PixelRenderTransform(Transform):
+    """A transform to call render on the parent environment and register the pixel observation in the tensordict.
+
+    This transform offers an alternative to the ``from_pixels`` syntactic sugar when instantiating an environment
+    that offers rendering is expensive, or when ``from_pixels`` is not implemented.
+    It can be used within a single environment or over batched environments alike.
+
+    Args:
+        out_keys (List[NestedKey] or Nested): List of keys where to register the pixel observations.
+        preproc (Callable, optional): a preproc function. Can be used to reshape the observation, or apply
+            any other transformation that makes it possible to register it in the output data.
+        as_non_tensor (bool, optional): if ``True``, the data will be written as a :class:`~tensordict.NonTensorData`
+            thereby relaxing the shape requirements. If not provided, it will be inferred automatically from the
+            input data type and shape.
+        render_method (str, optional): the name of the render method. Defaults to ``"render"``.
+        pass_tensordict (bool, optional): if ``True``, the input tensordict will be passed to the
+            render method. This enables rendering for stateless environments. Defaults to ``False``.
+        **kwargs: additional keyword arguments to pass to the render function (e.g. ``mode="rgb_array"``).
+
+    Examples:
+        >>> from torchrl.envs import GymEnv, check_env_specs, ParallelEnv, EnvCreator
+        >>> from torchrl.record.loggers import CSVLogger
+        >>> from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
+        >>>
+        >>> def make_env():
+        >>>     env = GymEnv("CartPole-v1", render_mode="rgb_array")
+        >>>     env = env.append_transform(PixelRenderTransform())
+        >>>     return env
+        >>>
+        >>> if __name__ == "__main__":
+        ...     logger = CSVLogger("dummy", video_format="mp4")
+        ...
+        ...     env = ParallelEnv(4, EnvCreator(make_env))
+        ...
+        ...     env = env.append_transform(VideoRecorder(logger=logger, tag="pixels_record"))
+        ...     env.rollout(3)
+        ...
+        ...     check_env_specs(env)
+        ...
+        ...     r = env.rollout(30)
+        ...     print(env)
+        ...     env.transform.dump()
+        ...     env.close()
+
+    This transform can also be used whenever a batched environment ``render()`` returns a single image:
+
+    Examples:
+        >>> from torchrl.envs import check_env_specs
+        >>> from torchrl.envs.libs.vmas import VmasEnv
+        >>> from torchrl.record.loggers import CSVLogger
+        >>> from torchrl.record.recorder import PixelRenderTransform, VideoRecorder
+        >>>
+        >>> env = VmasEnv(
+        ...     scenario="flocking",
+        ...     num_envs=32,
+        ...     continuous_actions=True,
+        ...     max_steps=200,
+        ...     device="cpu",
+        ...     seed=None,
+        ...     # Scenario kwargs
+        ...     n_agents=5,
+        ... )
+        >>>
+        >>> logger = CSVLogger("dummy", video_format="mp4")
+        >>>
+        >>> env = env.append_transform(PixelRenderTransform(mode="rgb_array", preproc=lambda x: x.copy()))
+        >>> env = env.append_transform(VideoRecorder(logger=logger, tag="pixels_record"))
+        >>>
+        >>> check_env_specs(env)
+        >>>
+        >>> r = env.rollout(30)
+        >>> env.transform[-1].dump()
+
+    The transform can be disabled using the :meth:`~torchrl.record.PixelRenderTransform.switch` method, which will
+    turn the rendering on if it's off or off if it's on (an argument can also be passed to control this behavior).
+    Since transforms are :class:`~torch.nn.Module` instances, :meth:`~torch.nn.Module.apply` can be used to control
+    this behavior:
+
+        >>> def switch(module):
+        ...     if isinstance(module, PixelRenderTransform):
+        ...         module.switch()
+        >>> env.apply(switch)
+
+    """
+
+    def __init__(
+        self,
+        out_keys: list[NestedKey] = None,
+        preproc: Callable[
+            [np.ndarray | torch.Tensor], np.ndarray | torch.Tensor
+        ] = None,
+        as_non_tensor: bool | None = None,
+        render_method: str = "render",
+        pass_tensordict: bool = False,
+        **kwargs,
+    ) -> None:
+        if out_keys is None:
+            out_keys = ["pixels"]
+        elif isinstance(out_keys, (str, tuple)):
+            out_keys = [out_keys]
+        if len(out_keys) != 1:
+            raise RuntimeError(
+                f"Expected one and only one out_key, got out_keys={out_keys}"
+            )
+        if preproc is not None and not _can_be_pickled(preproc):
+            preproc = CloudpickleWrapper(preproc)
+        self.preproc = preproc
+        self.as_non_tensor = as_non_tensor
+        self.kwargs = kwargs
+        self.render_method = render_method
+        self._enabled = True
+        self.pass_tensordict = pass_tensordict
+        super().__init__(in_keys=[], out_keys=out_keys)
+
+    def _reset(
+        self, tensordict: TensorDictBase, tensordict_reset: TensorDictBase
+    ) -> TensorDictBase:
+        return self._call(tensordict_reset)
+
+    def _call(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+        if not self._enabled:
+            return next_tensordict
+
+        method = getattr(self.parent, self.render_method)
+        if not self.pass_tensordict:
+            array = method(**self.kwargs)
+        else:
+            array = method(next_tensordict, **self.kwargs)
+
+        if self.preproc:
+            array = self.preproc(array)
+        if self.as_non_tensor is None:
+            if isinstance(array, list):
+                if isinstance(array[0], np.ndarray):
+                    array = np.asarray(array)
+                else:
+                    array = torch.as_tensor(array)
+            if (
+                array.ndim == 3
+                and array.shape[-1] == 3
+                and self.parent.batch_size != ()
+            ):
+                self.as_non_tensor = True
+            else:
+                self.as_non_tensor = False
+        if not self.as_non_tensor:
+            try:
+                next_tensordict.set(self.out_keys[0], array)
+            except Exception:
+                raise RuntimeError(
+                    f"An exception was raised while writing the rendered array "
+                    f"(shape={getattr(array, 'shape', None)}, dtype={getattr(array, 'dtype', None)}) in the tensordict with shape {next_tensordict.shape}. "
+                    f"Consider adapting your preproc function in {type(self).__name__}. You can also "
+                    f"pass keyword arguments to the render function of the parent environment, or save "
+                    f"this observation as a non-tensor data with as_non_tensor=True."
+                )
+        else:
+            next_tensordict.set_non_tensor(self.out_keys[0], array)
+        return next_tensordict
+
+    def transform_observation_spec(self, observation_spec: TensorSpec) -> TensorSpec:
+        # Adds the pixel observation spec by calling render on the parent env
+        switch = False
+        if not self.enabled:
+            switch = True
+            self.switch()
+        parent = self.parent
+        td_in = parent.reset()
+        self._call(td_in)
+        obs = td_in.get(self.out_keys[0])
+        if isinstance(obs, NonTensorData):
+            spec = NonTensor(device=obs.device, dtype=obs.dtype, shape=obs.shape)
+        else:
+            spec = Unbounded(device=obs.device, dtype=obs.dtype, shape=obs.shape)
+        observation_spec[self.out_keys[0]] = spec
+        if switch:
+            self.switch()
+        return observation_spec
+
+    def switch(self, mode: str | bool = None):
+        """Sets the transform on or off.
+
+        Args:
+            mode (str or bool, optional): if provided, sets the switch to the desired mode.
+                ``"on"``, ``"off"``, ``True`` and ``False`` are accepted values.
+                By default, ``switch`` sets the mode to the opposite of the current one.
+
+        """
+        if mode is None:
+            mode = not self._enabled
+        if not isinstance(mode, bool):
+            if mode not in ("on", "off"):
+                raise ValueError("mode must be either 'on' or 'off', or a boolean.")
+            mode = mode == "on"
+        self._enabled = mode
+
+    @property
+    def enabled(self) -> bool:
+        """Whether the recorder is enabled."""
+        return self._enabled
+
+    def set_container(self, container: Transform | EnvBase) -> None:
+        out = super().set_container(container)
+        if isinstance(self.parent, EnvBase):
+            # Start the env if needed
+            method = getattr(self.parent, self.render_method, None)
+            if method is None or not callable(method):
+                raise ValueError(
+                    f"The render method must exist and be a callable. Got render={method}."
+                )
+        return out

torchrl/services/__init__.py

@@ -0,0 +1,79 @@
+# 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.
+
+"""Distributed service registry for TorchRL.
+
+This module provides a service registry for managing distributed actors
+(tokenizers, replay buffers, etc.) that can be accessed across workers.
+
+Example:
+    >>> from torchrl.services import get_services
+    >>>
+    >>> # Worker 1: Register a tokenizer service
+    >>> services = get_services()
+    >>> services.register("tokenizer", TokenizerClass, num_cpus=1, num_gpus=0.1)
+    >>>
+    >>> # Worker 2: Access the same tokenizer
+    >>> services = get_services()
+    >>> tokenizer = services["tokenizer"]
+    >>> result = tokenizer.encode.remote(text)
+"""
+from __future__ import annotations
+
+from torchrl.services.base import ServiceBase
+from torchrl.services.ray_service import RayService
+
+__all__ = ["ServiceBase", "RayService", "get_services"]
+
+
+def get_services(backend: str = "ray", **init_kwargs) -> ServiceBase:
+    """Get a distributed service registry.
+
+    This function creates or retrieves a service registry for managing distributed
+    actors across workers. Services registered by one worker are immediately visible
+    to all other workers in the cluster.
+
+    Args:
+        backend: Service backend to use. Currently only "ray" is supported.
+        **init_kwargs: Backend-specific initialization arguments.
+            For Ray:
+
+            - ray_init_config (dict, optional): Arguments to pass to ray.init()
+            - namespace (str, optional): Ray namespace for service isolation.
+              Defaults to "torchrl_services".
+
+    Returns:
+        ServiceBase: A service registry instance.
+
+    Raises:
+        ValueError: If an unsupported backend is specified.
+        ImportError: If the required backend library is not installed.
+
+    Examples:
+        >>> # Basic usage - register and access services
+        >>> services = get_services()
+        >>> services.register("tokenizer", TokenizerClass, num_cpus=1)
+        >>> tokenizer = services["tokenizer"]
+        >>>
+        >>> # With custom Ray initialization
+        >>> services = get_services(
+        ...     backend="ray",
+        ...     ray_init_config={"address": "auto"},
+        ...     namespace="my_experiment"
+        ... )
+        >>>
+        >>> # Check if service exists
+        >>> if "tokenizer" in services:
+        ...     tokenizer = services["tokenizer"]
+        >>>
+        >>> # List all registered services
+        >>> service_names = services.list()
+    """
+    if backend != "ray":
+        raise ValueError(
+            f"Unsupported backend: {backend}. Currently only 'ray' is supported."
+        )
+
+    return RayService(**init_kwargs)