[Test] Fix error catches (#2982)

Author: vmoens

test/test_actors.py

@@ -79,7 +79,7 @@ def test_probabilistic_actor_nested_delta(log_prob_key, nested_dim=5, n_actions=
     if log_prob_key:
         assert td_out[log_prob_key].shape == (5,)
     else:
-        assert td_out["sample_log_prob"].shape == (5,)
+        assert td_out["data", "action_log_prob"].shape == (5,)
 
     policy = ProbabilisticActor(
         module=policy_module,
@@ -99,7 +99,7 @@ def test_probabilistic_actor_nested_delta(log_prob_key, nested_dim=5, n_actions=
     if log_prob_key:
         assert td_out[log_prob_key].shape == (5,)
     else:
-        assert td_out["sample_log_prob"].shape == (5,)
+        assert td_out["data", "action_log_prob"].shape == (5,)
 
 
 @pytest.mark.parametrize(
@@ -144,7 +144,7 @@ def test_probabilistic_actor_nested_normal(log_prob_key, nested_dim=5, n_actions
     if log_prob_key:
         assert td_out[log_prob_key].shape == (5,)
     else:
-        assert td_out["sample_log_prob"].shape == (5,)
+        assert td_out["data", "action_log_prob"].shape == (5,)
 
     policy = ProbabilisticActor(
         module=policy_module,
@@ -164,7 +164,7 @@ def test_probabilistic_actor_nested_normal(log_prob_key, nested_dim=5, n_actions
     if log_prob_key:
         assert td_out[log_prob_key].shape == (5,)
     else:
-        assert td_out["sample_log_prob"].shape == (5,)
+        assert td_out["data", "action_log_prob"].shape == (5,)
 
 
 class TestQValue:
@@ -867,7 +867,7 @@ def test_lmhead_actorvalueoperator(device):
 
     # check actor
     assert aco.module[1].in_keys == ["x"]
-    assert aco.module[1].out_keys == ["logits", "action", "sample_log_prob"]
+    assert aco.module[1].out_keys == ["logits", "action", "action_log_prob"]
     assert aco.module[1][0].module is base_model.lm_head
 
     # check critic

test/test_exploration.py

@@ -124,7 +124,9 @@ def test_egreedy_masked(self, module, eps_init, spec_class):
             {"observation": torch.zeros(*batch_size, action_size)},
             batch_size=batch_size,
         )
-        with pytest.raises(RuntimeError, match="Failed while executing module"):
+        with pytest.raises(
+            KeyError, match="Action mask key action_mask not found in TensorDict"
+        ):
             explorative_policy(td)
 
         torch.manual_seed(0)
@@ -163,9 +165,7 @@ def test_egreedy_masked(self, module, eps_init, spec_class):
         assert not (action[~action_mask] == 0).all()
         assert (masked_action[~action_mask] == 0).all()
 
-    def test_no_spec_error(
-        self,
-    ):
+    def test_no_spec_error(self):
         torch.manual_seed(0)
         action_size = 4
         batch_size = (3, 4, 2)
@@ -183,7 +183,10 @@ def test_no_spec_error(
             batch_size=batch_size,
         )
 
-        with pytest.raises(RuntimeError, match="Failed while executing module"):
+        with pytest.raises(
+            RuntimeError,
+            match="Failed while executing module|spec must be provided to the exploration wrapper",
+        ):
             explorative_policy(td)
 
     @pytest.mark.parametrize("module", [True, False])
@@ -200,7 +203,9 @@ def test_wrong_action_shape(self, module):
                 policy,
             )
         td = TensorDict({"observation": torch.zeros(10, 4)}, batch_size=[10])
-        with pytest.raises(RuntimeError, match="Failed while executing module"):
+        with pytest.raises(
+            ValueError, match="Action spec shape does not match the action shape"
+        ):
             explorative_policy(td)
 
 

torchrl/modules/tensordict_module/actors.py

@@ -2168,7 +2168,7 @@ def __init__(self, base_model):
             TensorDictModule(
                 base_model.transformer,
                 in_keys={"input_ids": "input_ids", "attention_mask": "attention_mask"},
-                out_keys=["x"],
+                out_keys=["x", "_"],
             ),
             TensorDictModule(lambda x: x[:, -1, :], in_keys=["x"], out_keys=["x"]),
         )