[BugFix] select_out_keys for Prob sequential

ghstack-source-id: a566ae225c54f07a680b4bf380b16d8e797f62ea
Pull Request resolved: #1103

tensordict/nn/probabilistic.py

@@ -283,7 +283,6 @@ def __init__(
         in_keys: NestedKey | List[NestedKey] | Dict[str, NestedKey],
         out_keys: NestedKey | List[NestedKey] | None = None,
         *,
-        default_interaction_mode: str | None = None,
         default_interaction_type: InteractionType = InteractionType.DETERMINISTIC,
         distribution_class: type = Delta,
         distribution_kwargs: dict | None = None,
@@ -332,11 +331,7 @@ def __init__(
             log_prob_key = "sample_log_prob"
         self.log_prob_key = log_prob_key
 
-        if default_interaction_mode is not None:
-            raise ValueError(
-                "default_interaction_mode is deprecated, use default_interaction_type instead."
-            )
-        self.default_interaction_type = default_interaction_type
+        self.default_interaction_type = InteractionType(default_interaction_type)
 
         if isinstance(distribution_class, str):
             distribution_class = distributions_maps.get(distribution_class.lower())
@@ -356,7 +351,7 @@ def get_dist(self, tensordict: TensorDictBase) -> D.Distribution:
             for dist_key, td_key in _zip_strict(self.dist_keys, self.in_keys):
                 if isinstance(dist_key, tuple):
                     dist_key = dist_key[-1]
-                dist_kwargs[dist_key] = tensordict.get(td_key)
+                dist_kwargs[dist_key] = tensordict.get(td_key, None)
             dist = self.distribution_class(
                 **dist_kwargs, **_dynamo_friendly_to_dict(self.distribution_kwargs)
             )
@@ -630,8 +625,24 @@ def forward(
         tensordict_out: TensorDictBase | None = None,
         **kwargs,
     ) -> TensorDictBase:
-        tensordict_out = self.get_dist_params(tensordict, tensordict_out, **kwargs)
-        return self.module[-1](tensordict_out, _requires_sample=self._requires_sample)
+        if (tensordict_out is None and self._select_before_return) or (
+            tensordict_out is not None
+        ):
+            tensordict_exec = tensordict.copy()
+        else:
+            tensordict_exec = tensordict
+        tensordict_exec = self.get_dist_params(tensordict_exec, **kwargs)
+        tensordict_exec = self.module[-1](
+            tensordict_exec, _requires_sample=self._requires_sample
+        )
+        if tensordict_out is not None:
+            result = tensordict_out
+            result.update(tensordict_exec, keys_to_update=self.out_keys)
+        else:
+            result = tensordict_exec
+            if self._select_before_return:
+                return tensordict.update(result, keys_to_update=self.out_keys)
+        return result
 
 
 def _dynamo_friendly_to_dict(data):

tensordict/nn/sequence.py

@@ -470,13 +470,13 @@ def forward(
         tensordict_out: TensorDictBase | None = None,
         **kwargs: Any,
     ) -> TensorDictBase:
-        if tensordict_out is None and self._select_before_return:
+        if (tensordict_out is None and self._select_before_return) or (
+            tensordict_out is not None
+        ):
             tensordict_exec = tensordict.copy()
         else:
             tensordict_exec = tensordict
         if not len(kwargs):
-            if tensordict_out is not None:
-                tensordict_exec = tensordict_exec.copy()
             for module in self.module:
                 tensordict_exec = self._run_module(module, tensordict_exec, **kwargs)
         else:

test/test_nn.py

@@ -890,6 +890,59 @@ def test_stateful_probabilistic_deprec(self, lazy):
         dist = tdmodule.get_dist(td)
         assert dist.rsample().shape[: td.ndimension()] == td.shape
 
+    @pytest.mark.parametrize("return_log_prob", [True, False])
+    @pytest.mark.parametrize("td_out", [True, False])
+    def test_probtdseq(self, return_log_prob, td_out):
+        mod = ProbabilisticTensorDictSequential(
+            TensorDictModule(lambda x: x + 2, in_keys=["a"], out_keys=["c"]),
+            TensorDictModule(lambda x: (x + 2, x), in_keys=["b"], out_keys=["d", "e"]),
+            ProbabilisticTensorDictModule(
+                in_keys={"loc": "d", "scale": "e"},
+                out_keys=["f"],
+                distribution_class=Normal,
+                return_log_prob=return_log_prob,
+                default_interaction_type="random",
+            ),
+        )
+        inp = TensorDict({"a": 0.0, "b": 1.0})
+        inp_clone = inp.clone()
+        if td_out:
+            out = TensorDict()
+        else:
+            out = None
+        out2 = mod(inp, tensordict_out=out)
+        assert not mod._select_before_return
+        if td_out:
+            assert out is out2
+        else:
+            assert out2 is inp
+        assert set(out2.keys()) - {"a", "b"} == set(mod.out_keys), (
+            td_out,
+            return_log_prob,
+        )
+
+        inp = inp_clone.clone()
+        mod.select_out_keys("f")
+        if td_out:
+            out = TensorDict()
+        else:
+            out = None
+        out2 = mod(inp, tensordict_out=out)
+        assert mod._select_before_return
+        if td_out:
+            assert out is out2
+        else:
+            assert out2 is inp
+        expected = {"f"}
+        if td_out:
+            assert set(out2.keys()) == set(mod.out_keys) == expected
+        else:
+            assert (
+                set(out2.keys()) - set(inp_clone.keys())
+                == set(mod.out_keys)
+                == expected
+            )
+
     @pytest.mark.parametrize("lazy", [True, False])
     def test_stateful_probabilistic(self, lazy):
         torch.manual_seed(0)