torchrl 0.11.0__cp314-cp314t-macosx_11_0_arm64.whl

torchrl/envs/libs/jumanji.py

@@ -0,0 +1,963 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import importlib.util
+
+import numpy as np
+import torch
+from packaging import version
+from tensordict import TensorDict, TensorDictBase
+from torchrl.envs.common import _EnvPostInit
+from torchrl.envs.utils import _classproperty
+
+_has_jumanji = importlib.util.find_spec("jumanji") is not None
+
+from torchrl.data.tensor_specs import (
+    Bounded,
+    Categorical,
+    Composite,
+    DEVICE_TYPING,
+    MultiCategorical,
+    MultiOneHot,
+    OneHot,
+    TensorSpec,
+    Unbounded,
+)
+from torchrl.data.utils import numpy_to_torch_dtype_dict
+from torchrl.envs.gym_like import GymLikeEnv
+
+from torchrl.envs.libs.jax_utils import (
+    _extract_spec,
+    _ndarray_to_tensor,
+    _object_to_tensordict,
+    _tensordict_to_object,
+    _tree_flatten,
+    _tree_reshape,
+)
+
+
+def _get_envs():
+    if not _has_jumanji:
+        raise ImportError("Jumanji is not installed in your virtual environment.")
+    import jumanji
+
+    return jumanji.registered_environments()
+
+
+def _jumanji_to_torchrl_spec_transform(
+    spec,
+    dtype: torch.dtype | None = None,
+    device: DEVICE_TYPING = None,
+    categorical_action_encoding: bool = True,
+) -> TensorSpec:
+    import jumanji
+
+    if isinstance(spec, jumanji.specs.DiscreteArray):
+        action_space_cls = Categorical if categorical_action_encoding else OneHot
+        if dtype is None:
+            dtype = numpy_to_torch_dtype_dict[spec.dtype]
+        return action_space_cls(spec.num_values, dtype=dtype, device=device)
+    if isinstance(spec, jumanji.specs.MultiDiscreteArray):
+        action_space_cls = (
+            MultiCategorical if categorical_action_encoding else MultiOneHot
+        )
+        if dtype is None:
+            dtype = numpy_to_torch_dtype_dict[spec.dtype]
+        return action_space_cls(
+            torch.as_tensor(np.asarray(spec.num_values)), dtype=dtype, device=device
+        )
+    elif isinstance(spec, jumanji.specs.BoundedArray):
+        shape = spec.shape
+        if dtype is None:
+            dtype = numpy_to_torch_dtype_dict[spec.dtype]
+        return Bounded(
+            shape=shape,
+            low=np.asarray(spec.minimum),
+            high=np.asarray(spec.maximum),
+            dtype=dtype,
+            device=device,
+        )
+    elif isinstance(spec, jumanji.specs.Array):
+        shape = spec.shape
+        if dtype is None:
+            dtype = numpy_to_torch_dtype_dict[spec.dtype]
+        if dtype in (torch.float, torch.double, torch.half):
+            return Unbounded(shape=shape, dtype=dtype, device=device)
+        else:
+            return Unbounded(shape=shape, dtype=dtype, device=device)
+    elif isinstance(spec, jumanji.specs.Spec) and hasattr(spec, "__dict__"):
+        new_spec = {}
+        for key, value in spec.__dict__.items():
+            if isinstance(value, jumanji.specs.Spec):
+                if key.endswith("_obs"):
+                    key = key[:-4]
+                if key.endswith("_spec"):
+                    key = key[:-5]
+                new_spec[key] = _jumanji_to_torchrl_spec_transform(
+                    value, dtype, device, categorical_action_encoding
+                )
+        return Composite(**new_spec)
+    else:
+        raise TypeError(f"Unsupported spec type {type(spec)}")
+
+
+class _JumanjiMakeRender(_EnvPostInit):
+    def __call__(self, *args, **kwargs):
+        instance = super().__call__(*args, **kwargs)
+        if instance.from_pixels:
+            return instance.make_render()
+        return instance
+
+
+class JumanjiWrapper(GymLikeEnv, metaclass=_JumanjiMakeRender):
+    """Jumanji's environment wrapper.
+
+    Jumanji offers a vectorized simulation framework based on Jax.
+    TorchRL's wrapper incurs some overhead for the jax-to-torch conversion,
+    but computational graphs can still be built on top of the simulated trajectories,
+    allowing for backpropagation through the rollout.
+
+    GitHub: https://github.com/instadeepai/jumanji
+
+    Doc: https://instadeepai.github.io/jumanji/
+
+    Paper: https://arxiv.org/abs/2306.09884
+
+    .. note:: For better performance, turn `jit` on when instantiating this class.
+        The `jit` attribute can also be flipped during code execution:
+
+            >>> env.jit = True # Used jit
+            >>> env.jit = False # eager
+
+    Args:
+        env (jumanji.env.Environment): the env to wrap.
+        categorical_action_encoding (bool, optional): if ``True``, categorical
+            specs will be converted to the TorchRL equivalent (:class:`torchrl.data.Categorical`),
+            otherwise a one-hot encoding will be used (:class:`torchrl.data.OneHot`).
+            Defaults to ``False``.
+
+    Keyword Args:
+        batch_size (torch.Size, optional): the batch size of the environment.
+            With ``jumanji``, this indicates the number of vectorized environments.
+            If the batch-size is empty, the environment is not batch-locked and an arbitrary number
+            of environments can be executed simultaneously.
+            Defaults to ``torch.Size([])``.
+
+                >>> import jumanji
+                >>> from torchrl.envs import JumanjiWrapper
+                >>> base_env = jumanji.make("Snake-v1")
+                >>> env = JumanjiWrapper(base_env)
+                >>> # Set the batch-size of the TensorDict instead of the env allows to control the number
+                >>> #  of envs being run simultaneously
+                >>> tdreset = env.reset(TensorDict(batch_size=[32]))
+                >>> # Execute a rollout until all envs are done or max steps is reached, whichever comes first
+                >>> rollout = env.rollout(100, break_when_all_done=True, auto_reset=False, tensordict=tdreset)
+
+        from_pixels (bool, optional): Whether the environment should render its output.
+            This will drastically impact the environment throughput. Only the first environment
+            will be rendered. See :meth:`~torchrl.envs.JumanjiWrapper.render` for more information.
+            Defaults to `False`.
+        frame_skip (int, optional): if provided, indicates for how many steps the
+            same action is to be repeated. The observation returned will be the
+            last observation of the sequence, whereas the reward will be the sum
+            of rewards across steps.
+        device (torch.device, optional): if provided, the device on which the data
+            is to be cast. Defaults to ``torch.device("cpu")``.
+        allow_done_after_reset (bool, optional): if ``True``, it is tolerated
+            for envs to be ``done`` just after :meth:`reset` is called.
+            Defaults to ``False``.
+        jit (bool, optional): whether the step and reset method should be wrapped in `jit`.
+            Defaults to ``False``.
+
+    Attributes:
+        available_envs: environments available to build
+
+    Examples:
+        >>> import jumanji
+        >>> from torchrl.envs import JumanjiWrapper
+        >>> base_env = jumanji.make("Snake-v1")
+        >>> env = JumanjiWrapper(base_env)
+        >>> env.set_seed(0)
+        >>> td = env.reset()
+        >>> td["action"] = env.action_spec.rand()
+        >>> td = env.step(td)
+        >>> print(td)
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
+                grid: Tensor(shape=torch.Size([12, 12, 5]), device=cpu, dtype=torch.float32, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        grid: Tensor(shape=torch.Size([12, 12, 5]), device=cpu, dtype=torch.float32, is_shared=False),
+                        reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
+                        state: TensorDict(
+                            fields={
+                                action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                                body: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False),
+                                body_state: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.int32, is_shared=False),
+                                fruit_position: TensorDict(
+                                    fields={
+                                        col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                        row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                                    batch_size=torch.Size([]),
+                                    device=cpu,
+                                    is_shared=False),
+                                head_position: TensorDict(
+                                    fields={
+                                        col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                        row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                                    batch_size=torch.Size([]),
+                                    device=cpu,
+                                    is_shared=False),
+                                key: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int32, is_shared=False),
+                                length: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                tail: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False)},
+                            batch_size=torch.Size([]),
+                            device=cpu,
+                            is_shared=False),
+                        step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([]),
+                    device=cpu,
+                    is_shared=False),
+                state: TensorDict(
+                    fields={
+                        action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                        body: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False),
+                        body_state: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.int32, is_shared=False),
+                        fruit_position: TensorDict(
+                            fields={
+                                col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                            batch_size=torch.Size([]),
+                            device=cpu,
+                            is_shared=False),
+                        head_position: TensorDict(
+                            fields={
+                                col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                            batch_size=torch.Size([]),
+                            device=cpu,
+                            is_shared=False),
+                        key: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int32, is_shared=False),
+                        length: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                        step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                        tail: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([]),
+                    device=cpu,
+                    is_shared=False),
+                step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([]),
+            device=cpu,
+            is_shared=False)
+        >>> print(env.available_envs)
+        ['Game2048-v1',
+         'Maze-v0',
+         'Cleaner-v0',
+         'CVRP-v1',
+         'MultiCVRP-v0',
+         'Minesweeper-v0',
+         'RubiksCube-v0',
+         'Knapsack-v1',
+         'Sudoku-v0',
+         'Snake-v1',
+         'TSP-v1',
+         'Connector-v2',
+         'MMST-v0',
+         'GraphColoring-v0',
+         'RubiksCube-partly-scrambled-v0',
+         'RobotWarehouse-v0',
+         'Tetris-v0',
+         'BinPack-v2',
+         'Sudoku-very-easy-v0',
+         'JobShop-v0']
+
+    To take advante of Jumanji, one usually executes multiple environments at the
+    same time.
+
+        >>> import jumanji
+        >>> from torchrl.envs import JumanjiWrapper
+        >>> base_env = jumanji.make("Snake-v1")
+        >>> env = JumanjiWrapper(base_env, batch_size=[10])
+        >>> env.set_seed(0)
+        >>> td = env.reset()
+        >>> td["action"] = env.action_spec.rand()
+        >>> td = env.step(td)
+
+    In the following example, we iteratively test different batch sizes
+    and report the execution time for a short rollout:
+
+    Examples:
+        >>> from torch.utils.benchmark import Timer
+        >>> for batch_size in [4, 16, 128]:
+        ...     timer = Timer(
+        ...     '''
+        ... env.rollout(100)
+        ... ''',
+        ... setup=f'''
+        ... from torchrl.envs import JumanjiWrapper
+        ... import jumanji
+        ... env = JumanjiWrapper(jumanji.make('Snake-v1'), batch_size=[{batch_size}])
+        ... env.set_seed(0)
+        ... env.rollout(2)
+        ... ''')
+        ...     print(batch_size, timer.timeit(number=10))
+        4
+        env.rollout(100)
+        setup: [...]
+        Median: 122.40 ms
+        2 measurements, 1 runs per measurement, 1 thread
+
+        16
+        env.rollout(100)
+        setup: [...]
+        Median: 134.39 ms
+        2 measurements, 1 runs per measurement, 1 thread
+
+        128
+        env.rollout(100)
+        setup: [...]
+        Median: 172.31 ms
+        2 measurements, 1 runs per measurement, 1 thread
+
+    """
+
+    git_url = "https://github.com/instadeepai/jumanji"
+    libname = "jumanji"
+
+    @_classproperty
+    def available_envs(cls):
+        if not _has_jumanji:
+            return []
+        return sorted(_get_envs())
+
+    @property
+    def lib(self):
+        import jumanji
+
+        if version.parse(jumanji.__version__) < version.parse("1.0.0"):
+            raise ImportError("jumanji version must be >= 1.0.0")
+        return jumanji
+
+    def __init__(
+        self,
+        env: jumanji.env.Environment = None,  # noqa: F821
+        categorical_action_encoding=True,
+        jit: bool = True,
+        **kwargs,
+    ):
+        if not _has_jumanji:
+            raise ImportError(
+                "jumanji is not installed or importing it failed. Consider checking your installation."
+            )
+        self.categorical_action_encoding = categorical_action_encoding
+        if env is not None:
+            kwargs["env"] = env
+        batch_locked = kwargs.pop("batch_locked", kwargs.get("batch_size") is not None)
+        super().__init__(**kwargs)
+        self._batch_locked = batch_locked
+        self.jit = jit
+
+    @property
+    def jit(self):
+        return self._jit
+
+    @jit.setter
+    def jit(self, value):
+        self._jit = value
+        if value:
+            import jax
+
+            self._env_reset = jax.jit(self._env.reset)
+            self._env_step = jax.jit(self._env.step)
+        else:
+            self._env_reset = self._env.reset
+            self._env_step = self._env.step
+
+    def _build_env(
+        self,
+        env,
+        _seed: int | None = None,
+        from_pixels: bool = False,
+        render_kwargs: dict | None = None,
+        pixels_only: bool = False,
+        camera_id: int | str = 0,
+        **kwargs,
+    ):
+        self.from_pixels = from_pixels
+        self.pixels_only = pixels_only
+
+        return env
+
+    def make_render(self):
+        """Returns a transformed environment that can be rendered.
+
+        Examples:
+            >>> from torchrl.envs import JumanjiEnv
+            >>> from torchrl.record import CSVLogger, VideoRecorder
+            >>>
+            >>> envname = JumanjiEnv.available_envs[-1]
+            >>> logger = CSVLogger("jumanji", video_format="mp4", video_fps=2)
+            >>> env = JumanjiEnv(envname, from_pixels=True)
+            >>>
+            >>> env = env.append_transform(
+            ...     VideoRecorder(logger=logger, in_keys=["pixels"], tag=envname)
+            ... )
+            >>> env.set_seed(0)
+            >>> r = env.rollout(100)
+            >>> env.transform.dump()
+
+        """
+        from torchrl.record import PixelRenderTransform
+
+        return self.append_transform(
+            PixelRenderTransform(
+                out_keys=["pixels"],
+                pass_tensordict=True,
+                as_non_tensor=bool(self.batch_size),
+                as_numpy=bool(self.batch_size),
+            )
+        )
+
+    def _make_state_example(self, env):
+        import jax
+        from jax import numpy as jnp
+
+        key = jax.random.PRNGKey(0)
+        keys = jax.random.split(key, self.batch_size.numel())
+        state, _ = jax.vmap(env.reset)(jnp.stack(keys))
+        state = _tree_reshape(state, self.batch_size)
+        return state
+
+    def _make_state_spec(self, env) -> TensorSpec:
+        import jax
+
+        key = jax.random.PRNGKey(0)
+        state, _ = env.reset(key)
+        state_dict = _object_to_tensordict(state, self.device, batch_size=())
+        state_spec = _extract_spec(state_dict)
+        return state_spec
+
+    def _make_action_spec(self, env) -> TensorSpec:
+        action_spec = _jumanji_to_torchrl_spec_transform(
+            env.action_spec,
+            device=self.device,
+            categorical_action_encoding=self.categorical_action_encoding,
+        )
+        action_spec = action_spec.expand(*self.batch_size, *action_spec.shape)
+        return action_spec
+
+    def _make_observation_spec(self, env) -> TensorSpec:
+        jumanji = self.lib
+
+        spec = env.observation_spec
+        new_spec = _jumanji_to_torchrl_spec_transform(spec, device=self.device)
+        if isinstance(spec, jumanji.specs.Array):
+            return Composite(observation=new_spec).expand(self.batch_size)
+        elif isinstance(spec, jumanji.specs.Spec):
+            return Composite(**{k: v for k, v in new_spec.items()}).expand(
+                self.batch_size
+            )
+        else:
+            raise TypeError(f"Unsupported spec type {type(spec)}")
+
+    def _make_reward_spec(self, env) -> TensorSpec:
+        reward_spec = _jumanji_to_torchrl_spec_transform(
+            env.reward_spec, device=self.device
+        )
+        if not len(reward_spec.shape):
+            reward_spec.shape = torch.Size([1])
+        return reward_spec.expand([*self.batch_size, *reward_spec.shape])
+
+    def _make_specs(self, env: jumanji.env.Environment) -> None:  # noqa: F821
+
+        # extract spec from jumanji definition
+        self.action_spec = self._make_action_spec(env)
+        self.observation_spec = self._make_observation_spec(env)
+        self.reward_spec = self._make_reward_spec(env)
+
+        # extract state spec from instance
+        state_spec = self._make_state_spec(env).expand(self.batch_size)
+        self.state_spec["state"] = state_spec
+        self.observation_spec["state"] = state_spec.clone()
+
+        # build state example for data conversion
+        self._state_example = self._make_state_example(env)
+
+    def _check_kwargs(self, kwargs: dict):
+        jumanji = self.lib
+        if "env" not in kwargs:
+            raise TypeError("Could not find environment key 'env' in kwargs.")
+        env = kwargs["env"]
+        if not isinstance(env, (jumanji.env.Environment,)):
+            raise TypeError("env is not of type 'jumanji.env.Environment'.")
+
+    def _init_env(self):
+        pass
+
+    @property
+    def key(self):
+        key = getattr(self, "_key", None)
+        if key is None:
+            raise RuntimeError(
+                "the env.key attribute wasn't found. Make sure to call `env.set_seed(seed)` before any interaction."
+            )
+        return key
+
+    @key.setter
+    def key(self, value):
+        self._key = value
+
+    def _set_seed(self, seed: int | None) -> None:
+        import jax
+
+        if seed is None:
+            raise Exception("Jumanji requires an integer seed.")
+        self.key = jax.random.PRNGKey(seed)
+
+    def read_state(self, state, batch_size=None):
+        state_dict = _object_to_tensordict(
+            state, self.device, self.batch_size if batch_size is None else batch_size
+        )
+        return self.state_spec["state"].encode(state_dict)
+
+    def read_obs(self, obs, batch_size=None):
+        from jax import numpy as jnp
+
+        if isinstance(obs, (list, jnp.ndarray, np.ndarray)):
+            obs_dict = _ndarray_to_tensor(obs).to(self.device)
+        else:
+            obs_dict = _object_to_tensordict(
+                obs, self.device, self.batch_size if batch_size is None else batch_size
+            )
+        return super().read_obs(obs_dict)
+
+    def render(
+        self,
+        tensordict,
+        matplotlib_backend: str | None = None,
+        as_numpy: bool = False,
+        **kwargs,
+    ):
+        """Renders the environment output given an input tensordict.
+
+        This method is intended to be called by the :class:`~torchrl.record.PixelRenderTransform`
+        created whenever `from_pixels=True` is selected.
+        To create an appropriate rendering transform, use a similar call as bellow:
+
+            >>> from torchrl.record import PixelRenderTransform
+            >>> matplotlib_backend = None # Change this value if a specific matplotlib backend has to be used.
+            >>> env = env.append_transform(
+            ...     PixelRenderTransform(out_keys=["pixels"], pass_tensordict=True, matplotlib_backend=matplotlib_backend)
+            ... )
+
+        This pipeline will write a `"pixels"` entry in your output tensordict.
+
+        Args:
+            tensordict (TensorDictBase): a tensordict containing a state to represent
+            matplotlib_backend (str, optional): the matplotlib backend
+            as_numpy (bool, optional): if ``False``, the np.ndarray will be converted to a torch.Tensor.
+                Defaults to ``False``.
+
+        """
+        import io
+
+        import jax
+        import jax.numpy as jnp
+        import jumanji
+
+        try:
+            import matplotlib
+            import matplotlib.pyplot as plt
+            import PIL
+            import torchvision.transforms.v2.functional
+        except ImportError as err:
+            raise ImportError(
+                "Rendering with Jumanji requires torchvision, matplotlib and PIL to be installed."
+            ) from err
+
+        if matplotlib_backend is not None:
+            matplotlib.use(matplotlib_backend)
+
+        # Get only one env
+        _state_example = self._state_example
+        while tensordict.ndim:
+            tensordict = tensordict[0]
+            _state_example = jax.tree_util.tree_map(
+                lambda x: jnp.take(x, 0, axis=0), _state_example
+            )
+        # Patch jumanji is_notebook
+        is_notebook = jumanji.environments.is_notebook
+        try:
+            jumanji.environments.is_notebook = lambda: False
+
+            isinteractive = plt.isinteractive()
+            plt.ion()
+            buf = io.BytesIO()
+            state = _tensordict_to_object(
+                tensordict.get("state"),
+                _state_example,
+                batch_size=tensordict.batch_size if not self.batch_locked else None,
+            )
+            self._env.render(state, **kwargs)
+            plt.savefig(buf, format="png")
+            buf.seek(0)
+            # Load the image into a PIL object.
+            img = PIL.Image.open(buf)
+            img_array = torchvision.transforms.v2.functional.pil_to_tensor(img)
+            if not isinteractive:
+                plt.ioff()
+            plt.close()
+            if not as_numpy:
+                return img_array[:3]
+            return img_array[:3].numpy().copy()
+        finally:
+            jumanji.environments.is_notebook = is_notebook
+
+    def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
+        import jax
+
+        if self.batch_locked:
+            batch_size = self.batch_size
+        else:
+            batch_size = tensordict.batch_size
+
+        # prepare inputs
+        state = _tensordict_to_object(
+            tensordict.get("state"),
+            self._state_example,
+            batch_size=tensordict.batch_size if not self.batch_locked else None,
+        )
+        action = self.read_action(tensordict.get("action"))
+
+        # flatten batch size into vector
+        state = _tree_flatten(state, batch_size)
+        action = _tree_flatten(action, batch_size)
+
+        # jax vectorizing map on env.step
+        state, timestep = jax.vmap(self._env_step)(state, action)
+
+        # reshape batch size from vector
+        state = _tree_reshape(state, batch_size)
+        timestep = _tree_reshape(timestep, batch_size)
+
+        # collect outputs
+        state_dict = self.read_state(state, batch_size=batch_size)
+        obs_dict = self.read_obs(timestep.observation, batch_size=batch_size)
+        reward = self.read_reward(np.asarray(timestep.reward))
+        done = timestep.step_type == self.lib.types.StepType.LAST
+        done = _ndarray_to_tensor(done).view(torch.bool).to(self.device)
+
+        # build results
+        tensordict_out = TensorDict(
+            source=obs_dict,
+            batch_size=tensordict.batch_size,
+            device=self.device,
+        )
+        tensordict_out.set("reward", reward)
+        tensordict_out.set("done", done)
+        tensordict_out.set("terminated", done)
+        # tensordict_out.set("terminated", done)
+        tensordict_out["state"] = state_dict
+
+        return tensordict_out
+
+    def _reset(
+        self, tensordict: TensorDictBase | None = None, **kwargs
+    ) -> TensorDictBase:
+        import jax
+        from jax import numpy as jnp
+
+        if self.batch_locked or tensordict is None:
+            numel = self.numel()
+            batch_size = self.batch_size
+        elif tensordict is not None:
+            numel = tensordict.numel()
+            batch_size = tensordict.batch_size
+
+        # generate random keys
+        self.key, *keys = jax.random.split(self.key, numel + 1)
+
+        # jax vectorizing map on env.reset
+        state, timestep = jax.vmap(self._env_reset)(jnp.stack(keys))
+
+        # reshape batch size from vector
+        state = _tree_reshape(state, batch_size)
+        timestep = _tree_reshape(timestep, batch_size)
+
+        # collect outputs
+        state_dict = self.read_state(state, batch_size=batch_size)
+        obs_dict = self.read_obs(timestep.observation, batch_size=batch_size)
+        if not self.batch_locked:
+            done_td = self.full_done_spec.zero(batch_size)
+        else:
+            done_td = self.full_done_spec.zero()
+
+        # build results
+        tensordict_out = TensorDict(
+            source=obs_dict,
+            batch_size=batch_size,
+            device=self.device,
+        )
+        tensordict_out.update(done_td)
+        tensordict_out["state"] = state_dict
+
+        return tensordict_out
+
+    def read_reward(self, reward):
+        """Reads the reward and maps it to the reward space.
+
+        Args:
+            reward (torch.Tensor or TensorDict): reward to be mapped.
+
+        """
+        if isinstance(reward, int) and reward == 0:
+            return self.reward_spec.zero()
+        if self.batch_locked:
+            reward = self.reward_spec.encode(reward, ignore_device=True)
+        else:
+            reward = torch.as_tensor(reward)
+            if not reward.ndim or (reward.shape[-1] != self.reward_spec.shape[-1]):
+                reward = reward.unsqueeze(-1)
+
+        if reward is None:
+            reward = torch.tensor(np.nan).expand(self.reward_spec.shape)
+
+        return reward
+
+    def _output_transform(self, step_outputs_tuple: tuple) -> tuple:
+        ...
+
+    def _reset_output_transform(self, reset_outputs_tuple: tuple) -> tuple:
+        ...
+
+
+class JumanjiEnv(JumanjiWrapper):
+    """Jumanji environment wrapper built with the environment name.
+
+    Jumanji offers a vectorized simulation framework based on Jax.
+    TorchRL's wrapper incurs some overhead for the jax-to-torch conversion,
+    but computational graphs can still be built on top of the simulated trajectories,
+    allowing for backpropagation through the rollout.
+
+    GitHub: https://github.com/instadeepai/jumanji
+
+    Doc: https://instadeepai.github.io/jumanji/
+
+    Paper: https://arxiv.org/abs/2306.09884
+
+    Args:
+        env_name (str): the name of the environment to wrap. Must be part of :attr:`~.available_envs`.
+        categorical_action_encoding (bool, optional): if ``True``, categorical
+            specs will be converted to the TorchRL equivalent (:class:`torchrl.data.Categorical`),
+            otherwise a one-hot encoding will be used (:class:`torchrl.data.OneHot`).
+            Defaults to ``False``.
+
+    Keyword Args:
+        from_pixels (bool, optional): Not yet supported.
+        frame_skip (int, optional): if provided, indicates for how many steps the
+            same action is to be repeated. The observation returned will be the
+            last observation of the sequence, whereas the reward will be the sum
+            of rewards across steps.
+        device (torch.device, optional): if provided, the device on which the data
+            is to be cast. Defaults to ``torch.device("cpu")``.
+        batch_size (torch.Size, optional): the batch size of the environment.
+            With ``jumanji``, this indicates the number of vectorized environments.
+            Defaults to ``torch.Size([])``.
+        allow_done_after_reset (bool, optional): if ``True``, it is tolerated
+            for envs to be ``done`` just after :meth:`reset` is called.
+            Defaults to ``False``.
+
+    Attributes:
+        available_envs: environments available to build
+
+    Examples:
+        >>> from torchrl.envs import JumanjiEnv
+        >>> env = JumanjiEnv("Snake-v1")
+        >>> env.set_seed(0)
+        >>> td = env.reset()
+        >>> td["action"] = env.action_spec.rand()
+        >>> td = env.step(td)
+        >>> print(td)
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
+                grid: Tensor(shape=torch.Size([12, 12, 5]), device=cpu, dtype=torch.float32, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                        done: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        grid: Tensor(shape=torch.Size([12, 12, 5]), device=cpu, dtype=torch.float32, is_shared=False),
+                        reward: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.float32, is_shared=False),
+                        state: TensorDict(
+                            fields={
+                                action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                                body: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False),
+                                body_state: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.int32, is_shared=False),
+                                fruit_position: TensorDict(
+                                    fields={
+                                        col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                        row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                                    batch_size=torch.Size([]),
+                                    device=cpu,
+                                    is_shared=False),
+                                head_position: TensorDict(
+                                    fields={
+                                        col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                        row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                                    batch_size=torch.Size([]),
+                                    device=cpu,
+                                    is_shared=False),
+                                key: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int32, is_shared=False),
+                                length: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                tail: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False)},
+                            batch_size=torch.Size([]),
+                            device=cpu,
+                            is_shared=False),
+                        step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([]),
+                    device=cpu,
+                    is_shared=False),
+                state: TensorDict(
+                    fields={
+                        action_mask: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.bool, is_shared=False),
+                        body: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False),
+                        body_state: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.int32, is_shared=False),
+                        fruit_position: TensorDict(
+                            fields={
+                                col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                            batch_size=torch.Size([]),
+                            device=cpu,
+                            is_shared=False),
+                        head_position: TensorDict(
+                            fields={
+                                col: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                                row: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False)},
+                            batch_size=torch.Size([]),
+                            device=cpu,
+                            is_shared=False),
+                        key: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int32, is_shared=False),
+                        length: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                        step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                        tail: Tensor(shape=torch.Size([12, 12]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([]),
+                    device=cpu,
+                    is_shared=False),
+                step_count: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int32, is_shared=False),
+                terminated: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([]),
+            device=cpu,
+            is_shared=False)
+        >>> print(env.available_envs)
+        ['Game2048-v1',
+         'Maze-v0',
+         'Cleaner-v0',
+         'CVRP-v1',
+         'MultiCVRP-v0',
+         'Minesweeper-v0',
+         'RubiksCube-v0',
+         'Knapsack-v1',
+         'Sudoku-v0',
+         'Snake-v1',
+         'TSP-v1',
+         'Connector-v2',
+         'MMST-v0',
+         'GraphColoring-v0',
+         'RubiksCube-partly-scrambled-v0',
+         'RobotWarehouse-v0',
+         'Tetris-v0',
+         'BinPack-v2',
+         'Sudoku-very-easy-v0',
+         'JobShop-v0']
+
+    To take advante of Jumanji, one usually executes multiple environments at the
+    same time.
+
+        >>> from torchrl.envs import JumanjiEnv
+        >>> env = JumanjiEnv("Snake-v1", batch_size=[10])
+        >>> env.set_seed(0)
+        >>> td = env.reset()
+        >>> td["action"] = env.action_spec.rand()
+        >>> td = env.step(td)
+
+    In the following example, we iteratively test different batch sizes
+    and report the execution time for a short rollout:
+
+    Examples:
+        >>> from torch.utils.benchmark import Timer
+        >>> for batch_size in [4, 16, 128]:
+        ...     timer = Timer(
+        ...     '''
+        ... env.rollout(100)
+        ... ''',
+        ... setup=f'''
+        ... from torchrl.envs import JumanjiEnv
+        ... env = JumanjiEnv('Snake-v1', batch_size=[{batch_size}])
+        ... env.set_seed(0)
+        ... env.rollout(2)
+        ... ''')
+        ...     print(batch_size, timer.timeit(number=10))
+        4 <torch.utils.benchmark.utils.common.Measurement object at 0x1fca91910>
+        env.rollout(100)
+        setup: [...]
+          Median: 122.40 ms
+          2 measurements, 1 runs per measurement, 1 thread
+        16 <torch.utils.benchmark.utils.common.Measurement object at 0x1ff9baee0>
+        env.rollout(100)
+        setup: [...]
+          Median: 134.39 ms
+          2 measurements, 1 runs per measurement, 1 thread
+        128 <torch.utils.benchmark.utils.common.Measurement object at 0x1ff9ba7c0>
+        env.rollout(100)
+        setup: [...]
+          Median: 172.31 ms
+          2 measurements, 1 runs per measurement, 1 thread
+    """
+
+    def __init__(self, env_name, **kwargs):
+        kwargs["env_name"] = env_name
+        super().__init__(**kwargs)
+
+    def _build_env(
+        self,
+        env_name: str,
+        **kwargs,
+    ) -> jumanji.env.Environment:  # noqa: F821
+        if not _has_jumanji:
+            raise ImportError(
+                f"jumanji not found, unable to create {env_name}. "
+                f"Consider installing jumanji. More info:"
+                f" {self.git_url}."
+            )
+        from_pixels = kwargs.pop("from_pixels", False)
+        pixels_only = kwargs.pop("pixels_only", True)
+        if kwargs:
+            raise ValueError(f"Extra kwargs are not supported by {type(self)}.")
+        self.wrapper_frame_skip = 1
+        env = self.lib.make(env_name, **kwargs)
+        return super()._build_env(env, pixels_only=pixels_only, from_pixels=from_pixels)
+
+    @property
+    def env_name(self):
+        return self._constructor_kwargs["env_name"]
+
+    def _check_kwargs(self, kwargs: dict):
+        if "env_name" not in kwargs:
+            raise TypeError("Expected 'env_name' to be part of kwargs")
+
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}(env={self.env_name}, batch_size={self.batch_size}, device={self.device})"