[Feature] Heterogeneous Environments compatibility

[matteobettini]

This introduces tests for het envs and key modifications to

• rollout
• step_mdp

[vmoens]

Thanks for this, really useful

[matteobettini]

So i have updated and added all the requested changes.
my takes on the new function are

• it should be public because any user which has some nested lazy stacks somewhere should use it after calling env.rollout(return_contiguous=False)
• my personal preference is that it should be in tensordict utils but the choice is yours

[vmoens]

So i have updated and added all the requested changes. my takes on the new function are

• it should be public because any user which has some nested lazy stacks somewhere should use it after calling env.rollout(return_contiguous=False)
• my personal preference is that it should be in tensordict utils but the choice is yours

let's put it in tensordict then!

[vmoens]

what's the overhead here?
Since this will likely be very frequent I'm a bit worried by increasing the control flow for every key

[vmoens]

For a regular td, this op means that we'll be accessing the tensor twice through get (once to get its shape, a second time to get it). Accessing it once should suffice.
Not I understand that we can't use get if the tensor is heterogeneous... A try/except may not be the right solution either since it'll be slow for heterogeneous envs. However regular envs won't be impacted

[vmoens]

I ran the benchmarks and this implementation is 10 - 35% slower than the current one. Any chance we can avoid taking that hit?

[matteobettini]

I am not sure what I can do here

• in the past i was using try blocks, but you made me move away from them
• an isinstance wont work since the key might be het even if the container is a regular td

so I can move back to try except if you prefer

[matteobettini]

lemme switch to try blocks

ideally tho this problem could be solved once and for all in tensordict, with an efficient way to understand if a key is het

[vmoens]

yep that's what I'm thinking too
If we have tensorstack we can even leave the block with try / except and it won't break (I guess)

[vmoens]

Just make sure you comment the try / excpet saying that it's temporary and for efficiency purpose

[matteobettini]

updated

[vmoens]

Het groene licht is gekomen