Repository: pytorch/tensordict

Issue:
[BUG] 'distribution_kwargs' are not correctly passed to dist in 'ProbabilisticTensorDictModule'
## Describe the bug

Hi All,

It seems to me that if I set custom arguments `distribution_kwargs` for a distribution with `ProbabilisticTensorDictModule`, these arguments only get assigned to the object (see `tensordict/tensordict/nn/probabilistic.py` line 215):
```python
self.distribution_kwargs = (
            distribution_kwargs if distribution_kwargs is not None else {}
        )
```
but later on, when the dist is fetched `def get_dist(self, tensordict: TensorDictBase) -> d.Distribution:`, they are not used (see `tensordict/tensordict/nn/probabilistic.py` line 225):

```python
dist_kwargs = {
    dist_key: tensordict[td_key]
    for dist_key, td_key in self.in_keys.items()
}
dist = self.distribution_class(**dist_kwargs)
```

In my case, this led to the issue that a custom min/max range for the dist didn't get set correctly as can be seen in the following small example.

## To Reproduce
```python
import torch
from tensordict.nn import ProbabilisticTensorDictModule
from tensordict import TensorDict
from torchrl.modules.distributions.continuous import TanhNormal

min_dist = torch.Tensor([-10, -10])
prob_module = ProbabilisticTensorDictModule(
    in_keys=["loc", "scale"],
    out_keys=["action"],
    distribution_class=TanhNormal,
    distribution_kwargs={
        "min": min_dist,
        "max": min_dist * -1,
    },
    return_log_prob=True,
)
print(f'prob_module.distribution_kwargs[\'min\']: {prob_module.distribution_kwargs["min"]}')

td = TensorDict(
    {
        'loc': torch.Tensor([0, 0]),
        'scale': torch.Tensor([0.1, 0.1]),

    }, [])

dist = prob_module.get_dist(td)
print(f'dist.min {dist.min}')

assert all(dist.min == min_dist)
```

## Expected behavior
`distribution_kwargs` should be used when calling `ProbabilisticTensorDictModule.get_dist()`.

## System info

tensordict.__version__ == 2023.01.09
numpy.__version__ == 1.23.2
sys.version == 3.8.15 (default, Oct 12 2022, 19:15:16) \n[GCC 11.2.0]
sys.platform == linux
torch.__version__ == 2.0.0.dev20230108+cu116

## Reason and Possible fixes
In `tensordict/tensordict/nn/probabilistic.py` line 229, just adding `**self.distribution_kwargs` should fix it:
```python
dist = self.distribution_class(**dist_kwargs, **self.distribution_kwargs)
```
## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)
- [x] I have read the [documentation](https://github.com/pytorch/rl/tree/main/docs/) (**required**)
- [x] I have provided a minimal working example to reproduce the bug (**required**)



Discussion:
Thanks for reporting! I'll fix that asap

Resolving patch:
diff --git a/tensordict/nn/probabilistic.py b/tensordict/nn/probabilistic.py
--- a/tensordict/nn/probabilistic.py
+++ b/tensordict/nn/probabilistic.py
@@ -226,7 +226,7 @@ def get_dist(self, tensordict: TensorDictBase) -> d.Distribution:
                 dist_key: tensordict[td_key]
                 for dist_key, td_key in self.in_keys.items()
             }
-            dist = self.distribution_class(**dist_kwargs)
+            dist = self.distribution_class(**dist_kwargs, **self.distribution_kwargs)
         except TypeError as err:
             if "an unexpected keyword argument" in str(err):
                 raise TypeError(


Test patch:
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
--- a/test/test_tensordictmodules.py
+++ b/test/test_tensordictmodules.py
@@ -90,6 +90,45 @@ def test_stateful_probabilistic_deprec(self, lazy, interaction_mode, out_keys):
         assert td.shape == torch.Size([3])
         assert td.get("out").shape == torch.Size([3, 4])
 
+    @pytest.mark.parametrize("out_keys", [["low"], ["low1"]])
+    @pytest.mark.parametrize("lazy", [True, False])
+    @pytest.mark.parametrize("max_dist", [1.0, 2.0])
+    @pytest.mark.parametrize("interaction_mode", ["mode", "random", None])
+    def test_stateful_probabilistic_kwargs(
+        self, lazy, interaction_mode, out_keys, max_dist
+    ):
+        torch.manual_seed(0)
+        if lazy:
+            net = nn.LazyLinear(4)
+        else:
+            net = nn.Linear(3, 4)
+
+        in_keys = ["in"]
+        net = TensorDictModule(module=net, in_keys=in_keys, out_keys=out_keys)
+
+        kwargs = {
+            "distribution_class": torch.distributions.Uniform,
+            "distribution_kwargs": {"high": max_dist},
+        }
+        if out_keys == ["low"]:
+            dist_in_keys = ["low"]
+        elif out_keys == ["low1"]:
+            dist_in_keys = {"low": "low1"}
+        else:
+            raise NotImplementedError
+
+        prob_module = ProbabilisticTensorDictModule(
+            in_keys=dist_in_keys, out_keys=["out"], **kwargs
+        )
+
+        tensordict_module = ProbabilisticTensorDictSequential(net, prob_module)
+
+        td = TensorDict({"in": torch.randn(3, 3)}, [3])
+        with set_interaction_mode(interaction_mode):
+            tensordict_module(td)
+        assert td.shape == torch.Size([3])
+        assert td.get("out").shape == torch.Size([3, 4])
+
     @pytest.mark.parametrize("out_keys", [["loc", "scale"], ["loc_1", "scale_1"]])
     @pytest.mark.parametrize("lazy", [True, False])
     @pytest.mark.parametrize("interaction_mode", ["mode", "random", None])