[BUG] Using sorted() on the keys() attribute of a CompositeSpec

[matteobettini]

Describe the bug

In a number of places in the code, the sorted() method is run on the keys() attribute of a CompositeSpec.

The bug was found when trying to create a ParallelEnv for Brax.

Here it fails as for nested specs the outputs of keys can be tuples.

env = ParallelEnv(1, lambda: BraxEnv(_get_envs()[0]))
env.rand_step()

The error is found is in

  File "/Users/Matteo/PycharmProjects/torchrl/torchrl/envs/vec_env.py", line 426, in _create_td
    self.env_input_keys = sorted(self.input_spec.keys())

but I suspect it might not be the only occurrance

To Reproduce

import torchrl

env = ParallelEnv(1, lambda: BraxEnv(_get_envs()[0]))
env.rand_step()

Traceback (most recent call last):
  File "/Users/Matteo/PycharmProjects/VectorizedMultiAgentSimulator/vmas/examples/torch_rl.py", line 30, in <module>
    env.rand_step()
  File "/Users/Matteo/PycharmProjects/torchrl/torchrl/envs/common.py", line 529, in rand_step
    return self.step(tensordict)
  File "/Users/Matteo/PycharmProjects/torchrl/torchrl/envs/common.py", line 341, in step
    tensordict_out = self._step(tensordict)
  File "/Users/Matteo/PycharmProjects/torchrl/torchrl/envs/vec_env.py", line 32, in decorated_fun
    self._create_td()
  File "/Users/Matteo/PycharmProjects/torchrl/torchrl/envs/vec_env.py", line 426, in _create_td
    self.env_input_keys = sorted(self.input_spec.keys())
TypeError: '<' not supported between instances of 'tuple' and 'str'