LossModule is not longer an abstract base class.

pytorch · vmoens · May 31, 2023 · May 22, 2023 · May 22, 2023 · May 23, 2023
commit 81c041336d2f1916dff0a3e128dc9a60c6c9dfaf
diff --git a/test/test_cost.py b/test/test_cost.py
@@ -8,7 +8,7 @@
 import operator
 import warnings
 from copy import deepcopy
-from dataclasses import asdict
+from dataclasses import asdict, dataclass
 
 from packaging import version as pack_version
 from tensordict.nn import InteractionType
@@ -6080,9 +6080,6 @@ def __init__(self, compare_against, expand_dim):
                     expand_dim=expand_dim,
                 )
 
-            def _forward_value_estimator_keys(self, **kwargs) -> None:
-                pass
-
         loss_module = MyLoss(compare_against=compare_against, expand_dim=expand_dim)
 
         for key in ["module.0.bias", "module.0.weight"]:
@@ -6104,6 +6101,36 @@ def _forward_value_estimator_keys(self, **kwargs) -> None:
         for key in ["module.1.bias", "module.1.weight"]:
             loss_module.module_b_params.flatten_keys()[key].requires_grad
 
+    def test_tensordict_keys(self):
+        """Test configurable tensordict key behavior with derived classes."""
+
+        class MyLoss(LossModule):
+            def __init__(self):
+                super().__init__()
+
+        loss_module = MyLoss()
+        with pytest.raises(AttributeError):
+            loss_module.set_keys()
+
+        class MyLoss2(MyLoss):
+            def _forward_value_estimator_keys(self, **kwargs) -> None:
+                pass
+
+        loss_module = MyLoss2()
+        assert loss_module.set_keys() is None
+        with pytest.raises(ValueError):
+            loss_module.set_keys(some_key="test")
+
+        class MyLoss3(MyLoss2):
+            @dataclass
+            class _AcceptedKeys:
+                some_key = "some_value"
+
+        loss_module = MyLoss3()
+        assert loss_module.tensor_keys.some_key == "some_value"
+        loss_module.set_keys(some_key="test")
+        assert loss_module.tensor_keys.some_key == "test"
+
 
 class TestUtils:
     @pytest.mark.parametrize("B", [None, (1, ), (4, ), (2, 2, ), (1, 2, 8, )])  # fmt: skip

diff --git a/torchrl/objectives/common.py b/torchrl/objectives/common.py
@@ -6,7 +6,6 @@
 from __future__ import annotations
 
 import warnings
-from abc import ABC, abstractmethod
 from copy import deepcopy
 from dataclasses import dataclass
 from typing import Iterator, List, Optional, Tuple, Union
@@ -40,7 +39,7 @@
     FUNCTORCH_ERROR = "functorch not installed. Consider installing functorch to use this functionality."
 
 
-class LossModule(nn.Module, ABC):
+class LossModule(nn.Module):
     """A parent class for RL losses.
 
     LossModule inherits from nn.Module. It is designed to read an input
@@ -59,6 +58,25 @@ class LossModule(nn.Module, ABC):
 
     By default, the forward method is always decorated with a
     gh :class:`torchrl.envs.ExplorationType.MODE`
+
+    To utilize the ability configuring the tensordict keys via
+    :meth:`~.set_keys()` a subclass must define an _AcceptedKeys dataclass.
+    This dataclass should include all keys that are intended to be configurable.
+    In addition, the subclass must implement the
+    :meth:._forward_value_estimator_keys() method. This function is crucial for
+    forwarding any altered tensordict keys to the underlying value_estimator.
+
+    Examples:
+        >>> class MyLoss(LossModule):
+        >>>     @dataclass
+        >>>     class _AcceptedKeys:
+        >>>         action = "action"
+        >>>
+        >>>     def _forward_value_estimator_keys(self, **kwargs) -> None:
+        >>>         pass
+        >>>
+        >>> loss = MyLoss()
+        >>> loss.set_keys(action="action2")
     """
 
     @dataclass
@@ -91,11 +109,6 @@ def __init__(self):
         self.value_type = self.default_value_estimator
         # self.register_forward_pre_hook(_parameters_to_tensordict)
 
-    @abstractmethod
-    def _forward_value_estimator_keys(self, **kwargs) -> None:
-        """Passes updated tensordict keys to the underlying value estimator."""
-        ...
-
     def _set_deprecated_ctor_keys(self, **kwargs) -> None:
         """Helper function to set a tensordict key from a constructor and raise a warning simultaneously."""
         for key, value in kwargs.items():
@@ -124,7 +137,15 @@ def set_keys(self, **kwargs) -> None:
             else:
                 setattr(self.tensor_keys, key, self.default_keys.key)
 
-        self._forward_value_estimator_keys(**kwargs)
+        try:
+            self._forward_value_estimator_keys(**kwargs)
+        except AttributeError:
+            raise AttributeError(
+                "To utilize `.set_keys(...)` for tensordict key configuration, the subclassed loss module "
+                "must define an _AcceptedKeys dataclass containing all keys intended for configuration. "
+                "Moreover, the subclass needs to implement `._forward_value_estimator_keys()` method to "
+                "facilitate forwarding of any modified tensordict keys to the underlying value_estimator."
+            )
 
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
         """It is designed to read an input TensorDict and return another tensordict with loss keys named "loss*".