[Doc] Document composite spec shapes for batched envs

docs/source/reference/data_specs.rst

@@ -26,6 +26,45 @@ TensorSpec classes define the shape, dtype, and domain of tensors in TorchRL.
     UnboundedContinuous
     UnboundedDiscrete
 
+.. _composite-spec-shapes:
+
+Composite shape vs leaf shape
+-----------------------------
+
+A :class:`~torchrl.data.Composite` has a ``shape`` that is the leading batch
+of every leaf, the same way a :class:`~tensordict.TensorDict` has a
+``batch_size``. Leaves carry that batch plus their feature size.
+
+Environment specs follow this layout. An unbatched env has composite shape
+``[]``; an image leaf is then just the feature, for example ``[3, 64, 64]``.
+A batched env with ``batch_size=[N]`` (as with
+:class:`~torchrl.envs.ParallelEnv`) gives every spec a leading ``[N, *]``:
+the composite is ``[N]`` and the image leaf is ``[N, 3, 64, 64]``. Nested
+composites can introduce extra leading dimensions for agent groups.
+
+.. code-block:: python
+
+    from torchrl.data import Bounded, Composite
+
+    # Unbatched env: empty composite shape, image leaf is feature-only.
+    observation_spec = Composite(
+        pixels=Bounded(low=0, high=255, shape=(3, 64, 64), dtype=torch.uint8),
+        shape=(),
+    )
+    assert observation_spec.shape == ()
+    assert observation_spec["pixels"].shape == (3, 64, 64)
+
+    # Batched env with N=2: every spec carries a leading [2].
+    batched_observation_spec = Composite(
+        pixels=Bounded(low=0, high=255, shape=(2, 3, 64, 64), dtype=torch.uint8),
+        shape=(2,),
+    )
+    assert batched_observation_spec.shape == (2,)
+    assert batched_observation_spec["pixels"].shape == (2, 3, 64, 64)
+
+See :ref:`env_spec_shapes` for the same examples in the env tutorial, and
+:ref:`Environment-spec-shapes` for the env-API statement of this rule.
+
 Supported PyTorch Operations
 ----------------------------
 

docs/source/reference/envs_api.rst

@@ -62,10 +62,28 @@ be avoided).
 Modifications of the specs such as `env.observation_spec = new_spec` are allowed: under the hood, TorchRL will erase
 the cache, unlock the specs, make the modification and relock the specs if the env was previously locked.
 
-Importantly, the environment spec shapes should contain the batch size, e.g.
-an environment with :attr:`env.batch_size` ``== torch.Size([4])`` should have
-an :attr:`env.action_spec` with shape :class:`torch.Size` ``([4, action_size])``.
-This is helpful when preallocation tensors, checking shape consistency etc.
+Importantly, the environment spec shapes should contain the batch size.
+
+.. _Environment-spec-shapes:
+
+The composite spec has the shape of the environment batch size (it can be
+empty) and the leaves have that size plus their feature size. Feature
+dimensions (an image, a board, a vector) belong on the leaf, not on the
+composite.
+
+* Unbatched env (``env.batch_size == torch.Size([])``): the composite
+  (``observation_spec``, ``full_observation_spec``, ...) has shape ``[]``.
+  An image leaf has the feature shape, for example ``[3, 64, 64]``.
+* Batched env (``env.batch_size == torch.Size([N])``, as with
+  :class:`~torchrl.envs.ParallelEnv`): *all* specs have a leading shape of
+  ``[N, *]``. The composite has shape ``[N]`` and the image leaf has shape
+  ``[N, 3, 64, 64]``. An :attr:`env.action_spec` has shape
+  :class:`torch.Size` ``([N, action_size])``.
+
+This is helpful when preallocating tensors, checking shape consistency, and
+writing custom :class:`~torchrl.envs.EnvBase` subclasses. Worked examples
+(including nested MARL groups) are in :ref:`env_spec_shapes`; the spec
+classes themselves are documented in :ref:`ref_specs`.
 
 .. _Environment-policy-arg:
 

torchrl/data/tensor_specs.py

@@ -5049,6 +5049,12 @@ class Composite(TensorSpec):
     the :class:`~tensordict.TensorDict` class. Like :class:`~tensordict.TensorDict`, it has a ``shape`` (akin to the
     ``TensorDict``'s ``batch_size``) and an optional ``device``.
 
+    Environment specs use that layout: the composite ``shape`` is the env batch
+    size (empty if the env is unbatched) and each leaf is that batch plus the
+    feature. An unbatched image observation is therefore composite shape ``[]``
+    with a ``[3, 64, 64]`` leaf; a batched env with ``N`` copies uses ``[N]``
+    and ``[N, 3, 64, 64]``. See :ref:`env_spec_shapes`.
+
     Args:
         *args: if an unnamed argument is passed, it must be a dictionary with keys
             matching the expected keys to be found in the :obj:`Composite` object.

tutorials/sphinx-tutorials/pendulum.py

@@ -431,10 +431,15 @@ def _reset(self, tensordict):
 # ^^^^^^^^^^^
 #
 # The environment specs leading dimensions must match the
-# environment batch-size. This is done to enforce that every component of an
+# environment batch-size. The composite spec has that batch shape (it can be
+# empty) and the leaves have that size plus their feature size: an unbatched
+# image env has composite shape ``[]`` and a ``[3, 64, 64]`` leaf, while a
+# batched env with ``batch_size=[N]`` uses ``[N]`` and ``[N, 3, 64, 64]``.
+# This is done to enforce that every component of an
 # environment (including its transforms) have an accurate representation of
 # the expected input and output shapes. This is something that should be
-# accurately coded in stateful settings.
+# accurately coded in stateful settings. Worked examples are in
+# :ref:`env_spec_shapes`.
 #
 # For non batch-locked environments, such as the one in our example (see below),
 # this is irrelevant as the environment batch size will most likely be empty.

tutorials/sphinx-tutorials/torchrl_envs.py

@@ -123,6 +123,76 @@
 # (e.g. Brax) this should also include a representation of the previous state,
 # or any other input to the environment (including inputs at reset time).
 #
+# .. _env_spec_shapes:
+#
+# Spec shapes and batch size
+# ------------------------------
+#
+# When you write a custom environment, the composite spec has the shape of
+# the environment batch size (it can be empty) and the leaves have that size
+# plus their feature size. Feature dimensions -- an image, a board of shape
+# ``[8, 8, 13]``, a vector -- belong on the leaf, not on the composite.
+# ``GymEnv`` and ``ParallelEnv`` follow the same rule; the snippets below
+# construct the specs a custom env would assign, without a new env class.
+# The spec API is documented in :ref:`ref_specs`, and how env specs relate
+# to ``batch_size`` in :ref:`Environment-spec-shapes`.
+#
+# Unbatched env, image observation: ``observation_spec`` has shape ``[]``.
+# Its leaf has the shape of the feature (for example ``[3, 64, 64]`` for an
+# image of 64 pixels width/height). The ``GymEnv`` created above is also
+# unbatched: its composite matches ``env.batch_size`` and the observation
+# leaf is just the 3-vector Pendulum returns.
+
+from torchrl.data import Bounded, Composite, Unbounded
+
+assert env.observation_spec.shape == env.batch_size == torch.Size([])
+assert env.observation_spec["observation"].shape == torch.Size([3])
+
+observation_spec = Composite(
+    pixels=Bounded(low=0, high=255, shape=(3, 64, 64), dtype=torch.uint8),
+    shape=(),
+)
+assert observation_spec.shape == torch.Size([])
+assert observation_spec["pixels"].shape == torch.Size([3, 64, 64])
+
+###############################################################################
+# Batched env: if the env has batch size ``[2]`` (as with a ``ParallelEnv``
+# of two workers), *all* its specs have a leading shape of ``[2, *]``. The
+# composite has shape ``[2]`` and the leaf has shape ``[2, 3, 64, 64]``.
+
+batched_observation_spec = Composite(
+    pixels=Bounded(low=0, high=255, shape=(2, 3, 64, 64), dtype=torch.uint8),
+    shape=(2,),
+)
+assert batched_observation_spec.shape == torch.Size([2])
+assert batched_observation_spec["pixels"].shape == torch.Size([2, 3, 64, 64])
+
+###############################################################################
+# Nested groups (MARL): the env itself can have an empty batch size while
+# inner composites carry a per-group agent dimension. Here ``agents1`` has
+# three members that output images and ``agents2`` has four members that
+# output a state vector of size 5.
+
+full_observation_spec = Composite(
+    agents1=Composite(
+        pixels=Bounded(low=0, high=255, shape=(3, 3, 64, 64), dtype=torch.uint8),
+        shape=(3,),
+    ),
+    agents2=Composite(
+        state=Unbounded(shape=(4, 5)),
+        shape=(4,),
+    ),
+    shape=(),
+)
+assert full_observation_spec.shape == torch.Size([])
+assert full_observation_spec["agents1"].shape == torch.Size([3])
+assert full_observation_spec["agents1", "pixels"].shape == torch.Size(
+    [3, 3, 64, 64]
+)
+assert full_observation_spec["agents2"].shape == torch.Size([4])
+assert full_observation_spec["agents2", "state"].shape == torch.Size([4, 5])
+
+###############################################################################
 # Seeding, resetting and steps
 # ------------------------------
 # The basic operations on an environment are (1) ``set_seed``, (2) ``reset``
@@ -241,6 +311,13 @@
 ###############################################################################
 
 data = env.reset()
+
+###############################################################################
+# The pixel ``GymEnv`` is unbatched: composite shape ``[]``, and the
+# ``"pixels"`` leaf is the image feature (Gym typically uses ``[H, W, C]``).
+
+assert env.observation_spec.shape == env.batch_size == torch.Size([])
+assert env.observation_spec["pixels"].shape == data.get("pixels").shape
 env.close()
 
 ###############################################################################
@@ -472,6 +549,12 @@ def env_make():
 
 parallel_env.reset()
 
+# The composite spec has the env batch ``[3]``; the observation leaf adds
+# the Pendulum feature size, so ``[3, 3]``.
+assert parallel_env.batch_size == torch.Size([3])
+assert parallel_env.observation_spec.shape == torch.Size([3])
+assert parallel_env.observation_spec["observation"].shape == torch.Size([3, 3])
+
 ###############################################################################
 # One can check that the parallel environment has the right batch size.
 # Conventionally, the first part of the ``batch_size`` indicates the batch,
@@ -637,6 +720,11 @@ def env_make(env_name):
 )
 data = parallel_env.reset()
 
+# Two pixel workers: composite shape ``[2]``, pixels leaf ``[2, 3, 64, 64]``.
+assert parallel_env.batch_size == torch.Size([2])
+assert parallel_env.observation_spec.shape == torch.Size([2])
+assert parallel_env.observation_spec["pixels"].shape == torch.Size([2, 3, 64, 64])
+
 plt.figure()
 plt.subplot(121)
 plt.imshow(data[0].get("pixels").permute(1, 2, 0).numpy())