Fixed runtime error for continuous single action envs

nikhilbarhate99 · Apr 10, 2021 · 2d6a390 · 2d6a390
1 parent b27f936
commit 2d6a390
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/PPO.py b/PPO.py
@@ -123,10 +123,12 @@ def evaluate(self, state, action):
             action_var = self.action_var.expand_as(action_mean)
             cov_mat = torch.diag_embed(action_var).to(device)
             dist = MultivariateNormal(action_mean, cov_mat)
+
+            action = action.reshape(-1, self.action_dim)
+
         else:
             action_probs = self.actor(state)
             dist = Categorical(action_probs)
-        action = action.reshape(-1, self.action_dim)
         action_logprobs = dist.log_prob(action)
         dist_entropy = dist.entropy()
         state_values = self.critic(state)