[Feature Request] Tutorial for custom env with complex shapes about rl HOT 1 OPEN

We can document this better.
In general the idea is that your composite spec has the shape of the batch size (it can be empty) and the leaves have that size plus their feature size.

Examples:
Your env has a single agent, no batch-size.
full_composite_spec has shape []
Its leaf has the shape of the feature (eg [3, 64, 64] if you have an image of 64 pixels width/height).

If you have a batched env with 2 agents it could have a batch size [2]. This is what will happen with a ParallelEnv for instance. All its specs will have a leading shape of [2, *], meaning that your full_composite_spec will have a shape of [2] and the leaf will have shape [2, 3, 64, 64].

Final case: your env has no batch size but it simulates several groups of agents (MARL setting). A first group named agents1 has 3 identical members and a second, agents2 has 4. The first outputs images from its steps and the second outputs a state vector of shape 5.

Here's how to build it:

full_observation_spec = CompositeSpec(
    agents1=CompositeSpec(pixels=SomeSpec(3, 3, 64, 64), shape=[3]),
    agents2=CompositeSpec(state=SomeOtherSpec(4, 5), shape=[4]), 
shape=[])

I hope that clarifies things a tiny bit

from rl.