fix(pu): add batch wrapper for alphazero ptree mcts

lzero/mcts/ptree/ptree_az.py

@@ -263,6 +263,46 @@ def get_next_action(
         # Return the selected action and the output probability of each action.
         return action, action_probs
 
+    def get_next_actions_batch(
+            self,
+            state_config_for_simulate_env_reset_list: List[Dict[str, Any]],
+            policy_forward_fn_batch: Callable,
+            temperature: float = 1.0,
+            sample: bool = True,
+            env_list: List[Type[BaseEnv]] = None,
+    ) -> List[Tuple[int, List[float]]]:
+        """
+        Overview:
+            Batch-compatible wrapper for the Python AlphaZero MCTS implementation.
+            The Python tree search is still executed sequentially, but this method
+            keeps the public interface aligned with the C++ AlphaZero MCTS.
+        """
+        if env_list is None:
+            env_list = [self.simulate_env for _ in state_config_for_simulate_env_reset_list]
+        assert len(state_config_for_simulate_env_reset_list) == len(env_list)
+
+        results = []
+        original_simulate_env = self.simulate_env
+        try:
+            for state_config, simulate_env in zip(state_config_for_simulate_env_reset_list, env_list):
+                self.simulate_env = simulate_env
+
+                def policy_forward_fn(env):
+                    return policy_forward_fn_batch([env])[0]
+
+                results.append(
+                    self.get_next_action(
+                        state_config_for_simulate_env_reset=state_config,
+                        policy_forward_fn=policy_forward_fn,
+                        temperature=temperature,
+                        sample=sample,
+                    )
+                )
+        finally:
+            self.simulate_env = original_simulate_env
+
+        return results
+
     def _simulate(self, node: Node, simulate_env: Type[BaseEnv], policy_forward_fn: Callable) -> None:
         """
         Overview:

lzero/mcts/ptree/ptree_az_sampled.py

@@ -282,6 +282,46 @@ def get_next_action(
         # Return the selected action and the output probability of each action.
         return action, action_probs
 
+    def get_next_actions_batch(
+            self,
+            state_config_for_env_reset_list: List[Dict[str, Any]],
+            policy_value_func_batch: Callable,
+            temperature: float = 1.0,
+            sample: bool = True,
+            env_list: List[Type[BaseEnv]] = None,
+    ) -> List[Tuple[int, List[float]]]:
+        """
+        Overview:
+            Batch-compatible wrapper for the Python Sampled AlphaZero MCTS
+            implementation. The Python tree search remains sequential, but the
+            public interface matches the C++ AlphaZero MCTS used by policy code.
+        """
+        if env_list is None:
+            env_list = [self.simulate_env for _ in state_config_for_env_reset_list]
+        assert len(state_config_for_env_reset_list) == len(env_list)
+
+        results = []
+        original_simulate_env = self.simulate_env
+        try:
+            for state_config, simulate_env in zip(state_config_for_env_reset_list, env_list):
+                self.simulate_env = simulate_env
+
+                def policy_value_func(env):
+                    return policy_value_func_batch([env])[0]
+
+                results.append(
+                    self.get_next_action(
+                        state_config_for_env_reset=state_config,
+                        policy_value_func=policy_value_func,
+                        temperature=temperature,
+                        sample=sample,
+                    )
+                )
+        finally:
+            self.simulate_env = original_simulate_env
+
+        return results
+
     def _simulate(self, node: Node, simulate_env: Type[BaseEnv], policy_value_func: Callable) -> None:
         """
         Overview:

lzero/mcts/tests/test_alphazero_ptree_batch.py

@@ -0,0 +1,69 @@
+import pytest
+from easydict import EasyDict
+
+from lzero.mcts.ptree.ptree_az import MCTS as AlphaZeroMCTS
+from lzero.mcts.ptree.ptree_az_sampled import MCTS as SampledAlphaZeroMCTS
+
+
+class DummyEnv:
+
+    def __init__(self, name):
+        self.name = name
+
+
+def _make_cfg():
+    return EasyDict(
+        dict(
+            max_moves=10,
+            num_simulations=1,
+            pb_c_base=1,
+            pb_c_init=1,
+            root_dirichlet_alpha=0.3,
+            root_noise_weight=0.25,
+            legal_actions=[0, 1],
+            action_space_size=2,
+            num_of_sampled_actions=2,
+            continuous_action_space=False,
+        )
+    )
+
+
+def _check_batch_wrapper(mcts_cls, monkeypatch):
+    env0 = DummyEnv('env0')
+    env1 = DummyEnv('env1')
+    mcts = mcts_cls(_make_cfg(), env0)
+    state_configs = [EasyDict(start_player_index=1, init_state='s0'), EasyDict(start_player_index=2, init_state='s1')]
+    calls = []
+
+    def fake_get_next_action(self, temperature=1.0, sample=True, **kwargs):
+        state_config = kwargs.get('state_config_for_simulate_env_reset', kwargs.get('state_config_for_env_reset'))
+        policy_fn = kwargs.get('policy_forward_fn', kwargs.get('policy_value_func'))
+        assert self.simulate_env in [env0, env1]
+        policy_result = policy_fn(self.simulate_env)
+        calls.append((state_config.init_state, self.simulate_env.name, policy_result, temperature, sample))
+        return len(calls), [0.0, 1.0]
+
+    monkeypatch.setattr(mcts_cls, 'get_next_action', fake_get_next_action)
+
+    def fake_policy_batch(env_list):
+        assert len(env_list) == 1
+        return [({'env': env_list[0].name}, 0.0)]
+
+    results = mcts.get_next_actions_batch(state_configs, fake_policy_batch, temperature=0.5, sample=False, env_list=[env0, env1])
+
+    assert results == [(1, [0.0, 1.0]), (2, [0.0, 1.0])]
+    assert calls == [
+        ('s0', 'env0', ({'env': 'env0'}, 0.0), 0.5, False),
+        ('s1', 'env1', ({'env': 'env1'}, 0.0), 0.5, False),
+    ]
+    assert mcts.simulate_env is env0
+
+
+@pytest.mark.unittest
+def test_alphazero_ptree_get_next_actions_batch(monkeypatch):
+    _check_batch_wrapper(AlphaZeroMCTS, monkeypatch)
+
+
+@pytest.mark.unittest
+def test_sampled_alphazero_ptree_get_next_actions_batch(monkeypatch):
+    _check_batch_wrapper(SampledAlphaZeroMCTS, monkeypatch)