Using dist.mode instead of logits.argmax (thu-ml#1066)

changed all the occurrences where an action is selected deterministically

- **from**: using the outputs of the actor network.
- **to**: using the mode of the PyTorch distribution.

---------

Co-authored-by: Arnau Jimenez <[email protected]>

tianshou/policy/modelfree/discrete_sac.py

@@ -108,7 +108,7 @@ def forward(  # type: ignore
         logits_BA, hidden_BH = self.actor(batch.obs, state=state, info=batch.info)
         dist = Categorical(logits=logits_BA)
         if self.deterministic_eval and not self.training:
-            act_B = dist.mode
+            act = dist.mode
         else:
             act_B = dist.sample()
         return Batch(logits=logits_BA, act=act_B, state=hidden_BH, dist=dist)

tianshou/policy/modelfree/pg.py

@@ -198,7 +198,7 @@ def forward(
         dist = self.dist_fn(action_dist_input_BD)
 
         if self.deterministic_eval and not self.training:
-            act_B = dist.mode
+            act = dist.mode
         else:
             act_B = dist.sample()
         # act is of dimension BA in continuous case and of dimension B in discrete

tianshou/policy/modelfree/redq.py

@@ -151,13 +151,14 @@ def forward(  # type: ignore
         state: dict | Batch | np.ndarray | None = None,
         **kwargs: Any,
     ) -> Batch:
-        (loc_B, scale_B), h_BH = self.actor(batch.obs, state=state, info=batch.info)
-        dist = Independent(Normal(loc_B, scale_B), 1)
+        loc_scale, h = self.actor(batch.obs, state=state, info=batch.info)
+        loc, scale = loc_scale
+        dist = Independent(Normal(loc, scale), 1)
         if self.deterministic_eval and not self.training:
-            act_B = dist.mode
+            act = dist.mode
         else:
-            act_B = dist.rsample()
-        log_prob = dist.log_prob(act_B).unsqueeze(-1)
+            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.
         # in appendix C to get some understanding of this equation.

tianshou/policy/modelfree/sac.py

@@ -176,7 +176,7 @@ def forward(  # type: ignore
         (loc_B, scale_B), hidden_BH = self.actor(batch.obs, state=state, info=batch.info)
         dist = Independent(Normal(loc=loc_B, scale=scale_B), 1)
         if self.deterministic_eval and not self.training:
-            act_B = dist.mode
+            act = dist.mode
         else:
             act_B = dist.rsample()
         log_prob = dist.log_prob(act_B).unsqueeze(-1)