[Feature] Deterministic sample for Masked one-hot

ghstack-source-id: 27787eab47324c5af152f706d81687e71b5b9803 Pull Request resolved: #2440
pytorch · Sep 17, 2024 · e294c68 · e294c68 · github-actions · Sep 17, 2024
1 parent 0a410ff
commit e294c68
Showing 1 changed file with 11 additions and 0 deletions.
diff --git a/torchrl/modules/distributions/discrete.py b/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))