Remove cast to long

Author: araffin

stable_baselines3/a2c/a2c.py

@@ -144,8 +144,8 @@ def train(self) -> None:
         for rollout_data in self.rollout_buffer.get(batch_size=None):
             actions = rollout_data.actions
             if isinstance(self.action_space, spaces.Discrete):
-                # Convert discrete action from float to long
-                actions = actions.long().flatten()
+                # Flatten discrete actions for correct computation of log prob
+                actions = actions.flatten()
 
             values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions)
             values = values.flatten()

stable_baselines3/ppo/ppo.py

@@ -207,9 +207,8 @@ def train(self) -> None:
             for rollout_data in self.rollout_buffer.get(self.batch_size):
                 actions = rollout_data.actions
                 if isinstance(self.action_space, spaces.Discrete):
-                    # Convert discrete action from float to long
-                    actions = rollout_data.actions.long().flatten()
-
+                    # Flatten discrete actions for correct computation of log prob
+                    actions = rollout_data.actions.flatten()
                 values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions)
                 values = values.flatten()
                 # Normalize advantage