Remove current_tensordict references (#229)

Author: vmoens

test/mocking_classes.py

@@ -129,9 +129,14 @@ class DiscreteActionVecMockEnv(_MockEnv):
     )
     action_spec = OneHotDiscreteTensorSpec(7)
     reward_spec = UnboundedContinuousTensorSpec()
+
     from_pixels = False
 
     out_key = "observation"
+    _out_key = "observation_orig"
+    input_spec = CompositeSpec(
+        **{_out_key: observation_spec["next_observation"], "action": action_spec}
+    )
 
     def _get_in_obs(self, obs):
         return obs
@@ -145,6 +150,7 @@ def _reset(self, tensordict: _TensorDict) -> _TensorDict:
         tensordict = tensordict.select().set(
             "next_" + self.out_key, self._get_out_obs(state)
         )
+        tensordict = tensordict.set("next_" + self._out_key, self._get_out_obs(state))
         tensordict.set("done", torch.zeros(*tensordict.shape, 1, dtype=torch.bool))
         return tensordict
 
@@ -157,12 +163,12 @@ def _step(
         assert (a.sum(-1) == 1).all()
         assert not self.is_done, "trying to execute step in done env"
 
-        obs = (
-            self._get_in_obs(self.current_tensordict.get(self.out_key))
-            + a / self.maxstep
-        )
+        obs = self._get_in_obs(tensordict.get(self._out_key)) + a / self.maxstep
         tensordict = tensordict.select()  # empty tensordict
+
         tensordict.set("next_" + self.out_key, self._get_out_obs(obs))
+        tensordict.set("next_" + self._out_key, self._get_out_obs(obs))
+
         done = torch.isclose(obs, torch.ones_like(obs) * (self.counter + 1))
         reward = done.any(-1).unsqueeze(-1)
         # set done to False
@@ -182,6 +188,10 @@ class ContinuousActionVecMockEnv(_MockEnv):
     from_pixels = False
 
     out_key = "observation"
+    _out_key = "observation_orig"
+    input_spec = CompositeSpec(
+        **{_out_key: observation_spec["next_observation"], "action": action_spec}
+    )
 
     def _get_in_obs(self, obs):
         return obs
@@ -193,9 +203,9 @@ def _reset(self, tensordict: _TensorDict) -> _TensorDict:
         self.counter += 1
         self.step_count = 0
         state = torch.zeros(self.size) + self.counter
-        tensordict = tensordict.select().set(
-            "next_" + self.out_key, self._get_out_obs(state)
-        )
+        tensordict = tensordict.select()
+        tensordict.set("next_" + self.out_key, self._get_out_obs(state))
+        tensordict.set("next_" + self._out_key, self._get_out_obs(state))
         tensordict.set("done", torch.zeros(*tensordict.shape, 1, dtype=torch.bool))
         return tensordict
 
@@ -208,11 +218,12 @@ def _step(
         a = tensordict.get("action")
         assert not self.is_done, "trying to execute step in done env"
 
-        obs = self._obs_step(
-            self._get_in_obs(self.current_tensordict.get(self.out_key)), a
-        )
+        obs = self._obs_step(self._get_in_obs(tensordict.get(self._out_key)), a)
         tensordict = tensordict.select()  # empty tensordict
+
         tensordict.set("next_" + self.out_key, self._get_out_obs(obs))
+        tensordict.set("next_" + self._out_key, self._get_out_obs(obs))
+
         done = torch.isclose(obs, torch.ones_like(obs) * (self.counter + 1))
         reward = done.any(-1).unsqueeze(-1)
         done = done.all(-1).unsqueeze(-1)
@@ -251,6 +262,10 @@ class DiscreteActionConvMockEnv(DiscreteActionVecMockEnv):
     from_pixels = True
 
     out_key = "pixels"
+    _out_key = "pixels_orig"
+    input_spec = CompositeSpec(
+        **{_out_key: observation_spec["next_pixels"], "action": action_spec}
+    )
 
     def _get_out_obs(self, obs):
         obs = torch.diag_embed(obs, 0, -2, -1).unsqueeze(0)
@@ -287,6 +302,10 @@ class ContinuousActionConvMockEnv(ContinuousActionVecMockEnv):
     from_pixels = True
 
     out_key = "pixels"
+    _out_key = "pixels_orig"
+    input_spec = CompositeSpec(
+        **{_out_key: observation_spec["next_pixels"], "action": action_spec}
+    )
 
     def _get_out_obs(self, obs):
         obs = torch.diag_embed(obs, 0, -2, -1).unsqueeze(0)

test/test_collector.py

@@ -274,7 +274,7 @@ def env_fn():
     ccollector.shutdown()
 
 
-@pytest.mark.parametrize("num_env", [3, 1])
+@pytest.mark.parametrize("num_env", [1, 3])
 @pytest.mark.parametrize("env_name", ["conv", "vec"])
 def test_collector_consistency(num_env, env_name, seed=100):
     if num_env == 1:
@@ -320,9 +320,9 @@ def env_fn(seed):
         device="cpu",
         pin_memory=False,
     )
-    collector = iter(collector)
-    b1 = next(collector)
-    b2 = next(collector)
+    collector_iter = iter(collector)
+    b1 = next(collector_iter)
+    b2 = next(collector_iter)
     with pytest.raises(AssertionError):
         assert_allclose_td(b1, b2)
 
@@ -334,6 +334,7 @@ def env_fn(seed):
     ), f"got batch_size {rollout1a.batch_size} and {b1.batch_size}"
 
     assert_allclose_td(rollout1a, b1.select(*rollout1a.keys()))
+    collector.shutdown()
 
 
 @pytest.mark.parametrize("num_env", [1, 3])