feature(yzj): add multi-agent and structured observation env (GoBigger)

lzero/entry/__init__.py

@@ -3,4 +3,6 @@
 from .train_muzero import train_muzero
 from .eval_muzero import eval_muzero
 from .eval_muzero_with_gym_env import eval_muzero_with_gym_env
-from .train_muzero_with_gym_env import train_muzero_with_gym_env
+from .train_muzero_with_gym_env import train_muzero_with_gym_env
+from .train_muzero_gobigger import train_muzero_gobigger
+from .eval_muzero_gobigger import eval_muzero_gobigger

lzero/entry/eval_muzero_gobigger.py

@@ -0,0 +1,116 @@
+import logging
+import os
+from functools import partial
+from typing import Optional, Tuple
+import numpy as np
+import torch
+from tensorboardX import SummaryWriter
+import copy
+
+from ding.config import compile_config
+from ding.envs import create_env_manager
+from ding.envs import get_vec_env_setting
+from ding.policy import create_policy
+from ding.utils import set_pkg_seed
+from ding.worker import BaseLearner
+from lzero.worker import GoBiggerMuZeroEvaluator
+
+
+def eval_muzero_gobigger(
+        input_cfg: Tuple[dict, dict],
+        seed: int = 0,
+        model: Optional[torch.nn.Module] = None,
+        model_path: Optional[str] = None,
+) -> 'Policy':  # noqa
+    """
+    Overview:
+        The eval entry for GoBigger MCTS+RL algorithms, including MuZero, EfficientZero, Sampled EfficientZero.
+    Arguments:
+        - input_cfg (:obj:`Tuple[dict, dict]`): Config in dict type.
+            ``Tuple[dict, dict]`` type means [user_config, create_cfg].
+        - seed (:obj:`int`): Random seed.
+        - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module.
+        - model_path (:obj:`Optional[str]`): The pretrained model path, which should
+            point to the ckpt file of the pretrained model, and an absolute path is recommended.
+            In LightZero, the path is usually something like ``exp_name/ckpt/ckpt_best.pth.tar``.
+    Returns:
+        - policy (:obj:`Policy`): Converged policy.
+    """
+    cfg, create_cfg = input_cfg
+    assert create_cfg.policy.type in ['gobigger_efficientzero', 'gobigger_muzero', 'gobigger_sampled_efficientzero'], \
+        "train_muzero entry now only support the following algo.: 'gobigger_efficientzero', 'gobigger_muzero', 'gobigger_sampled_efficientzero'"
+
+    if create_cfg.policy.type == 'gobigger_efficientzero':
+        from lzero.mcts import GoBiggerEfficientZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'gobigger_muzero':
+        from lzero.mcts import GoBiggerMuZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'gobigger_sampled_efficientzero':
+        from lzero.mcts import GoBiggerSampledEfficientZeroGameBuffer as GameBuffer
+
+    if cfg.policy.cuda and torch.cuda.is_available():
+        cfg.policy.device = 'cuda'
+    else:
+        cfg.policy.device = 'cpu'
+
+    cfg = compile_config(cfg, seed=seed, env=None, auto=True, create_cfg=create_cfg, save_cfg=True)
+    # Create main components: env, policy
+    env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env)
+    collector_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in collector_env_cfg])
+    evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg])
+
+    env_cfg = copy.deepcopy(evaluator_env_cfg[0])
+    env_cfg.contain_raw_obs = True
+    vsbot_evaluator_env_cfg = [env_cfg for _ in range(len(evaluator_env_cfg))]
+    vsbot_evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in vsbot_evaluator_env_cfg])
+
+    collector_env.seed(cfg.seed)
+    evaluator_env.seed(cfg.seed, dynamic_seed=False)
+    vsbot_evaluator_env.seed(cfg.seed, dynamic_seed=False)
+    set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)
+
+    policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval'])
+
+    # load pretrained model
+    if model_path is not None:
+        policy.learn_mode.load_state_dict(torch.load(model_path, map_location=cfg.policy.device))
+
+    # Create worker components: learner, collector, evaluator, replay buffer, commander.
+    tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial'))
+    learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name)
+
+    # ==============================================================
+    # MCTS+RL algorithms related core code
+    # ==============================================================
+    policy_config = cfg.policy
+    evaluator = GoBiggerMuZeroEvaluator(
+        eval_freq=cfg.policy.eval_freq,
+        n_evaluator_episode=cfg.env.n_evaluator_episode,
+        stop_value=cfg.env.stop_value,
+        env=evaluator_env,
+        policy=policy.eval_mode,
+        tb_logger=tb_logger,
+        exp_name=cfg.exp_name,
+        policy_config=policy_config
+    )
+    vsbot_evaluator = GoBiggerMuZeroEvaluator(
+        eval_freq=cfg.policy.eval_freq,
+        n_evaluator_episode=cfg.env.n_evaluator_episode,
+        stop_value=cfg.env.stop_value,
+        env=vsbot_evaluator_env,
+        policy=policy.eval_mode,
+        tb_logger=tb_logger,
+        exp_name=cfg.exp_name,
+        policy_config=policy_config,
+        instance_name='vsbot_evaluator'
+    )
+    # ==============================================================
+    # Main loop
+    # ==============================================================
+    # Learner's before_run hook.
+    learner.call_hook('before_run')
+    # ==============================================================
+    # eval trained model
+    # ==============================================================
+    _, reward_sp = evaluator.eval(learner.save_checkpoint, learner.train_iter)
+    _, reward_vsbot = vsbot_evaluator.eval_vsbot(learner.save_checkpoint, learner.train_iter)
+    return reward_sp, reward_vsbot

lzero/entry/train_muzero_gobigger.py

@@ -0,0 +1,196 @@
+import logging
+import os
+from functools import partial
+from typing import Optional, Tuple
+
+import torch
+from ding.config import compile_config
+from ding.envs import create_env_manager
+from ding.envs import get_vec_env_setting
+from ding.policy import create_policy
+from ding.utils import set_pkg_seed
+from ding.worker import BaseLearner
+from tensorboardX import SummaryWriter
+import copy
+from ding.rl_utils import get_epsilon_greedy_fn
+from lzero.entry.utils import log_buffer_memory_usage, random_collect
+from lzero.policy import visit_count_temperature
+from lzero.worker import GoBiggerMuZeroCollector, GoBiggerMuZeroEvaluator
+
+
+def train_muzero_gobigger(
+        input_cfg: Tuple[dict, dict],
+        seed: int = 0,
+        model: Optional[torch.nn.Module] = None,
+        model_path: Optional[str] = None,
+        max_train_iter: Optional[int] = int(1e10),
+        max_env_step: Optional[int] = int(1e10),
+) -> 'Policy':  # noqa
+    """
+    Overview:
+        The train entry for GoBigger MCTS+RL algorithms, including MuZero, EfficientZero, Sampled EfficientZero.
+    Arguments:
+        - input_cfg (:obj:`Tuple[dict, dict]`): Config in dict type.
+            ``Tuple[dict, dict]`` type means [user_config, create_cfg].
+        - seed (:obj:`int`): Random seed.
+        - model (:obj:`Optional[torch.nn.Module]`): Instance of torch.nn.Module.
+        - model_path (:obj:`Optional[str]`): The pretrained model path, which should
+            point to the ckpt file of the pretrained model, and an absolute path is recommended.
+            In LightZero, the path is usually something like ``exp_name/ckpt/ckpt_best.pth.tar``.
+        - max_train_iter (:obj:`Optional[int]`): Maximum policy update iterations in training.
+        - max_env_step (:obj:`Optional[int]`): Maximum collected environment interaction steps.
+    Returns:
+        - policy (:obj:`Policy`): Converged policy.
+    """
+
+    cfg, create_cfg = input_cfg
+    assert create_cfg.policy.type in ['gobigger_efficientzero', 'gobigger_muzero', 'gobigger_sampled_efficientzero'], \
+        "train_muzero entry now only support the following algo.: 'gobigger_efficientzero', 'gobigger_muzero', 'gobigger_sampled_efficientzero'"
+
+    if create_cfg.policy.type == 'gobigger_efficientzero':
+        from lzero.mcts import GoBiggerEfficientZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'gobigger_muzero':
+        from lzero.mcts import GoBiggerMuZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'gobigger_sampled_efficientzero':
+        from lzero.mcts import GoBiggerSampledEfficientZeroGameBuffer as GameBuffer
+
+    if cfg.policy.cuda and torch.cuda.is_available():
+        cfg.policy.device = 'cuda'
+    else:
+        cfg.policy.device = 'cpu'
+
+    cfg = compile_config(cfg, seed=seed, env=None, auto=True, create_cfg=create_cfg, save_cfg=True)
+    # Create main components: env, policy
+    env_fn, collector_env_cfg, evaluator_env_cfg = get_vec_env_setting(cfg.env)
+    collector_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in collector_env_cfg])
+    evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in evaluator_env_cfg])
+
+    env_cfg = copy.deepcopy(evaluator_env_cfg[0])
+    env_cfg.contain_raw_obs = True
+    vsbot_evaluator_env_cfg = [env_cfg for _ in range(len(evaluator_env_cfg))]
+    vsbot_evaluator_env = create_env_manager(cfg.env.manager, [partial(env_fn, cfg=c) for c in vsbot_evaluator_env_cfg])
+
+    collector_env.seed(cfg.seed)
+    evaluator_env.seed(cfg.seed, dynamic_seed=False)
+    vsbot_evaluator_env.seed(cfg.seed, dynamic_seed=False)
+    set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)
+
+    policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval'])
+
+    # load pretrained model
+    if model_path is not None:
+        policy.learn_mode.load_state_dict(torch.load(model_path, map_location=cfg.policy.device))
+
+    # Create worker components: learner, collector, evaluator, replay buffer, commander.
+    tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial'))
+    learner = BaseLearner(cfg.policy.learn.learner, policy.learn_mode, tb_logger, exp_name=cfg.exp_name)
+
+    # ==============================================================
+    # MCTS+RL algorithms related core code
+    # ==============================================================
+    policy_config = cfg.policy
+    batch_size = policy_config.batch_size
+    # specific game buffer for MCTS+RL algorithms
+    replay_buffer = GameBuffer(policy_config)
+    collector = GoBiggerMuZeroCollector(
+        collect_print_freq=cfg.collect.collector.collect_print_freq,
+        env=collector_env,
+        policy=policy.collect_mode,
+        tb_logger=tb_logger,
+        exp_name=cfg.exp_name,
+        policy_config=policy_config
+    )
+    evaluator = GoBiggerMuZeroEvaluator(
+        eval_freq=cfg.policy.eval_freq,
+        n_evaluator_episode=cfg.env.n_evaluator_episode,
+        stop_value=cfg.env.stop_value,
+        env=evaluator_env,
+        policy=policy.eval_mode,
+        tb_logger=tb_logger,
+        exp_name=cfg.exp_name,
+        policy_config=policy_config
+    )
+
+    vsbot_evaluator = GoBiggerMuZeroEvaluator(
+        eval_freq=cfg.policy.eval_freq,
+        n_evaluator_episode=cfg.env.n_evaluator_episode,
+        stop_value=cfg.env.stop_value,
+        env=vsbot_evaluator_env,
+        policy=policy.eval_mode,
+        tb_logger=tb_logger,
+        exp_name=cfg.exp_name,
+        policy_config=policy_config,
+        instance_name='vsbot_evaluator'
+    )
+
+    # ==============================================================
+    # Main loop
+    # ==============================================================
+    # Learner's before_run hook.
+    learner.call_hook('before_run')
+    if cfg.policy.random_collect_episode_num > 0:
+        random_collect(cfg.policy, policy, collector, collector_env, replay_buffer)
+        # reset the random_collect_episode_num to 0
+        cfg.policy.random_collect_episode_num = 0
+
+    while True:
+        log_buffer_memory_usage(learner.train_iter, replay_buffer, tb_logger)
+        collect_kwargs = {}
+        # set temperature for visit count distributions according to the train_iter,
+        # please refer to Appendix D in MuZero paper for details.
+        collect_kwargs['temperature'] = visit_count_temperature(
+            policy_config.manual_temperature_decay,
+            policy_config.fixed_temperature_value,
+            policy_config.threshold_training_steps_for_final_temperature,
+            trained_steps=learner.train_iter
+        )
+        if policy_config.eps.eps_greedy_exploration_in_collect:
+            epsilon_greedy_fn = get_epsilon_greedy_fn(
+                start=policy_config.eps.start,
+                end=policy_config.eps.end,
+                decay=policy_config.eps.decay,
+                type_=policy_config.eps.type
+            )
+            collect_kwargs['epsilon'] = epsilon_greedy_fn(collector.envstep)
+        else:
+            collect_kwargs['epsilon'] = 0.0
+        # Evaluate policy performance.
+        if evaluator.should_eval(learner.train_iter):
+            stop, reward = evaluator.eval(learner.save_checkpoint, learner.train_iter, collector.envstep)
+            stop, reward = vsbot_evaluator.eval_vsbot(learner.save_checkpoint, learner.train_iter, collector.envstep)
+            if stop:
+                break
+
+        # Collect data by default config n_sample/n_episode.
+        new_data = collector.collect(train_iter=learner.train_iter, policy_kwargs=collect_kwargs)
+        # save returned new_data collected by the collector
+        replay_buffer.push_game_segments(new_data)
+        # remove the oldest data if the replay buffer is full.
+        replay_buffer.remove_oldest_data_to_fit()
+
+        # Learn policy from collected data.
+        for i in range(cfg.policy.update_per_collect):
+            # Learner will train ``update_per_collect`` times in one iteration.
+            if replay_buffer.get_num_of_transitions() > batch_size:
+                train_data = replay_buffer.sample(batch_size, policy)
+            else:
+                logging.warning(
+                    f'The data in replay_buffer is not sufficient to sample a mini-batch: '
+                    f'batch_size: {batch_size}, '
+                    f'{replay_buffer} '
+                    f'continue to collect now ....'
+                )
+                break
+
+            # The core train steps for MCTS+RL algorithms.
+            log_vars = learner.train(train_data, collector.envstep)
+
+            if cfg.policy.use_priority:
+                replay_buffer.update_priority(train_data, log_vars[0]['value_priority_orig'])
+
+        if collector.envstep >= max_env_step or learner.train_iter >= max_train_iter:
+            break
+
+    # Learner's after_run hook.
+    learner.call_hook('after_run')
+    return policy

lzero/entry/utils.py

@@ -1,9 +1,48 @@
 import os
+from typing import Optional, Callable
 
 import psutil
+from easydict import EasyDict
 from pympler.asizeof import asizeof
 from tensorboardX import SummaryWriter
 
+from lzero.policy.gobigger_random_policy import GoBiggerRandomPolicy
+
+
+def random_collect(
+        policy_cfg: 'EasyDict',  # noqa
+        policy: 'Policy',  # noqa
+        collector: 'ISerialCollector',  # noqa
+        collector_env: 'BaseEnvManager',  # noqa
+        replay_buffer: 'IBuffer',  # noqa
+        postprocess_data_fn: Optional[Callable] = None
+) -> None:  # noqa
+    assert policy_cfg.random_collect_episode_num > 0
+
+    random_policy = GoBiggerRandomPolicy(cfg=policy_cfg)
+    # set the policy to random policy
+    collector.reset_policy(random_policy.collect_mode)
+
+    collect_kwargs = {}
+    # set temperature for visit count distributions according to the train_iter,
+    # please refer to Appendix D in MuZero paper for details.
+    collect_kwargs['temperature'] = 1
+    collect_kwargs['epsilon'] = 0.0
+
+    # Collect data by default config n_sample/n_episode.
+    new_data = collector.collect(train_iter=0, policy_kwargs=collect_kwargs)
+
+    if postprocess_data_fn is not None:
+        new_data = postprocess_data_fn(new_data)
+
+    # save returned new_data collected by the collector
+    replay_buffer.push_game_segments(new_data)
+    # remove the oldest data if the replay buffer is full.
+    replay_buffer.remove_oldest_data_to_fit()
+
+    # restore the policy
+    collector.reset_policy(policy.collect_mode)
+
 
 def log_buffer_memory_usage(train_iter: int, buffer: "GameBuffer", writer: SummaryWriter) -> None:
     """

lzero/mcts/buffer/__init__.py

@@ -1,4 +1,7 @@
 from .game_buffer_muzero import MuZeroGameBuffer
 from .game_buffer_efficientzero import EfficientZeroGameBuffer
 from .game_buffer_sampled_efficientzero import SampledEfficientZeroGameBuffer
+from .gobigger_game_buffer_muzero import GoBiggerMuZeroGameBuffer
+from .gobigger_game_buffer_efficientzero import GoBiggerEfficientZeroGameBuffer
+from .gobigger_game_buffer_sampled_efficientzero import GoBiggerSampledEfficientZeroGameBuffer
 from .game_buffer_gumbel_muzero import GumbelMuZeroGameBuffer