amend

pytorch · Mar 25, 2024 · 236f41a · 236f41a
1 parent 7c5e618
commit 236f41a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/test/test_transforms.py b/test/test_transforms.py
@@ -10266,6 +10266,7 @@ def test_multistep_transform(self):
             rollout = env.rollout(
                 2, auto_reset=False, tensordict=td, break_when_any_done=False
             ).contiguous()
+            assert "reward" not in rollout.keys()
             out = t._inv_call(rollout)
             td = rollout[..., -1]["next"]
             if out is not None:

diff --git a/torchrl/envs/utils.py b/torchrl/envs/utils.py
@@ -141,6 +141,11 @@ def __init__(
             self.keys_from_root += self.state_keys
         else:
             self.exclude_from_root += self.state_keys
+
+        reset_keys = {_replace_last(key, "_reset") for key in self.done_keys}
+        self.exclude_from_root += list(reset_keys)
+        self.exclude_from_root += list(self.reward_keys)
+
         self.exclude_from_root = self._repr_key_list_as_tree(self.exclude_from_root)
         self.keys_from_root = self._repr_key_list_as_tree(self.keys_from_root)
         self.keys_from_next = self._repr_key_list_as_tree(self.keys_from_next)