Update stable_behavior_policy.py

Author: wuxiyang1996

nova/stable_behavior_policy.py

@@ -186,7 +186,11 @@ def learn(self, batch, t_env):
             # => [n_thread, num_history, max_history_len, max_vehicle_num, obs_dim]
             agent_history = history[:, :, i]
             agent_behavior_latent = behavior_latent[:, :, i]
-            mask = 1 - agent_terminate[:, :, i, 0]
+            
+            if self.args.env == "MPE":
+                mask = 1 - agent_terminate[:, :, i, 0]
+            else:
+                mask = agent_terminate[:, :, i, 0]
 
             # Prev latent / Updated latent: [n_thread, max_vehicle_num, latent_dim]
             latent = torch.zeros((n_thread, max_vehicle_num, self.latent_dim)).to(self.device)