[Feature] Deterministic sample for Masked one-hot

ghstack-source-id: 27787ea
Pull Request resolved: #2440

Author: Vincent Moens

torchrl/modules/distributions/discrete.py

@@ -389,6 +389,17 @@ def sample(
     ) -> torch.Tensor:
         ...
 
+    @property
+    def deterministic_sample(self):
+        return self.mode
+
+    @property
+    def mode(self) -> torch.Tensor:
+        if hasattr(self, "logits"):
+            return (self.logits == self.logits.max(-1, True)[0]).to(torch.long)
+        else:
+            return (self.probs == self.probs.max(-1, True)[0]).to(torch.long)
+
     def log_prob(self, value: torch.Tensor) -> torch.Tensor:
         return super().log_prob(value.argmax(dim=-1))