polish(pu): optimize the implementation of transformation from action_mask to legal_actions (#466)

puyuan1996 · web-flow · commit 556b2ec5451c · 2026-01-26T12:20:41.000+08:00
diff --git a/lzero/mcts/buffer/game_buffer_efficientzero.py b/lzero/mcts/buffer/game_buffer_efficientzero.py
@@ -185,7 +185,7 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A
             game_segment_batch_size, to_play_segment, action_mask_segment, pos_in_game_segment_list
         )
 
-        legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+        legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         # ==============================================================
         # EfficientZero related core code
@@ -344,7 +344,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
             game_segment_batch_size, to_play_segment, action_mask_segment, pos_in_game_segment_list
         )
 
-        legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+        legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
         with torch.no_grad():
             policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)
             # split a full batch into slices of mini_infer_size: to save the GPU memory for more GPU actors
diff --git a/lzero/mcts/buffer/game_buffer_muzero.py b/lzero/mcts/buffer/game_buffer_muzero.py
@@ -597,7 +597,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self._cfg.model.num_of_sampled_actions)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         with torch.no_grad():
             policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)
@@ -755,7 +755,7 @@ def _compute_target_policy_non_reanalyzed(
                 [-1 for _ in range(self.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
             
         with torch.no_grad():
             policy_index = 0
diff --git a/lzero/mcts/buffer/game_buffer_rezero_ez.py b/lzero/mcts/buffer/game_buffer_rezero_ez.py
@@ -158,7 +158,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         with torch.no_grad():
             policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)
diff --git a/lzero/mcts/buffer/game_buffer_rezero_mz.py b/lzero/mcts/buffer/game_buffer_rezero_mz.py
@@ -230,7 +230,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         with torch.no_grad():
             policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)
diff --git a/lzero/mcts/buffer/game_buffer_sampled_efficientzero.py b/lzero/mcts/buffer/game_buffer_sampled_efficientzero.py
@@ -272,7 +272,7 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A
                 [-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         batch_target_values, batch_value_prefixs = [], []
         with torch.no_grad():
@@ -452,7 +452,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         with torch.no_grad():
             policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)
diff --git a/lzero/mcts/buffer/game_buffer_sampled_muzero.py b/lzero/mcts/buffer/game_buffer_sampled_muzero.py
@@ -272,7 +272,7 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A
                 [-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         batch_target_values, batch_rewards = [], []
         with torch.no_grad():
@@ -437,7 +437,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         with torch.no_grad():
             policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)
diff --git a/lzero/mcts/buffer/game_buffer_sampled_unizero.py b/lzero/mcts/buffer/game_buffer_sampled_unizero.py
@@ -485,7 +485,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self._cfg.model.num_of_sampled_actions)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         # NOTE: TODO
         model.world_model.reanalyze_phase = True
@@ -658,7 +658,7 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A
                 [-1 for _ in range(self._cfg.model.num_of_sampled_actions)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         batch_target_values, batch_rewards = [], []
         with torch.no_grad():
diff --git a/lzero/mcts/buffer/game_buffer_unizero.py b/lzero/mcts/buffer/game_buffer_unizero.py
@@ -431,7 +431,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:
                 [-1 for _ in range(self.action_space_size)] for _ in range(transition_batch_size)
             ]
         else:
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]
 
         # NOTE: check the effect of reanalyze_phase
         model.world_model.reanalyze_phase = True
diff --git a/lzero/mcts/tests/eval_tree_speed.py b/lzero/mcts/tests/eval_tree_speed.py
@@ -116,7 +116,7 @@ def ptree_func(policy_config, num_simulations):
         assert len(action_mask[0]) == action_space_size
 
         action_num = [int(np.array(action_mask[i]).sum()) for i in range(env_nums)]
-        legal_actions_list = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(env_nums)]
+        legal_actions_list = [np.nonzero(action_mask[j])[0].tolist() for j in range(env_nums)]
         to_play = [np.random.randint(1, 3) for i in range(env_nums)]
         assert len(to_play) == batch_size
         # ============================================ptree=====================================#
@@ -212,7 +212,7 @@ def ctree_func(policy_config, num_simulations):
         assert len(action_mask[0]) == action_space_size
 
         action_num = [int(np.array(action_mask[i]).sum()) for i in range(env_nums)]
-        legal_actions_list = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(env_nums)]
+        legal_actions_list = [np.nonzero(action_mask[j])[0].tolist() for j in range(env_nums)]
         to_play = [np.random.randint(1, 3) for i in range(env_nums)]
         assert len(to_play) == batch_size
         # ============================================ctree=====================================#
diff --git a/lzero/mcts/tests/test_mcts_ctree.py b/lzero/mcts/tests/test_mcts_ctree.py
@@ -143,7 +143,7 @@ def recurrent_inference(self, latent_states, reward_hidden_states, actions=None)
 action_num = [
     int(np.array(action_mask[i]).sum()) for i in range(env_nums)
 ]  # [3, 3, 5, 4, 3, 3, 6, 6, 3, 6, 6, 5, 2, 5, 1, 4]
-legal_actions_list = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(env_nums)]
+legal_actions_list = [np.nonzero(action_mask[j])[0].tolist() for j in range(env_nums)]
 # legal_actions_list =
 # [[3, 5, 6], [0, 3, 6], [0, 1, 4, 6, 8], [0, 3, 4, 5],
 # [2, 5, 8], [1, 2, 4], [0, 2, 3, 4, 7, 8], [0, 1, 2, 3, 4, 8],
diff --git a/lzero/mcts/tests/test_mcts_ptree.py b/lzero/mcts/tests/test_mcts_ptree.py
@@ -121,7 +121,7 @@ def recurrent_inference(self, hidden_states, reward_hidden_states, actions):
 action_num = [
     int(np.array(action_mask[i]).sum()) for i in range(env_nums)
 ]  # [3, 3, 5, 4, 3, 3, 6, 6, 3, 6, 6, 5, 2, 5, 1, 4]
-legal_actions_list = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(env_nums)]
+legal_actions_list = [np.nonzero(action_mask[j])[0].tolist() for j in range(env_nums)]
 # legal_actions_list =
 # [[3, 5, 6], [0, 3, 6], [0, 1, 4, 6, 8], [0, 3, 4, 5],
 # [2, 5, 8], [1, 2, 4], [0, 2, 3, 4, 7, 8], [0, 1, 2, 3, 4, 8],
diff --git a/lzero/policy/efficientzero.py b/lzero/policy/efficientzero.py
@@ -592,7 +592,7 @@ def _forward_collect(
             )
             policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
 
             if not self._cfg.collect_with_pure_policy:
                 # collect with MCTS guided with policy.
@@ -711,7 +711,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: Union[in
                 )
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape（B, A）
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # cpp mcts_tree
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/gumbel_muzero.py b/lzero/policy/gumbel_muzero.py
@@ -546,7 +546,7 @@ def _forward_collect(
             latent_state_roots = latent_state_roots.detach().cpu().numpy()
             policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             # the only difference between collect and eval is the dirichlet noise
             noises = [
                 np.random.dirichlet([self._cfg.root_dirichlet_alpha] * int(sum(action_mask[j]))
@@ -655,7 +655,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: List = [
                 latent_state_roots = latent_state_roots.detach().cpu().numpy()
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape（B, A）
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # cpp mcts_tree
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/muzero.py b/lzero/policy/muzero.py
@@ -742,7 +742,7 @@ def _forward_collect(
             latent_state_roots = latent_state_roots.detach().cpu().numpy()
             policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             if not self._cfg.collect_with_pure_policy:
                 # the only difference between collect and eval is the dirichlet noise
                 noises = [
@@ -895,7 +895,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: List = [
                 latent_state_roots = latent_state_roots.detach().cpu().numpy()
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape（B, A）
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # cpp mcts_tree
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/muzero_multitask.py b/lzero/policy/muzero_multitask.py
@@ -711,7 +711,7 @@ def _forward_collect(
             latent_state_roots = latent_state_roots.detach().cpu().numpy()
             policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             if not self._cfg.collect_with_pure_policy:
                 # The only difference between collect and eval is the dirichlet noise.
                 noises = [
@@ -849,7 +849,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: int = -1
                 latent_state_roots = latent_state_roots.detach().cpu().numpy()
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape (B, A)
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # C++ MCTS tree.
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/muzero_rnn_full_obs.py b/lzero/policy/muzero_rnn_full_obs.py
@@ -588,7 +588,7 @@ def _forward_collect(
             world_model_latent_history_roots = world_model_latent_history_roots.detach().cpu().numpy()
             policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             # the only difference between collect and eval is the dirichlet noise.
             noises = [
                 np.random.dirichlet([self._cfg.root_dirichlet_alpha] * int(sum(action_mask[j]))
@@ -715,7 +715,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: List = [
                 world_model_latent_history_roots = world_model_latent_history_roots.detach().cpu().numpy()
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape（B, A）
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # cpp mcts_tree
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/sampled_unizero_multitask.py b/lzero/policy/sampled_unizero_multitask.py
@@ -684,7 +684,7 @@ def _forward_collect(
 
             # 2. Prepare MCTS roots.
             if not self._cfg.model.continuous_action_space:
-                legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+                legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             else:
                 legal_actions = [[-1] * self._cfg.model.world_model_cfg.num_of_sampled_actions for _ in range(active_collect_env_num)]
 
@@ -796,7 +796,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: int = -1
 
             # 2. Prepare MCTS roots without noise for deterministic evaluation.
             if not self._cfg.model.continuous_action_space:
-                legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+                legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             else:
                 legal_actions = [[-1] * self._cfg.model.world_model_cfg.num_of_sampled_actions for _ in range(active_eval_env_num)]
 
diff --git a/lzero/policy/stochastic_muzero.py b/lzero/policy/stochastic_muzero.py
@@ -625,7 +625,7 @@ def _forward_collect(
                 latent_state_roots = latent_state_roots.detach().cpu().numpy()
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             # the only difference between collect and eval is the dirichlet noise
             noises = [
                 np.random.dirichlet([self._cfg.root_dirichlet_alpha] * int(sum(action_mask[j]))
@@ -715,7 +715,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: List = [
                 latent_state_roots = latent_state_roots.detach().cpu().numpy()
                 policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape（B, A）
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # cpp mcts_tree
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/unizero.py b/lzero/policy/unizero.py
@@ -1397,7 +1397,7 @@ def _forward_collect(
             latent_state_roots = latent_state_roots.detach().cpu().numpy()
             policy_logits = policy_logits.detach().cpu().numpy().tolist()
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_collect_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_collect_env_num)]
             # the only difference between collect and eval is the dirichlet noise
             noises = [
                 np.random.dirichlet([self._cfg.root_dirichlet_alpha] * int(sum(action_mask[j]))
@@ -1554,7 +1554,7 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: int = -1
             latent_state_roots = latent_state_roots.detach().cpu().numpy()
             policy_logits = policy_logits.detach().cpu().numpy().tolist()  # list shape（B, A）
 
-            legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(active_eval_env_num)]
+            legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(active_eval_env_num)]
             if self._cfg.mcts_ctree:
                 # cpp mcts_tree
                 roots = MCTSCtree.roots(active_eval_env_num, legal_actions)
diff --git a/lzero/policy/unizero_multitask.py b/lzero/policy/unizero_multitask.py

Original file line number	Diff line number	Diff line change
`@@ -185,7 +185,7 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A`
`185`	`185`	`game_segment_batch_size, to_play_segment, action_mask_segment, pos_in_game_segment_list`
`186`	`186`	`)`
`187`	`187`
`188`		`- legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]`
	`188`	`+ legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]`
`189`	`189`
`190`	`190`	`# ==============================================================`
`191`	`191`	`# EfficientZero related core code`
`@@ -344,7 +344,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:`
`344`	`344`	`game_segment_batch_size, to_play_segment, action_mask_segment, pos_in_game_segment_list`
`345`	`345`	`)`
`346`	`346`
`347`		`- legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]`
	`347`	`+ legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]`
`348`	`348`	`with torch.no_grad():`
`349`	`349`	`policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)`
`350`	`350`	`# split a full batch into slices of mini_infer_size: to save the GPU memory for more GPU actors`
Original file line number	Diff line number	Diff line change
`@@ -158,7 +158,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:`
`158`	`158`	`[-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)`
`159`	`159`	`]`
`160`	`160`	`else:`
`161`		`- legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]`
	`161`	`+ legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]`
`162`	`162`
`163`	`163`	`with torch.no_grad():`
`164`	`164`	`policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)`
Original file line number	Diff line number	Diff line change
`@@ -230,7 +230,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:`
`230`	`230`	`[-1 for _ in range(self._cfg.model.action_space_size)] for _ in range(transition_batch_size)`
`231`	`231`	`]`
`232`	`232`	`else:`
`233`		`- legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]`
	`233`	`+ legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]`
`234`	`234`
`235`	`235`	`with torch.no_grad():`
`236`	`236`	`policy_obs_list = prepare_observation(policy_obs_list, self._cfg.model.model_type)`
Original file line number	Diff line number	Diff line change
`@@ -431,7 +431,7 @@ def _compute_target_policy_reanalyzed(self, policy_re_context: List[Any], model:`
`431`	`431`	`[-1 for _ in range(self.action_space_size)] for _ in range(transition_batch_size)`
`432`	`432`	`]`
`433`	`433`	`else:`
`434`		`- legal_actions = [[i for i, x in enumerate(action_mask[j]) if x == 1] for j in range(transition_batch_size)]`
	`434`	`+ legal_actions = [np.nonzero(action_mask[j])[0].tolist() for j in range(transition_batch_size)]`
`435`	`435`
`436`	`436`	`# NOTE: check the effect of reanalyze_phase`
`437`	`437`	`model.world_model.reanalyze_phase = True`