Repository: pytorch/rl

Issue:
[Feature Request] Make DQN compatible with nn.Module (not wrapped in a QValue specialized network)
## Motivation

For DQNLoss to work out of the box, it'd be nice to support nn.Module (and not necessarily specialized TensorDictModule instances).

It should be fairly easy: we can already easily wrap a nn.Module in a TensorDictModule assuming that the nn.Module has only one output (see [here](https://github.com/pytorch/rl/blob/487838844b1cca2c4c200cf10b0e10a5d322de31/torchrl/collectors/collectors.py#L165-L204)). In this example, we check if the policy (in our case: the value network) is some sort of `TensorDictModule`. If not, we wrap it in that.

We could do the same with DQN: if the value network is an nn.Module that just returns a single value tensor, but does not expect a `TensorDict` as input, we should wrap it in a [`QValueActor`](https://github.com/pytorch/rl/blob/487838844b1cca2c4c200cf10b0e10a5d322de31/torchrl/modules/tensordict_module/actors.py#L409-L446) and the trick is done.




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
@@ -5,6 +5,7 @@
 
 from __future__ import annotations
 
+import inspect
 import warnings
 from copy import deepcopy
 from textwrap import indent
@@ -14,6 +15,7 @@
     List,
     Optional,
     Sequence,
+    Type,
     Union,
 )
 
@@ -603,3 +605,83 @@ def __getattr__(self, name: str) -> Any:
 
     def forward(self, *args, **kwargs):
         return self.td_module.forward(*args, **kwargs)
+
+
+def is_tensordict_compatible(module: Union[TensorDictModule, nn.Module]):
+    sig = inspect.signature(module.forward)
+
+    if isinstance(module, TensorDictModule) or (
+        len(sig.parameters) == 1
+        and hasattr(module, "in_keys")
+        and hasattr(module, "out_keys")
+    ):
+        # if the module is a TensorDictModule or takes a single argument and defines
+        # in_keys and out_keys then we assume it can already deal with TensorDict input
+        # to forward and we return True
+        return True
+    elif not hasattr(module, "in_keys") and not hasattr(module, "out_keys"):
+        # if it's not a TensorDictModule, and in_keys and out_keys are not defined then
+        # we assume no TensorDict compatibility and will try to wrap it.
+        return False
+
+    # if in_keys or out_keys were defined but module is not a TensorDictModule or
+    # accepts multiple arguments then it's likely the user is trying to do something
+    # that will have undetermined behaviour, we raise an error
+    raise TypeError(
+        "Received a module that defines in_keys or out_keys and also expects multiple "
+        "arguments to module.forward. If the module is compatible with TensorDict, it "
+        "should take a single argument of type TensorDict to module.forward and define "
+        "both in_keys and out_keys. Alternatively, module.forward can accept "
+        "arbitrarily many tensor inputs and leave in_keys and out_keys undefined and "
+        "TorchRL will attempt to automatically wrap the module with a TensorDictModule."
+    )
+
+
+def ensure_tensordict_compatible(
+    module: Union[
+        FunctionalModule, FunctionalModuleWithBuffers, TensorDictModule, nn.Module
+    ],
+    in_keys: Optional[Iterable[str]] = None,
+    out_keys: Optional[Iterable[str]] = None,
+    safe: bool = False,
+    wrapper_type: Optional[Type] = TensorDictModule,
+):
+    """Checks and ensures an object with forward method is TensorDict compatible."""
+    if is_tensordict_compatible(module):
+        if in_keys is not None and set(in_keys) != set(module.in_keys):
+            raise TypeError(
+                f"Arguments to module.forward ({set(module.in_keys)}) doesn't match "
+                f"with the expected TensorDict in_keys ({set(in_keys)})."
+            )
+        if out_keys is not None and set(module.out_keys) != set(out_keys):
+            raise TypeError(
+                f"Outputs of module.forward ({set(module.out_keys)}) doesn't match "
+                f"with the expected TensorDict out_keys ({set(out_keys)})."
+            )
+        # return module itself if it's already tensordict compatible
+        return module
+
+    if not isinstance(module, nn.Module):
+        raise TypeError(
+            "Argument to ensure_tensordict_compatible should be either "
+            "a TensorDictModule or an nn.Module"
+        )
+
+    sig = inspect.signature(module.forward)
+    if in_keys is not None and set(sig.parameters) != set(in_keys):
+        raise TypeError(
+            "Arguments to module.forward are incompatible with entries in "
+            "env.observation_spec. If you want TorchRL to automatically "
+            "wrap your module with a TensorDictModule then the arguments "
+            "to module must correspond one-to-one with entries in "
+            "in_keys. For more complex behaviour and more control you can "
+            "consider writing your own TensorDictModule."
+        )
+
+    # TODO: Check whether out_keys match (at least in number) if they are provided.
+    kwargs = {}
+    if in_keys is not None:
+        kwargs["in_keys"] = in_keys
+    if out_keys is not None:
+        kwargs["out_keys"] = out_keys
+    return wrapper_type(module, **kwargs)
diff --git a/torchrl/objectives/dqn.py b/torchrl/objectives/dqn.py
--- a/torchrl/objectives/dqn.py
+++ b/torchrl/objectives/dqn.py
@@ -3,7 +3,10 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
+from typing import Union
+
 import torch
+from torch import nn
 
 from torchrl.data import TensorDict
 from torchrl.envs.utils import step_mdp
@@ -11,6 +14,7 @@
     DistributionalQValueActor,
     QValueActor,
 )
+from torchrl.modules.tensordict_module.common import ensure_tensordict_compatible
 from ..data.tensordict.tensordict import TensorDictBase
 from .common import LossModule
 from .utils import distance_loss, next_state_value
@@ -20,7 +24,7 @@ class DQNLoss(LossModule):
     """The DQN Loss class.
 
     Args:
-        value_network (ProbabilisticTDModule): a Q value operator.
+        value_network (QValueActor or nn.Module): a Q value operator.
         gamma (scalar): a discount factor for return computation.
         loss_function (str): loss function for the value discrepancy. Can be one of "l1", "l2" or "smooth_l1".
         delay_value (bool, optional): whether to duplicate the value network into a new target value network to
@@ -30,7 +34,7 @@ class DQNLoss(LossModule):
 
     def __init__(
         self,
-        value_network: QValueActor,
+        value_network: Union[QValueActor, nn.Module],
         gamma: float,
         loss_function: str = "l2",
         priority_key: str = "td_error",
@@ -39,6 +43,11 @@ def __init__(
 
         super().__init__()
         self.delay_value = delay_value
+
+        value_network = ensure_tensordict_compatible(
+            module=value_network, wrapper_type=QValueActor
+        )
+
         self.convert_to_functional(
             value_network,
             "value_network",
@@ -46,10 +55,7 @@ def __init__(
         )
 
         self.value_network_in_keys = value_network.in_keys
-        if not isinstance(value_network, QValueActor):
-            raise TypeError(
-                f"DQNLoss requires value_network to be of QValueActor dtype, got {type(value_network)}"
-            )
+
         self.register_buffer("gamma", torch.tensor(gamma))
         self.loss_function = loss_function
         self.priority_key = priority_key
@@ -138,7 +144,7 @@ class DistributionalDQNLoss(LossModule):
     https://arxiv.org/pdf/1707.06887.pdf
 
     Args:
-        value_network (DistributionalQValueActor): the distributional Q
+        value_network (DistributionalQValueActor or nn.Module): the distributional Q
             value operator.
         gamma (scalar): a discount factor for return computation.
         delay_value (bool): whether to duplicate the value network into a new target value network to create double DQN
@@ -146,7 +152,7 @@ class DistributionalDQNLoss(LossModule):
 
     def __init__(
         self,
-        value_network: DistributionalQValueActor,
+        value_network: Union[DistributionalQValueActor, nn.Module],
         gamma: float,
         priority_key: str = "td_error",
         delay_value: bool = False,
@@ -155,12 +161,11 @@ def __init__(
         self.register_buffer("gamma", torch.tensor(gamma))
         self.priority_key = priority_key
         self.delay_value = delay_value
-        if not isinstance(value_network, DistributionalQValueActor):
-            raise TypeError(
-                "Expected value_network to be of type "
-                "DistributionalQValueActor "
-                f"but got {type(value_network)}"
-            )
+
+        value_network = ensure_tensordict_compatible(
+            module=value_network, wrapper_type=DistributionalQValueActor
+        )
+
         self.convert_to_functional(
             value_network,
             "value_network",


Test patch:
diff --git a/test/test_cost.py b/test/test_cost.py
--- a/test/test_cost.py
+++ b/test/test_cost.py
@@ -115,7 +115,13 @@ class TestDQN:
     seed = 0
 
     def _create_mock_actor(
-        self, action_spec_type, batch=2, obs_dim=3, action_dim=4, device="cpu"
+        self,
+        action_spec_type,
+        batch=2,
+        obs_dim=3,
+        action_dim=4,
+        device="cpu",
+        is_nn_module=False,
     ):
         # Actor
         if action_spec_type == "one_hot":
@@ -130,6 +136,8 @@ def _create_mock_actor(
             raise ValueError(f"Wrong {action_spec_type}")
 
         module = nn.Linear(obs_dim, action_dim)
+        if is_nn_module:
+            return module.to(device)
         actor = QValueActor(
             spec=CompositeSpec(
                 action=action_spec, action_value=None, chosen_action_value=None
@@ -147,6 +155,7 @@ def _create_mock_distributional_actor(
         atoms=5,
         vmin=1,
         vmax=5,
+        is_nn_module=False,
     ):
         # Actor
         if action_spec_type == "mult_one_hot":
@@ -159,6 +168,11 @@ def _create_mock_distributional_actor(
             raise ValueError(f"Wrong {action_spec_type}")
         support = torch.linspace(vmin, vmax, atoms, dtype=torch.float)
         module = MLP(obs_dim, (atoms, action_dim))
+        # TODO: Fails tests with
+        # TypeError: __init__() missing 1 required keyword-only argument: 'support'
+        # DistributionalQValueActor initializer expects additional inputs.
+        # if is_nn_module:
+        #     return module
         actor = DistributionalQValueActor(
             spec=CompositeSpec(action=action_spec, action_value=None),
             module=module,
@@ -260,10 +274,11 @@ def _create_seq_mock_data_dqn(
     @pytest.mark.parametrize(
         "action_spec_type", ("nd_bounded", "one_hot", "categorical")
     )
-    def test_dqn(self, delay_value, device, action_spec_type):
+    @pytest.mark.parametrize("is_nn_module", (False, True))
+    def test_dqn(self, delay_value, device, action_spec_type, is_nn_module):
         torch.manual_seed(self.seed)
         actor = self._create_mock_actor(
-            action_spec_type=action_spec_type, device=device
+            action_spec_type=action_spec_type, device=device, is_nn_module=is_nn_module
         )
         td = self._create_mock_data_dqn(
             action_spec_type=action_spec_type, device=device
@@ -359,12 +374,13 @@ def test_dqn_batcher(self, n, delay_value, device, action_spec_type, gamma=0.9):
     @pytest.mark.parametrize(
         "action_spec_type", ("mult_one_hot", "one_hot", "categorical")
     )
+    @pytest.mark.parametrize("is_nn_module", (False, True))
     def test_distributional_dqn(
-        self, atoms, delay_value, device, action_spec_type, gamma=0.9
+        self, atoms, delay_value, device, action_spec_type, is_nn_module, gamma=0.9
     ):
         torch.manual_seed(self.seed)
         actor = self._create_mock_distributional_actor(
-            action_spec_type=action_spec_type, atoms=atoms
+            action_spec_type=action_spec_type, atoms=atoms, is_nn_module=is_nn_module
         ).to(device)
 
         td = self._create_mock_data_dqn(
diff --git a/test/test_tensordictmodules.py b/test/test_tensordictmodules.py
--- a/test/test_tensordictmodules.py
+++ b/test/test_tensordictmodules.py
@@ -17,6 +17,10 @@
 )
 from torchrl.envs.utils import set_exploration_mode
 from torchrl.modules import NormalParamWrapper, TanhNormal, TensorDictModule
+from torchrl.modules.tensordict_module.common import (
+    is_tensordict_compatible,
+    ensure_tensordict_compatible,
+)
 from torchrl.modules.tensordict_module.probabilistic import (
     ProbabilisticTensorDictModule,
 )
@@ -1775,3 +1779,110 @@ def test_subsequence_weight_update(self):
     if __name__ == "__main__":
         args, unknown = argparse.ArgumentParser().parse_known_args()
         pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)
+
+
+def test_is_tensordict_compatible():
+    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)
+
+    td_module = TensorDictModule(
+        MultiHeadLinear(5, 4, 3, 2),
+        in_keys=["in_1", "in_2"],
+        out_keys=["out_1", "out_2"],
+    )
+    assert is_tensordict_compatible(td_module)
+
+    class MockCompatibleModule(nn.Module):
+        def __init__(self, in_keys, out_keys):
+            self.in_keys = in_keys
+            self.out_keys = out_keys
+
+        def forward(self, tensordict):
+            pass
+
+    compatible_nn_module = MockCompatibleModule(
+        in_keys=["in_1", "in_2"],
+        out_keys=["out_1", "out_2"],
+    )
+    assert is_tensordict_compatible(compatible_nn_module)
+
+    class MockIncompatibleModuleNoKeys(nn.Module):
+        def forward(self, input):
+            pass
+
+    incompatible_nn_module_no_keys = MockIncompatibleModuleNoKeys()
+    assert not is_tensordict_compatible(incompatible_nn_module_no_keys)
+
+    class MockIncompatibleModuleMultipleArgs(nn.Module):
+        def __init__(self, in_keys, out_keys):
+            self.in_keys = in_keys
+            self.out_keys = out_keys
+
+        def forward(self, input_1, input_2):
+            pass
+
+    incompatible_nn_module_multi_args = MockIncompatibleModuleMultipleArgs(
+        in_keys=["in_1", "in_2"],
+        out_keys=["out_1", "out_2"],
+    )
+    with pytest.raises(TypeError):
+        is_tensordict_compatible(incompatible_nn_module_multi_args)
+
+
+def test_ensure_tensordict_compatible():
+    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)
+
+    td_module = TensorDictModule(
+        MultiHeadLinear(5, 4, 3, 2),
+        in_keys=["in_1", "in_2"],
+        out_keys=["out_1", "out_2"],
+    )
+    ensured_module = ensure_tensordict_compatible(td_module)
+    assert ensured_module is td_module
+    with pytest.raises(TypeError):
+        ensure_tensordict_compatible(td_module, in_keys=["input"])
+    with pytest.raises(TypeError):
+        ensure_tensordict_compatible(td_module, out_keys=["output"])
+
+    class NonNNModule:
+        def __init__(self):
+            pass
+
+        def forward(self, x):
+            pass
+
+    non_nn_module = NonNNModule()
+    with pytest.raises(TypeError):
+        ensure_tensordict_compatible(non_nn_module)
+
+    class ErrorNNModule(nn.Module):
+        def forward(self, in_1, in_2):
+            pass
+
+    error_nn_module = ErrorNNModule()
+    with pytest.raises(TypeError):
+        ensure_tensordict_compatible(error_nn_module, in_keys=["input"])
+
+    nn_module = MultiHeadLinear(5, 4, 3, 2)
+    ensured_module = ensure_tensordict_compatible(
+        nn_module,
+        in_keys=["x"],
+        out_keys=["out_1", "out_2", "out_3"],
+    )
+    assert set(ensured_module.in_keys) == {"x"}
+    assert isinstance(ensured_module, TensorDictModule)