Repository: pytorch/rl

Issue:
[Feature Request] Make `params` and `buffers` explicit kwargs in `TensorDictModule.forward` and sub-classes
## Motivation

For functional evaluation, we pass the module params (and possibly buffers) together with the inputs of the module (stored in a TensorDict) using 
```python
def forward(self, tensordict, **kwargs):
    params = kwargs.get("params", None)
    buffers = kwargs.get("buffers", None)
    ...
```

However, this breaks when using [higher](https://github.com/facebookresearch/higher). 

We should have default values for these instead of using `get("params", None)` etc.

The changes should be compatible with: `TensorDictModule`, `TensorDictSequence` and `ProbabilisticTensorDict`.


Discussion:

Resolving patch:
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
@@ -267,13 +267,12 @@ def _write_to_tensordict(
             tensordict_out.set(_out_key, _tensor)
         return tensordict_out
 
-    def _make_vmap(self, kwargs, n_input):
+    def _make_vmap(self, buffers, kwargs, n_input):
         if "vmap" in kwargs and kwargs["vmap"]:
             if not isinstance(kwargs["vmap"], (tuple, bool)):
                 raise RuntimeError(
                     "vmap argument must be a boolean or a tuple of dim expensions."
                 )
-            _buffers = "buffers" in kwargs
             # if vmap is a tuple, we make sure the number of inputs after params and buffers match
             if isinstance(kwargs["vmap"], (tuple, list)):
                 err_msg = f"the vmap argument had {len(kwargs['vmap'])} elements, but the module has {len(self.in_keys)} inputs"
@@ -306,12 +305,18 @@ def _make_vmap(self, kwargs, n_input):
                     )
             else:
                 _vmap = (
-                    (0, 0, *(None,) * n_input) if _buffers else (0, *(None,) * n_input)
+                    (0, 0, *(None,) * n_input)
+                    if buffers is not None
+                    else (0, *(None,) * n_input)
                 )
             return _vmap
 
     def _call_module(
-        self, tensors: Sequence[Tensor], **kwargs
+        self,
+        tensors: Sequence[Tensor],
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        **kwargs,
     ) -> Union[Tensor, Sequence[Tensor]]:
         err_msg = "Did not find the {0} keyword argument to be used with the functional module. Check it was passed to the TensorDictModule method."
         if isinstance(
@@ -323,37 +328,33 @@ def _call_module(
                 rlFunctionalModuleWithBuffers,
             ),
         ):
-            _vmap = self._make_vmap(kwargs, len(tensors))
+            _vmap = self._make_vmap(buffers, kwargs, len(tensors))
             if _vmap:
                 module = vmap(self.module, _vmap)
             else:
                 module = self.module
 
         if isinstance(self.module, (FunctionalModule, rlFunctionalModule)):
-            if "params" not in kwargs:
+            if params is None:
                 raise KeyError(err_msg.format("params"))
             kwargs_pruned = {
-                key: item
-                for key, item in kwargs.items()
-                if key not in ("params", "vmap")
+                key: item for key, item in kwargs.items() if key not in ("vmap")
             }
-            out = module(kwargs["params"], *tensors, **kwargs_pruned)
+            out = module(params, *tensors, **kwargs_pruned)
             return out
 
         elif isinstance(
             self.module, (FunctionalModuleWithBuffers, rlFunctionalModuleWithBuffers)
         ):
-            if "params" not in kwargs:
+            if params is None:
                 raise KeyError(err_msg.format("params"))
-            if "buffers" not in kwargs:
+            if buffers is None:
                 raise KeyError(err_msg.format("buffers"))
 
             kwargs_pruned = {
-                key: item
-                for key, item in kwargs.items()
-                if key not in ("params", "buffers", "vmap")
+                key: item for key, item in kwargs.items() if key not in ("vmap")
             }
-            out = module(kwargs["params"], kwargs["buffers"], *tensors, **kwargs_pruned)
+            out = module(params, buffers, *tensors, **kwargs_pruned)
             return out
         else:
             out = self.module(*tensors, **kwargs)
@@ -363,10 +364,12 @@ def forward(
         self,
         tensordict: TensorDictBase,
         tensordict_out: Optional[TensorDictBase] = None,
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
         **kwargs,
     ) -> TensorDictBase:
         tensors = tuple(tensordict.get(in_key, None) for in_key in self.in_keys)
-        tensors = self._call_module(tensors, **kwargs)
+        tensors = self._call_module(tensors, params=params, buffers=buffers, **kwargs)
         if not isinstance(tensors, tuple):
             tensors = (tensors,)
         tensordict_out = self._write_to_tensordict(
diff --git a/torchrl/modules/tensordict_module/probabilistic.py b/torchrl/modules/tensordict_module/probabilistic.py
--- a/torchrl/modules/tensordict_module/probabilistic.py
+++ b/torchrl/modules/tensordict_module/probabilistic.py
@@ -6,7 +6,7 @@
 import re
 from copy import deepcopy
 from textwrap import indent
-from typing import Sequence, Union, Type, Optional, Tuple
+from typing import List, Sequence, Union, Type, Optional, Tuple
 
 from torch import Tensor
 from torch import distributions as d
@@ -193,8 +193,14 @@ def __init__(
         self.cache_dist = cache_dist if hasattr(distribution_class, "update") else False
         self.return_log_prob = return_log_prob
 
-    def _call_module(self, tensordict: TensorDictBase, **kwargs) -> TensorDictBase:
-        return self.module(tensordict, **kwargs)
+    def _call_module(
+        self,
+        tensordict: TensorDictBase,
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        **kwargs,
+    ) -> TensorDictBase:
+        return self.module(tensordict, params=params, buffers=buffers, **kwargs)
 
     def make_functional_with_buffers(self, clone: bool = True, native: bool = False):
         module_params = self.parameters(recurse=False)
@@ -218,6 +224,8 @@ def get_dist(
         self,
         tensordict: TensorDictBase,
         tensordict_out: Optional[TensorDictBase] = None,
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
         **kwargs,
     ) -> Tuple[d.Distribution, TensorDictBase]:
         interaction_mode = exploration_mode()
@@ -225,7 +233,11 @@ def get_dist(
             interaction_mode = self.default_interaction_mode
         with set_exploration_mode(interaction_mode):
             tensordict_out = self._call_module(
-                tensordict, tensordict_out=tensordict_out, **kwargs
+                tensordict,
+                tensordict_out=tensordict_out,
+                params=params,
+                buffers=buffers,
+                **kwargs,
             )
         dist = self.build_dist_from_params(tensordict_out)
         return dist, tensordict_out
@@ -256,11 +268,17 @@ def forward(
         self,
         tensordict: TensorDictBase,
         tensordict_out: Optional[TensorDictBase] = None,
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
         **kwargs,
     ) -> TensorDictBase:
 
         dist, tensordict_out = self.get_dist(
-            tensordict, tensordict_out=tensordict_out, **kwargs
+            tensordict,
+            tensordict_out=tensordict_out,
+            params=params,
+            buffers=buffers,
+            **kwargs,
         )
         if self._requires_sample:
             out_tensors = self._dist_sample(dist, interaction_mode=exploration_mode())
diff --git a/torchrl/modules/tensordict_module/sequence.py b/torchrl/modules/tensordict_module/sequence.py
--- a/torchrl/modules/tensordict_module/sequence.py
+++ b/torchrl/modules/tensordict_module/sequence.py
@@ -6,7 +6,7 @@
 from __future__ import annotations
 
 from copy import copy, deepcopy
-from typing import List, Iterable, Union, Tuple
+from typing import List, Iterable, Optional, Union, Tuple
 
 _has_functorch = False
 try:
@@ -262,17 +262,24 @@ def select_subsequence(
 
         return TensorDictSequence(*modules)
 
-    def _run_module(self, module, tensordict, **kwargs):
+    def _run_module(
+        self,
+        module,
+        tensordict,
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        **kwargs,
+    ):
         tensordict_keys = set(tensordict.keys())
         if not self.partial_tolerant or all(
             key in tensordict_keys for key in module.in_keys
         ):
-            tensordict = module(tensordict, **kwargs)
+            tensordict = module(tensordict, params=params, buffers=buffers, **kwargs)
         elif self.partial_tolerant and isinstance(tensordict, LazyStackedTensorDict):
             for sub_td in tensordict.tensordicts:
                 tensordict_keys = set(sub_td.keys())
                 if all(key in tensordict_keys for key in module.in_keys):
-                    module(sub_td, **kwargs)
+                    module(sub_td, params=params, buffers=buffers, **kwargs)
             tensordict._update_valid_keys()
         return tensordict
 
@@ -280,62 +287,52 @@ def forward(
         self,
         tensordict: TensorDictBase,
         tensordict_out=None,
+        params: Optional[Union[TensorDictBase, List[Tensor]]] = None,
+        buffers: Optional[Union[TensorDictBase, List[Tensor]]] = None,
         **kwargs,
     ) -> TensorDictBase:
-
-        if "params" in kwargs and "buffers" in kwargs:
-            params = kwargs["params"]
-            buffers = kwargs["buffers"]
+        if params is not None and buffers is not None:
             if isinstance(params, TensorDictBase):
                 # TODO: implement sorted values and items
                 param_splits = list(zip(*sorted(list(params.items()))))[1]
                 buffer_splits = list(zip(*sorted(list(buffers.items()))))[1]
             else:
-                param_splits = self._split_param(kwargs["params"], "params")
-                buffer_splits = self._split_param(kwargs["buffers"], "buffers")
-            kwargs_pruned = {
-                key: item
-                for key, item in kwargs.items()
-                if key not in ("params", "buffers")
-            }
+                param_splits = self._split_param(params, "params")
+                buffer_splits = self._split_param(buffers, "buffers")
             for i, (module, param, buffer) in enumerate(
                 zip(self.module, param_splits, buffer_splits)
             ):
-                if "vmap" in kwargs_pruned and i > 0:
+                if "vmap" in kwargs and i > 0:
                     # the tensordict is already expended
-                    if not isinstance(kwargs_pruned["vmap"], tuple):
-                        kwargs_pruned["vmap"] = (0, 0, *(0,) * len(module.in_keys))
+                    if not isinstance(kwargs["vmap"], tuple):
+                        kwargs["vmap"] = (0, 0, *(0,) * len(module.in_keys))
                     else:
-                        kwargs_pruned["vmap"] = (
-                            *kwargs_pruned["vmap"][:2],
+                        kwargs["vmap"] = (
+                            *kwargs["vmap"][:2],
                             *(0,) * len(module.in_keys),
                         )
                 tensordict = self._run_module(
-                    module, tensordict, params=param, buffers=buffer, **kwargs_pruned
+                    module, tensordict, params=param, buffers=buffer, **kwargs
                 )
 
-        elif "params" in kwargs:
-            params = kwargs["params"]
+        elif params is not None:
             if isinstance(params, TensorDictBase):
                 # TODO: implement sorted values and items
                 param_splits = list(zip(*sorted(list(params.items()))))[1]
             else:
-                param_splits = self._split_param(kwargs["params"], "params")
-            kwargs_pruned = {
-                key: item for key, item in kwargs.items() if key not in ("params",)
-            }
+                param_splits = self._split_param(params, "params")
             for i, (module, param) in enumerate(zip(self.module, param_splits)):
-                if "vmap" in kwargs_pruned and i > 0:
+                if "vmap" in kwargs and i > 0:
                     # the tensordict is already expended
-                    if not isinstance(kwargs_pruned["vmap"], tuple):
-                        kwargs_pruned["vmap"] = (0, *(0,) * len(module.in_keys))
+                    if not isinstance(kwargs["vmap"], tuple):
+                        kwargs["vmap"] = (0, *(0,) * len(module.in_keys))
                     else:
-                        kwargs_pruned["vmap"] = (
-                            *kwargs_pruned["vmap"][:1],
+                        kwargs["vmap"] = (
+                            *kwargs["vmap"][:1],
                             *(0,) * len(module.in_keys),
                         )
                 tensordict = self._run_module(
-                    module, tensordict, params=param, **kwargs_pruned
+                    module, tensordict, params=param, **kwargs
                 )
 
         elif not len(kwargs):


Test patch: