[Feature] Allow for actions and rewards to be in the reset tensordict (…

…pytorch#458)
maxweissenbacher · Sep 16, 2022 · cd90398 · cd90398
1 parent 8c313c8
commit cd90398
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/torchrl/envs/common.py b/torchrl/envs/common.py
@@ -367,8 +367,8 @@ def reset(
             tensordict_reset = step_tensordict(
                 tensordict_reset,
                 exclude_done=False,
-                exclude_reward=True,
-                exclude_action=True,
+                exclude_reward=False,  # some policies may need reward and action at reset time
+                exclude_action=False,
             )
         if tensordict is not None:
             tensordict.update(tensordict_reset)