[BE] Make better logits in cost tests

ghstack-source-id: be9ea92b3f3d2592e426eaeaff7b81e50472cf16
Pull Request resolved: #2775

Author: vmoens

test/test_cost.py

@@ -4710,7 +4710,7 @@ def _create_mock_actor(
     ):
         # Actor
         action_spec = OneHot(action_dim)
-        net = nn.Sequential(nn.Linear(obs_dim, 2 * action_dim), NormalParamExtractor())
+        net = nn.Linear(obs_dim, action_dim)
         module = TensorDictModule(net, in_keys=[observation_key], out_keys=["logits"])
         actor = ProbabilisticActor(
             spec=action_spec,
@@ -11388,7 +11388,7 @@ def _create_mock_actor(self, batch=2, obs_dim=3, action_dim=4, device="cpu"):
         action_spec = Bounded(
             -torch.ones(action_dim), torch.ones(action_dim), (action_dim,)
         )
-        net = nn.Sequential(nn.Linear(obs_dim, 2 * action_dim), NormalParamExtractor())
+        net = nn.Linear(obs_dim, action_dim)
         module = TensorDictModule(net, in_keys=["observation"], out_keys=["param"])
         actor = ProbabilisticActor(
             module=module,
@@ -12632,7 +12632,7 @@ def _create_mock_actor(
     ):
         # Actor
         action_spec = OneHot(action_dim)
-        net = nn.Sequential(nn.Linear(obs_dim, 2 * action_dim), NormalParamExtractor())
+        net = nn.Linear(obs_dim, action_dim)
         module = TensorDictModule(net, in_keys=[observation_key], out_keys=["logits"])
         actor = ProbabilisticActor(
             spec=action_spec,
@@ -12729,8 +12729,7 @@ def _create_mock_common_layer_setup(
         common = Mod(common_net, in_keys=["obs"], out_keys=["hidden"])
         actor = ProbSeq(
             common,
-            Mod(actor_net, in_keys=["hidden"], out_keys=["param"]),
-            Mod(NormalParamExtractor(), in_keys=["param"], out_keys=["logits"]),
+            Mod(actor_net, in_keys=["hidden"], out_keys=["logits"]),
             ProbMod(
                 in_keys=["logits"],
                 out_keys=["action"],