[QUESTION] Using LSTMs with vectorized envs without ParallelEnv

[ADebor]

Hi there,

As mentioned here, I'm trying to use torchrl with NVIDIA Orbit to train an agent in parallel robot environments. I tried to draw inspiration from your recent IsaacGymEnv class to create a simple OrbitEnv class, inheriting from GymEnv directly as Orbit environments are registered in gym (made sense to me, but I may be wrong). I'm thus able to create a torchrl environment and add transforms, but I get in trouble when trying to use ("parallel") RNNs and thus hidden states.

Since Orbit environments are vectorized environments, I only create one torchrl environment wrapping the orbit one and I set the batch_size equal to the number of environments in the Orbit vectorized one. As I want to use LSTMs, I add the make_tensordict_primer() transform to my environment. If I then try to reset my environment, I get the following error:

RuntimeError: batch dimension mismatch, got self.batch_size=torch.Size([2]) and value.shape[:self.batch_dims]=torch.Size([1]) with value tensor([[0., 0., 0., 0., 0.]])

where 2 is the number of parallel environments and the value tensor is a hidden state one (of length 5). I don't get what the problem is exactly, but it looks like the batch size is not taken into account when it comes to hidden states. When trying with a basic gym environment, I wrap my environment within a ParallelEnv instance and I have no issue, but in this case, using ParallelEnv causes a problem since (as I get it) torchrl tries to create multiple parallel and separate environments, which seems to be in conflict with Orbit where all environments are in the same scene in Isaac Sim.

Can you see a way to properly use LSTMs in torchrl while using a vectorized environment from Orbit? Don't hesitate to tell me if the question or context isn't clear. I apologize if I'm out of line here, I'm maybe missing something important in the way torchrl is supposed to be used.

---

Here is the complete stack trace:

Traceback (most recent call last):
  File "test.py", line 61, in <module>
    td = env.reset()
  File "/home/adebor/isaacsim_ws/rl/torchrl/envs/common.py", line 949, in reset
    tensordict_reset = self._reset(tensordict, **kwargs)
  File "/home/adebor/isaacsim_ws/rl/torchrl/envs/transforms/transforms.py", line 651, in _reset
    out_tensordict = self.transform.reset(out_tensordict)
  File "/home/adebor/isaacsim_ws/rl/torchrl/envs/transforms/transforms.py", line 889, in reset
    tensordict = t.reset(tensordict)
  File "/home/adebor/isaacsim_ws/rl/torchrl/envs/transforms/transforms.py", line 3016, in reset
    tensordict.set(key, value)
  File "/home/adebor/isaacsim_ws/tensordict/tensordict/tensordict.py", line 761, in set
    return self._set_tuple(key, item, inplace=inplace, validated=False)
  File "/home/adebor/isaacsim_ws/tensordict/tensordict/tensordict.py", line 4122, in _set_tuple
    return self._set_str(key[0], value, inplace=inplace, validated=validated)
  File "/home/adebor/isaacsim_ws/tensordict/tensordict/tensordict.py", line 4093, in _set_str
    value = self._validate_value(value, check_shape=True)
  File "/home/adebor/isaacsim_ws/tensordict/tensordict/tensordict.py", line 1732, in _validate_value
    f"batch dimension mismatch, got self.batch_size"
RuntimeError: batch dimension mismatch, got self.batch_size=torch.Size([2]) and value.shape[:self.batch_dims]=torch.Size([1]) with value tensor([[0., 0., 0., 0., 0.]])

---

Versions:
torchrl==0.1.1
torch==1.13.1

[btx0424]

It would be clearer if you could provide the specs of your OrbitEnv before and after the transformation.

Also, it seems that the LSTMModule can not properly work with vectorized environments in general because the current implementation of TensordictPremier.reset would always reset all the hidden states.

[ADebor]

I'm not sure these are the specs you mention, but here's what I get:

• before any transformation

OrbitEnv(env=Isaac-Cartpole-v0, batch_size=torch.Size([2]), device=cuda:0)

• after all the transforms except the make_tensordict_primer one

TransformedEnv(
    env=OrbitEnv(env=Isaac-Cartpole-v0, batch_size=torch.Size([2]), device=cuda:0),
    transform=Compose(
            InitTracker(keys=[]),
            DoubleToFloat(in_keys=['observation'], out_keys=['observation'], in_keys_inv=[], out_keys_inv=[]),
            StepCounter(keys=[]),
            ObservationNorm(keys=['observation'])))

• after the make_tensordict_primer one

TransformedEnv(
    env=OrbitEnv(env=Isaac-Cartpole-v0, batch_size=torch.Size([2]), device=cuda:0),
    transform=Compose(
            InitTracker(keys=[]),
            DoubleToFloat(in_keys=['observation'], out_keys=['observation'], in_keys_inv=[], out_keys_inv=[]),
            StepCounter(keys=[]),
            ObservationNorm(keys=['observation']),
            TensorDictPrimer(primers={('recurrent_state_h',): UnboundedContinuousTensorSpec(
                shape=torch.Size([1, 5]),
                space=None,
                device=cpu,
                dtype=torch.float32,
                domain=continuous), ('recurrent_state_c',): UnboundedContinuousTensorSpec(
                shape=torch.Size([1, 5]),
                space=None,
                device=cpu,
                dtype=torch.float32,
                domain=continuous)}, default_value=0.0, random=False)))

Concerning the LSTMModule-Vectorized envs issue, I thought this had been fixed here.

[vmoens]

I would be interested in seeing the what the OrbitEnv looks like @ADebor. It would be cool to integrate that in torchrl!
(and easier to debug :p )

[btx0424]

Hi there, I am glad to give it a try or share some experiences if you like. What I personally have been doing is very similar to GymWrapper. Some tricky differences prevent me from directly treating it as a Vectorized Gym environment. Meanwhile, Orbit's behavior upon reset is like early Gym versions, i.e., no terminal observations are returned. It could be a problem in some cases. I will try reaching out to the authors shortly to see what they plan to do about it.

[ADebor]

Hi there, sorry for the late reply. I had to work on other stuff for a while, and still have. I'll get back to you as soon as I'm able to look at the code again. Since Orbit just merged their dev branch, I'm not sure my previous attempts are still compatible with this new version. Plus, my OrbitEnv was a very basic and naive class, and only small tests were conducted on my side. I'd be thrilled to see what you did @btx0424, thanks for the help.

[Cadene]

@btx0424 @ADebor Thanks for your effort on porting Orbit. Is there a PR somewhere or an issue to track progress? Would be happy to see how you did it and maybe contribute :)

[yu-fz]

I believe the batch dimension mismatch you see here is actually unrelated to the use of LSTMs and its associated make_tensordict_primer() transform. I was able to successfully call env.reset() on a GymEnv wrapping an IsaacLab env with an LSTMModule. It seems like the error comes from the way you are calling the GymEnv constructor in your wrapper when you pass the batch_size as an argument, it should work if you just use take the default value.