[Feature Request] Tutorial for custom env with complex shapes

🚀 The feature, motivation and pitch

I am totally unable to create a EnvBase subclass, where the *_spec attribute have complex shapes.

For example, I have a state with shape (8,8,13), what shape/batch size should i give to the observation_spec? If I have an action of shape (8,18), what value to the batch size of action_spec?

        self.action_spec = BoundedTensorSpec(
            minimum=0,
            maximum=1,
            shape=action_space.size(),  # `action_space` is a (N,) one-hot tensor
            dtype=self.dtype,
        )

        observation_spec = BoundedTensorSpec(
            low=0,
            high=1,
            shape=state.size(),  # `state` is a 8x8x13 tensor
            dtype=self.dtype,
        )
        self.observation_spec = CompositeSpec(observation=observation_spec)

Would this work?

Solution

A tutorial where the shape attribute/argument is better explored will suffice. Just give examples of all edge cases, for how to use shape.