[BUG?] How to handle next with custom environment and check_env_specs() about rl HOT 2 CLOSED

        self.observation_spec = CompositeSpec({
            "observation": CompositeSpec({
                "positions": BoundedTensorSpec(
                    low=0.0,
                    high=1.0,
                    shape=torch.Size([NUM_BLOCKS, 2]),
                    dtype=torch.float32
                ),
                "sizes": BoundedTensorSpec(
                    low=0.1,
                    high=1.0,
                    shape=torch.Size([NUM_BLOCKS, 2]),
                dtype=torch.float32
                )
            }),
            "state": CompositeSpec({
                "distance_from_center": UnboundedContinuousTensorSpec(
                    shape=torch.Size([NUM_BLOCKS]),
                    dtype=torch.float32
                ),
            }),
        })

import torch
from torchrl.envs import EnvBase
from torchrl.envs.utils import check_env_specs
from torchrl.data import BoundedTensorSpec, CompositeSpec, UnboundedContinuousTensorSpec
from tensordict import TensorDict

COUNTER = 0
NUM_BLOCKS = 4

class BlockArrangementEnv(EnvBase):

    def __init__(self):
        super().__init__()

        self.observation_spec = CompositeSpec({
            "observation": CompositeSpec({
                "positions": BoundedTensorSpec(
                    low=0.0,
                    high=1.0,
                    shape=torch.Size([NUM_BLOCKS, 2]),
                    dtype=torch.float32
                ),
                "sizes": BoundedTensorSpec(
                    low=0.1,
                    high=1.0,
                    shape=torch.Size([NUM_BLOCKS, 2]),
                dtype=torch.float32
                )
            }),
            "state": CompositeSpec({
                "distance_from_center": UnboundedContinuousTensorSpec(
                    shape=torch.Size([NUM_BLOCKS]),
                    dtype=torch.float32
                ),
            }),
        })

        self.full_state_spec = CompositeSpec({
            "state": CompositeSpec({
                "distance_from_center": UnboundedContinuousTensorSpec(
                    shape=torch.Size([NUM_BLOCKS]),
                    dtype=torch.float32
                ),
            }),
            "context": UnboundedContinuousTensorSpec(
                    shape=torch.Size([NUM_BLOCKS]),
                    dtype=torch.float32
                )
        })

        self.action_spec = CompositeSpec({
            "action": CompositeSpec({
                "index": BoundedTensorSpec(
                    low=0,
                    high=NUM_BLOCKS - 1,
                    shape=torch.Size([1]),
                    dtype=torch.int
                ),
                "delta": BoundedTensorSpec(
                    low=-1.0,
                    high=1.0,
                    shape=torch.Size([2]),
                    dtype=torch.float32
                )
            })
        })

        self.reward_spec = UnboundedContinuousTensorSpec(
            shape=torch.Size([NUM_BLOCKS]),
            dtype=torch.float32
        )


    def _reset(self, td):
        global COUNTER
        COUNTER += 1
        return TensorDict({
            "observation": {
                "positions": torch.rand([NUM_BLOCKS, 2]),
                "sizes": torch.FloatTensor(NUM_BLOCKS, 2).uniform_(0.1, 1.0),
            },
            "state": {
                "distance_from_center": torch.rand([NUM_BLOCKS]),
            },
            "context": self.full_state_spec["context"].zero() + COUNTER,
        }, batch_size=[])

    def _step(self, td, **kwargs):
        return TensorDict({
            "observation": {
                "positions": torch.rand([NUM_BLOCKS, 2]),
                "sizes": torch.FloatTensor(NUM_BLOCKS, 2).uniform_(0.1, 1.0),
            },
            "state": {
                "distance_from_center": torch.rand([NUM_BLOCKS]),
            },
            "reward": torch.rand([NUM_BLOCKS]),
            "done": torch.tensor(False)
        }, batch_size=[])

    def _set_seed(self, seed):
        pass


env = BlockArrangementEnv()
check_env_specs(env)
assert (env.rollout(3)["context"] == COUNTER).all()