diff --git a/lzero/mcts/ptree/ptree_az.py b/lzero/mcts/ptree/ptree_az.py index 58143c481..63d9b6c6f 100644 --- a/lzero/mcts/ptree/ptree_az.py +++ b/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: diff --git a/lzero/mcts/ptree/ptree_az_sampled.py b/lzero/mcts/ptree/ptree_az_sampled.py index 7ebfefd05..85adde3a0 100644 --- a/lzero/mcts/ptree/ptree_az_sampled.py +++ b/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: diff --git a/lzero/mcts/tests/test_alphazero_ptree_batch.py b/lzero/mcts/tests/test_alphazero_ptree_batch.py new file mode 100644 index 000000000..f1003cc62 --- /dev/null +++ b/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)