[Feature] Heterogeneous Environments compatibility #1411
Conversation
Signed-off-by: Matteo Bettini <matbet@meta.com>
Signed-off-by: Matteo Bettini <matbet@meta.com>
Signed-off-by: Matteo Bettini <matbet@meta.com>
Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
Co-authored-by: Vincent Moens <vincentmoens@gmail.com>
# Conflicts: # test/test_specs.py # torchrl/data/tensor_specs.py # torchrl/data/utils.py
|
So i have updated and added all the requested changes.
|
let's put it in tensordict then! |
torchrl/envs/utils.py
Outdated
| shape = source.get_item_shape(key) | ||
| if -1 not in shape: |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
There was a problem hiding this comment.
I ran the benchmarks and this implementation is 10 - 35% slower than the current one. Any chance we can avoid taking that hit?
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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)
There was a problem hiding this comment.
Just make sure you comment the try / excpet saying that it's temporary and for efficiency purpose
Signed-off-by: Matteo Bettini <matbet@meta.com>
|
updated |
Signed-off-by: Matteo Bettini <matbet@meta.com> Co-authored-by: vmoens <vincentmoens@gmail.com>
This introduces tests for het envs and key modifications to