Update

[ghstack-poisoned]
pytorch · Sep 17, 2024 · 3d7a5c4 · 3d7a5c4
1 parent 200e5b6
commit 3d7a5c4
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))