[Refactor] Better init for CatFrames buffers + removing default init values (#874)

Author: vmoens

examples/dreamer/dreamer_utils.py

@@ -90,7 +90,7 @@ def make_env_transforms(
         if cfg.grayscale:
             env.append_transform(GrayScale())
         env.append_transform(FlattenObservation(0, -3))
-        env.append_transform(CatFrames(N=cfg.catframes, in_keys=["pixels"]))
+        env.append_transform(CatFrames(N=cfg.catframes, in_keys=["pixels"], dim=-3))
         if stats is None:
             obs_stats = {
                 "loc": torch.zeros(env.observation_spec["pixels"].shape),

test/test_helpers.py

@@ -89,15 +89,6 @@ def dreamer_constructor_fixture():
     sys.path.pop()
 
 
-def _assert_keys_match(td, expeceted_keys):
-    td_keys = list(td.keys())
-    d = set(td_keys) - set(expeceted_keys)
-    assert len(d) == 0, f"{d} is in tensordict but unexpected: {td.keys()}"
-    d = set(expeceted_keys) - set(td_keys)
-    assert len(d) == 0, f"{d} is expected but not in tensordict: {td.keys()}"
-    assert len(td_keys) == len(expeceted_keys)
-
-
 @pytest.mark.skipif(not _has_gym, reason="No gym library found")
 @pytest.mark.skipif(not _has_tv, reason="No torchvision library found")
 @pytest.mark.skipif(not _has_hydra, reason="No hydra library found")
@@ -152,16 +143,20 @@ def test_dqn_maker(
         else:
             actor(td)
 
-        expected_keys = ["done", "action", "action_value"]
+        expected_keys = [
+            "done",
+            "action",
+            "action_value",
+        ]
         if from_pixels:
-            expected_keys += ["pixels", "pixels_orig"]
+            expected_keys += ["pixels", "pixels_orig", "_reset"]
         else:
             expected_keys += ["observation_orig", "observation_vector"]
 
         if not distributional:
             expected_keys += ["chosen_action_value"]
         try:
-            _assert_keys_match(td, expected_keys)
+            assert set(td.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -217,15 +212,15 @@ def test_ddpg_maker(device, from_pixels, gsde, exploration):
                 actor(td)
         expected_keys = ["done", "action", "param"]
         if from_pixels:
-            expected_keys += ["pixels", "hidden", "pixels_orig"]
+            expected_keys += ["pixels", "hidden", "pixels_orig", "_reset"]
         else:
             expected_keys += ["observation_vector", "observation_orig"]
 
         if cfg.gSDE:
             expected_keys += ["scale", "loc", "_eps_gSDE"]
 
         try:
-            _assert_keys_match(td, expected_keys)
+            assert set(td.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -245,7 +240,7 @@ def test_ddpg_maker(device, from_pixels, gsde, exploration):
             value(td)
         expected_keys += ["state_action_value"]
         try:
-            _assert_keys_match(td, expected_keys)
+            assert set(td.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -359,8 +354,12 @@ def test_ppo_maker(
             else:
                 actor(td_clone)
 
+        if from_pixels:
+            # for CatFrames
+            expected_keys += ["_reset"]
+
         try:
-            _assert_keys_match(td_clone, expected_keys)
+            assert set(td_clone.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -386,6 +385,9 @@ def test_ppo_maker(
             "pixels_orig" if len(from_pixels) else "observation_orig",
             "state_value",
         ]
+        if from_pixels:
+            # for CatFrames
+            expected_keys += ["_reset"]
         if shared_mapping:
             expected_keys += ["hidden"]
         if len(gsde):
@@ -398,7 +400,7 @@ def test_ppo_maker(
         else:
             value(td_clone)
         try:
-            _assert_keys_match(td_clone, expected_keys)
+            assert set(td_clone.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -495,6 +497,9 @@ def test_a2c_maker(
             "action",
             "sample_log_prob",
         ]
+        if from_pixels:
+            # for CatFrames
+            expected_keys += ["_reset"]
         if action_space == "continuous":
             expected_keys += ["loc", "scale"]
         else:
@@ -514,7 +519,7 @@ def test_a2c_maker(
                 actor(td_clone)
 
         try:
-            _assert_keys_match(td_clone, expected_keys)
+            assert set(td_clone.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -540,6 +545,9 @@ def test_a2c_maker(
             "pixels_orig" if len(from_pixels) else "observation_orig",
             "state_value",
         ]
+        if from_pixels:
+            # for CatFrames
+            expected_keys += ["_reset"]
         if shared_mapping:
             expected_keys += ["hidden"]
         if len(gsde):
@@ -552,7 +560,7 @@ def test_a2c_maker(
         else:
             value(td_clone)
         try:
-            _assert_keys_match(td_clone, expected_keys)
+            assert set(td_clone.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -631,6 +639,9 @@ def test_sac_make(device, gsde, tanh_loc, from_pixels, exploration):
             "loc",
             "scale",
         ]
+        if from_pixels:
+            # for CatFrames
+            expected_keys += ["_reset"]
         if len(gsde):
             expected_keys += ["_eps_gSDE"]
 
@@ -643,7 +654,7 @@ def test_sac_make(device, gsde, tanh_loc, from_pixels, exploration):
                 torch.testing.assert_close(td_clone.get("action"), tsf_loc)
 
         try:
-            _assert_keys_match(td_clone, expected_keys)
+            assert set(td_clone.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -667,7 +678,7 @@ def test_sac_make(device, gsde, tanh_loc, from_pixels, exploration):
             expected_keys += ["_eps_gSDE"]
 
         try:
-            _assert_keys_match(td_clone, expected_keys)
+            assert set(td_clone.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -687,7 +698,7 @@ def test_sac_make(device, gsde, tanh_loc, from_pixels, exploration):
             expected_keys += ["_eps_gSDE"]
 
         try:
-            _assert_keys_match(td, expected_keys)
+            assert set(td.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -756,12 +767,12 @@ def test_redq_make(device, from_pixels, gsde, exploration):
         if len(gsde):
             expected_keys += ["_eps_gSDE"]
         if from_pixels:
-            expected_keys += ["hidden", "pixels", "pixels_orig"]
+            expected_keys += ["hidden", "pixels", "pixels_orig", "_reset"]
         else:
             expected_keys += ["observation_vector", "observation_orig"]
 
         try:
-            _assert_keys_match(td, expected_keys)
+            assert set(td.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -786,11 +797,11 @@ def test_redq_make(device, from_pixels, gsde, exploration):
         if len(gsde):
             expected_keys += ["_eps_gSDE"]
         if from_pixels:
-            expected_keys += ["hidden", "pixels", "pixels_orig"]
+            expected_keys += ["hidden", "pixels", "pixels_orig", "_reset"]
         else:
             expected_keys += ["observation_vector", "observation_orig"]
         try:
-            _assert_keys_match(td, expected_keys)
+            assert set(td.keys()) == set(expected_keys)
         except AssertionError:
             proof_environment.close()
             raise
@@ -861,6 +872,7 @@ def test_dreamer_make(device, tanh_loc, exploration, dreamer_constructor_fixture
             "state",
             ("next", "reco_pixels"),
             "next",
+            "_reset",
         }
         assert set(out.keys(True)) == expected_keys
 

test/test_rb.py

@@ -798,7 +798,7 @@ def test_insert_transform():
     pytest.param(partial(SqueezeTransform, squeeze_dim=-1), id="SqueezeTransform"),
     GrayScale,
     pytest.param(partial(ObservationNorm, loc=1, scale=2), id="ObservationNorm"),
-    CatFrames,
+    pytest.param(partial(CatFrames, dim=-3, N=4), id="CatFrames"),
     pytest.param(partial(RewardScaling, loc=1, scale=2), id="RewardScaling"),
     DoubleToFloat,
     VecNorm,

test/test_transforms.py

@@ -878,8 +878,6 @@ def test_time_max_pool(self, T, seq_len, device):
             tensor_list.append(torch.rand(batch, nodes).to(device))
         max_vals, _ = torch.max(torch.stack(tensor_list[-T:]), dim=0)
 
-        print(f"max vals: {max_vals}")
-
         for i in range(seq_len):
             env_td = TensorDict(
                 {
@@ -946,7 +944,11 @@ def test_totensorimage(self, keys, batch, device):
     @pytest.mark.parametrize("device", get_available_devices())
     def test_compose(self, keys, batch, device, nchannels=1, N=4):
         torch.manual_seed(0)
-        t1 = CatFrames(in_keys=keys, N=4)
+        t1 = CatFrames(
+            in_keys=keys,
+            N=4,
+            dim=-3,
+        )
         t2 = FiniteTensorDictCheck()
         compose = Compose(t1, t2)
         dont_touch = torch.randn(*batch, nchannels, 16, 16, device=device)
@@ -1287,7 +1289,11 @@ def test_catframes_transform_observation_spec(self):
         key1 = "first key"
         key2 = "second key"
         keys = [key1, key2]
-        cat_frames = CatFrames(N=N, in_keys=keys)
+        cat_frames = CatFrames(
+            N=N,
+            in_keys=keys,
+            dim=-3,
+        )
         mins = [0, 0.5]
         maxes = [0.5, 1]
         observation_spec = CompositeSpec(
@@ -1321,31 +1327,50 @@ def test_catframes_transform_observation_spec(self):
                 )
 
     @pytest.mark.parametrize("device", get_available_devices())
+    @pytest.mark.parametrize("batch_size", [(), (1,), (1, 2)])
     @pytest.mark.parametrize("d", range(1, 4))
-    def test_catframes_buffer_check_latest_frame(self, device, d):
+    @pytest.mark.parametrize("dim", [-3, -2, 1])
+    @pytest.mark.parametrize("N", [2, 4])
+    def test_catframes_buffer_check_latest_frame(self, device, d, batch_size, dim, N):
         key1 = "first key"
         key2 = "second key"
-        N = 4
         keys = [key1, key2]
-        key1_tensor = torch.ones(1, d, 3, 3, device=device) * 2
-        key2_tensor = torch.ones(1, d, 3, 3, device=device)
+        extra_d = (3,) * (-dim - 1)
+        key1_tensor = torch.ones(*batch_size, d, *extra_d, device=device) * 2
+        key2_tensor = torch.ones(*batch_size, d, *extra_d, device=device)
         key_tensors = [key1_tensor, key2_tensor]
-        td = TensorDict(dict(zip(keys, key_tensors)), [1], device=device)
-        cat_frames = CatFrames(N=N, in_keys=keys)
+        td = TensorDict(dict(zip(keys, key_tensors)), batch_size, device=device)
+        if dim > 0:
+            with pytest.raises(
+                ValueError, match="dim must be > 0 to accomodate for tensordict"
+            ):
+                cat_frames = CatFrames(N=N, in_keys=keys, dim=dim)
+            return
+        cat_frames = CatFrames(N=N, in_keys=keys, dim=dim)
 
         tdclone = cat_frames(td.clone())
         latest_frame = tdclone.get(key2)
 
-        assert latest_frame.shape[1] == N * d
-        assert (latest_frame[0, :-d] == 0).all()
-        assert (latest_frame[0, -d:] == 1).all()
+        assert latest_frame.shape[dim] == N * d
+        slices = (slice(None),) * (-dim - 1)
+        index1 = (Ellipsis, slice(None, -d), *slices)
+        index2 = (Ellipsis, slice(-d, None), *slices)
+        assert (latest_frame[index1] == 0).all()
+        assert (latest_frame[index2] == 1).all()
+        v1 = latest_frame[index1]
 
         tdclone = cat_frames(td.clone())
         latest_frame = tdclone.get(key2)
 
-        assert latest_frame.shape[1] == N * d
-        assert (latest_frame[0, : -2 * d] == 0).all()
-        assert (latest_frame[0, -2 * d :] == 1).all()
+        assert latest_frame.shape[dim] == N * d
+        index1 = (Ellipsis, slice(None, -2 * d), *slices)
+        index2 = (Ellipsis, slice(-2 * d, None), *slices)
+        assert (latest_frame[index1] == 0).all()
+        assert (latest_frame[index2] == 1).all()
+        v2 = latest_frame[index1]
+
+        # we don't want the same tensor to be returned twice, but they're all copies of the same buffer
+        assert v1 is not v2
 
     @pytest.mark.parametrize("device", get_available_devices())
     def test_catframes_reset(self, device):
@@ -1357,19 +1382,20 @@ def test_catframes_reset(self, device):
         key2_tensor = torch.randn(1, 1, 3, 3, device=device)
         key_tensors = [key1_tensor, key2_tensor]
         td = TensorDict(dict(zip(keys, key_tensors)), [1], device=device)
-        cat_frames = CatFrames(N=N, in_keys=keys)
+        cat_frames = CatFrames(N=N, in_keys=keys, dim=-3)
 
-        cat_frames(td)
+        cat_frames(td.clone())
         buffer = getattr(cat_frames, f"_cat_buffers_{key1}")
 
-        passed_back_td = cat_frames.reset(td)
+        tdc = td.clone()
+        passed_back_td = cat_frames.reset(tdc)
+        assert "_reset" in tdc.keys()
 
-        assert td is passed_back_td
-        assert (0 == buffer).all()
+        assert tdc is passed_back_td
+        assert (buffer == 0).all()
 
-        _ = cat_frames._call(td)
-        assert (0 == buffer[..., :-1, :, :]).all()
-        assert (0 != buffer[..., -1:, :, :]).all()
+        _ = cat_frames._call(tdc)
+        assert (buffer != 0).all()
 
     @pytest.mark.parametrize("device", get_available_devices())
     def test_finitetensordictcheck(self, device):
@@ -1691,7 +1717,7 @@ def test_append(self):
         (key,) = itertools.islice(obs_spec.keys(), 1)
 
         env = TransformedEnv(env)
-        env.append_transform(CatFrames(N=4, cat_dim=-1, in_keys=[key]))
+        env.append_transform(CatFrames(N=4, dim=-1, in_keys=[key]))
         assert isinstance(env.transform, Compose)
         assert len(env.transform) == 1
         obs_spec = env.observation_spec
@@ -1715,7 +1741,7 @@ def test_insert(self):
         assert env._observation_spec is not None
         assert env._reward_spec is not None
 
-        env.insert_transform(0, CatFrames(N=4, cat_dim=-1, in_keys=[key]))
+        env.insert_transform(0, CatFrames(N=4, dim=-1, in_keys=[key]))
 
         # transformed envs do not have spec after insert -- they need to be computed
         assert env._input_spec is None
@@ -1762,7 +1788,7 @@ def test_insert(self):
         assert env._observation_spec is None
         assert env._reward_spec is None
 
-        env.insert_transform(-5, CatFrames(N=4, cat_dim=-1, in_keys=[key]))
+        env.insert_transform(-5, CatFrames(N=4, dim=-1, in_keys=[key]))
         assert isinstance(env.transform, Compose)
         assert len(env.transform) == 6
 
@@ -2441,7 +2467,7 @@ def test_select(self):
     pytest.param(partial(SqueezeTransform, squeeze_dim=-1), id="SqueezeTransform"),
     GrayScale,
     ObservationNorm,
-    CatFrames,
+    pytest.param(partial(CatFrames, dim=-3, N=4), id="CatFrames"),
     pytest.param(partial(RewardScaling, loc=1, scale=2), id="RewardScaling"),
     FiniteTensorDictCheck,
     DoubleToFloat,