[BugFix] step_mdp nested keys
#1339
Conversation
Signed-off-by: Matteo Bettini <matbet@meta.com>
facebook-github-bot
added
the
CLA Signed
Signed-off-by: Matteo Bettini <matbet@meta.com>
Signed-off-by: Matteo Bettini <matbet@meta.com>
There was a problem hiding this comment.
The tests look good, but I suspect the implementation of step_mdp will be drastically slower than the one we have now.
Given the number of issues opened by users lately regarding environments overhead compared to gym, I'm reluctant to any change in step_mdp that makes it slower than what it already is.
Let's wait until we merge the 2 benchmark PRs (the one with step_mdp and the other to run benchmarks in PRs), then it'll be easier to iterate over this
torchrl/envs/utils.py
Outdated
| if isinstance(done_key, tuple) and len(done_key) == 1: | ||
| done_key = done_key[0] | ||
| if isinstance(reward_key, tuple) and len(reward_key) == 1: | ||
| reward_key = reward_key[0] | ||
| if isinstance(action_key, tuple) and len(action_key) == 1: | ||
| action_key = action_key[0] |
There was a problem hiding this comment.
this is expensive let's not do that if we can avoid it
We could blend this in unravel_keys in tensordict (which is coded in c++) if it's there for efficiency purposes
torchrl/envs/utils.py
Outdated
| if not exclude_action: | ||
| out._set(action_key, tensordict.get(action_key)) |
There was a problem hiding this comment.
Here if action has to be kepts we're removing it and adding it back, which is expensive, let's avoid that
There was a problem hiding this comment.
we are not removing it, the exclusion before is on "next"
torchrl/envs/utils.py
Outdated
| out = tensordict.get("next").clone(False) | ||
| excluded = set() | ||
| excluded = {action_key} | ||
| if exclude_done: | ||
| excluded.add(done_key) | ||
| if exclude_reward: | ||
| excluded.add(reward_key) | ||
| if len(excluded): |
There was a problem hiding this comment.
that will never happen anymore, so we'll always be calling exclude, which is expensive. Let's avoid that
torchrl/envs/utils.py
Outdated
| # out.update(tensordict.select(*td_keys)) | ||
| for key in td_keys: | ||
| out._set(key, tensordict.get(key)) | ||
| excluded = set.union(excluded, set(out.keys(True, True))) |
There was a problem hiding this comment.
set on all the keys is super expensive, i'd avoid it if we can
Signed-off-by: Matteo Bettini <matbet@meta.com>
Signed-off-by: Matteo Bettini <matbet@meta.com>
Signed-off-by: Matteo Bettini <matbet@meta.com>
torchrl/envs/utils.py
Outdated
| td_next = tensordict.get("next") | ||
|
|
||
| td_keys = td.keys(True, True) | ||
| td_next_keys = td_next.keys(True, True) |
There was a problem hiding this comment.
These are is the only time we traverse.
Basically we just visit every key in tensordict
torchrl/envs/utils.py
Outdated
|
|
||
| # Set the keys from root | ||
| if not exclude_action: | ||
| _set_key(dest=out, source=tensordict, key=action_key) |
There was a problem hiding this comment.
here we do the action separate just to not have another if later, perf is the same
torchrl/envs/utils.py
Outdated
| td = tensordict.exclude("next") | ||
| td_next = tensordict.get("next") |
torchrl/envs/utils.py
Outdated
| excluded = { | ||
| done_key if exclude_done else None, | ||
| reward_key if exclude_reward else None, | ||
| } |
There was a problem hiding this comment.
Isn't weird to have None in the set?
Should we not build it iteratively?
There was a problem hiding this comment.
i can do it, i thought preformance was similar and it was more readable, but i ll change it
torchrl/envs/utils.py
Outdated
| if isinstance(key, tuple) and len(key) > 1: # Setting the batch_sizes | ||
| for k in range(1, len(key)): | ||
| dest[key[:k]].batch_size = source[key[:k]].batch_size |
There was a problem hiding this comment.
if two keys share the same root, we'll set the batch-time multiple times, which is time consuming
There was a problem hiding this comment.
yeah I know, i ll improve it
Signed-off-by: Matteo Bettini <matbet@meta.com>
Signed-off-by: Matteo Bettini <matbet@meta.com>
No description provided.