amend

pytorch · vmoens · Oct 5, 2023 · Jul 28, 2023 · Jul 28, 2023 · Aug 10, 2023
commit 858da29bfab4bbbf0ca219d122a248da8e209565
diff --git a/torchrl/envs/vec_env.py b/torchrl/envs/vec_env.py
@@ -632,7 +632,7 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
                 {}, batch_size=self.shared_tensordict_parent.shape, device=self.device
             )
             for key in self._selected_reset_keys:
-                if key != "_reset":
+                if key != ("_reset",):
                     _set_single_key(self.shared_tensordict_parent, out, key, clone=True)
             return out
         else:
@@ -879,7 +879,7 @@ def _reset(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
                 {}, batch_size=self.shared_tensordict_parent.shape, device=self.device
             )
             for key in self._selected_reset_keys:
-                if key != "_reset":
+                if key != ("_reset",):
                     _set_single_key(self.shared_tensordict_parent, out, key, clone=True)
             return out
         else:

diff --git a/tutorials/sphinx-tutorials/pendulum.py b/tutorials/sphinx-tutorials/pendulum.py
@@ -241,16 +241,13 @@ def _step(tensordict):
     new_th = th + new_thdot * dt
     reward = -costs.view(*tensordict.shape, 1)
     done = torch.zeros_like(reward, dtype=torch.bool)
-    # The output must be written in a ``"next"`` entry
     out = TensorDict(
         {
-            "next": {
-                "th": new_th,
-                "thdot": new_thdot,
-                "params": tensordict["params"],
-                "reward": reward,
-                "done": done,
-            }
+            "th": new_th,
+            "thdot": new_thdot,
+            "params": tensordict["params"],
+            "reward": reward,
+            "done": done,
         },
         tensordict.shape,
     )