Repository: pytorch/rl

Issue:
[Feature Request] Exclude "_" keys from the output tensordict in TensorDictModule
## Describe the bug

We sometimes want to ignore outputs of a NN in TensorDictModule.
To do this, we can rename the key to "_". However, this will still be written out.
We should avoid writing tensors names "_".

Feature request: 
- The specs should follow (no spec for "_" keys)
- We should test that this works (test models with multiple, i.e. > 2, outputs, where at least 2 are marked as undesired with a "_" key)


Discussion:

Resolving patch:
diff --git a/torchrl/modules/distributions/continuous.py b/torchrl/modules/distributions/continuous.py
--- a/torchrl/modules/distributions/continuous.py
+++ b/torchrl/modules/distributions/continuous.py
@@ -101,7 +101,7 @@ def _inverse(self, y: torch.Tensor) -> torch.Tensor:
 
 
 class NormalParamWrapper(nn.Module):
-    """A wrapper for normal distirbution parameters.
+    """A wrapper for normal distribution parameters.
 
     Args:
         operator (nn.Module): operator whose output will be transformed_in in location and scale parameters
diff --git a/torchrl/modules/tensordict_module/common.py b/torchrl/modules/tensordict_module/common.py
--- a/torchrl/modules/tensordict_module/common.py
+++ b/torchrl/modules/tensordict_module/common.py
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import warnings
 from copy import deepcopy
 from textwrap import indent
 from typing import (
@@ -94,7 +95,7 @@ class TensorDictModule(nn.Module):
         in_keys (iterable of str): keys to be read from input tensordict and passed to the module. If it
             contains more than one element, the values will be passed in the order given by the in_keys iterable.
         out_keys (iterable of str): keys to be written to the input tensordict. The length of out_keys must match the
-            number of tensors returned by the embedded module.
+            number of tensors returned by the embedded module. Using "_" as a key avoid writing tensor to output.
         spec (TensorSpec): specs of the output tensor. If the module outputs multiple output tensors,
             spec characterize the space of the first output tensor.
         safe (bool): if True, the value of the output is checked against the input spec. Out-of-domain sampling can
@@ -186,6 +187,11 @@ def __init__(
         self.in_keys = in_keys
         _check_all_str(self.in_keys)
 
+        if "_" in in_keys:
+            warnings.warn(
+                'key "_" is for ignoring output, it should not be used in input keys'
+            )
+
         if spec is not None and not isinstance(spec, TensorSpec):
             raise TypeError("spec must be a TensorSpec subclass")
         elif spec is not None and not isinstance(spec, CompositeSpec):
@@ -195,6 +201,9 @@ def __init__(
                     "Consider using a CompositeSpec object or no spec at all."
                 )
             spec = CompositeSpec(**{self.out_keys[0]: spec})
+        elif spec is not None and isinstance(spec, CompositeSpec):
+            if "_" in spec.keys():
+                warnings.warn('got a spec with key "_": it will be ignored')
         elif spec is None:
             spec = CompositeSpec()
 
@@ -259,13 +268,13 @@ def _write_to_tensordict(
             and vmap
             and (isinstance(vmap, bool) or vmap[-1] is None)
         ):
-            #
             dim = tensors[0].shape[0]
             tensordict_out = tensordict.expand(dim, *tensordict.batch_size).contiguous()
         elif tensordict_out is None:
             tensordict_out = tensordict
         for _out_key, _tensor in zip(out_keys, tensors):
-            tensordict_out.set(_out_key, _tensor)
+            if _out_key != "_":
+                tensordict_out.set(_out_key, _tensor)
         return tensordict_out
 
     def _make_vmap(self, buffers, kwargs, n_input):


Test patch:
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
--- a/test/test_tensordictmodules.py
+++ b/test/test_tensordictmodules.py
@@ -24,6 +24,69 @@
 
 
 class TestTDModule:
+    def test_multiple_output(self):
+        class MultiHeadLinear(nn.Module):
+            def __init__(self, in_1, out_1, out_2, out_3):
+                super().__init__()
+                self.linear_1 = nn.Linear(in_1, out_1)
+                self.linear_2 = nn.Linear(in_1, out_2)
+                self.linear_3 = nn.Linear(in_1, out_3)
+
+            def forward(self, x):
+                return self.linear_1(x), self.linear_2(x), self.linear_3(x)
+
+        tensordict_module = TensorDictModule(
+            MultiHeadLinear(5, 4, 3, 2),
+            in_keys=["input"],
+            out_keys=["out_1", "out_2", "out_3"],
+        )
+        td = TensorDict({"input": torch.randn(3, 5)}, batch_size=[3])
+        td = tensordict_module(td)
+        assert td.shape == torch.Size([3])
+        assert "input" in td.keys()
+        assert "out_1" in td.keys()
+        assert "out_2" in td.keys()
+        assert "out_3" in td.keys()
+        assert td.get("out_3").shape == torch.Size([3, 2])
+
+        # Using "_" key to ignore some output
+        tensordict_module = TensorDictModule(
+            MultiHeadLinear(5, 4, 3, 2),
+            in_keys=["input"],
+            out_keys=["_", "_", "out_3"],
+        )
+        td = TensorDict({"input": torch.randn(3, 5)}, batch_size=[3])
+        td = tensordict_module(td)
+        assert td.shape == torch.Size([3])
+        assert "input" in td.keys()
+        assert "out_3" in td.keys()
+        assert "_" not in td.keys()
+        assert td.get("out_3").shape == torch.Size([3, 2])
+
+    def test_spec_key_warning(self):
+        class MultiHeadLinear(nn.Module):
+            def __init__(self, in_1, out_1, out_2):
+                super().__init__()
+                self.linear_1 = nn.Linear(in_1, out_1)
+                self.linear_2 = nn.Linear(in_1, out_2)
+
+            def forward(self, x):
+                return self.linear_1(x), self.linear_2(x)
+
+        spec_dict = {
+            "_": NdUnboundedContinuousTensorSpec((4,)),
+            "out_2": NdUnboundedContinuousTensorSpec((3,)),
+        }
+
+        # warning due to "_" in spec keys
+        with pytest.warns(UserWarning, match='got a spec with key "_"'):
+            tensordict_module = TensorDictModule(
+                MultiHeadLinear(5, 4, 3),
+                in_keys=["input"],
+                out_keys=["_", "out_2"],
+                spec=CompositeSpec(**spec_dict),
+            )
+
     @pytest.mark.parametrize("safe", [True, False])
     @pytest.mark.parametrize("spec_type", [None, "bounded", "unbounded"])
     @pytest.mark.parametrize("lazy", [True, False])
@@ -773,6 +836,16 @@ def test_vmap_probabilistic_laterconstruct(self, safe, spec_type):
 
 
 class TestTDSequence:
+    def test_in_key_warning(self):
+        with pytest.warns(UserWarning, match='key "_" is for ignoring output'):
+            tensordict_module = TensorDictModule(
+                nn.Linear(3, 4), in_keys=["_"], out_keys=["out1"]
+            )
+        with pytest.warns(UserWarning, match='key "_" is for ignoring output'):
+            tensordict_module = TensorDictModule(
+                nn.Linear(3, 4), in_keys=["_", "key2"], out_keys=["out1"]
+            )
+
     def test_key_exclusion(self):
         module1 = TensorDictModule(
             nn.Linear(3, 4), in_keys=["key1", "key2"], out_keys=["foo1"]