[Feature] map_names for composite dists

tensordict/_torch_func.py

@@ -46,7 +46,7 @@
     from torch.utils._pytree import tree_flatten
 
     def tree_leaves(pytree):
-        """torch 2.0 compatible version of tree_leaves."""
+        """Torch 2.0 compatible version of tree_leaves."""
         return tree_flatten(pytree)[0]
 
 

tensordict/nn/distributions/composite.py

@@ -2,11 +2,11 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
 
 import torch
 from tensordict import TensorDict, TensorDictBase
-from tensordict._tensordict import unravel_keys
-from tensordict.utils import NestedKey
+from tensordict.utils import NestedKey, unravel_key, unravel_keys
 from torch import distributions as d
 
 
@@ -27,6 +27,9 @@ class CompositeDistribution(d.Distribution):
             will match the names of the samples in the tensordict.
 
     Keyword Arguments:
+        name_map (Dict[NestedKey, NestedKey]]): a dictionary representing where each
+            sample should be written. If not provided, the key names from ``distribution_map``
+            will be used.
         extra_kwargs (Dict[NestedKey, Dict]): a possibly incomplete dictionary of
             extra keyword arguments for the distributions to be built.
 
@@ -61,18 +64,41 @@ class CompositeDistribution(d.Distribution):
             is_shared=False)
     """
 
-    def __init__(self, params: TensorDictBase, distribution_map, *, extra_kwargs=None):
+    def __init__(
+        self,
+        params: TensorDictBase,
+        distribution_map: dict,
+        *,
+        name_map: dict | None = None,
+        extra_kwargs=None,
+    ):
         self._batch_shape = params.shape
         if extra_kwargs is None:
             extra_kwargs = {}
         dists = {}
+        if name_map is not None:
+            name_map = {
+                unravel_key(key): unravel_key(other_key)
+                for key, other_key in name_map.items()
+            }
         for name, dist_class in distribution_map.items():
+            name_unravel = unravel_key(name)
+            if name_map:
+                try:
+                    write_name = unravel_key(name_map.get(name, name_unravel))
+                except KeyError:
+                    raise KeyError(
+                        f"Failed to retrieve the key {name} from the name_map with keys {name_map.keys()}."
+                    )
+            else:
+                write_name = name_unravel
+            name = name_unravel
             dist_params = params.get(name, None)
             kwargs = extra_kwargs.get(name, {})
             if dist_params is None:
                 raise KeyError
             dist = dist_class(**dist_params, **kwargs)
-            dists[name] = dist
+            dists[write_name] = dist
         self.dists = dists
 
     def sample(self, shape=None) -> TensorDictBase:

test/test_nn.py

@@ -3123,6 +3123,43 @@ def test_sample(self):
         sample = dist.sample((4,))
         assert sample.shape == torch.Size((4,) + params.shape)
 
+    def test_sample_named(self):
+        params = TensorDict(
+            {
+                "cont": {"loc": torch.randn(3, 4), "scale": torch.rand(3, 4)},
+                ("nested", "disc"): {"logits": torch.randn(3, 10)},
+            },
+            [3],
+        )
+        dist = CompositeDistribution(
+            params,
+            distribution_map={
+                "cont": distributions.Normal,
+                ("nested", "disc"): distributions.Categorical,
+            },
+            name_map={
+                "cont": ("sample", "cont"),
+                (("nested",), "disc"): ("sample", "disc"),
+            },
+        )
+        sample = dist.sample()
+        assert sample.shape == params.shape
+        sample = dist.sample((4,))
+        assert sample.shape == torch.Size((4,) + params.shape)
+        assert sample["sample", "cont"].shape == torch.Size(
+            (
+                4,
+                3,
+                4,
+            )
+        )
+        assert sample["sample", "disc"].shape == torch.Size(
+            (
+                4,
+                3,
+            )
+        )
+
     def test_rsample(self):
         params = TensorDict(
             {
@@ -3231,7 +3268,8 @@ def test_icdf(self):
         "interaction", [InteractionType.MODE, InteractionType.MEAN]
     )
     @pytest.mark.parametrize("return_log_prob", [True, False])
-    def test_prob_module(self, interaction, return_log_prob):
+    @pytest.mark.parametrize("map_names", [True, False])
+    def test_prob_module(self, interaction, return_log_prob, map_names):
         params = TensorDict(
             {
                 "params": {
@@ -3253,31 +3291,46 @@ def test_prob_module(self, interaction, return_log_prob):
             "cont": distributions.Normal,
             ("nested", "cont"): distributions.Normal,
         }
+        distribution_kwargs = {"distribution_map": distribution_map}
+        if map_names:
+            distribution_kwargs.update(
+                {
+                    "name_map": {
+                        "cont": ("sample", "cont"),
+                        ("nested", "cont"): ("sample", "nested", "cont"),
+                    }
+                }
+            )
         module = ProbabilisticTensorDictModule(
             in_keys=in_keys,
             out_keys=out_keys,
             distribution_class=CompositeDistribution,
-            distribution_kwargs={"distribution_map": distribution_map},
+            distribution_kwargs=distribution_kwargs,
             default_interaction_type=interaction,
             return_log_prob=return_log_prob,
         )
         sample = module(params)
+        key_logprob0 = ("sample", "cont_log_prob") if map_names else "cont_log_prob"
+        key_logprob1 = (
+            ("sample", "nested", "cont_log_prob")
+            if map_names
+            else ("nested", "cont_log_prob")
+        )
         if return_log_prob:
-            assert "cont_log_prob" in sample.keys()
-            assert ("nested", "cont_log_prob") in sample.keys(True)
+            assert key_logprob0 in sample
+            assert key_logprob1 in sample
         sample_clone = sample.clone()
         lp = module.log_prob(sample_clone)
         if return_log_prob:
             torch.testing.assert_close(
                 lp,
-                sample.get("cont_log_prob").sum(-1)
-                + sample.get(("nested", "cont_log_prob")).sum(-1),
+                sample.get(key_logprob0).sum(-1) + sample.get(key_logprob1).sum(-1),
             )
         else:
             torch.testing.assert_close(
                 lp,
-                sample_clone.get("cont_log_prob").sum(-1)
-                + sample_clone.get(("nested", "cont_log_prob")).sum(-1),
+                sample_clone.get(key_logprob0).sum(-1)
+                + sample_clone.get(key_logprob1).sum(-1),
             )
 
     @pytest.mark.parametrize(