[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))