Using dist.mode instead of logits.argmax. More compact.

tianshou/policy/modelfree/discrete_sac.py

@@ -109,7 +109,7 @@ def forward(  # type: ignore
         logits, hidden = self.actor(batch.obs, state=state, info=batch.info)
         dist = Categorical(logits=logits)
         if self.deterministic_eval and not self.training:
-            act = logits.argmax(axis=-1)
+            act = dist.mode
         else:
             act = dist.sample()
         return Batch(logits=logits, act=act, state=hidden, dist=dist)

tianshou/policy/modelfree/pg.py

@@ -158,19 +158,6 @@ def process_fn(
         batch: BatchWithReturnsProtocol
         return batch
 
-    def _get_deterministic_action(self, logits: torch.Tensor) -> torch.Tensor:
-        if self.action_type == "discrete":
-            return logits.argmax(-1)
-        if self.action_type == "continuous":
-            # assume that the mode of the distribution is the first element
-            # of the actor's output (the "logits")
-            return logits[0]
-        raise RuntimeError(
-            f"Unknown action type: {self.action_type}. "
-            f"This should not happen and might be a bug."
-            f"Supported action types are: 'discrete' and 'continuous'.",
-        )
-
     def forward(
         self,
         batch: ObsBatchProtocol,
@@ -198,7 +185,7 @@ def forward(
 
         # in this case, the dist is unused!
         if self.deterministic_eval and not self.training:
-            act = self._get_deterministic_action(logits)
+            act = dist.mode
         else:
             act = dist.sample()
         result = Batch(logits=logits, act=act, state=hidden, dist=dist)

tianshou/policy/modelfree/redq.py

@@ -153,7 +153,10 @@ def forward(  # type: ignore
         loc_scale, h = self.actor(batch.obs, state=state, info=batch.info)
         loc, scale = loc_scale
         dist = Independent(Normal(loc, scale), 1)
-        act = loc if self.deterministic_eval and not self.training else dist.rsample()
+        if self.deterministic_eval and not self.training:
+            act = dist.mode
+        else:
+            act = dist.rsample()
         log_prob = dist.log_prob(act).unsqueeze(-1)
         # apply correction for Tanh squashing when computing logprob from Gaussian
         # You can check out the original SAC paper (arXiv 1801.01290): Eq 21.

tianshou/policy/modelfree/sac.py

@@ -177,7 +177,7 @@ def forward(  # type: ignore
         assert isinstance(logits, tuple)
         dist = Independent(Normal(*logits), 1)
         if self.deterministic_eval and not self.training:
-            act = logits[0]
+            act = dist.mode
         else:
             act = dist.rsample()
         log_prob = dist.log_prob(act).unsqueeze(-1)