[Deprecation] Enact deprecations

ghstack-source-id: 690a9f6
Pull Request resolved: #2917

Author: Vincent Moens

test/mocking_classes.py

@@ -1045,6 +1045,7 @@ def forward(self, observation, action):
 class CountingEnvCountPolicy(TensorDictModuleBase):
     def __init__(self, action_spec: TensorSpec, action_key: NestedKey = "action"):
         super().__init__()
+        assert not isinstance(action_spec, Composite)
         self.action_spec = action_spec
         self.action_key = action_key
         self.in_keys = []
@@ -1411,11 +1412,13 @@ def __init__(
                 },
                 shape=self.batch_size,
             )
-            self.action_spec = Composite(
+            action_spec = self.full_action_spec[self.action_key]
+            assert not isinstance(action_spec, Composite)
+            self.full_action_spec = Composite(
                 {
                     "data": Composite(
                         {
-                            "action": self.action_spec.unsqueeze(-1).expand(
+                            "action": action_spec.unsqueeze(-1).expand(
                                 *self.batch_size, self.nested_dim, 1
                             )
                         },

test/test_actors.py

@@ -218,7 +218,10 @@ def test_nested_keys(self, nested_action, batch_size, nested_dim=5):
             nest_obs_action=nested_action, batch_size=batch_size, nested_dim=nested_dim
         )
         action_spec = env._input_spec["full_action_spec"]
-        leaf_action_spec = env.action_spec
+        if nested_action:
+            leaf_action_spec = env.full_action_spec[env.action_keys[0]]
+        else:
+            leaf_action_spec = env.action_spec
 
         space_str, spec = _process_action_space_spec(None, action_spec)
         assert spec == action_spec

test/test_collector.py

@@ -2311,7 +2311,9 @@ def test_multi_collector_nested_env_consistency(self, seed=1):
         torch.manual_seed(seed)
         env_fn = lambda: TransformedEnv(NestedCountingEnv(), InitTracker())
         env = NestedCountingEnv()
-        policy = CountingEnvCountPolicy(env.action_spec, env.action_key)
+        policy = CountingEnvCountPolicy(
+            env.full_action_spec[env.action_key], env.action_key
+        )
 
         ccollector = MultiaSyncDataCollector(
             create_env_fn=[env_fn],
@@ -2377,7 +2379,9 @@ def test_collector_nested_env_combinations(
             nest_obs_action=nested_obs_action,
         )
         torch.manual_seed(seed)
-        policy = CountingEnvCountPolicy(env.action_spec, env.action_key)
+        policy = CountingEnvCountPolicy(
+            env.full_action_spec[env.action_key], env.action_key
+        )
         ccollector = SyncDataCollector(
             create_env_fn=env,
             policy=policy,
@@ -2404,7 +2408,9 @@ def test_nested_env_dims(self, batch_size, nested_dim=5, frames_per_batch=20):
         env = NestedCountingEnv(batch_size=batch_size, nested_dim=nested_dim)
         env_fn = lambda: NestedCountingEnv(batch_size=batch_size, nested_dim=nested_dim)
         torch.manual_seed(0)
-        policy = CountingEnvCountPolicy(env.action_spec, env.action_key)
+        policy = CountingEnvCountPolicy(
+            env.full_action_spec[env.action_key], env.action_key
+        )
         policy(env.reset())
         ccollector = SyncDataCollector(
             create_env_fn=env_fn,

test/test_env.py

@@ -87,6 +87,7 @@
     pytest.mark.filterwarnings(
         "ignore:Got multiple backends for torchrl.data.replay_buffers.storages"
     ),
+    pytest.mark.filterwarnings("ignore:unclosed file"),
 ]
 
 gym_version = None
@@ -1241,6 +1242,7 @@ def env_make():
             td_serial = env_serial.rollout(max_steps=50)
         finally:
             env_serial.close(raise_if_closed=False)
+            gc.collect()
 
         try:
             env_parallel = maybe_fork_ParallelEnv(
@@ -1256,6 +1258,7 @@ def env_make():
             assert_allclose_td(td_serial, td_parallel)
         finally:
             env_parallel.close(raise_if_closed=False)
+            gc.collect()
 
     @pytest.mark.skipif(not _has_dmc, reason="no dm_control")
     def test_multitask(self, maybe_fork_ParallelEnv):
@@ -2809,18 +2812,23 @@ def test_nested_env(self, envclass):
         else:
             raise NotImplementedError
         reset = env.reset()
-        with pytest.warns(
-            DeprecationWarning, match="non-trivial"
-        ) if envclass == "NestedCountingEnv" else contextlib.nullcontext():
+        if envclass == "NestedCountingEnv":
+            assert isinstance(env.reward_spec, Composite)
+        else:
             assert not isinstance(env.reward_spec, Composite)
         for done_key in env.done_keys:
             assert (
                 env.full_done_spec[done_key]
                 == env.output_spec[("full_done_spec", *_unravel_key_to_tuple(done_key))]
             )
-        with pytest.warns(
-            DeprecationWarning, match="non-trivial"
-        ) if envclass == "NestedCountingEnv" else contextlib.nullcontext():
+        if envclass == "NestedCountingEnv":
+            assert (
+                env.full_reward_spec[env.reward_key]
+                == env.output_spec[
+                    ("full_reward_spec", *_unravel_key_to_tuple(env.reward_key))
+                ]
+            )
+        else:
             assert (
                 env.reward_spec
                 == env.output_spec[

test/test_exploration.py

@@ -391,7 +391,9 @@ def test_nested(
             exploratory_policy = TensorDictSequential(
                 policy,
                 OrnsteinUhlenbeckProcessModule(
-                    spec=action_spec, action_key=env.action_key, is_init_key=is_init_key
+                    spec=action_spec.clone(),
+                    action_key=env.action_key,
+                    is_init_key=is_init_key,
                 ).to(device),
             )
         else:

test/test_libs.py

@@ -2833,7 +2833,7 @@ def test_collector(
 
         env = maybe_fork_ParallelEnv(n_workers, env_fun)
 
-        n_actions_per_agent = env.action_spec.shape[-1]
+        n_actions_per_agent = env.full_action_spec[env.action_key].shape[-1]
         n_observations_per_agent = env.observation_spec["agents", "observation"].shape[
             -1
         ]
@@ -2845,7 +2845,7 @@ def test_collector(
             ),
             in_keys=[("agents", "observation")],
             out_keys=[env.action_key],
-            spec=env.action_spec,
+            spec=env.full_action_spec[env.action_key],
             safe=True,
         )
         ccollector = SyncDataCollector(
@@ -4207,7 +4207,7 @@ def test_parallel_env(self, maybe_fork_ParallelEnv):
     def test_collector(self):
         env = SMACv2Env(map_name="MMM2", seed=0, categorical_actions=True)
         in_feats = env.observation_spec["agents", "observation"].shape[-1]
-        out_feats = env.action_spec.space.n
+        out_feats = env.full_action_spec[env.action_key].space.n
 
         module = TensorDictModule(
             nn.Linear(in_feats, out_feats),

test/test_loggers.py

@@ -31,7 +31,7 @@
         version.parse(torchvision.__version__).base_version
     )
 else:
-    TORCHVISION_VERSION = version.parse("0.0.1").base_version
+    TORCHVISION_VERSION = version.parse("0.0.1")
 
 if _has_tb:
     from tensorboard.backend.event_processing.event_accumulator import EventAccumulator

test/test_specs.py

@@ -820,12 +820,6 @@ def test_create_composite_nested(shape, device):
 class TestLock:
     @pytest.mark.parametrize("recurse", [None, True, False])
     def test_lock(self, recurse):
-        catch_warn = (
-            pytest.warns(DeprecationWarning, match="recurse")
-            if recurse is None
-            else contextlib.nullcontext()
-        )
-
         shape = [3, 4, 5]
         spec = Composite(
             a=Composite(b=Composite(shape=shape[:3], device="cpu"), shape=shape[:2]),
@@ -834,14 +828,13 @@ def test_lock(self, recurse):
         spec["a"] = spec["a"].clone()
         spec["a", "b"] = spec["a", "b"].clone()
         assert not spec.locked
-        with catch_warn:
-            spec.lock_(recurse=recurse)
+        spec.lock_(recurse=recurse)
         assert spec.locked
         with pytest.raises(RuntimeError, match="Cannot modify a locked Composite."):
             spec["a"] = spec["a"].clone()
         with pytest.raises(RuntimeError, match="Cannot modify a locked Composite."):
             spec.set("a", spec["a"].clone())
-        if recurse:
+        if recurse in (None, True):
             assert spec["a"].locked
             with pytest.raises(RuntimeError, match="Cannot modify a locked Composite."):
                 spec["a"].set("b", spec["a", "b"].clone())
@@ -851,8 +844,7 @@ def test_lock(self, recurse):
             assert not spec["a"].locked
             spec["a", "b"] = spec["a", "b"].clone()
             spec["a"].set("b", spec["a", "b"].clone())
-        with catch_warn:
-            spec.unlock_(recurse=recurse)
+        spec.unlock_(recurse=recurse)
         spec["a"] = spec["a"].clone()
         spec["a", "b"] = spec["a", "b"].clone()
         spec["a"].set("b", spec["a", "b"].clone())