torchrl 0.11.0__cp314-cp314-macosx_11_0_arm64.whl

torchrl/envs/common.py

@@ -0,0 +1,4241 @@
+# 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 re
+import warnings
+import weakref
+from collections.abc import Callable, Iterator, Sequence
+from copy import deepcopy
+from functools import partial, wraps
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn as nn
+from tensordict import (
+    is_tensor_collection,
+    LazyStackedTensorDict,
+    TensorDictBase,
+    unravel_key,
+)
+from tensordict.base import _is_leaf_nontensor, NO_DEFAULT
+from tensordict.utils import is_non_tensor, NestedKey
+from torchrl._utils import (
+    _ends_with,
+    _make_ordinal_device,
+    _replace_last,
+    implement_for,
+    prod,
+    seed_generator,
+)
+
+from torchrl.data.tensor_specs import (
+    Categorical,
+    Composite,
+    NonTensor,
+    TensorSpec,
+    Unbounded,
+)
+from torchrl.data.utils import DEVICE_TYPING
+from torchrl.envs.utils import (
+    _make_compatible_policy,
+    _repr_by_depth,
+    _StepMDP,
+    _terminated_or_truncated,
+    _update_during_reset,
+    check_env_specs as check_env_specs_func,
+    get_available_libraries,
+)
+
+LIBRARIES = get_available_libraries()
+
+
+def _tensor_to_np(t):
+    return t.detach().cpu().numpy()
+
+
+dtype_map = {
+    torch.float: np.float32,
+    torch.double: np.float64,
+    torch.bool: bool,
+}
+
+
+def _maybe_unlock(func):
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        is_locked = self.is_spec_locked
+        try:
+            if is_locked:
+                self.set_spec_lock_(False)
+            result = func(self, *args, **kwargs)
+        finally:
+            if is_locked:
+                if not hasattr(self, "_cache"):
+                    self._cache = {}
+                self._cache.clear()
+                self.set_spec_lock_(True)
+        return result
+
+    return wrapper
+
+
+def _cache_value(func):
+    """Caches the result of the decorated function in env._cache dictionary."""
+    func_name = func.__name__
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if not self.is_spec_locked:
+            return func(self, *args, **kwargs)
+        result = self.__dict__.setdefault("_cache", {}).get(func_name, NO_DEFAULT)
+        if result is NO_DEFAULT:
+            result = func(self, *args, **kwargs)
+            self.__dict__.setdefault("_cache", {})[func_name] = result
+        return result
+
+    return wrapper
+
+
+def _clear_cache_when_set(func):
+    """A decorator for EnvBase methods that should clear the caches when called."""
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        # if there's no cache we'll just recompute the value
+        if "_cache" not in self.__dict__:
+            self._cache = {}
+        else:
+            self._cache.clear()
+        result = func(self, *args, **kwargs)
+        self._cache.clear()
+        return result
+
+    return wrapper
+
+
+class EnvMetaData:
+    """A class for environment meta-data storage and passing in multiprocessed settings."""
+
+    def __init__(
+        self,
+        *,
+        tensordict: TensorDictBase,
+        specs: Composite,
+        batch_size: torch.Size,
+        env_str: str,
+        device: torch.device,
+        batch_locked: bool,
+        device_map: dict,
+    ):
+        self.device = device
+        self.tensordict = tensordict
+        self.specs = specs
+        self.batch_size = batch_size
+        self.env_str = env_str
+        self.batch_locked = batch_locked
+        self.device_map = device_map
+        self.has_dynamic_specs = _has_dynamic_specs(specs)
+
+    @property
+    def tensordict(self) -> TensorDictBase:
+        td = self._tensordict.copy()
+        if td.device != self.device:
+            if self.device is None:
+                return td.clear_device_()
+            else:
+                return td.to(self.device)
+        return td
+
+    @property
+    def specs(self):
+        return self._specs.to(self.device)
+
+    @tensordict.setter
+    def tensordict(self, value: TensorDictBase):
+        self._tensordict = value.to("cpu")
+
+    @specs.setter
+    def specs(self, value: Composite):
+        self._specs = value.to("cpu")
+
+    @staticmethod
+    def metadata_from_env(env) -> EnvMetaData:
+        tensordict = env.fake_tensordict().clone()
+
+        for done_key in env.done_keys:
+            tensordict.set(
+                _replace_last(done_key, "_reset"),
+                torch.zeros_like(tensordict.get(("next", done_key))),
+            )
+
+        specs = env.specs.to("cpu")
+
+        batch_size = env.batch_size
+        try:
+            env_str = str(env)
+        except Exception:
+            env_str = f"{env.__class__.__name__}()"
+        device = env.device
+        specs = specs.to("cpu")
+        batch_locked = env.batch_locked
+        # we need to save the device map, as the tensordict will be placed on cpu
+        device_map = {}
+
+        def fill_device_map(name, val, device_map=device_map):
+            device_map[name] = val.device
+
+        tensordict.named_apply(fill_device_map, nested_keys=True, filter_empty=True)
+        return EnvMetaData(
+            tensordict=tensordict,
+            specs=specs,
+            batch_size=batch_size,
+            env_str=env_str,
+            device=device,
+            batch_locked=batch_locked,
+            device_map=device_map,
+        )
+
+    def expand(self, *size: int) -> EnvMetaData:
+        tensordict = self.tensordict.expand(*size).clone()
+        batch_size = torch.Size(list(size))
+        return EnvMetaData(
+            tensordict=tensordict,
+            specs=self.specs.expand(*size),
+            batch_size=batch_size,
+            env_str=self.env_str,
+            device=self.device,
+            batch_locked=self.batch_locked,
+            device_map=self.device_map,
+        )
+
+    def clone(self):
+        return EnvMetaData(
+            tensordict=self.tensordict.clone(),
+            specs=self.specs.clone(),
+            batch_size=torch.Size([*self.batch_size]),
+            env_str=deepcopy(self.env_str),
+            device=self.device,
+            batch_locked=self.batch_locked,
+            device_map=self.device_map,
+        )
+
+    def to(self, device: DEVICE_TYPING) -> EnvMetaData:
+        if device is not None:
+            device = _make_ordinal_device(torch.device(device))
+            device_map = {key: device for key in self.device_map}
+        tensordict = self.tensordict.contiguous().to(device)
+        specs = self.specs.to(device)
+        return EnvMetaData(
+            tensordict=tensordict,
+            specs=specs,
+            batch_size=self.batch_size,
+            env_str=self.env_str,
+            device=device,
+            batch_locked=self.batch_locked,
+            device_map=device_map,
+        )
+
+    def __getitem__(self, item):
+        from tensordict.utils import _getitem_batch_size
+
+        return EnvMetaData(
+            tensordict=self.tensordict[item],
+            specs=self.specs[item],
+            batch_size=_getitem_batch_size(self.batch_size, item),
+            env_str=self.env_str,
+            device=self.device,
+            batch_locked=self.batch_locked,
+            device_map=self.device_map,
+        )
+
+
+class _EnvPostInit(abc.ABCMeta):
+    def __call__(cls, *args, **kwargs):
+        spec_locked = kwargs.pop("spec_locked", True)
+        auto_reset = kwargs.pop("auto_reset", False)
+        auto_reset_replace = kwargs.pop("auto_reset_replace", True)
+        instance: EnvBase = super().__call__(*args, **kwargs)
+        if "_cache" not in instance.__dict__:
+            instance._cache = {}
+
+        if spec_locked:
+            instance.input_spec.lock_(recurse=True)
+            instance.output_spec.lock_(recurse=True)
+        instance._is_spec_locked = spec_locked
+
+        # we create the done spec by adding a done/terminated entry if one is missing
+        instance._create_done_specs()
+        # we access lazy attributed to make sure they're built properly.
+        # This isn't done in `__init__` because we don't know if super().__init__
+        # will be called before or after the specs, batch size etc are set.
+        _ = instance.done_spec
+        _ = instance.reward_keys
+        # _ = instance.action_keys
+        _ = instance.state_spec
+        if auto_reset:
+            from torchrl.envs.transforms.transforms import (
+                AutoResetEnv,
+                AutoResetTransform,
+            )
+
+            return AutoResetEnv(
+                instance, AutoResetTransform(replace=auto_reset_replace)
+            )
+
+        done_keys = set(instance.full_done_spec.keys(True, True))
+        obs_keys = set(instance.full_observation_spec.keys(True, True))
+        reward_keys = set(instance.full_reward_spec.keys(True, True))
+        # state_keys can match obs_keys so we don't test that
+        action_keys = set(instance.full_action_spec.keys(True, True))
+        state_keys = set(instance.full_state_spec.keys(True, True))
+        total_set = set()
+        for keyset in (done_keys, obs_keys, reward_keys):
+            if total_set.intersection(keyset):
+                raise RuntimeError(
+                    f"The set of keys of one spec collides (culprit: {total_set.intersection(keyset)}) with another."
+                )
+            total_set = total_set.union(keyset)
+        total_set = set()
+        for keyset in (state_keys, action_keys):
+            if total_set.intersection(keyset):
+                raise RuntimeError(
+                    f"The set of keys of one spec collides (culprit: {total_set.intersection(keyset)}) with another."
+                )
+            total_set = total_set.union(keyset)
+        return instance
+
+
+class EnvBase(nn.Module, metaclass=_EnvPostInit):
+    """Abstract environment parent class.
+
+    Keyword Args:
+        device (torch.device): The device of the environment. Deviceless environments
+            are allowed (device=None). If not ``None``, all specs will be cast
+            on that device and it is expected that all inputs and outputs will
+            live on that device.
+            Defaults to ``None``.
+        batch_size (torch.Size or equivalent, optional): batch-size of the environment.
+            Corresponds to the leading dimension of all the input and output
+            tensordicts the environment reads and writes. Defaults to an empty batch-size.
+        run_type_checks (bool, optional): If ``True``, type-checks will occur
+            at every reset and every step. Defaults to ``False``.
+        allow_done_after_reset (bool, optional): if ``True``, an environment can
+            be done after a call to :meth:`reset` is made. Defaults to ``False``.
+        spec_locked (bool, optional): if ``True``, the specs are locked and can only be
+            modified if :meth:`~torchrl.envs.EnvBase.set_spec_lock_` is called.
+
+            .. note:: The locking is achieved by the `EnvBase` metaclass. It does not appear in the
+                `__init__` method and is included in the keyword arguments strictly for type-hinting purpose.
+
+            .. seealso:: :ref:`Locking environment specs <Environment-lock>`.
+
+            Defaults to ``True``.
+        auto_reset (bool, optional): if ``True``, the env is assumed to reset automatically
+            when done. Defaults to ``False``.
+
+            .. note:: The auto-resetting is achieved by the `EnvBase` metaclass. It does not appear in the
+                `__init__` method and is included in the keyword arguments strictly for type-hinting purpose.
+
+            .. seealso:: The :ref:`auto-resetting environments API <autoresetting_envs>` section in the API
+                documentation.
+
+    Attributes:
+        done_spec (Composite): equivalent to ``full_done_spec`` as all
+            ``done_specs`` contain at least a ``"done"`` and a ``"terminated"`` entry
+        action_spec (TensorSpec): the spec of the action. Links to the spec of the leaf
+            action if only one action tensor is to be expected. Otherwise links to
+            ``full_action_spec``.
+        observation_spec (Composite): equivalent to ``full_observation_spec``.
+        reward_spec (TensorSpec): the spec of the reward. Links to the spec of the leaf
+            reward if only one reward tensor is to be expected. Otherwise links to
+            ``full_reward_spec``.
+        state_spec (Composite): equivalent to ``full_state_spec``.
+        full_done_spec (Composite): a composite spec such that ``full_done_spec.zero()``
+            returns a tensordict containing only the leaves encoding the done status of the
+            environment.
+        full_action_spec (Composite): a composite spec such that ``full_action_spec.zero()``
+            returns a tensordict containing only the leaves encoding the action of the
+            environment.
+        full_observation_spec (Composite): a composite spec such that ``full_observation_spec.zero()``
+            returns a tensordict containing only the leaves encoding the observation of the
+            environment.
+        full_reward_spec (Composite): a composite spec such that ``full_reward_spec.zero()``
+            returns a tensordict containing only the leaves encoding the reward of the
+            environment.
+        full_state_spec (Composite): a composite spec such that ``full_state_spec.zero()``
+            returns a tensordict containing only the leaves encoding the inputs (actions
+            excluded) of the environment.
+        batch_size (torch.Size): The batch-size of the environment.
+        device (torch.device): the device where the input/outputs of the environment
+            are to be expected. Can be ``None``.
+        is_spec_locked (bool): returns ``True`` if the specs are locked. See the :attr:`spec_locked`
+            argument above.
+
+    Methods:
+        step (TensorDictBase -> TensorDictBase): step in the environment
+        reset (TensorDictBase, optional -> TensorDictBase): reset the environment
+        set_seed (int -> int): sets the seed of the environment
+        rand_step (TensorDictBase, optional -> TensorDictBase): random step given the action spec
+        rollout (Callable, ... -> TensorDictBase): executes a rollout in the environment with the given policy (or random
+            steps if no policy is provided)
+
+    Examples:
+        >>> from torchrl.envs import EnvBase
+        >>> class CounterEnv(EnvBase):
+        ...     def __init__(self, batch_size=(), device=None, **kwargs):
+        ...         self.observation_spec = Composite(
+        ...             count=Unbounded(batch_size, device=device, dtype=torch.int64))
+        ...         self.action_spec = Unbounded(batch_size, device=device, dtype=torch.int8)
+        ...         # done spec and reward spec are set automatically
+        ...     def _step(self, tensordict):
+        ...
+        >>> from torchrl.envs.libs.gym import GymEnv
+        >>> env = GymEnv("Pendulum-v1")
+        >>> env.batch_size  # how many envs are run at once
+        torch.Size([])
+        >>> env.input_spec
+        Composite(
+            full_state_spec: None,
+            full_action_spec: Composite(
+                action: BoundedContinuous(
+                    shape=torch.Size([1]),
+                    space=ContinuousBox(
+                        low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                        high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                    device=cpu,
+                    dtype=torch.float32,
+                    domain=continuous), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+        >>> env.action_spec
+        BoundedContinuous(
+            shape=torch.Size([1]),
+            space=ContinuousBox(
+                low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+            device=cpu,
+            dtype=torch.float32,
+            domain=continuous)
+        >>> env.observation_spec
+        Composite(
+            observation: BoundedContinuous(
+                shape=torch.Size([3]),
+                space=ContinuousBox(
+                    low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
+                    high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous), device=cpu, shape=torch.Size([]))
+        >>> env.reward_spec
+        UnboundedContinuous(
+            shape=torch.Size([1]),
+            space=None,
+            device=cpu,
+            dtype=torch.float32,
+            domain=continuous)
+        >>> env.done_spec
+        Categorical(
+            shape=torch.Size([1]),
+            space=DiscreteBox(n=2),
+            device=cpu,
+            dtype=torch.bool,
+            domain=discrete)
+        >>> # the output_spec contains all the expected outputs
+        >>> env.output_spec
+        Composite(
+            full_reward_spec: Composite(
+                reward: UnboundedContinuous(
+                    shape=torch.Size([1]),
+                    space=None,
+                    device=cpu,
+                    dtype=torch.float32,
+                    domain=continuous), device=cpu, shape=torch.Size([])),
+            full_observation_spec: Composite(
+                observation: BoundedContinuous(
+                    shape=torch.Size([3]),
+                    space=ContinuousBox(
+                        low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
+                        high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
+                    device=cpu,
+                    dtype=torch.float32,
+                    domain=continuous), device=cpu, shape=torch.Size([])),
+            full_done_spec: Composite(
+                done: Categorical(
+                    shape=torch.Size([1]),
+                    space=DiscreteBox(n=2),
+                    device=cpu,
+                    dtype=torch.bool,
+                    domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+    .. note:: Learn more about dynamic specs and environments :ref:`here <dynamic_envs>`.
+    """
+
+    _batch_size: torch.Size | None
+    _device: torch.device | None
+    _is_spec_locked: bool = False
+
+    def __init__(
+        self,
+        *,
+        device: DEVICE_TYPING | None = None,
+        batch_size: tuple | torch.Size | None = None,
+        run_type_checks: bool = False,
+        allow_done_after_reset: bool = False,
+        spec_locked: bool = True,
+        auto_reset: bool = False,
+    ):
+        if "_cache" not in self.__dict__:
+            self._cache = {}
+        super().__init__()
+
+        self.__dict__.setdefault("_batch_size", None)
+        self.__dict__.setdefault("_device", None)
+
+        if batch_size is not None:
+            # we want an error to be raised if we pass batch_size but
+            # it's already been set
+            batch_size = self.batch_size = torch.Size(batch_size)
+        else:
+            batch_size = torch.Size(())
+
+        if device is not None:
+            device = self.__dict__["_device"] = _make_ordinal_device(
+                torch.device(device)
+            )
+
+        output_spec = self.__dict__.get("_output_spec")
+        if output_spec is None:
+            output_spec = self.__dict__["_output_spec"] = Composite(
+                shape=batch_size, device=device
+            )
+        elif self._output_spec.device != device and device is not None:
+            self.__dict__["_output_spec"] = self.__dict__["_output_spec"].to(
+                self.device
+            )
+        input_spec = self.__dict__.get("_input_spec")
+        if input_spec is None:
+            input_spec = self.__dict__["_input_spec"] = Composite(
+                shape=batch_size, device=device
+            )
+        elif self._input_spec.device != device and device is not None:
+            self.__dict__["_input_spec"] = self.__dict__["_input_spec"].to(self.device)
+
+        output_spec.unlock_(recurse=True)
+        input_spec.unlock_(recurse=True)
+        if "full_observation_spec" not in output_spec:
+            output_spec["full_observation_spec"] = Composite(batch_size=batch_size)
+        if "full_done_spec" not in output_spec:
+            output_spec["full_done_spec"] = Composite(batch_size=batch_size)
+        if "full_reward_spec" not in output_spec:
+            output_spec["full_reward_spec"] = Composite(batch_size=batch_size)
+        if "full_state_spec" not in input_spec:
+            input_spec["full_state_spec"] = Composite(batch_size=batch_size)
+        if "full_action_spec" not in input_spec:
+            input_spec["full_action_spec"] = Composite(batch_size=batch_size)
+
+        if "is_closed" not in self.__dir__():
+            self.is_closed = True
+        self._run_type_checks = run_type_checks
+        self._allow_done_after_reset = allow_done_after_reset
+
+    _collector: weakref.ReferenceType[
+        LLMCollector  # noqa: F821 # type: ignore
+    ] | None = None
+
+    def register_collector(self, collector: BaseCollector):  # noqa: F821 # type: ignore
+        """Registers a collector with the environment.
+
+        Args:
+            collector (BaseCollector): The collector to register.
+        """
+        self._collector = weakref.ref(collector)
+
+    @property
+    def collector(self) -> BaseCollector | None:  # noqa: F821 # type: ignore
+        """Returns the collector associated with the container, if it exists."""
+        return self._collector() if self._collector is not None else None
+
+    def set_spec_lock_(self, mode: bool = True) -> EnvBase:
+        """Locks or unlocks the environment's specs.
+
+        Args:
+            mode (bool): Whether to lock (`True`) or unlock (`False`) the specs. Defaults to `True`.
+
+        Returns:
+            EnvBase: The environment instance itself.
+
+        .. seealso:: :ref:`Locking environment specs <Environment-lock>`.
+
+        """
+        output_spec = self.__dict__.get("_output_spec")
+        input_spec = self.__dict__.get("_input_spec")
+        if mode:
+            if output_spec is not None:
+                output_spec.lock_(recurse=True)
+            if input_spec is not None:
+                input_spec.lock_(recurse=True)
+        else:
+            self._cache.clear()
+            if output_spec is not None:
+                output_spec.unlock_(recurse=True)
+            if input_spec is not None:
+                input_spec.unlock_(recurse=True)
+        self.__dict__["_is_spec_locked"] = mode
+        return self
+
+    @property
+    def is_spec_locked(self):
+        """Gets whether the environment's specs are locked.
+
+        This property can be modified directly.
+
+        Returns:
+            bool: True if the specs are locked, False otherwise.
+
+        .. seealso:: :ref:`Locking environment specs <Environment-lock>`.
+
+        """
+        return self.__dict__.get("_is_spec_locked", False)
+
+    @is_spec_locked.setter
+    def is_spec_locked(self, value: bool):
+        self.set_spec_lock_(value)
+
+    def auto_specs_(
+        self,
+        policy: Callable[[TensorDictBase], TensorDictBase],
+        *,
+        tensordict: TensorDictBase | None = None,
+        action_key: NestedKey | list[NestedKey] = "action",
+        done_key: NestedKey | list[NestedKey] | None = None,
+        observation_key: NestedKey | list[NestedKey] = "observation",
+        reward_key: NestedKey | list[NestedKey] = "reward",
+    ):
+        """Automatically sets the specifications (specs) of the environment based on a random rollout using a given policy.
+
+        This method performs a rollout using the provided policy to infer the input and output specifications of the environment.
+        It updates the environment's specs for actions, observations, rewards, and done signals based on the data collected
+        during the rollout.
+
+        Args:
+            policy (Callable[[TensorDictBase], TensorDictBase]):
+                A callable policy that takes a `TensorDictBase` as input and returns a `TensorDictBase` as output.
+                This policy is used to perform the rollout and determine the specs.
+
+        Keyword Args:
+            tensordict (TensorDictBase, optional):
+                An optional `TensorDictBase` instance to be used as the initial state for the rollout.
+                If not provided, the environment's `reset` method will be called to obtain the initial state.
+            action_key (NestedKey or List[NestedKey], optional):
+                The key(s) used to identify actions in the `TensorDictBase`. Defaults to "action".
+            done_key (NestedKey or List[NestedKey], optional):
+                The key(s) used to identify done signals in the `TensorDictBase`. Defaults to ``None``, which will
+                attempt to use ["done", "terminated", "truncated"] as potential keys.
+            observation_key (NestedKey or List[NestedKey], optional):
+                The key(s) used to identify observations in the `TensorDictBase`. Defaults to "observation".
+            reward_key (NestedKey or List[NestedKey], optional):
+                The key(s) used to identify rewards in the `TensorDictBase`. Defaults to "reward".
+
+        Returns:
+            EnvBase: The environment instance with updated specs.
+
+        Raises:
+            RuntimeError: If there are keys in the output specs that are not accounted for in the provided keys.
+        """
+        if self.batch_locked or tensordict is None:
+            batch_size = self.batch_size
+        else:
+            batch_size = tensordict.batch_size
+        if tensordict is None:
+            tensordict = self.reset()
+
+        # Input specs
+        tensordict.update(policy(tensordict))
+        step_0 = self.step(tensordict.copy())
+        tensordict2 = step_0.get("next").copy()
+        step_1 = self.step(policy(tensordict2).copy())
+        nexts_0: TensorDictBase = step_0.pop("next")
+        nexts_1: TensorDictBase = step_1.pop("next")
+
+        input_spec_stack = {}
+        tensordict.apply(
+            partial(_tensor_to_spec, stack=input_spec_stack),
+            tensordict2,
+            named=True,
+            nested_keys=True,
+            is_leaf=_is_leaf_nontensor,
+        )
+        input_spec = Composite(input_spec_stack, batch_size=batch_size)
+        if not self.batch_locked and batch_size != self.batch_size:
+            while input_spec.shape:
+                input_spec = input_spec[0]
+        if isinstance(action_key, NestedKey):
+            action_key = [action_key]
+        full_action_spec = input_spec.separates(*action_key, default=None)
+
+        # Output specs
+
+        output_spec_stack = {}
+        nexts_0.apply(
+            partial(_tensor_to_spec, stack=output_spec_stack),
+            nexts_1,
+            named=True,
+            nested_keys=True,
+            is_leaf=_is_leaf_nontensor,
+        )
+
+        output_spec = Composite(output_spec_stack, batch_size=batch_size)
+        if not self.batch_locked and batch_size != self.batch_size:
+            while output_spec.shape:
+                output_spec = output_spec[0]
+
+        if done_key is None:
+            done_key = ["done", "terminated", "truncated"]
+        full_done_spec = output_spec.separates(*done_key, default=None)
+        if full_done_spec is not None:
+            self.full_done_spec = full_done_spec
+
+        if isinstance(reward_key, NestedKey):
+            reward_key = [reward_key]
+        full_reward_spec = output_spec.separates(*reward_key, default=None)
+
+        if isinstance(observation_key, NestedKey):
+            observation_key = [observation_key]
+        full_observation_spec = output_spec.separates(*observation_key, default=None)
+        if not output_spec.is_empty(recurse=True):
+            raise RuntimeError(
+                f"Keys {list(output_spec.keys(True, True))} are unaccounted for. "
+                f"Make sure you have passed all the leaf names to the auto_specs_ method."
+            )
+
+        if full_action_spec is not None:
+            self.full_action_spec = full_action_spec
+        if full_done_spec is not None:
+            self.full_done_spec = full_done_spec
+        if full_observation_spec is not None:
+            self.full_observation_spec = full_observation_spec
+        if full_reward_spec is not None:
+            self.full_reward_spec = full_reward_spec
+        full_state_spec = input_spec
+        self.full_state_spec = full_state_spec
+
+        return self
+
+    def check_env_specs(self, *args, **kwargs):
+        kwargs.setdefault("return_contiguous", not self._has_dynamic_specs)
+        return check_env_specs_func(self, *args, **kwargs)
+
+    check_env_specs.__doc__ = check_env_specs_func.__doc__
+
+    def cardinality(self, tensordict: TensorDictBase | None = None) -> int:
+        """The cardinality of the action space.
+
+        By default, this is just a wrapper around :meth:`env.action_space.cardinality <~torchrl.data.TensorSpec.cardinality>`.
+
+        This class is useful when the action spec is variable:
+
+        - The number of actions can be undefined, e.g., ``Categorical(n=-1)``;
+        - The action cardinality may depend on the action mask;
+        - The shape can be dynamic, as in ``Unbound(shape=(-1))``.
+
+        In these cases, the :meth:`cardinality` should be overwritten,
+
+        Args:
+            tensordict (TensorDictBase, optional): a tensordict containing the data required to compute the cardinality.
+
+        """
+        return self.full_action_spec.cardinality()
+
+    def configure_parallel(
+        self,
+        *,
+        use_buffers: bool | None = None,
+        shared_memory: bool | None = None,
+        memmap: bool | None = None,
+        mp_start_method: str | None = None,
+        num_threads: int | None = None,
+        num_sub_threads: int | None = None,
+        non_blocking: bool | None = None,
+        daemon: bool | None = None,
+    ) -> EnvBase:
+        """Configure parallel execution parameters.
+
+        This method allows configuring parameters for parallel environment
+        execution before the environment is started. It is only effective
+        on :class:`~torchrl.envs.BatchedEnvBase` and its subclasses.
+
+        Args:
+            use_buffers (bool, optional): whether communication between workers should
+                occur via circular preallocated memory buffers.
+            shared_memory (bool, optional): whether the returned tensordict will be
+                placed in shared memory.
+            memmap (bool, optional): whether the returned tensordict will be placed
+                in memory map.
+            mp_start_method (str, optional): the multiprocessing start method.
+            num_threads (int, optional): number of threads for this process.
+            num_sub_threads (int, optional): number of threads of the subprocesses.
+            non_blocking (bool, optional): if ``True``, device moves will be done using
+                the ``non_blocking=True`` option.
+            daemon (bool, optional): whether the processes should be daemonized.
+
+        Returns:
+            self: Returns self for method chaining.
+
+        Raises:
+            NotImplementedError: If called on an environment that does not support
+                parallel configuration.
+            RuntimeError: If called after the environment has already started.
+
+        Example:
+            >>> env = DMControlEnv("cheetah", "run", num_envs=4)
+            >>> env.configure_parallel(use_buffers=True, num_threads=2)
+            >>> env.reset()  # Environment starts here, configure_parallel no longer effective
+
+        """
+        raise NotImplementedError(
+            f"{type(self).__name__} does not support configure_parallel. "
+            "This method is only available on BatchedEnvBase and its subclasses."
+        )
+
+    @classmethod
+    def make_parallel(
+        cls,
+        create_env_fn,
+        *,
+        num_envs: int = 1,
+        create_env_kwargs: dict | Sequence[dict] | None = None,
+        pin_memory: bool = False,
+        share_individual_td: bool | None = None,
+        shared_memory: bool = True,
+        memmap: bool = False,
+        policy_proof: Callable | None = None,
+        device: DEVICE_TYPING | None = None,
+        allow_step_when_done: bool = False,
+        num_threads: int | None = None,
+        num_sub_threads: int = 1,
+        serial_for_single: bool = False,
+        non_blocking: bool = False,
+        mp_start_method: str | None = None,
+        use_buffers: bool | None = None,
+        consolidate: bool = True,
+        daemon: bool = False,
+        **parallel_kwargs,
+    ) -> EnvBase:
+        """Factory method to create a ParallelEnv from an environment creator.
+
+        This method provides a convenient way to create parallel environments
+        with the same signature as :class:`~torchrl.envs.ParallelEnv`.
+
+        Args:
+            create_env_fn (callable): A callable that creates an environment instance.
+            num_envs (int, optional): Number of parallel environments. Defaults to 1.
+            create_env_kwargs (dict or list of dicts, optional): kwargs to be used
+                with the environments being created.
+            pin_memory (bool, optional): Whether to pin memory. Defaults to False.
+            share_individual_td (bool, optional): if ``True``, a different tensordict
+                is created for every process/worker and a lazy stack is returned.
+            shared_memory (bool, optional): whether the returned tensordict will be
+                placed in shared memory. Defaults to True.
+            memmap (bool, optional): whether the returned tensordict will be placed
+                in memory map. Defaults to False.
+            policy_proof (callable, optional): if provided, it'll be used to get
+                the list of tensors to return through step() and reset() methods.
+            device (str, int, torch.device, optional): The device of the batched
+                environment.
+            allow_step_when_done (bool, optional): Allow stepping when done.
+                Defaults to False.
+            num_threads (int, optional): number of threads for this process.
+            num_sub_threads (int, optional): number of threads of the subprocesses.
+                Defaults to 1.
+            serial_for_single (bool, optional): if ``True``, creating a parallel
+                environment with a single worker will return a SerialEnv instead.
+                Defaults to False.
+            non_blocking (bool, optional): if ``True``, device moves will be done
+                using the ``non_blocking=True`` option. Defaults to False.
+            mp_start_method (str, optional): the multiprocessing start method.
+            use_buffers (bool, optional): whether communication between workers
+                should occur via circular preallocated memory buffers.
+            consolidate (bool, optional): Whether to consolidate tensordicts.
+                Defaults to True.
+            daemon (bool, optional): whether the processes should be daemonized.
+                Defaults to False.
+            **parallel_kwargs: Additional keyword arguments passed to ParallelEnv.
+
+        Returns:
+            EnvBase: A ParallelEnv (or SerialEnv if serial_for_single=True and num_envs=1).
+
+        """
+        from torchrl.envs import ParallelEnv
+
+        return ParallelEnv(
+            num_workers=num_envs,
+            create_env_fn=create_env_fn,
+            create_env_kwargs=create_env_kwargs,
+            pin_memory=pin_memory,
+            share_individual_td=share_individual_td,
+            shared_memory=shared_memory,
+            memmap=memmap,
+            policy_proof=policy_proof,
+            device=device,
+            allow_step_when_done=allow_step_when_done,
+            num_threads=num_threads,
+            num_sub_threads=num_sub_threads,
+            serial_for_single=serial_for_single,
+            non_blocking=non_blocking,
+            mp_start_method=mp_start_method,
+            use_buffers=use_buffers,
+            consolidate=consolidate,
+            daemon=daemon,
+            **parallel_kwargs,
+        )
+
+    @classmethod
+    def __new__(cls, *args, _inplace_update=False, _batch_locked=True, **kwargs):
+        # inplace update will write tensors in-place on the provided tensordict.
+        # This is risky, especially if gradients need to be passed (in-place copy
+        # for tensors that are part of computational graphs will result in an error).
+        # It can also lead to inconsistencies when calling rollout.
+        cls._inplace_update = _inplace_update
+        cls._batch_locked = _batch_locked
+        cls._device = None
+        # cached in_keys to be excluded from update when calling step
+        cls._cache_in_keys = None
+
+        # We may assign _input_spec to the cls, but it must be assigned to the instance
+        # we pull it off, and place it back where it belongs
+        _input_spec = None
+        if hasattr(cls, "_input_spec"):
+            _input_spec = cls._input_spec.clone()
+            delattr(cls, "_input_spec")
+        _output_spec = None
+        if hasattr(cls, "_output_spec"):
+            _output_spec = cls._output_spec.clone()
+            delattr(cls, "_output_spec")
+        env = super().__new__(cls)
+        if _input_spec is not None:
+            env.__dict__["_input_spec"] = _input_spec
+        if _output_spec is not None:
+            env.__dict__["_output_spec"] = _output_spec
+        return env
+
+        return super().__new__(cls)
+
+    def __setattr__(self, key, value):
+        if key in (
+            "_input_spec",
+            "_observation_spec",
+            "_action_spec",
+            "_reward_spec",
+            "_output_spec",
+            "_state_spec",
+            "_done_spec",
+        ):
+            raise AttributeError(
+                "To set an environment spec, please use `env.observation_spec = obs_spec` (without the leading"
+                " underscore)."
+            )
+        super().__setattr__(key, value)
+
+    @property
+    def batch_locked(self) -> bool:
+        """Whether the environment can be used with a batch size different from the one it was initialized with or not.
+
+        If True, the env needs to be used with a tensordict having the same batch size as the env.
+        batch_locked is an immutable property.
+        """
+        return self._batch_locked
+
+    @batch_locked.setter
+    def batch_locked(self, value: bool) -> None:
+        raise RuntimeError("batch_locked is a read-only property")
+
+    @property
+    def run_type_checks(self) -> bool:
+        return self._run_type_checks
+
+    @run_type_checks.setter
+    def run_type_checks(self, run_type_checks: bool) -> None:
+        self._run_type_checks = run_type_checks
+
+    @property
+    def batch_size(self) -> torch.Size:
+        """Number of envs batched in this environment instance organised in a `torch.Size()` object.
+
+        Environment may be similar or different but it is assumed that they have little if
+        not no interactions between them (e.g., multi-task or batched execution
+        in parallel).
+
+        """
+        _batch_size = self.__dict__.get("_batch_size")
+        if _batch_size is None:
+            _batch_size = self._batch_size = torch.Size([])
+        return _batch_size
+
+    @batch_size.setter
+    @_maybe_unlock
+    def batch_size(self, value: torch.Size) -> None:
+        self._batch_size = torch.Size(value)
+        if (
+            hasattr(self, "output_spec")
+            and self.output_spec.shape[: len(value)] != value
+        ):
+            self.output_spec.shape = value
+        if hasattr(self, "input_spec") and self.input_spec.shape[: len(value)] != value:
+            self.input_spec.shape = value
+
+    @property
+    def shape(self):
+        """Equivalent to :attr:`~.batch_size`."""
+        return self.batch_size
+
+    @property
+    def device(self) -> torch.device:
+        device = self.__dict__.get("_device")
+        return device
+
+    @device.setter
+    def device(self, value: torch.device) -> None:
+        device = self.__dict__.get("_device")
+        if device is None:
+            self.__dict__["_device"] = value
+            return
+        raise RuntimeError("device cannot be set. Call env.to(device) instead.")
+
+    def ndimension(self):
+        return len(self.batch_size)
+
+    @property
+    def ndim(self):
+        return self.ndimension()
+
+    def append_transform(
+        self,
+        transform: Transform | Callable[[TensorDictBase], TensorDictBase],  # noqa: F821
+    ) -> torchrl.envs.TransformedEnv:  # noqa
+        """Returns a transformed environment where the callable/transform passed is applied.
+
+        Args:
+            transform (Transform or Callable[[TensorDictBase], TensorDictBase]): the transform to apply
+                to the environment.
+
+        Examples:
+            >>> from torchrl.envs import GymEnv
+            >>> import torch
+            >>> env = GymEnv("CartPole-v1")
+            >>> loc = 0.5
+            >>> scale = 1.0
+            >>> transform = lambda data: data.set("observation", (data.get("observation") - loc)/scale)
+            >>> env = env.append_transform(transform=transform)
+            >>> print(env)
+            TransformedEnv(
+                env=GymEnv(env=CartPole-v1, batch_size=torch.Size([]), device=cpu),
+                transform=_CallableTransform(keys=[]))
+
+        """
+        from torchrl.envs.transforms.transforms import TransformedEnv
+
+        return TransformedEnv(self, transform)
+
+    # Parent specs: input and output spec.
+    @property
+    def input_spec(self) -> TensorSpec:
+        """Input spec.
+
+        The composite spec containing all specs for data input to the environments.
+
+        It contains:
+
+        - "full_action_spec": the spec of the input actions
+        - "full_state_spec": the spec of all other environment inputs
+
+        This attribute is locked and should be read-only.
+        Instead, to set the specs contained in it, use the respective properties.
+
+        Examples:
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> env = GymEnv("Pendulum-v1")
+            >>> env.input_spec
+            Composite(
+                full_state_spec: None,
+                full_action_spec: Composite(
+                    action: BoundedContinuous(
+                        shape=torch.Size([1]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+
+        """
+        input_spec = self.__dict__.get("_input_spec")
+        if input_spec is None:
+            is_locked = self.is_spec_locked
+            if is_locked:
+                self.set_spec_lock_(False)
+            input_spec = Composite(
+                full_state_spec=None,
+                shape=self.batch_size,
+                device=self.device,
+            )
+            self.__dict__["_input_spec"] = input_spec
+            if is_locked:
+                self.set_spec_lock_(True)
+        return input_spec
+
+    @input_spec.setter
+    def input_spec(self, value: TensorSpec) -> None:
+        raise RuntimeError("input_spec is protected.")
+
+    @property
+    def output_spec(self) -> TensorSpec:
+        """Output spec.
+
+        The composite spec containing all specs for data output from the environments.
+
+        It contains:
+
+        - "full_reward_spec": the spec of reward
+        - "full_done_spec": the spec of done
+        - "full_observation_spec": the spec of all other environment outputs
+
+        This attribute is locked and should be read-only.
+        Instead, to set the specs contained in it, use the respective properties.
+
+        Examples:
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> env = GymEnv("Pendulum-v1")
+            >>> env.output_spec
+            Composite(
+                full_reward_spec: Composite(
+                    reward: UnboundedContinuous(
+                        shape=torch.Size([1]),
+                        space=None,
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous), device=cpu, shape=torch.Size([])),
+                full_observation_spec: Composite(
+                    observation: BoundedContinuous(
+                        shape=torch.Size([3]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous), device=cpu, shape=torch.Size([])),
+                full_done_spec: Composite(
+                    done: Categorical(
+                        shape=torch.Size([1]),
+                        space=DiscreteBox(n=2),
+                        device=cpu,
+                        dtype=torch.bool,
+                        domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+
+        """
+        output_spec = self.__dict__.get("_output_spec")
+        if output_spec is None:
+            is_locked = self.is_spec_locked
+            if is_locked:
+                self.set_spec_lock_(False)
+            output_spec = Composite(
+                shape=self.batch_size,
+                device=self.device,
+            )
+            self.__dict__["_output_spec"] = output_spec
+            if is_locked:
+                self.set_spec_lock_(True)
+        return output_spec
+
+    @output_spec.setter
+    def output_spec(self, value: TensorSpec) -> None:
+        raise RuntimeError("output_spec is protected.")
+
+    @property
+    @_cache_value
+    def action_keys(self) -> list[NestedKey]:
+        """The action keys of an environment.
+
+        By default, there will only be one key named "action".
+
+        Keys are sorted by depth in the data tree.
+        """
+        keys = self.full_action_spec.keys(True, True)
+        keys = sorted(keys, key=_repr_by_depth)
+        return keys
+
+    @property
+    @_cache_value
+    def state_keys(self) -> list[NestedKey]:
+        """The state keys of an environment.
+
+        By default, there will only be one key named "state".
+
+        Keys are sorted by depth in the data tree.
+        """
+        state_keys = self.__dict__.get("_state_keys")
+        if state_keys is not None:
+            return state_keys
+        keys = self.input_spec["full_state_spec"].keys(True, True)
+        keys = sorted(keys, key=_repr_by_depth)
+        self.__dict__["_state_keys"] = keys
+        return keys
+
+    @property
+    def action_key(self) -> NestedKey:
+        """The action key of an environment.
+
+        By default, this will be "action".
+
+        If there is more than one action key in the environment, this function will raise an exception.
+        """
+        if len(self.action_keys) > 1:
+            raise KeyError(
+                "action_key requested but more than one key present in the environment"
+            )
+        return self.action_keys[0]
+
+    # Action spec: action specs belong to input_spec
+    @property
+    def action_spec(self) -> TensorSpec:
+        """The ``action`` spec.
+
+        The ``action_spec`` is always stored as a composite spec.
+
+        If the action spec is provided as a simple spec, this will be returned.
+
+            >>> env.action_spec = Unbounded(1)
+            >>> env.action_spec
+            UnboundedContinuous(
+                shape=torch.Size([1]),
+                space=ContinuousBox(
+                    low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                    high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous)
+
+        If the action spec is provided as a composite spec and contains only one leaf,
+        this function will return just the leaf.
+
+            >>> env.action_spec = Composite({"nested": {"action": Unbounded(1)}})
+            >>> env.action_spec
+            UnboundedContinuous(
+                shape=torch.Size([1]),
+                space=ContinuousBox(
+                    low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                    high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous)
+
+        If the action spec is provided as a composite spec and has more than one leaf,
+        this function will return the whole spec.
+
+            >>> env.action_spec = Composite({"nested": {"action": Unbounded(1), "another_action": Categorical(1)}})
+            >>> env.action_spec
+            Composite(
+                nested: Composite(
+                    action: UnboundedContinuous(
+                        shape=torch.Size([1]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    another_action: Categorical(
+                        shape=torch.Size([]),
+                        space=DiscreteBox(n=1),
+                        device=cpu,
+                        dtype=torch.int64,
+                        domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+        To retrieve the full spec passed, use:
+
+            >>> env.input_spec["full_action_spec"]
+
+        This property is mutable.
+
+        Examples:
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> env = GymEnv("Pendulum-v1")
+            >>> env.action_spec
+            BoundedContinuous(
+                shape=torch.Size([1]),
+                space=ContinuousBox(
+                    low=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True),
+                    high=Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, contiguous=True)),
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous)
+        """
+        try:
+            action_spec = self.input_spec["full_action_spec"]
+        except (KeyError, AttributeError):
+            raise KeyError("Failed to find the action_spec.")
+
+        if len(self.action_keys) > 1:
+            return action_spec
+        else:
+            if len(self.action_keys) == 1 and self.action_keys[0] != "action":
+                return action_spec
+            try:
+                return action_spec[self.action_key]
+            except KeyError:
+                # the key may have changed
+                raise KeyError(
+                    "The action_key attribute seems to have changed. "
+                    "This occurs when a action_spec is updated without "
+                    "calling `env.action_spec = new_spec`. "
+                    "Make sure you rely on this  type of command "
+                    "to set the action and other specs."
+                )
+
+    @action_spec.setter
+    @_maybe_unlock
+    def action_spec(self, value: TensorSpec) -> None:
+        device = self.input_spec._device
+        if not hasattr(value, "shape"):
+            raise TypeError(
+                f"action_spec of type {type(value)} do not have a shape attribute."
+            )
+        if value.shape[: len(self.batch_size)] != self.batch_size:
+            raise ValueError(
+                f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size}). "
+                "Please use `env.action_spec_unbatched = value` to set unbatched versions instead."
+            )
+
+        if not isinstance(value, Composite):
+            value = Composite(
+                action=value.to(device), shape=self.batch_size, device=device
+            )
+
+        self.input_spec["full_action_spec"] = value.to(device)
+
+    @property
+    def full_action_spec(self) -> Composite:
+        """The full action spec.
+
+        ``full_action_spec`` is a :class:`~torchrl.data.Composite`` instance
+        that contains all the action entries.
+
+        Examples:
+            >>> from torchrl.envs import BraxEnv
+            >>> for envname in BraxEnv.available_envs:
+            ...     break
+            >>> env = BraxEnv(envname)
+            >>> env.full_action_spec
+            Composite(
+                action: BoundedContinuous(
+                    shape=torch.Size([8]),
+                    space=ContinuousBox(
+                        low=Tensor(shape=torch.Size([8]), device=cpu, dtype=torch.float32, contiguous=True),
+                        high=Tensor(shape=torch.Size([8]), device=cpu, dtype=torch.float32, contiguous=True)),
+                    device=cpu,
+                    dtype=torch.float32,
+                    domain=continuous), device=cpu, shape=torch.Size([]))
+
+        """
+        full_action_spec = self.input_spec.get("full_action_spec", None)
+        if full_action_spec is None:
+            is_locked = self.is_spec_locked
+            if is_locked:
+                self.set_spec_lock_(False)
+            full_action_spec = Composite(shape=self.batch_size, device=self.device)
+            self.input_spec["full_action_spec"] = full_action_spec
+            if is_locked:
+                self.set_spec_lock_(True)
+        return full_action_spec
+
+    @full_action_spec.setter
+    def full_action_spec(self, spec: Composite) -> None:
+        self.action_spec = spec
+
+    # Reward spec
+    @property
+    @_cache_value
+    def reward_keys(self) -> list[NestedKey]:
+        """The reward keys of an environment.
+
+        By default, there will only be one key named "reward".
+
+        Keys are sorted by depth in the data tree.
+        """
+        reward_keys = sorted(self.full_reward_spec.keys(True, True), key=_repr_by_depth)
+        return reward_keys
+
+    @property
+    @_cache_value
+    def observation_keys(self) -> list[NestedKey]:
+        """The observation keys of an environment.
+
+        By default, there will only be one key named "observation".
+
+        Keys are sorted by depth in the data tree.
+        """
+        observation_keys = sorted(
+            self.full_observation_spec.keys(True, True), key=_repr_by_depth
+        )
+        return observation_keys
+
+    @property
+    @_cache_value
+    def _observation_keys_step_mdp(self) -> list[NestedKey]:
+        """The observation keys of an environment that are static under step_mdp (i.e. to be copied as-is during step_mdp)."""
+        observation_keys_leaves = sorted(
+            self.full_observation_spec.keys(True, True, step_mdp_static_only=True),
+            key=_repr_by_depth,
+        )
+        return observation_keys_leaves
+
+    @property
+    @_cache_value
+    def _state_keys_step_mdp(self) -> list[NestedKey]:
+        """The state keys of an environment that are static under step_mdp (i.e. to be copied as-is during step_mdp)."""
+        state_keys_leaves = sorted(
+            self.full_state_spec.keys(True, True, step_mdp_static_only=True),
+            key=_repr_by_depth,
+        )
+        return state_keys_leaves
+
+    @property
+    @_cache_value
+    def _action_keys_step_mdp(self) -> list[NestedKey]:
+        """The action keys of an environment that are static under step_mdp (i.e. to be copied as-is during step_mdp)."""
+        action_keys_leaves = sorted(
+            self.full_action_spec.keys(True, True, step_mdp_static_only=True),
+            key=_repr_by_depth,
+        )
+        return action_keys_leaves
+
+    @property
+    @_cache_value
+    def _done_keys_step_mdp(self) -> list[NestedKey]:
+        """The done keys of an environment that are static under step_mdp (i.e. to be copied as-is during step_mdp)."""
+        done_keys_leaves = sorted(
+            self.full_done_spec.keys(True, True, step_mdp_static_only=True),
+            key=_repr_by_depth,
+        )
+        return done_keys_leaves
+
+    @property
+    @_cache_value
+    def _reward_keys_step_mdp(self) -> list[NestedKey]:
+        """The reward keys of an environment that are static under step_mdp (i.e. to be copied as-is during step_mdp)."""
+        reward_keys_leaves = sorted(
+            self.full_reward_spec.keys(True, True, step_mdp_static_only=True),
+            key=_repr_by_depth,
+        )
+        return reward_keys_leaves
+
+    @property
+    def reward_key(self):
+        """The reward key of an environment.
+
+        By default, this will be "reward".
+
+        If there is more than one reward key in the environment, this function will raise an exception.
+        """
+        if len(self.reward_keys) > 1:
+            raise KeyError(
+                "reward_key requested but more than one key present in the environment"
+            )
+        return self.reward_keys[0]
+
+    # Reward spec: reward specs belong to output_spec
+    @property
+    def reward_spec(self) -> TensorSpec:
+        """The ``reward`` spec.
+
+        The ``reward_spec`` is always stored as a composite spec.
+
+        If the reward spec is provided as a simple spec, this will be returned.
+
+            >>> env.reward_spec = Unbounded(1)
+            >>> env.reward_spec
+            UnboundedContinuous(
+                shape=torch.Size([1]),
+                space=ContinuousBox(
+                    low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                    high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous)
+
+        If the reward spec is provided as a composite spec and contains only one leaf,
+        this function will return just the leaf.
+
+            >>> env.reward_spec = Composite({"nested": {"reward": Unbounded(1)}})
+            >>> env.reward_spec
+            UnboundedContinuous(
+                shape=torch.Size([1]),
+                space=ContinuousBox(
+                    low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                    high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous)
+
+        If the reward spec is provided as a composite spec and has more than one leaf,
+        this function will return the whole spec.
+
+            >>> env.reward_spec = Composite({"nested": {"reward": Unbounded(1), "another_reward": Categorical(1)}})
+            >>> env.reward_spec
+            Composite(
+                nested: Composite(
+                    reward: UnboundedContinuous(
+                        shape=torch.Size([1]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous),
+                    another_reward: Categorical(
+                        shape=torch.Size([]),
+                        space=DiscreteBox(n=1),
+                        device=cpu,
+                        dtype=torch.int64,
+                        domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+        To retrieve the full spec passed, use:
+
+            >>> env.output_spec["full_reward_spec"]
+
+        This property is mutable.
+
+        Examples:
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> env = GymEnv("Pendulum-v1")
+            >>> env.reward_spec
+            UnboundedContinuous(
+                shape=torch.Size([1]),
+                space=None,
+                device=cpu,
+                dtype=torch.float32,
+                domain=continuous)
+        """
+        try:
+            reward_spec = self.output_spec["full_reward_spec"]
+        except (KeyError, AttributeError):
+            # populate the "reward" entry
+            # this will be raised if there is not full_reward_spec (unlikely) or no reward_key
+            # Since output_spec is lazily populated with an empty composite spec for
+            # reward_spec, the second case is much more likely to occur.
+            self.reward_spec = Unbounded(
+                shape=(*self.batch_size, 1),
+                device=self.device,
+            )
+            reward_spec = self.output_spec["full_reward_spec"]
+
+        reward_keys = self.reward_keys
+        if len(reward_keys) > 1 or not len(reward_keys):
+            return reward_spec
+        else:
+            if len(self.reward_keys) == 1 and self.reward_keys[0] != "reward":
+                return reward_spec
+            return reward_spec[self.reward_keys[0]]
+
+    @reward_spec.setter
+    @_maybe_unlock
+    def reward_spec(self, value: TensorSpec) -> None:
+        device = self.output_spec._device
+        if not hasattr(value, "shape"):
+            raise TypeError(
+                f"reward_spec of type {type(value)} do not have a shape " f"attribute."
+            )
+        if value.shape[: len(self.batch_size)] != self.batch_size:
+            raise ValueError(
+                f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size}). "
+                "Please use `env.reward_spec_unbatched = value` to set unbatched versions instead."
+            )
+        if not isinstance(value, Composite):
+            value = Composite(
+                reward=value.to(device), shape=self.batch_size, device=device
+            )
+        for leaf in value.values(True, True):
+            if len(leaf.shape) == 0:
+                raise RuntimeError(
+                    "the reward_spec's leaves shape cannot be empty (this error"
+                    " usually comes from trying to set a reward_spec"
+                    " with a null number of dimensions. Try using a multidimensional"
+                    " spec instead, for instance with a singleton dimension at the tail)."
+                )
+        self.output_spec["full_reward_spec"] = value.to(device)
+
+    @property
+    def full_reward_spec(self) -> Composite:
+        """The full reward spec.
+
+        ``full_reward_spec`` is a :class:`~torchrl.data.Composite`` instance
+        that contains all the reward entries.
+
+        Examples:
+            >>> import gymnasium
+            >>> from torchrl.envs import GymWrapper, TransformedEnv, RenameTransform
+            >>> base_env = GymWrapper(gymnasium.make("Pendulum-v1"))
+            >>> env = TransformedEnv(base_env, RenameTransform("reward", ("nested", "reward")))
+            >>> env.full_reward_spec
+            Composite(
+                nested: Composite(
+                    reward: UnboundedContinuous(
+                        shape=torch.Size([1]),
+                        space=ContinuousBox(
+                            low=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True),
+                            high=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, contiguous=True)),
+                        device=cpu,
+                        dtype=torch.float32,
+                        domain=continuous), device=None, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+        """
+        try:
+            return self.output_spec["full_reward_spec"]
+        except KeyError:
+            # populate the "reward" entry
+            # this will be raised if there is not full_reward_spec (unlikely) or no reward_key
+            # Since output_spec is lazily populated with an empty composite spec for
+            # reward_spec, the second case is much more likely to occur.
+            self.reward_spec = Unbounded(
+                shape=(*self.batch_size, 1),
+                device=self.device,
+            )
+            return self.output_spec["full_reward_spec"]
+
+    @full_reward_spec.setter
+    @_maybe_unlock
+    def full_reward_spec(self, spec: Composite) -> None:
+        self.reward_spec = spec.to(self.device) if self.device is not None else spec
+
+    # done spec
+    @property
+    @_cache_value
+    def done_keys(self) -> list[NestedKey]:
+        """The done keys of an environment.
+
+        By default, there will only be one key named "done".
+
+        Keys are sorted by depth in the data tree.
+        """
+        done_keys = sorted(self.full_done_spec.keys(True, True), key=_repr_by_depth)
+        return done_keys
+
+    @property
+    def done_key(self):
+        """The done key of an environment.
+
+        By default, this will be "done".
+
+        If there is more than one done key in the environment, this function will raise an exception.
+        """
+        done_keys = self.done_keys
+        if len(done_keys) > 1:
+            raise KeyError(
+                "done_key requested but more than one key present in the environment"
+            )
+        return done_keys[0]
+
+    @property
+    def full_done_spec(self) -> Composite:
+        """The full done spec.
+
+        ``full_done_spec`` is a :class:`~torchrl.data.Composite`` instance
+        that contains all the done entries.
+        It can be used to generate fake data with a structure that mimics the
+        one obtained at runtime.
+
+        Examples:
+            >>> import gymnasium
+            >>> from torchrl.envs import GymWrapper
+            >>> env = GymWrapper(gymnasium.make("Pendulum-v1"))
+            >>> env.full_done_spec
+            Composite(
+                done: Categorical(
+                    shape=torch.Size([1]),
+                    space=DiscreteBox(n=2),
+                    device=cpu,
+                    dtype=torch.bool,
+                    domain=discrete),
+                truncated: Categorical(
+                    shape=torch.Size([1]),
+                    space=DiscreteBox(n=2),
+                    device=cpu,
+                    dtype=torch.bool,
+                    domain=discrete), device=cpu, shape=torch.Size([]))
+
+        """
+        return self.output_spec["full_done_spec"]
+
+    @full_done_spec.setter
+    @_maybe_unlock
+    def full_done_spec(self, spec: Composite) -> None:
+        self.done_spec = spec.to(self.device) if self.device is not None else spec
+
+    # Done spec: done specs belong to output_spec
+    @property
+    def done_spec(self) -> TensorSpec:
+        """The ``done`` spec.
+
+        The ``done_spec`` is always stored as a composite spec.
+
+        If the done spec is provided as a simple spec, this will be returned.
+
+            >>> env.done_spec = Categorical(2, dtype=torch.bool)
+            >>> env.done_spec
+            Categorical(
+                shape=torch.Size([]),
+                space=DiscreteBox(n=2),
+                device=cpu,
+                dtype=torch.bool,
+                domain=discrete)
+
+        If the done spec is provided as a composite spec and contains only one leaf,
+        this function will return just the leaf.
+
+            >>> env.done_spec = Composite({"nested": {"done": Categorical(2, dtype=torch.bool)}})
+            >>> env.done_spec
+            Categorical(
+                shape=torch.Size([]),
+                space=DiscreteBox(n=2),
+                device=cpu,
+                dtype=torch.bool,
+                domain=discrete)
+
+        If the done spec is provided as a composite spec and has more than one leaf,
+        this function will return the whole spec.
+
+            >>> env.done_spec = Composite({"nested": {"done": Categorical(2, dtype=torch.bool), "another_done": Categorical(2, dtype=torch.bool)}})
+            >>> env.done_spec
+            Composite(
+                nested: Composite(
+                    done: Categorical(
+                        shape=torch.Size([]),
+                        space=DiscreteBox(n=2),
+                        device=cpu,
+                        dtype=torch.bool,
+                        domain=discrete),
+                    another_done: Categorical(
+                        shape=torch.Size([]),
+                        space=DiscreteBox(n=2),
+                        device=cpu,
+                        dtype=torch.bool,
+                        domain=discrete), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+        To always retrieve the full spec passed, use:
+
+            >>> env.output_spec["full_done_spec"]
+
+        This property is mutable.
+
+        Examples:
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> env = GymEnv("Pendulum-v1")
+            >>> env.done_spec
+            Categorical(
+                shape=torch.Size([1]),
+                space=DiscreteBox(n=2),
+                device=cpu,
+                dtype=torch.bool,
+                domain=discrete)
+        """
+        done_spec = self.output_spec["full_done_spec"]
+        return done_spec
+
+    @_maybe_unlock
+    def _create_done_specs(self):
+        """Reads through the done specs and makes it so that it's complete.
+
+        If the done_specs contain only a ``"done"`` entry, a similar ``"terminated"`` entry is created.
+        Same goes if only ``"terminated"`` key is present.
+
+        If none of ``"done"`` and ``"terminated"`` can be found and the spec is not
+        empty, nothing is changed.
+
+        """
+        try:
+            full_done_spec = self.output_spec["full_done_spec"]
+        except KeyError:
+            full_done_spec = Composite(
+                shape=self.output_spec.shape, device=self.output_spec.device
+            )
+            full_done_spec["done"] = Categorical(
+                n=2,
+                shape=(*full_done_spec.shape, 1),
+                dtype=torch.bool,
+                device=self.device,
+            )
+            full_done_spec["terminated"] = Categorical(
+                n=2,
+                shape=(*full_done_spec.shape, 1),
+                dtype=torch.bool,
+                device=self.device,
+            )
+            self.output_spec["full_done_spec"] = full_done_spec
+            return
+
+        def check_local_done(spec):
+            shape = None
+            for key, item in list(
+                spec.items()
+            ):  # list to avoid error due to in-loop changes
+                # in the case where the spec is non-empty and there is no done and no terminated, we do nothing
+                if key == "done" and "terminated" not in spec.keys():
+                    spec["terminated"] = item.clone()
+                elif key == "terminated" and "done" not in spec.keys():
+                    spec["done"] = item.clone()
+                elif isinstance(item, Composite):
+                    check_local_done(item)
+                else:
+                    if shape is None:
+                        shape = item.shape
+                        continue
+                    # checks that all shape match
+                    if shape != item.shape:
+                        raise ValueError(
+                            f"All shapes should match in done_spec {spec} (shape={shape}, key={key})."
+                        )
+
+            # if the spec is empty, we need to add a done and terminated manually
+            if spec.is_empty():
+                spec["done"] = Categorical(
+                    n=2, shape=(*spec.shape, 1), dtype=torch.bool, device=self.device
+                )
+                spec["terminated"] = Categorical(
+                    n=2, shape=(*spec.shape, 1), dtype=torch.bool, device=self.device
+                )
+
+        if_locked = self.is_spec_locked
+        if if_locked:
+            self.is_spec_locked = False
+        check_local_done(full_done_spec)
+        self.output_spec["full_done_spec"] = full_done_spec
+        if if_locked:
+            self.is_spec_locked = True
+        return
+
+    @done_spec.setter
+    @_maybe_unlock
+    def done_spec(self, value: TensorSpec) -> None:
+        device = self.output_spec.device
+        if not hasattr(value, "shape"):
+            raise TypeError(
+                f"done_spec of type {type(value)} do not have a shape " f"attribute."
+            )
+        if value.shape[: len(self.batch_size)] != self.batch_size:
+            raise ValueError(
+                f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size})."
+            )
+        if not isinstance(value, Composite):
+            value = Composite(
+                done=value.to(device),
+                terminated=value.to(device),
+                shape=self.batch_size,
+                device=device,
+            )
+        for leaf in value.values(True, True):
+            if len(leaf.shape) == 0:
+                raise RuntimeError(
+                    "the done_spec's leaves shape cannot be empty (this error"
+                    " usually comes from trying to set a reward_spec"
+                    " with a null number of dimensions. Try using a multidimensional"
+                    " spec instead, for instance with a singleton dimension at the tail)."
+                )
+        self.output_spec["full_done_spec"] = value.to(device)
+        self._create_done_specs()
+
+    # observation spec: observation specs belong to output_spec
+    @property
+    def observation_spec(self) -> Composite:
+        """Observation spec.
+
+        Must be a :class:`torchrl.data.Composite` instance.
+        The keys listed in the spec are directly accessible after reset and step.
+
+        In TorchRL, even though they are not properly speaking "observations"
+        all info, states, results of transforms etc. outputs from the environment are stored in the
+        ``observation_spec``.
+
+        Therefore, ``"observation_spec"`` should be thought as
+        a generic data container for environment outputs that are not done or reward data.
+
+        Examples:
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> env = GymEnv("Pendulum-v1")
+            >>> env.observation_spec
+            Composite(
+                observation: BoundedContinuous(
+                    shape=torch.Size([3]),
+                    space=ContinuousBox(
+                        low=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True),
+                        high=Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, contiguous=True)),
+                    device=cpu,
+                    dtype=torch.float32,
+                    domain=continuous), device=cpu, shape=torch.Size([]))
+
+        """
+        observation_spec = self.output_spec.get("full_observation_spec", default=None)
+        if observation_spec is None:
+            is_locked = self.is_spec_locked
+            if is_locked:
+                self.set_spec_lock_(False)
+            observation_spec = Composite(shape=self.batch_size, device=self.device)
+            self.output_spec["full_observation_spec"] = observation_spec
+            if is_locked:
+                self.set_spec_lock_(True)
+
+        return observation_spec
+
+    @observation_spec.setter
+    @_maybe_unlock
+    def observation_spec(self, value: TensorSpec) -> None:
+        if not isinstance(value, Composite):
+            value = Composite(
+                observation=value,
+                device=self.device,
+                batch_size=self.output_spec.batch_size,
+            )
+        elif value.shape[: len(self.batch_size)] != self.batch_size:
+            raise ValueError(
+                f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size})."
+            )
+        if value.shape[: len(self.batch_size)] != self.batch_size:
+            raise ValueError(
+                f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size})."
+            )
+        device = self.output_spec._device
+        self.output_spec["full_observation_spec"] = (
+            value.to(device) if device is not None else value
+        )
+
+    @property
+    def full_observation_spec(self) -> Composite:
+        return self.observation_spec
+
+    @full_observation_spec.setter
+    @_maybe_unlock
+    def full_observation_spec(self, spec: Composite):
+        self.observation_spec = spec
+
+    # state spec: state specs belong to input_spec
+    @property
+    def state_spec(self) -> Composite:
+        """State spec.
+
+        Must be a :class:`torchrl.data.Composite` instance.
+        The keys listed here should be provided as input alongside actions to the environment.
+
+        In TorchRL, even though they are not properly speaking "state"
+        all inputs to the environment that are not actions are stored in the
+        ``state_spec``.
+
+        Therefore, ``"state_spec"`` should be thought as
+        a generic data container for environment inputs that are not action data.
+
+        Examples:
+            >>> from torchrl.envs import BraxEnv
+            >>> for envname in BraxEnv.available_envs:
+            ...     break
+            >>> env = BraxEnv(envname)
+            >>> env.state_spec
+            Composite(
+                state: Composite(
+                    pipeline_state: Composite(
+                        q: UnboundedContinuous(
+                            shape=torch.Size([15]),
+                            space=None,
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                [...], device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+
+        """
+        state_spec = self.input_spec["full_state_spec"]
+        if state_spec is None:
+            is_locked = self.is_spec_locked
+            if is_locked:
+                self.set_spec_lock_(False)
+            state_spec = Composite(shape=self.batch_size, device=self.device)
+            self.input_spec["full_state_spec"] = state_spec
+            if is_locked:
+                self.set_spec_lock_(True)
+        return state_spec
+
+    @state_spec.setter
+    @_maybe_unlock
+    def state_spec(self, value: Composite) -> None:
+        if value is None:
+            self.input_spec["full_state_spec"] = Composite(
+                device=self.device, shape=self.batch_size
+            )
+        else:
+            device = self.input_spec.device
+            if not isinstance(value, Composite):
+                raise TypeError("The type of an state_spec must be Composite.")
+            elif value.shape[: len(self.batch_size)] != self.batch_size:
+                raise ValueError(
+                    f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size})."
+                )
+            if value.shape[: len(self.batch_size)] != self.batch_size:
+                raise ValueError(
+                    f"The value of spec.shape ({value.shape}) must match the env batch size ({self.batch_size})."
+                )
+            self.input_spec["full_state_spec"] = (
+                value.to(device) if device is not None else value
+            )
+
+    @property
+    def full_state_spec(self) -> Composite:
+        """The full state spec.
+
+        ``full_state_spec`` is a :class:`~torchrl.data.Composite`` instance
+        that contains all the state entries (ie, the input data that is not action).
+
+        Examples:
+            >>> from torchrl.envs import BraxEnv
+            >>> for envname in BraxEnv.available_envs:
+            ...     break
+            >>> env = BraxEnv(envname)
+            >>> env.full_state_spec
+            Composite(
+                state: Composite(
+                    pipeline_state: Composite(
+                        q: UnboundedContinuous(
+                            shape=torch.Size([15]),
+                            space=None,
+                            device=cpu,
+                            dtype=torch.float32,
+                            domain=continuous),
+                [...], device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([])), device=cpu, shape=torch.Size([]))
+
+        """
+        return self.state_spec
+
+    @full_state_spec.setter
+    @_maybe_unlock
+    def full_state_spec(self, spec: Composite) -> None:
+        self.state_spec = spec
+
+    # Single-env specs can be used to remove the batch size from the spec
+    @property
+    def batch_dims(self) -> int:
+        """Number of batch dimensions of the env."""
+        return len(self.batch_size)
+
+    def _make_single_env_spec(self, spec: TensorSpec) -> TensorSpec:
+        if not self.batch_dims:
+            return spec
+        idx = tuple(0 for _ in range(self.batch_dims))
+        return spec[idx]
+
+    @property
+    def full_action_spec_unbatched(self) -> Composite:
+        """Returns the action spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.full_action_spec)
+
+    @full_action_spec_unbatched.setter
+    @_maybe_unlock
+    def full_action_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.full_action_spec = spec
+
+    @property
+    def action_spec_unbatched(self) -> TensorSpec:
+        """Returns the action spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.action_spec)
+
+    @action_spec_unbatched.setter
+    @_maybe_unlock
+    def action_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.action_spec = spec
+
+    @property
+    def full_observation_spec_unbatched(self) -> Composite:
+        """Returns the observation spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.full_observation_spec)
+
+    @full_observation_spec_unbatched.setter
+    @_maybe_unlock
+    def full_observation_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.full_observation_spec = spec
+
+    @property
+    def observation_spec_unbatched(self) -> Composite:
+        """Returns the observation spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.observation_spec)
+
+    @observation_spec_unbatched.setter
+    @_maybe_unlock
+    def observation_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.observation_spec = spec
+
+    @property
+    def full_reward_spec_unbatched(self) -> Composite:
+        """Returns the reward spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.full_reward_spec)
+
+    @full_reward_spec_unbatched.setter
+    @_maybe_unlock
+    def full_reward_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.full_reward_spec = spec
+
+    @property
+    def reward_spec_unbatched(self) -> TensorSpec:
+        """Returns the reward spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.reward_spec)
+
+    @reward_spec_unbatched.setter
+    @_maybe_unlock
+    def reward_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.reward_spec = spec
+
+    @property
+    def full_done_spec_unbatched(self) -> Composite:
+        """Returns the done spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.full_done_spec)
+
+    @full_done_spec_unbatched.setter
+    @_maybe_unlock
+    def full_done_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.full_done_spec = spec
+
+    @property
+    def done_spec_unbatched(self) -> TensorSpec:
+        """Returns the done spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.done_spec)
+
+    @done_spec_unbatched.setter
+    @_maybe_unlock
+    def done_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.done_spec = spec
+
+    @property
+    def output_spec_unbatched(self) -> Composite:
+        """Returns the output spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.output_spec)
+
+    @output_spec_unbatched.setter
+    @_maybe_unlock
+    def output_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.output_spec = spec
+
+    @property
+    def input_spec_unbatched(self) -> Composite:
+        """Returns the input spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.input_spec)
+
+    @input_spec_unbatched.setter
+    @_maybe_unlock
+    def input_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.input_spec = spec
+
+    @property
+    def full_state_spec_unbatched(self) -> Composite:
+        """Returns the state spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.full_state_spec)
+
+    @full_state_spec_unbatched.setter
+    @_maybe_unlock
+    def full_state_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.full_state_spec = spec
+
+    @property
+    def state_spec_unbatched(self) -> TensorSpec:
+        """Returns the state spec of the env as if it had no batch dimensions."""
+        return self._make_single_env_spec(self.state_spec)
+
+    @state_spec_unbatched.setter
+    @_maybe_unlock
+    def state_spec_unbatched(self, spec: Composite):
+        spec = spec.expand(self.batch_size + spec.shape)
+        self.state_spec = spec
+
+    def _skip_tensordict(self, tensordict: TensorDictBase) -> TensorDictBase:
+        # Creates a "skip" tensordict, ie a placeholder for when a step is skipped
+        next_tensordict = self.full_done_spec.zero()
+        next_tensordict.update(self.full_observation_spec.zero())
+        next_tensordict.update(self.full_reward_spec.zero())
+
+        # Copy the data from tensordict in `next`
+        keys = set()
+
+        def select_and_clone(name, x, y):
+            keys.add(name)
+            if y is not None:
+                if y.device == x.device:
+                    return x.clone()
+                return x.to(y.device)
+
+        result = tensordict._fast_apply(
+            select_and_clone,
+            next_tensordict,
+            device=self.device,
+            default=None,
+            filter_empty=True,
+            is_leaf=_is_leaf_nontensor,
+            named=True,
+            nested_keys=True,
+        )
+        result.update(next_tensordict.exclude(*keys).filter_empty_())
+        return result
+
+    def step(self, tensordict: TensorDictBase) -> TensorDictBase:
+        """Makes a step in the environment.
+
+        Step accepts a single argument, tensordict, which usually carries an 'action' key which indicates the action
+        to be taken.
+        Step will call an out-place private method, _step, which is the method to be re-written by EnvBase subclasses.
+
+        Args:
+            tensordict (TensorDictBase): Tensordict containing the action to be taken.
+                If the input tensordict contains a ``"next"`` entry, the values contained in it
+                will prevail over the newly computed values. This gives a mechanism
+                to override the underlying computations.
+
+        Returns:
+            the input tensordict, modified in place with the resulting observations, done state and reward
+            (+ others if needed).
+
+        """
+        # sanity check
+        self._assert_tensordict_shape(tensordict)
+        partial_steps = tensordict.pop("_step", None)
+
+        next_tensordict = None
+
+        if partial_steps is not None:
+            tensordict_batch_size = None
+            if not self.batch_locked:
+                # Batched envs have their own way of dealing with this - batched envs that are not batched-locked may fail here
+                if partial_steps.all():
+                    partial_steps = None
+                else:
+                    tensordict_batch_size = tensordict.batch_size
+                    partial_steps = partial_steps.view(tensordict_batch_size)
+                    tensordict = tensordict[partial_steps]
+            else:
+                if not partial_steps.any():
+                    next_tensordict = self._skip_tensordic(tensordict)
+                else:
+                    # trust that the _step can handle this!
+                    tensordict.set("_step", partial_steps)
+            if tensordict_batch_size is None:
+                tensordict_batch_size = self.batch_size
+
+        next_preset = tensordict.get("next", None)
+
+        if next_tensordict is None:
+            next_tensordict = self._step(tensordict)
+            next_tensordict = self._step_proc_data(next_tensordict)
+        if next_preset is not None:
+            # tensordict could already have a "next" key
+            # this could be done more efficiently by not excluding but just passing
+            # the necessary keys
+            next_tensordict.update(
+                next_preset.exclude(*next_tensordict.keys(True, True))
+            )
+        tensordict.set("next", next_tensordict)
+        if partial_steps is not None and tensordict_batch_size != self.batch_size:
+            result = tensordict.new_zeros(tensordict_batch_size)
+
+            if tensordict_batch_size == tensordict.batch_size:
+
+                def select_and_clone(x, y):
+                    if y is not None:
+                        if x.device == y.device:
+                            return x.clone()
+                        return x.to(y.device)
+
+                result.update(
+                    tensordict._fast_apply(
+                        select_and_clone,
+                        result,
+                        device=result.device,
+                        filter_empty=True,
+                        default=None,
+                        batch_size=result.batch_size,
+                        is_leaf=_is_leaf_nontensor,
+                    )
+                )
+            if partial_steps.any():
+                result[partial_steps] = tensordict
+            return result
+        return tensordict
+
+    @classmethod
+    def _complete_done(
+        cls, done_spec: Composite, data: TensorDictBase
+    ) -> TensorDictBase:
+        """Completes the data structure at step time to put missing done keys."""
+        # by default, if a done key is missing, it is assumed that it is False
+        # except in 2 cases: (1) there is a "done" but no "terminated" or (2)
+        # there is a "terminated" but no "done".
+        if done_spec.ndim:
+            leading_dim = data.shape[: -done_spec.ndim]
+        else:
+            leading_dim = data.shape
+        vals = {}
+        i = -1
+        for i, (key, item) in enumerate(done_spec.items()):  # noqa: B007
+            val = data.get(key, None)
+            if isinstance(item, Composite):
+                if val is not None:
+                    cls._complete_done(item, val)
+                continue
+            shape = (*leading_dim, *item.shape)
+            if val is not None:
+                if val.shape != shape:
+                    val = val.reshape(shape)
+                    data.set(key, val)
+                vals[key] = val
+
+        if len(vals) < i + 1:
+            # complete missing dones: we only want to do that if we don't have enough done values
+            data_keys = set(data.keys())
+            done_spec_keys = set(done_spec.keys())
+            for key, item in done_spec.items(False, True):
+                val = vals.get(key, None)
+                if (
+                    key == "done"
+                    and val is not None
+                    and "terminated" in done_spec_keys
+                    and "terminated" not in data_keys
+                ):
+                    if "truncated" in data_keys:
+                        raise RuntimeError(
+                            "Cannot infer the value of terminated when only done and truncated are present."
+                        )
+                    data.set("terminated", val)
+                    data_keys.add("terminated")
+                elif (
+                    key == "terminated"
+                    and val is not None
+                    and "done" in done_spec_keys
+                    and "done" not in data_keys
+                ):
+                    if "truncated" in data_keys:
+                        val = val | data.get("truncated")
+                    data.set("done", val)
+                    data_keys.add("done")
+                elif val is None and key not in data_keys:
+                    # we must keep this here: we only want to fill with 0s if we're sure
+                    # done should not be copied to terminated or terminated to done
+                    # in this case, just fill with 0s
+                    data.set(key, item.zero(leading_dim))
+        return data
+
+    def _step_proc_data(self, next_tensordict_out):
+        batch_size = self.batch_size
+        dims = len(batch_size)
+        leading_batch_size = (
+            next_tensordict_out.batch_size[:-dims]
+            if dims
+            else next_tensordict_out.shape
+        )
+        for reward_key in self.reward_keys:
+            expected_reward_shape = torch.Size(
+                [
+                    *leading_batch_size,
+                    *self.output_spec["full_reward_spec"][reward_key].shape,
+                ]
+            )
+            # If the reward has a variable shape, we don't want to perform this check
+            if all(s > 0 for s in expected_reward_shape):
+                reward = next_tensordict_out.get(reward_key)
+                actual_reward_shape = reward.shape
+                if actual_reward_shape != expected_reward_shape:
+                    reward = reward.view(expected_reward_shape)
+                    next_tensordict_out.set(reward_key, reward)
+
+        self._complete_done(self.full_done_spec, next_tensordict_out)
+
+        if self.run_type_checks:
+            for key, spec in self.observation_spec.items():
+                obs = next_tensordict_out.get(key)
+                spec.type_check(obs)
+
+            for reward_key in self.reward_keys:
+                if (
+                    next_tensordict_out.get(reward_key).dtype
+                    is not self.output_spec[
+                        unravel_key(("full_reward_spec", reward_key))
+                    ].dtype
+                ):
+                    raise TypeError(
+                        f"expected reward.dtype to be {self.output_spec[unravel_key(('full_reward_spec',reward_key))]} "
+                        f"but got {next_tensordict_out.get(reward_key).dtype}"
+                    )
+
+            for done_key in self.done_keys:
+                if (
+                    next_tensordict_out.get(done_key).dtype
+                    is not self.output_spec["full_done_spec", done_key].dtype
+                ):
+                    raise TypeError(
+                        f"expected done.dtype to be {self.output_spec['full_done_spec', done_key].dtype} but got {next_tensordict_out.get(done_key).dtype}"
+                    )
+        return next_tensordict_out
+
+    def _get_in_keys_to_exclude(self, tensordict):
+        if self._cache_in_keys is None:
+            self._cache_in_keys = list(
+                set(self.input_spec.keys(True)).intersection(
+                    tensordict.keys(True, True)
+                )
+            )
+        return self._cache_in_keys
+
+    @classmethod
+    def register_gym(
+        cls,
+        id: str,
+        *,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        backend: str | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        """Registers an environment in gym(nasium).
+
+        This method is designed with the following scopes in mind:
+
+        - Incorporate a TorchRL-first environment in a framework that uses Gym;
+        - Incorporate another environment (eg, DeepMind Control, Brax, Jumanji, ...)
+          in a framework that uses Gym.
+
+        Args:
+            id (str): the name of the environment. Should follow the
+                `gym naming convention <https://www.gymlibrary.dev/content/environment_creation/#registering-envs>`_.
+
+        Keyword Args:
+            entry_point (callable, optional): the entry point to build the environment.
+                If none is passed, the parent class will be used as entry point.
+                Typically, this is used to register an environment that does not
+                necessarily inherit from the base being used:
+
+                    >>> from torchrl.envs import DMControlEnv
+                    >>> DMControlEnv.register_gym("DMC-cheetah-v0", env_name="cheetah", task="run")
+                    >>> # equivalently
+                    >>> EnvBase.register_gym("DMC-cheetah-v0", entry_point=DMControlEnv, env_name="cheetah", task="run")
+
+            transform (torchrl.envs.Transform): a transform (or list of transforms
+                within a :class:`torchrl.envs.Compose` instance) to be used with the env.
+                This arg can be passed during a call to :func:`~gym.make` (see
+                example below).
+            info_keys (List[NestedKey], optional): if provided, these keys will
+                be used to build the info dictionary and will be excluded from
+                the observation keys.
+                This arg can be passed during a call to :func:`~gym.make` (see
+                example below).
+
+                .. warning::
+                  It may be the case that using ``info_keys`` makes a spec empty
+                  because the content has been moved to the info dictionary.
+                  Gym does not like empty ``Dict`` in the specs, so this empty
+                  content should be removed with :class:`~torchrl.envs.transforms.RemoveEmptySpecs`.
+
+            backend (str, optional): the backend. Can be either `"gym"` or `"gymnasium"`
+                or any other backend compatible with :class:`~torchrl.envs.libs.gym.set_gym_backend`.
+            to_numpy (bool, optional): if ``True``, the result of calls to `step` and
+                `reset` will be mapped to numpy arrays. Defaults to ``False``
+                (results are tensors).
+                This arg can be passed during a call to :func:`~gym.make` (see
+                example below).
+            reward_threshold (:obj:`float`, optional): [Gym kwarg] The reward threshold
+                considered to have learnt an environment.
+            nondeterministic (bool, optional): [Gym kwarg If the environment is nondeterministic
+                (even with knowledge of the initial seed and all actions). Defaults to
+                ``False``.
+            max_episode_steps (int, optional): [Gym kwarg] The maximum number
+                of episodes steps before truncation. Used by the Time Limit wrapper.
+            order_enforce (bool, optional): [Gym >= 0.14] Whether the order
+                enforcer wrapper should be applied to ensure users run functions
+                in the correct order.
+                Defaults to ``True``.
+            autoreset (bool, optional): [Gym >= 0.14 and <1.0.0] Whether the autoreset wrapper
+                should be added such that reset does not need to be called.
+                Defaults to ``False``.
+            disable_env_checker: [Gym >= 0.14] Whether the environment
+                checker should be disabled for the environment. Defaults to ``False``.
+            apply_api_compatibility: [Gym >= 0.26 and <1.0.0] If to apply the `StepAPICompatibility` wrapper.
+                Defaults to ``False``.
+            **kwargs: arbitrary keyword arguments which are passed to the environment constructor.
+
+        .. note::
+            TorchRL's environment do not have the concept of an ``"info"`` dictionary,
+            as ``TensorDict`` offers all the storage requirements deemed necessary
+            in most training settings. Still, you can use the ``info_keys`` argument to
+            have a fine grained control over what is deemed to be considered
+            as an observation and what should be seen as info.
+
+        Examples:
+            >>> # Register the "cheetah" env from DMControl with the "run" task
+            >>> from torchrl.envs import DMControlEnv
+            >>> import torch
+            >>> DMControlEnv.register_gym("DMC-cheetah-v0", to_numpy=False, backend="gym", env_name="cheetah", task_name="run")
+            >>> import gym
+            >>> envgym = gym.make("DMC-cheetah-v0")
+            >>> envgym.seed(0)
+            >>> torch.manual_seed(0)
+            >>> envgym.reset()
+            ({'position': tensor([-0.0855,  0.0215, -0.0881, -0.0412, -0.1101,  0.0080,  0.0254,  0.0424],
+                   dtype=torch.float64), 'velocity': tensor([ 1.9609e-02, -1.9776e-04, -1.6347e-03,  3.3842e-02,  2.5338e-02,
+                     3.3064e-02,  1.0381e-04,  7.6656e-05,  1.0204e-02],
+                   dtype=torch.float64)}, {})
+            >>> envgym.step(envgym.action_space.sample())
+            ({'position': tensor([-0.0833,  0.0275, -0.0612, -0.0770, -0.1256,  0.0082,  0.0186,  0.0476],
+                   dtype=torch.float64), 'velocity': tensor([ 0.2221,  0.2256,  0.5930,  2.6937, -3.5865, -1.5479,  0.0187, -0.6825,
+                     0.5224], dtype=torch.float64)}, tensor([0.0018], dtype=torch.float64), tensor([False]), tensor([False]), {})
+            >>> # same environment with observation stacked
+            >>> from torchrl.envs import CatTensors
+            >>> envgym = gym.make("DMC-cheetah-v0", transform=CatTensors(in_keys=["position", "velocity"], out_key="observation"))
+            >>> envgym.reset()
+            ({'observation': tensor([-0.1005,  0.0335, -0.0268,  0.0133, -0.0627,  0.0074, -0.0488, -0.0353,
+                    -0.0075, -0.0069,  0.0098, -0.0058,  0.0033, -0.0157, -0.0004, -0.0381,
+                    -0.0452], dtype=torch.float64)}, {})
+            >>> # same environment with numpy observations
+            >>> envgym = gym.make("DMC-cheetah-v0", transform=CatTensors(in_keys=["position", "velocity"], out_key="observation"), to_numpy=True)
+            >>> envgym.reset()
+            ({'observation': array([-0.11355747,  0.04257728,  0.00408397,  0.04155852, -0.0389733 ,
+                   -0.01409826, -0.0978704 , -0.08808327,  0.03970837,  0.00535434,
+                   -0.02353762,  0.05116226,  0.02788907,  0.06848346,  0.05154399,
+                    0.0371798 ,  0.05128025])}, {})
+            >>> # If gymnasium is installed, we can register the environment there too.
+            >>> DMControlEnv.register_gym("DMC-cheetah-v0", to_numpy=False, backend="gymnasium", env_name="cheetah", task_name="run")
+            >>> import gymnasium
+            >>> envgym = gymnasium.make("DMC-cheetah-v0")
+            >>> envgym.seed(0)
+            >>> torch.manual_seed(0)
+            >>> envgym.reset()
+            ({'position': tensor([-0.0855,  0.0215, -0.0881, -0.0412, -0.1101,  0.0080,  0.0254,  0.0424],
+                   dtype=torch.float64), 'velocity': tensor([ 1.9609e-02, -1.9776e-04, -1.6347e-03,  3.3842e-02,  2.5338e-02,
+                     3.3064e-02,  1.0381e-04,  7.6656e-05,  1.0204e-02],
+                   dtype=torch.float64)}, {})
+
+        .. note::
+            This feature also works for stateless environments (eg, :class:`~torchrl.envs.BraxEnv`).
+
+                >>> import gymnasium
+                >>> import torch
+                >>> from tensordict import TensorDict
+                >>> from torchrl.envs import BraxEnv, SelectTransform
+                >>>
+                >>> # get action for dydactic purposes
+                >>> env = BraxEnv("ant", batch_size=[2])
+                >>> env.set_seed(0)
+                >>> torch.manual_seed(0)
+                >>> td = env.rollout(10)
+                >>>
+                >>> actions = td.get("action")
+                >>>
+                >>> # register env
+                >>> env.register_gym("Brax-Ant-v0", env_name="ant", batch_size=[2], info_keys=["state"])
+                >>> gym_env = gymnasium.make("Brax-Ant-v0")
+                >>> gym_env.seed(0)
+                >>> torch.manual_seed(0)
+                >>>
+                >>> gym_env.reset()
+                >>> obs = []
+                >>> for i in range(10):
+                ...     obs, reward, terminated, truncated, info = gym_env.step(td[..., i].get("action"))
+
+
+        """
+        from torchrl.envs.libs.gym import gym_backend, set_gym_backend
+
+        if backend is None:
+            backend = gym_backend()
+
+        with set_gym_backend(backend):
+            return cls._register_gym(
+                id=id,
+                entry_point=entry_point,
+                transform=transform,
+                info_keys=info_keys,
+                to_numpy=to_numpy,
+                reward_threshold=reward_threshold,
+                nondeterministic=nondeterministic,
+                max_episode_steps=max_episode_steps,
+                order_enforce=order_enforce,
+                autoreset=autoreset,
+                disable_env_checker=disable_env_checker,
+                apply_api_compatibility=apply_api_compatibility,
+                **kwargs,
+            )
+
+    _GYM_UNRECOGNIZED_KWARG = (
+        "The keyword argument {} is not compatible with gym version {}"
+    )
+
+    @implement_for("gym", "0.26", None, class_method=True)
+    def _register_gym(
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gym
+        from torchrl.envs.libs._gym_utils import _TorchRLGymWrapper
+
+        if entry_point is None:
+            entry_point = cls
+        entry_point = partial(
+            _TorchRLGymWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        return gym.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+            order_enforce=order_enforce,
+            autoreset=bool(autoreset),
+            disable_env_checker=disable_env_checker,
+            apply_api_compatibility=apply_api_compatibility,
+        )
+
+    @implement_for("gym", "0.25", "0.26", class_method=True)
+    def _register_gym(  # noqa: F811
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gym
+
+        if apply_api_compatibility is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "apply_api_compatibility", gym.__version__
+                )
+            )
+        from torchrl.envs.libs._gym_utils import _TorchRLGymWrapper
+
+        if entry_point is None:
+            entry_point = cls
+        entry_point = partial(
+            _TorchRLGymWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        return gym.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+            order_enforce=order_enforce,
+            autoreset=bool(autoreset),
+            disable_env_checker=disable_env_checker,
+        )
+
+    @implement_for("gym", "0.24", "0.25", class_method=True)
+    def _register_gym(  # noqa: F811
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gym
+
+        if apply_api_compatibility is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "apply_api_compatibility", gym.__version__
+                )
+            )
+        if disable_env_checker is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "disable_env_checker", gym.__version__
+                )
+            )
+        from torchrl.envs.libs._gym_utils import _TorchRLGymWrapper
+
+        if entry_point is None:
+            entry_point = cls
+        entry_point = partial(
+            _TorchRLGymWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        return gym.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+            order_enforce=order_enforce,
+            autoreset=bool(autoreset),
+        )
+
+    @implement_for("gym", "0.21", "0.24", class_method=True)
+    def _register_gym(  # noqa: F811
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gym
+
+        if apply_api_compatibility is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "apply_api_compatibility", gym.__version__
+                )
+            )
+        if disable_env_checker is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "disable_env_checker", gym.__version__
+                )
+            )
+        if autoreset is not None:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format("autoreset", gym.__version__)
+            )
+        from torchrl.envs.libs._gym_utils import _TorchRLGymWrapper
+
+        if entry_point is None:
+            entry_point = cls
+        entry_point = partial(
+            _TorchRLGymWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        return gym.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+            order_enforce=order_enforce,
+            autoreset=bool(autoreset),
+        )
+
+    @implement_for("gym", None, "0.21", class_method=True)
+    def _register_gym(  # noqa: F811
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gym
+        from torchrl.envs.libs._gym_utils import _TorchRLGymWrapper
+
+        if order_enforce is not True:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format("order_enforce", gym.__version__)
+            )
+        if disable_env_checker is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "disable_env_checker", gym.__version__
+                )
+            )
+        if autoreset is not None:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format("autoreset", gym.__version__)
+            )
+        if apply_api_compatibility is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "apply_api_compatibility", gym.__version__
+                )
+            )
+        if entry_point is None:
+            entry_point = cls
+        entry_point = partial(
+            _TorchRLGymWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        return gym.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+        )
+
+    @implement_for("gymnasium", None, "1.0.0", class_method=True)
+    def _register_gym(  # noqa: F811
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gymnasium
+        from torchrl.envs.libs._gym_utils import _TorchRLGymnasiumWrapper
+
+        if entry_point is None:
+            entry_point = cls
+
+        entry_point = partial(
+            _TorchRLGymnasiumWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        return gymnasium.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+            order_enforce=order_enforce,
+            autoreset=bool(autoreset),
+            disable_env_checker=disable_env_checker,
+            apply_api_compatibility=apply_api_compatibility,
+        )
+
+    @implement_for("gymnasium", "1.1.0", class_method=True)
+    def _register_gym(  # noqa: F811
+        cls,
+        id,
+        entry_point: Callable | None = None,
+        transform: Transform | None = None,  # noqa: F821
+        info_keys: list[NestedKey] | None = None,
+        to_numpy: bool = False,
+        reward_threshold: float | None = None,
+        nondeterministic: bool = False,
+        max_episode_steps: int | None = None,
+        order_enforce: bool = True,
+        autoreset: bool | None = None,
+        disable_env_checker: bool = False,
+        apply_api_compatibility: bool = False,
+        **kwargs,
+    ):
+        import gymnasium
+        from torchrl.envs.libs._gym_utils import _TorchRLGymnasiumWrapper
+
+        if autoreset is not None:
+            raise TypeError(
+                f"the autoreset argument is deprecated in gymnasium>=1.0. Got autoreset={autoreset}"
+            )
+        if entry_point is None:
+            entry_point = cls
+
+        entry_point = partial(
+            _TorchRLGymnasiumWrapper,
+            entry_point=entry_point,
+            info_keys=info_keys,
+            to_numpy=to_numpy,
+            transform=transform,
+            **kwargs,
+        )
+        if apply_api_compatibility is not False:
+            raise TypeError(
+                cls._GYM_UNRECOGNIZED_KWARG.format(
+                    "apply_api_compatibility", gymnasium.__version__
+                )
+            )
+        return gymnasium.register(
+            id=id,
+            entry_point=entry_point,
+            reward_threshold=reward_threshold,
+            nondeterministic=nondeterministic,
+            max_episode_steps=max_episode_steps,
+            order_enforce=order_enforce,
+            disable_env_checker=disable_env_checker,
+        )
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError(
+            "EnvBase.forward is not implemented. If you ended here during a call to `ParallelEnv(...)`, please use "
+            "a constructor such as `ParallelEnv(num_env, lambda env=env: env)` instead. "
+            "Batched envs require constructors because environment instances may not always be serializable."
+        )
+
+    @abc.abstractmethod
+    def _step(
+        self,
+        tensordict: TensorDictBase,
+    ) -> TensorDictBase:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
+        raise NotImplementedError
+
+    def reset(
+        self,
+        tensordict: TensorDictBase | None = None,
+        **kwargs,
+    ) -> TensorDictBase:
+        """Resets the environment.
+
+        As for step and _step, only the private method :obj:`_reset` should be overwritten by EnvBase subclasses.
+
+        Args:
+            tensordict (TensorDictBase, optional): tensordict to be used to contain the resulting new observation.
+                In some cases, this input can also be used to pass argument to the reset function.
+            kwargs (optional): other arguments to be passed to the native
+                reset function.
+
+        Returns:
+            a tensordict (or the input tensordict, if any), modified in place with the resulting observations.
+
+        .. note:: `reset` should not be overwritten by :class:`~torchrl.envs.EnvBase` subclasses. The method to
+            modify is :meth:`~torchrl.envs.EnvBase._reset`.
+
+        """
+        if tensordict is not None:
+            self._assert_tensordict_shape(tensordict)
+
+        select_reset_only = kwargs.pop("select_reset_only", False)
+        if select_reset_only and tensordict is not None:
+            # When making rollouts with step_and_maybe_reset, it can happen that a tensordict has
+            # keys that are used by reset to optionally set the reset state (eg, the fen in chess). If that's the
+            # case and we don't throw them away here, reset will just be a no-op (put the env in the state reached
+            # during the previous step).
+            # Therefore, maybe_reset tells reset to temporarily hide the non-reset keys.
+            # To make step_and_maybe_reset handle custom reset states, some version of TensorDictPrimer should be used.
+            tensordict_reset = self._reset(
+                tensordict.select(*self.reset_keys, strict=False), **kwargs
+            )
+        else:
+            tensordict_reset = self._reset(tensordict, **kwargs)
+        # We assume that this is done properly
+        # if reset.device != self.device:
+        #     reset = reset.to(self.device, non_blocking=True)
+        if tensordict_reset is tensordict:
+            raise RuntimeError(
+                "EnvBase._reset should return outplace changes to the input "
+                "tensordict. Consider emptying the TensorDict first (e.g. tensordict.empty()) "
+                "inside _reset before writing new tensors onto this new instance."
+            )
+        if not isinstance(tensordict_reset, TensorDictBase):
+            raise RuntimeError(
+                f"env._reset returned an object of type {type(tensordict_reset)} but a TensorDict was expected."
+            )
+        return self._reset_proc_data(tensordict, tensordict_reset)
+
+    def _reset_proc_data(self, tensordict, tensordict_reset):
+        self._complete_done(self.full_done_spec, tensordict_reset)
+        self._reset_check_done(tensordict, tensordict_reset)
+        if tensordict is not None:
+            return _update_during_reset(tensordict_reset, tensordict, self.reset_keys)
+        return tensordict_reset
+
+    def _reset_check_done(self, tensordict, tensordict_reset):
+        """Checks the done status after reset.
+
+        If _reset signals were passed, we check that the env is not done for these
+        indices.
+
+        We also check that the input tensordict contained ``"done"``s if the
+        reset is partial and incomplete.
+
+        """
+        # we iterate over (reset_key, (done_key, truncated_key)) and check that all
+        # values where reset was true now have a done set to False.
+        # If no reset was present, all done and truncated must be False
+
+        # Once we checked a root, we don't check its leaves - so keep track of the roots. Fortunately, we sort the done
+        #  keys in the done_keys_group from root to leaf
+        prefix_complete = set()
+        for reset_key, done_key_group in zip(self.reset_keys, self.done_keys_groups):
+            skip = False
+            if isinstance(reset_key, tuple):
+                for i in range(len(reset_key) - 1):
+                    if reset_key[:i] in prefix_complete:
+                        skip = True
+                        break
+            if skip:
+                continue
+            reset_value = (
+                tensordict.get(reset_key, default=None)
+                if tensordict is not None
+                else None
+            )
+            prefix_complete.add(() if isinstance(reset_key, str) else reset_key[:-1])
+            if reset_value is not None:
+                for done_key in done_key_group:
+                    done_val = tensordict_reset.get(done_key)
+                    if (
+                        done_val.any()
+                        and done_val[reset_value].any()
+                        and not self._allow_done_after_reset
+                    ):
+                        raise RuntimeError(
+                            f"Env done entry '{done_key}' was (partially) True after reset on specified '_reset' dimensions. This is not allowed."
+                        )
+                    if (
+                        done_key not in tensordict.keys(True)
+                        and done_val[~reset_value].any()
+                    ):
+                        warnings.warn(
+                            f"A partial `'_reset'` key has been passed to `reset` ({reset_key}), "
+                            f"but the corresponding done_key ({done_key}) wasn't present in the input "
+                            f"tensordict. "
+                            f"This is discouraged, since the input tensordict should contain "
+                            f"all the data not being reset."
+                        )
+                        # we set the done val to tensordict, to make sure that
+                        # _update_during_reset does not pad the value
+                        tensordict.set(done_key, done_val)
+            elif not self._allow_done_after_reset:
+                for done_key in done_key_group:
+                    if tensordict_reset.get(done_key).any():
+                        raise RuntimeError(
+                            f"The done entry '{done_key}' was (partially) True after a call to reset() in env {self}."
+                        )
+
+    def numel(self) -> int:
+        return prod(self.batch_size)
+
+    def set_seed(
+        self, seed: int | None = None, static_seed: bool = False
+    ) -> int | None:
+        """Sets the seed of the environment and returns the next seed to be used (which is the input seed if a single environment is present).
+
+        Args:
+            seed (int): seed to be set. The seed is set only locally in the environment. To handle the global seed,
+                see :func:`~torch.manual_seed`.
+            static_seed (bool, optional): if ``True``, the seed is not incremented.
+                Defaults to False
+
+        Returns:
+            integer representing the "next seed": i.e. the seed that should be
+            used for another environment if created concomitantly to this environment.
+
+        """
+        self._set_seed(seed)
+        if seed is not None and not static_seed:
+            new_seed = seed_generator(seed)
+            seed = new_seed
+        return seed
+
+    @abc.abstractmethod
+    def _set_seed(self, seed: int | None) -> None:
+        raise NotImplementedError
+
+    def set_state(self):
+        raise NotImplementedError
+
+    def _assert_tensordict_shape(self, tensordict: TensorDictBase) -> None:
+        if (
+            self.batch_locked or self.batch_size != ()
+        ) and tensordict.batch_size != self.batch_size:
+            raise RuntimeError(
+                f"Expected a tensordict with shape==env.batch_size, "
+                f"got {tensordict.batch_size} and {self.batch_size}"
+            )
+
+    def all_actions(self, tensordict: TensorDictBase | None = None) -> TensorDictBase:
+        """Generates all possible actions from the action spec.
+
+        This only works in environments with fully discrete actions.
+
+        Args:
+            tensordict (TensorDictBase, optional): If given, :meth:`~.reset`
+                is called with this tensordict.
+
+        Returns:
+            a tensordict object with the "action" entry updated with a batch of
+            all possible actions. The actions are stacked together in the
+            leading dimension.
+        """
+        if tensordict is not None:
+            self.reset(tensordict)
+
+        return self.full_action_spec.enumerate(use_mask=True)
+
+    def rand_action(self, tensordict: TensorDictBase | None = None):
+        """Performs a random action given the action_spec attribute.
+
+        Args:
+            tensordict (TensorDictBase, optional): tensordict where the resulting action should be written.
+
+        Returns:
+            a tensordict object with the "action" entry updated with a random
+            sample from the action-spec.
+
+        """
+        shape = torch.Size([])
+        if not self.batch_locked:
+            if not self.batch_size and tensordict is not None:
+                # if we can't infer the batch-size from the env, take it from tensordict
+                shape = tensordict.shape
+            elif not self.batch_size:
+                # if tensordict wasn't provided, we assume empty batch size
+                shape = torch.Size([])
+            elif tensordict.shape != self.batch_size:
+                # if tensordict is not None and the env has a batch size, their shape must match
+                raise RuntimeError(
+                    "The input tensordict and the env have a different batch size: "
+                    f"env.batch_size={self.batch_size} and tensordict.batch_size={tensordict.shape}. "
+                    f"Non batch-locked environment require the env batch-size to be either empty or to"
+                    f" match the tensordict one."
+                )
+        # We generate the action from the full_action_spec
+        r = self.input_spec["full_action_spec"].rand(shape)
+        if tensordict is None:
+            return r
+        tensordict.update(r)
+        return tensordict
+
+    def rand_step(self, tensordict: TensorDictBase | None = None) -> TensorDictBase:
+        """Performs a random step in the environment given the action_spec attribute.
+
+        Args:
+            tensordict (TensorDictBase, optional): tensordict where the resulting info should be written.
+
+        Returns:
+            a tensordict object with the new observation after a random step in the environment. The action will
+            be stored with the "action" key.
+
+        """
+        tensordict = self.rand_action(tensordict)
+        return self.step(tensordict)
+
+    @property
+    def specs(self) -> Composite:
+        """Returns a Composite container where all the environment are present.
+
+        This feature allows one to create an environment, retrieve all of the specs in a single data container and then
+        erase the environment from the workspace.
+
+        """
+        return Composite(
+            output_spec=self.output_spec,
+            input_spec=self.input_spec,
+            shape=self.batch_size,
+        )
+
+    @property
+    @_cache_value
+    def _has_dynamic_specs(self) -> bool:
+        return _has_dynamic_specs(self.specs)
+
+    def rollout(
+        self,
+        max_steps: int,
+        policy: Callable[[TensorDictBase], TensorDictBase] | None = None,
+        callback: Callable[[TensorDictBase, ...], Any] | None = None,
+        *,
+        auto_reset: bool = True,
+        auto_cast_to_device: bool = False,
+        break_when_any_done: bool | None = None,
+        break_when_all_done: bool | None = None,
+        return_contiguous: bool | None = False,
+        tensordict: TensorDictBase | None = None,
+        set_truncated: bool = False,
+        out=None,
+        trust_policy: bool = False,
+        storing_device: DEVICE_TYPING | None = None,
+    ) -> TensorDictBase:
+        """Executes a rollout in the environment.
+
+        The function will return as soon as any of the contained environments
+        reaches any of the done states.
+
+        Args:
+            max_steps (int): maximum number of steps to be executed. The actual number of steps can be smaller if
+                the environment reaches a done state before max_steps have been executed.
+            policy (callable, optional): callable to be called to compute the desired action.
+                If no policy is provided, actions will be called using :obj:`env.rand_step()`.
+                The policy can be any callable that reads either a tensordict or
+                the entire sequence of observation entries __sorted as__ the ``env.observation_spec.keys()``.
+                Defaults to `None`.
+            callback (Callable[[TensorDict], Any], optional): function to be called at each iteration with the given
+                TensorDict. Defaults to ``None``. The output of ``callback`` will not be collected, it is the user
+                responsibility to save any result within the callback call if data needs to be carried over beyond
+                the call to ``rollout``.
+
+        Keyword Args:
+            auto_reset (bool, optional): if ``True``, the contained environments will be reset before starting the
+                rollout. If ``False``, then the rollout will continue from a previous state, which requires the
+                ``tensordict`` argument to be passed with the previous rollout. Default is ``True``.
+            auto_cast_to_device (bool, optional): if ``True``, the device of the tensordict is automatically cast to the
+                policy device before the policy is used. Default is ``False``.
+            break_when_any_done (bool): if ``True``, break when any of the contained environments reaches any of the
+                done states. If ``False``, then the done environments are reset automatically. Default is ``True``.
+
+                .. seealso:: The :ref:`Partial resets <ref_partial_resets>` of the documentation gives more
+                    information about partial resets.
+
+            break_when_all_done (bool, optional): if ``True``, break if all of the contained environments reach any
+                of the done states. If ``False``, break if at least one environment reaches any of the done states.
+                Default is ``False``.
+
+                .. seealso:: The :ref:`Partial steps <ref_partial_steps>` of the documentation gives more
+                    information about partial resets.
+
+            return_contiguous (bool): if False, a LazyStackedTensorDict will be returned. Default is `True` if
+                the env does not have dynamic specs, otherwise `False`.
+            tensordict (TensorDict, optional): if ``auto_reset`` is False, an initial
+                tensordict must be provided. Rollout will check if this tensordict has done flags and reset the
+                environment in those dimensions (if needed).
+                This normally should not occur if ``tensordict`` is the output of a reset, but can occur
+                if ``tensordict`` is the last step of a previous rollout.
+                A ``tensordict`` can also be provided when ``auto_reset=True`` if metadata need to be passed
+                to the ``reset`` method, such as a batch-size or a device for stateless environments.
+            set_truncated (bool, optional): if ``True``, ``"truncated"`` and ``"done"`` keys will be set to
+                ``True`` after completion of the rollout. If no ``"truncated"`` is found within the
+                ``done_spec``, an exception is raised.
+                Truncated keys can be set through ``env.add_truncated_keys``.
+                Defaults to ``False``.
+            trust_policy (bool, optional): if ``True``, a non-TensorDictModule policy will be trusted to be
+                assumed to be compatible with the collector. This defaults to ``True`` for CudaGraphModules
+                and ``False`` otherwise.
+            storing_device (Device, optional): if provided, the tensordict will be stored on this device.
+                Defaults to ``None``.
+
+        Returns:
+            TensorDict object containing the resulting trajectory.
+
+        The data returned will be marked with a "time" dimension name for the last
+        dimension of the tensordict (at the ``env.ndim`` index).
+
+        ``rollout`` is quite handy to display what the data structure of the
+        environment looks like.
+
+        Examples:
+            >>> # Using rollout without a policy
+            >>> from torchrl.envs.libs.gym import GymEnv
+            >>> from torchrl.envs.transforms import TransformedEnv, StepCounter
+            >>> env = TransformedEnv(GymEnv("Pendulum-v1"), StepCounter(max_steps=20))
+            >>> rollout = env.rollout(max_steps=1000)
+            >>> print(rollout)
+            TensorDict(
+                fields={
+                    action: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                    done: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    next: TensorDict(
+                        fields={
+                            done: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            observation: Tensor(shape=torch.Size([20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                            reward: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                            step_count: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                            truncated: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                        batch_size=torch.Size([20]),
+                        device=cpu,
+                        is_shared=False),
+                    observation: Tensor(shape=torch.Size([20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                    step_count: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                    truncated: Tensor(shape=torch.Size([20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                batch_size=torch.Size([20]),
+                device=cpu,
+                is_shared=False)
+            >>> print(rollout.names)
+            ['time']
+            >>> # with envs that contain more dimensions
+            >>> from torchrl.envs import SerialEnv
+            >>> env = SerialEnv(3, lambda: TransformedEnv(GymEnv("Pendulum-v1"), StepCounter(max_steps=20)))
+            >>> rollout = env.rollout(max_steps=1000)
+            >>> print(rollout)
+            TensorDict(
+                fields={
+                    action: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                    done: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    next: TensorDict(
+                        fields={
+                            done: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            observation: Tensor(shape=torch.Size([3, 20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                            reward: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                            step_count: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                            truncated: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                        batch_size=torch.Size([3, 20]),
+                        device=cpu,
+                        is_shared=False),
+                    observation: Tensor(shape=torch.Size([3, 20, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                    step_count: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                    truncated: Tensor(shape=torch.Size([3, 20, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                batch_size=torch.Size([3, 20]),
+                device=cpu,
+                is_shared=False)
+            >>> print(rollout.names)
+            [None, 'time']
+
+        Using a policy (a regular :class:`~torch.nn.Module` or a :class:`~tensordict.nn.TensorDictModule`)
+        is also easy:
+
+        Examples:
+            >>> from torch import nn
+            >>> env = GymEnv("CartPole-v1", categorical_action_encoding=True)
+            >>> class ArgMaxModule(nn.Module):
+            ...     def forward(self, values):
+            ...         return values.argmax(-1)
+            >>> n_obs = env.observation_spec["observation"].shape[-1]
+            >>> n_act = env.action_spec.n
+            >>> # A deterministic policy
+            >>> policy = nn.Sequential(
+            ...     nn.Linear(n_obs, n_act),
+            ...     ArgMaxModule())
+            >>> env.rollout(max_steps=10, policy=policy)
+            TensorDict(
+                fields={
+                    action: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False),
+                    done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    next: TensorDict(
+                        fields={
+                            done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                            reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                            terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                        batch_size=torch.Size([10]),
+                        device=cpu,
+                        is_shared=False),
+                    observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                    terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                batch_size=torch.Size([10]),
+                device=cpu,
+                is_shared=False)
+            >>> # Under the hood, rollout will wrap the policy in a TensorDictModule
+            >>> # To speed things up we can do that ourselves
+            >>> from tensordict.nn import TensorDictModule
+            >>> policy = TensorDictModule(policy, in_keys=list(env.observation_spec.keys()), out_keys=["action"])
+            >>> env.rollout(max_steps=10, policy=policy)
+            TensorDict(
+                fields={
+                    action: Tensor(shape=torch.Size([10]), device=cpu, dtype=torch.int64, is_shared=False),
+                    done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    next: TensorDict(
+                        fields={
+                            done: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                            reward: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                            terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                        batch_size=torch.Size([10]),
+                        device=cpu,
+                        is_shared=False),
+                    observation: Tensor(shape=torch.Size([10, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                    terminated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    truncated: Tensor(shape=torch.Size([10, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                batch_size=torch.Size([10]),
+                device=cpu,
+                is_shared=False)
+
+
+        In some instances, contiguous tensordict cannot be obtained because
+        they cannot be stacked. This can happen when the data returned at
+        each step may have a different shape, or when different environments
+        are executed together. In that case, ``return_contiguous=False``
+        will cause the returned tensordict to be a lazy stack of tensordicts:
+
+        Examples of non-contiguous rollout:
+            >>> rollout = env.rollout(4, return_contiguous=False)
+            >>> print(rollout)
+            LazyStackedTensorDict(
+                fields={
+                    action: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                    done: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    next: LazyStackedTensorDict(
+                        fields={
+                            done: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            observation: Tensor(shape=torch.Size([3, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                            reward: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                            step_count: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                            truncated: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                        batch_size=torch.Size([3, 4]),
+                        device=cpu,
+                        is_shared=False),
+                    observation: Tensor(shape=torch.Size([3, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False),
+                    step_count: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                    truncated: Tensor(shape=torch.Size([3, 4, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                batch_size=torch.Size([3, 4]),
+                device=cpu,
+                is_shared=False)
+                >>> print(rollout.names)
+                [None, 'time']
+
+        Rollouts can be used in a loop to emulate data collection.
+        To do so, you need to pass as input the last tensordict coming from the previous rollout after calling
+        :func:`~torchrl.envs.utils.step_mdp` on it.
+
+        Examples of data collection rollouts:
+            >>> from torchrl.envs import GymEnv, step_mdp
+            >>> env = GymEnv("CartPole-v1")
+            >>> epochs = 10
+            >>> input_td = env.reset()
+            >>> for i in range(epochs):
+            ...     rollout_td = env.rollout(
+            ...         max_steps=100,
+            ...         break_when_any_done=False,
+            ...         auto_reset=False,
+            ...         tensordict=input_td,
+            ...     )
+            ...     input_td = step_mdp(
+            ...         rollout_td[..., -1],
+            ...     )
+
+        """
+        if break_when_any_done is None:  # True by default
+            if break_when_all_done:  # all overrides
+                break_when_any_done = False
+            else:
+                break_when_any_done = True
+        if break_when_all_done is None:
+            # There is no case where break_when_all_done is True by default
+            break_when_all_done = False
+        if break_when_all_done and break_when_any_done:
+            raise TypeError(
+                "Cannot have both break_when_all_done and break_when_any_done True at the same time."
+            )
+        if return_contiguous is None:
+            return_contiguous = not self._has_dynamic_specs
+        if policy is not None:
+            policy = _make_compatible_policy(
+                policy,
+                self.observation_spec,
+                env=self,
+                fast_wrap=True,
+                trust_policy=trust_policy,
+            )
+            if auto_cast_to_device:
+                try:
+                    policy_device = next(policy.parameters()).device
+                except (StopIteration, AttributeError):
+                    policy_device = None
+            else:
+                policy_device = None
+        else:
+            policy = self.rand_action
+            policy_device = None
+
+        env_device = self.device
+
+        if auto_reset:
+            tensordict = self.reset(tensordict)
+        elif tensordict is None:
+            raise RuntimeError("tensordict must be provided when auto_reset is False")
+        else:
+            tensordict = self.maybe_reset(tensordict)
+
+        kwargs = {
+            "tensordict": tensordict,
+            "auto_cast_to_device": auto_cast_to_device,
+            "max_steps": max_steps,
+            "policy": policy,
+            "policy_device": policy_device,
+            "env_device": env_device,
+            "storing_device": None
+            if storing_device is None
+            else torch.device(storing_device),
+            "callback": callback,
+        }
+        if break_when_any_done or break_when_all_done:
+            tensordicts = self._rollout_stop_early(
+                break_when_all_done=break_when_all_done,
+                break_when_any_done=break_when_any_done,
+                **kwargs,
+            )
+        else:
+            tensordicts = self._rollout_nonstop(**kwargs)
+        batch_size = self.batch_size if tensordict is None else tensordict.batch_size
+        if return_contiguous:
+            try:
+                out_td = torch.stack(tensordicts, len(batch_size), out=out)
+            except RuntimeError as err:
+                if (
+                    re.match(
+                        "The shapes of the tensors to stack is incompatible", str(err)
+                    )
+                    and self._has_dynamic_specs
+                ):
+                    raise RuntimeError(
+                        "The environment specs are dynamic. Call rollout with return_contiguous=False."
+                    )
+                if re.match(
+                    "The sets of keys in the tensordicts to stack are exclusive",
+                    str(err),
+                ):
+                    for reward_key in self.reward_keys:
+                        if any(reward_key in td for td in tensordicts):
+                            raise RuntimeError(
+                                "The reward key was present in the root tensordict of at least one of the tensordicts to stack. "
+                                "The likely cause is that your environment returns a reward during a call to `reset`, which is not allowed. "
+                                "To fix this, you should return the reward in the `step` method but not in during `reset`. If you need a reward "
+                                "to be returned during `reset`, submit an issue on github."
+                            )
+                raise
+        else:
+            out_td = LazyStackedTensorDict.maybe_dense_stack(
+                tensordicts, len(batch_size), out=out
+            )
+        if set_truncated:
+            found_truncated = False
+            for key in self.done_keys:
+                if _ends_with(key, "truncated"):
+                    val = out_td.get(("next", key))
+                    done = out_td.get(("next", _replace_last(key, "done")))
+                    val[(slice(None),) * (out_td.ndim - 1) + (-1,)] = True
+                    out_td.set(("next", key), val)
+                    out_td.set(("next", _replace_last(key, "done")), val | done)
+                    found_truncated = True
+            if not found_truncated:
+                raise RuntimeError(
+                    "set_truncated was set to True but no truncated key could be found. "
+                    "Make sure a 'truncated' entry was set in the environment "
+                    "full_done_keys using `env.add_truncated_keys()`."
+                )
+
+        out_td.refine_names(..., "time")
+        return out_td
+
+    @_maybe_unlock
+    def add_truncated_keys(self) -> EnvBase:
+        """Adds truncated keys to the environment."""
+        i = 0
+        for key in self.done_keys:
+            i += 1
+            truncated_key = _replace_last(key, "truncated")
+            self.full_done_spec[truncated_key] = self.full_done_spec[key].clone()
+        if i == 0:
+            raise KeyError(f"Couldn't find done keys. done_spec={self.full_done_specs}")
+
+        return self
+
+    def step_mdp(self, next_tensordict: TensorDictBase) -> TensorDictBase:
+        """Advances the environment state by one step using the provided `next_tensordict`.
+
+        This method updates the environment's state by transitioning from the current
+        state to the next, as defined by the `next_tensordict`. The resulting tensordict
+        includes updated observations and any other relevant state information, with
+        keys managed according to the environment's specifications.
+
+        Internally, this method utilizes a precomputed :class:`~torchrl.envs.utils._StepMDP` instance to efficiently
+        handle the transition of state, observation, action, reward, and done keys. The
+        :class:`~torchrl.envs.utils._StepMDP` class optimizes the process by precomputing the keys to include and
+        exclude, reducing runtime overhead during repeated calls. The :class:`~torchrl.envs.utils._StepMDP` instance
+        is created with `exclude_action=False`, meaning that action keys are retained in
+        the root tensordict.
+
+        Args:
+            next_tensordict (TensorDictBase): A tensordict containing the state of the
+                environment at the next time step. This tensordict should include keys
+                for observations, actions, rewards, and done flags, as defined by the
+                environment's specifications.
+
+        Returns:
+            TensorDictBase: A new tensordict representing the environment state after
+            advancing by one step.
+
+        .. note:: The method ensures that the environment's key specifications are validated
+              against the provided `next_tensordict`, issuing warnings if discrepancies
+              are found.
+
+        .. note:: This method is designed to work efficiently with environments that have
+              consistent key specifications, leveraging the `_StepMDP` class to minimize
+              overhead.
+
+        Example:
+            >>> from torchrl.envs import GymEnv
+            >>> env = GymEnv("Pendulum-1")
+            >>> data = env.reset()
+            >>> for i in range(10):
+            ...     # compute action
+            ...     env.rand_action(data)
+            ...     # Perform action
+            ...     next_data = env.step(reset_data)
+            ...     data = env.step_mdp(next_data)
+        """
+        return self._step_mdp(next_tensordict)
+
+    @property
+    @_cache_value
+    def _step_mdp(self) -> Callable[[TensorDictBase], TensorDictBase]:
+        return _StepMDP(self, exclude_action=False)
+
+    def _rollout_stop_early(
+        self,
+        *,
+        break_when_any_done,
+        break_when_all_done,
+        tensordict,
+        auto_cast_to_device,
+        max_steps,
+        policy,
+        policy_device,
+        env_device,
+        storing_device,
+        callback,
+    ):
+        # Get the sync func
+        if auto_cast_to_device:
+            sync_func = _get_sync_func(policy_device, env_device)
+        tensordicts = []
+        partial_steps = True
+        for i in range(max_steps):
+            if auto_cast_to_device:
+                if policy_device is not None:
+                    tensordict = tensordict.to(policy_device, non_blocking=True)
+                    sync_func()
+                else:
+                    tensordict.clear_device_()
+            # In case policy(..) does not modify in-place - no-op for TensorDict and related
+            tensordict.update(policy(tensordict))
+            if auto_cast_to_device:
+                if env_device is not None:
+                    tensordict = tensordict.to(env_device, non_blocking=True)
+                    sync_func()
+                else:
+                    tensordict.clear_device_()
+            tensordict = self.step(tensordict)
+            if storing_device is None or tensordict.device == storing_device:
+                td_append = tensordict.copy()
+            else:
+                td_append = tensordict.to(storing_device)
+            if break_when_all_done:
+                if partial_steps is not True and not partial_steps.all():
+                    # At least one step is partial
+                    td_append.pop("_step", None)
+                    td_append = torch.where(
+                        partial_steps.view(td_append.shape), td_append, tensordicts[-1]
+                    )
+
+            tensordicts.append(td_append)
+
+            if i == max_steps - 1:
+                # we don't truncate as one could potentially continue the run
+                break
+            tensordict = self._step_mdp(tensordict)
+
+            if break_when_any_done:
+                # done and truncated are in done_keys
+                # We read if any key is done.
+                any_done = _terminated_or_truncated(
+                    tensordict,
+                    full_done_spec=self.output_spec["full_done_spec"],
+                    key=None,
+                )
+                if any_done:
+                    break
+            else:
+                # Write the '_step' entry, indicating which step is to be undertaken
+                _terminated_or_truncated(
+                    tensordict,
+                    full_done_spec=self.output_spec["full_done_spec"],
+                    key="_neg_step",
+                    write_full_false=False,
+                )
+                # This is what differentiates _step and _reset: we need to flip _step False -> True
+                partial_step_curr = tensordict.pop("_neg_step", None)
+                if partial_step_curr is not None:
+                    partial_step_curr = ~partial_step_curr
+                    partial_steps = partial_steps & partial_step_curr
+                if partial_steps is not True:
+                    if not partial_steps.any():
+                        break
+                    # Write the final _step entry
+                    tensordict.set("_step", partial_steps)
+
+            if callback is not None:
+                callback(self, tensordict)
+        return tensordicts
+
+    def _rollout_nonstop(
+        self,
+        *,
+        tensordict,
+        auto_cast_to_device,
+        max_steps,
+        policy,
+        policy_device,
+        env_device,
+        storing_device,
+        callback,
+    ):
+        if auto_cast_to_device:
+            sync_func = _get_sync_func(policy_device, env_device)
+        tensordicts = []
+        tensordict_ = tensordict
+        for i in range(max_steps):
+            if auto_cast_to_device:
+                if policy_device is not None:
+                    tensordict_ = tensordict_.to(policy_device, non_blocking=True)
+                    sync_func()
+                else:
+                    tensordict_.clear_device_()
+            # In case policy(..) does not modify in-place - no-op for TensorDict and related
+            tensordict_.update(policy(tensordict_))
+            if auto_cast_to_device:
+                if env_device is not None:
+                    tensordict_ = tensordict_.to(env_device, non_blocking=True)
+                    sync_func()
+                else:
+                    tensordict_.clear_device_()
+            if i == max_steps - 1:
+                tensordict = self.step(tensordict_)
+            else:
+                tensordict, tensordict_ = self.step_and_maybe_reset(tensordict_)
+            if storing_device is None or tensordict.device == storing_device:
+                tensordicts.append(tensordict)
+            else:
+                tensordicts.append(tensordict.to(storing_device))
+            if i == max_steps - 1:
+                # we don't truncate as one could potentially continue the run
+                break
+            if callback is not None:
+                callback(self, tensordict)
+
+        return tensordicts
+
+    def step_and_maybe_reset(
+        self, tensordict: TensorDictBase
+    ) -> tuple[TensorDictBase, TensorDictBase]:
+        """Runs a step in the environment and (partially) resets it if needed.
+
+        Args:
+            tensordict (TensorDictBase): an input data structure for the :meth:`step`
+                method.
+
+        This method allows to easily code non-stopping rollout functions.
+
+        Examples:
+            >>> from torchrl.envs import ParallelEnv, GymEnv
+            >>> def rollout(env, n):
+            ...     data_ = env.reset()
+            ...     result = []
+            ...     for i in range(n):
+            ...         data, data_ = env.step_and_maybe_reset(data_)
+            ...         result.append(data)
+            ...     return torch.stack(result)
+            >>> env = ParallelEnv(2, lambda: GymEnv("CartPole-v1"))
+            >>> print(rollout(env, 2))
+            TensorDict(
+                fields={
+                    done: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    next: TensorDict(
+                        fields={
+                            done: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            observation: Tensor(shape=torch.Size([2, 2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                            reward: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.float32, is_shared=False),
+                            terminated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                            truncated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                        batch_size=torch.Size([2, 2]),
+                        device=cpu,
+                        is_shared=False),
+                    observation: Tensor(shape=torch.Size([2, 2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
+                    terminated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                    truncated: Tensor(shape=torch.Size([2, 2, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                batch_size=torch.Size([2, 2]),
+                device=cpu,
+                is_shared=False)
+        """
+        if tensordict.device != self.device:
+            tensordict = tensordict.to(self.device)
+        tensordict = self.step(tensordict)
+        # done and truncated are in done_keys
+        # We read if any key is done.
+        tensordict_ = self._step_mdp(tensordict)
+        # if self._post_step_mdp_hooks is not None:
+        # tensordict_ = self._post_step_mdp_hooks(tensordict_)
+        tensordict_ = self.maybe_reset(tensordict_)
+        return tensordict, tensordict_
+
+    # _post_step_mdp_hooks: Callable[[TensorDictBase], TensorDictBase] | None = None
+
+    @property
+    @_cache_value
+    def _simple_done(self):
+        key_set = set(self.full_done_spec.keys())
+
+        _simple_done = "done" in key_set and "terminated" in key_set
+        return _simple_done
+
+    def any_done(self, tensordict: TensorDictBase) -> bool:
+        """Checks if the tensordict is in a "done" state (or if an element of the batch is).
+
+        Writes the result under the `"_reset"` entry.
+
+        Returns: a bool indicating whether there is an element in the tensordict that is marked
+            as done.
+
+        .. note:: The tensordict passed should be a `"next"` tensordict or equivalent -- i.e., it should not
+            contain a `"next"` value.
+
+        """
+        if self._simple_done:
+            done = tensordict._get_str("done", default=None)
+            if done is not None:
+                any_done = done.any()
+            else:
+                any_done = False
+            if any_done:
+                tensordict._set_str(
+                    "_reset",
+                    done.clone(),
+                    validated=True,
+                    inplace=False,
+                    non_blocking=False,
+                )
+        else:
+            any_done = _terminated_or_truncated(
+                tensordict,
+                full_done_spec=self.output_spec["full_done_spec"],
+                key="_reset",
+            )
+        return any_done
+
+    def maybe_reset(self, tensordict: TensorDictBase) -> TensorDictBase:
+        """Checks the done keys of the input tensordict and, if needed, resets the environment where it is done.
+
+        Args:
+            tensordict (TensorDictBase): a tensordict coming from the output of :func:`~torchrl.envs.utils.step_mdp`.
+
+        Returns:
+            A tensordict that is identical to the input where the environment was
+            not reset and contains the new reset data where the environment was reset.
+
+        """
+        any_done = self.any_done(tensordict)
+        if any_done:
+            tensordict = self.reset(tensordict, select_reset_only=True)
+        return tensordict
+
+    def empty_cache(self):
+        """Erases all the cached values.
+
+        For regular envs, the key lists (reward, done etc) are cached, but in some cases
+        they may change during the execution of the code (eg, when adding a transform).
+
+        """
+        self._cache.clear()
+
+    @property
+    @_cache_value
+    def reset_keys(self) -> list[NestedKey]:
+        """Returns a list of reset keys.
+
+        Reset keys are keys that indicate partial reset, in batched, multitask or multiagent
+        settings. They are structured as ``(*prefix, "_reset")`` where ``prefix`` is
+        a (possibly empty) tuple of strings pointing to a tensordict location
+        where a done state can be found.
+
+        Keys are sorted by depth in the data tree.
+        """
+        reset_keys = sorted(
+            (
+                _replace_last(done_key, "_reset")
+                for (done_key, *_) in self.done_keys_groups
+            ),
+            key=_repr_by_depth,
+        )
+        return reset_keys
+
+    @property
+    def _filtered_reset_keys(self):
+        """Returns only the effective reset keys, discarding nested resets if they're not being used."""
+        reset_keys = self.reset_keys
+        result = []
+
+        def _root(key):
+            if isinstance(key, str):
+                return ()
+            return key[:-1]
+
+        roots = []
+        for reset_key in reset_keys:
+            cur_root = _root(reset_key)
+            for root in roots:
+                if cur_root[: len(root)] == root:
+                    break
+            else:
+                roots.append(cur_root)
+                result.append(reset_key)
+        return result
+
+    @property
+    @_cache_value
+    def done_keys_groups(self):
+        """A list of done keys, grouped as the reset keys.
+
+        This is a list of lists. The outer list has the length of reset keys, the
+        inner lists contain the done keys (eg, done and truncated) that can
+        be read to determine a reset when it is absent.
+        """
+        # done keys, sorted as reset keys
+        done_keys_group = []
+        roots = set()
+        fds = self.full_done_spec
+        for done_key in self.done_keys:
+            root_name = done_key[:-1] if isinstance(done_key, tuple) else ()
+            root = fds[root_name] if root_name else fds
+            n = len(roots)
+            roots.add(root_name)
+            if len(roots) - n:
+                done_keys_group.append(
+                    [
+                        unravel_key(root_name + (key,))
+                        for key in root.keys(include_nested=False, leaves_only=True)
+                    ]
+                )
+        return done_keys_group
+
+    def _select_observation_keys(self, tensordict: TensorDictBase) -> Iterator[str]:
+        for key in tensordict.keys():
+            if key.rfind("observation") >= 0:
+                yield key
+
+    def close(self, *, raise_if_closed: bool = True):
+        self.is_closed = True
+
+    def __del__(self):
+        # if del occurs before env has been set up, we don't want a recursion
+        # error
+        if "is_closed" in self.__dict__ and not self.is_closed:
+            try:
+                self.close()
+            except Exception:
+                # a TypeError will typically be raised if the env is deleted when the program ends.
+                # In the future, insignificant changes to the close method may change the error type.
+                # We excplicitely assume that any error raised during closure in
+                # __del__ will not affect the program.
+                pass
+
+    @_maybe_unlock
+    def to(self, device: DEVICE_TYPING) -> EnvBase:
+        device = _make_ordinal_device(torch.device(device))
+        if device == self.device:
+            return self
+        self.__dict__["_input_spec"] = self.input_spec.to(device)
+        self.__dict__["_output_spec"] = self.output_spec.to(device)
+        self._device = device
+        return super().to(device)
+
+    def fake_tensordict(self) -> TensorDictBase:
+        """Returns a fake tensordict with key-value pairs that match in shape, device and dtype what can be expected during an environment rollout."""
+        state_spec = self.state_spec
+        observation_spec = self.observation_spec
+        action_spec = self.input_spec["full_action_spec"]
+        # instantiates reward_spec if needed
+        _ = self.full_reward_spec
+        reward_spec = self.output_spec["full_reward_spec"]
+        full_done_spec = self.output_spec["full_done_spec"]
+
+        fake_obs = observation_spec.zero()
+        fake_reward = reward_spec.zero()
+        fake_done = full_done_spec.zero()
+        fake_state = state_spec.zero()
+        fake_action = action_spec.zero()
+
+        if any(
+            isinstance(val, LazyStackedTensorDict) for val in fake_action.values(True)
+        ):
+            fake_input = fake_action.update(fake_state)
+        else:
+            fake_input = fake_state.update(fake_action)
+
+        # the input and output key may match, but the output prevails
+        # Hence we generate the input, and override using the output
+        fake_in_out = fake_input.update(fake_obs)
+
+        next_output = fake_obs.clone()
+        next_output.update(fake_reward)
+        next_output.update(fake_done)
+        fake_in_out.update(fake_done.clone())
+        if "next" not in fake_in_out.keys():
+            fake_in_out.set("next", next_output)
+        else:
+            fake_in_out.get("next").update(next_output)
+
+        fake_in_out.batch_size = self.batch_size
+        fake_in_out = fake_in_out.to(self.device)
+        return fake_in_out
+
+
+class _EnvWrapper(EnvBase):
+    """Abstract environment wrapper class.
+
+    Unlike EnvBase, _EnvWrapper comes with a :obj:`_build_env` private method that will be called upon instantiation.
+    Interfaces with other libraries should be coded using _EnvWrapper.
+
+    It is possible to directly query attributed from the nested environment it its name does not conflict with
+    an attribute of the wrapper:
+        >>> env = SomeWrapper(...)
+        >>> custom_attribute0 = env._env.custom_attribute
+        >>> custom_attribute1 = env.custom_attribute
+        >>> assert custom_attribute0 is custom_attribute1  # should return True
+
+    """
+
+    git_url: str = ""
+    available_envs: dict[str, Any] = {}
+    libname: str = ""
+
+    def __init__(
+        self,
+        *args,
+        device: DEVICE_TYPING = None,
+        batch_size: torch.Size | None = None,
+        allow_done_after_reset: bool = False,
+        spec_locked: bool = True,
+        **kwargs,
+    ):
+        super().__init__(
+            device=device,
+            batch_size=batch_size,
+            allow_done_after_reset=allow_done_after_reset,
+            spec_locked=spec_locked,
+        )
+        if len(args):
+            raise ValueError(
+                "`_EnvWrapper.__init__` received a non-empty args list of arguments. "
+                "Make sure only keywords arguments are used when calling `super().__init__`."
+            )
+
+        frame_skip = kwargs.pop("frame_skip", 1)
+        if not isinstance(frame_skip, int):
+            raise ValueError(f"frame_skip must be an integer, got {frame_skip}")
+        self.frame_skip = frame_skip
+        # this value can be changed if frame_skip is passed during env construction
+        self.wrapper_frame_skip = frame_skip
+
+        self._constructor_kwargs = kwargs
+        self._check_kwargs(kwargs)
+        self._convert_actions_to_numpy = kwargs.pop("convert_actions_to_numpy", True)
+        self._env = self._build_env(**kwargs)  # writes the self._env attribute
+        self._make_specs(self._env)  # writes the self._env attribute
+        self.is_closed = False
+        self._init_env()  # runs all the steps to have a ready-to-use env
+
+    def _sync_device(self):
+        sync_func = self.__dict__.get("_sync_device_val")
+        if sync_func is None:
+            device = self.device
+            if device.type != "cuda":
+                if torch.cuda.is_available():
+                    self._sync_device_val = torch.cuda.synchronize
+                elif torch.backends.mps.is_available():
+                    self._sync_device_val = torch.cuda.synchronize
+                elif device.type == "cpu":
+                    self._sync_device_val = _do_nothing
+            else:
+                self._sync_device_val = _do_nothing
+            return self._sync_device
+        return sync_func
+
+    @abc.abstractmethod
+    def _check_kwargs(self, kwargs: dict):
+        raise NotImplementedError
+
+    def __getattr__(self, attr: str) -> Any:
+        if attr in self.__dir__():
+            return self.__getattribute__(
+                attr
+            )  # make sure that appropriate exceptions are raised
+
+        elif attr.startswith("__"):
+            raise AttributeError(
+                "passing built-in private methods is "
+                f"not permitted with type {type(self)}. "
+                f"Got attribute {attr}."
+            )
+
+        elif "_env" in self.__dir__():
+            env = self.__getattribute__("_env")
+            return getattr(env, attr)
+        super().__getattr__(attr)
+
+        raise AttributeError(
+            f"The env wasn't set in {self.__class__.__name__}, cannot access {attr}"
+        )
+
+    @abc.abstractmethod
+    def _init_env(self) -> int | None:
+        """Runs all the necessary steps such that the environment is ready to use.
+
+        This step is intended to ensure that a seed is provided to the environment (if needed) and that the environment
+        is reset (if needed). For instance, DMControl envs require the env to be reset before being used, but Gym envs
+        don't.
+
+        Returns:
+            the resulting seed
+
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _build_env(self, **kwargs) -> gym.Env:  # noqa: F821
+        """Creates an environment from the target library and stores it with the `_env` attribute.
+
+        When overwritten, this function should pass all the required kwargs to the env instantiation method.
+
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def _make_specs(self, env: gym.Env) -> None:  # noqa: F821
+        raise NotImplementedError
+
+    def close(self, *, raise_if_closed: bool = True) -> None:
+        """Closes the contained environment if possible."""
+        self.is_closed = True
+        try:
+            self._env.close()
+        except AttributeError:
+            pass
+
+
+def make_tensordict(
+    env: _EnvWrapper,
+    policy: Callable[[TensorDictBase, ...], TensorDictBase] | None = None,
+) -> TensorDictBase:
+    """Returns a zeroed-tensordict with fields matching those required for a full step (action selection and environment step) in the environment.
+
+    Args:
+        env (_EnvWrapper): environment defining the observation, action and reward space;
+        policy (Callable, optional): policy corresponding to the environment.
+
+    """
+    with torch.no_grad():
+        tensordict = env.reset()
+        if policy is not None:
+            tensordict.update(policy(tensordict))
+        else:
+            tensordict.set("action", env.action_spec.rand(), inplace=False)
+        tensordict = env.step(tensordict)
+        return tensordict.zero_()
+
+
+def _get_sync_func(policy_device, env_device):
+    if torch.cuda.is_available():
+        # Look for a specific device
+        if policy_device is not None and policy_device.type == "cuda":
+            if env_device is None or env_device.type == "cuda":
+                return torch.cuda.synchronize
+            return partial(torch.cuda.synchronize, device=policy_device)
+        if env_device is not None and env_device.type == "cuda":
+            if policy_device is None:
+                return torch.cuda.synchronize
+            return partial(torch.cuda.synchronize, device=env_device)
+        return torch.cuda.synchronize
+    if torch.backends.mps.is_available():
+        return torch.mps.synchronize
+    return _do_nothing
+
+
+def _do_nothing():
+    return
+
+
+def _has_dynamic_specs(spec: Composite):
+    from tensordict.base import _NESTED_TENSORS_AS_LISTS
+
+    return any(
+        any(s == -1 for s in spec.shape)
+        for spec in spec.values(True, True, is_leaf=_NESTED_TENSORS_AS_LISTS)
+    )
+
+
+def _tensor_to_spec(name, leaf, leaf_compare=None, *, stack):
+    if not (isinstance(leaf, torch.Tensor) or is_tensor_collection(leaf)):
+        stack[name] = NonTensor(shape=())
+        return
+    elif is_non_tensor(leaf):
+        stack[name] = NonTensor(shape=leaf.shape)
+        return
+    shape = leaf.shape
+    if leaf_compare is not None:
+        shape_compare = leaf_compare.shape
+        shape = [s0 if s0 == s1 else -1 for s0, s1 in zip(shape, shape_compare)]
+    stack[name] = Unbounded(shape, device=leaf.device, dtype=leaf.dtype)