amend

pytorch · Mar 8, 2024 · a00dd12 · a00dd12
1 parent dbfffff
commit a00dd12
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 2 deletions.
diff --git a/torchrl/envs/common.py b/torchrl/envs/common.py
@@ -2680,7 +2680,7 @@ def maybe_reset(self, tensordict: TensorDictBase) -> TensorDictBase:
         """
         if self._simple_done:
             done = tensordict._get_str("done", default=None)
-            tensordict._set_str("done", done.clone(), validated=True, inplace=False)
+            tensordict._set_str("_reset", done.clone(), validated=True, inplace=False)
             any_done = done.any()
         else:
             any_done = _terminated_or_truncated(

diff --git a/torchrl/envs/transforms/transforms.py b/torchrl/envs/transforms/transforms.py
@@ -733,7 +733,6 @@ def input_spec(self) -> TensorSpec:
     def _step(self, tensordict: TensorDictBase) -> TensorDictBase:
         # No need to clone here because inv does it already
         # tensordict = tensordict.clone(False)
-
         next_preset = tensordict.get("next", None)
         tensordict_in = self.transform.inv(tensordict)
         next_tensordict = self.base_env._step(tensordict_in)