[Feature] DETERMINISTIC interaction mode (#824)

Co-authored-by: Matteo Bettini <[email protected]>

tensordict/nn/distributions/continuous.py

@@ -264,3 +264,7 @@ def mode(self) -> torch.Tensor:
     @property
     def mean(self) -> torch.Tensor:
         return self.param
+
+    @property
+    def deterministic_sample(self) -> torch.Tensor:
+        return self.param

tensordict/nn/distributions/discrete.py

@@ -63,6 +63,8 @@ def mode(self) -> torch.Tensor:
         else:
             return (self.probs == self.probs.max(-1, True)[0]).to(torch.long)
 
+    determnistic_sample = mode
+
     def sample(
         self,
         sample_shape: torch.Size | Sequence[int] | None = None,

tensordict/nn/probabilistic.py

@@ -28,10 +28,24 @@
 
 
 class InteractionType(Enum):
+    """A list of possible interaction types with a distribution.
+
+    MODE, MEDIAN and MEAN point to the property / attribute with the same name.
+    RANDOM points to ``rsample()`` if that method exists or ``sample()`` if not.
+
+    DETERMINISTIC can be used as a generic fallback if ``MEAN`` or ``MODE`` are not guaranteed to
+    be analytically tractable. In such cases, a rude deterministic estimate can be used
+    in some cases even if it lacks a true algebraic meaning.
+    This value will trigger a query to the ``deterministic_sample`` attribute in the distribution
+    and if it does not exist, the ``mean`` will be used.
+
+    """
+
     MODE = auto()
     MEDIAN = auto()
     MEAN = auto()
     RANDOM = auto()
+    DETERMINISTIC = auto()
 
     @classmethod
     def from_str(cls, type_str: str) -> InteractionType:
@@ -452,9 +466,27 @@ def _dist_sample(
         if interaction_type is None:
             interaction_type = self.default_interaction_type
 
+        if interaction_type is InteractionType.DETERMINISTIC:
+            try:
+                return dist.deterministic_sample
+            except AttributeError:
+                try:
+                    return dist.mean
+                except AttributeError:
+                    raise NotImplementedError(
+                        f"method {type(dist)}.deterministic_sample is not implemented."
+                    )
+                finally:
+                    warnings.warn(
+                        "deterministic_sample wasn't found when queried. "
+                        f"{type(self).__name__} is falling back on mean instead. "
+                        f"For better code quality and efficiency, make sure to either "
+                        f"provide a distribution with a deterministic_sample attribute or "
+                        f"to change the InteractionMode to the desired value.",
+                        category=UserWarning,
+                    )
+
         if interaction_type is InteractionType.MODE:
-            if hasattr(dist, "get_mode"):
-                return dist.get_mode()
             try:
                 return dist.mode
             except AttributeError:
@@ -463,8 +495,6 @@ def _dist_sample(
                 )
 
         elif interaction_type is InteractionType.MEDIAN:
-            if hasattr(dist, "get_median"):
-                return dist.get_median()
             try:
                 return dist.median
             except AttributeError:
@@ -473,8 +503,6 @@ def _dist_sample(
                 )
 
         elif interaction_type is InteractionType.MEAN:
-            if hasattr(dist, "get_mean"):
-                return dist.get_mean()
             try:
                 return dist.mean
             except (AttributeError, NotImplementedError):