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

lzero/entry/__init__.py

@@ -3,4 +3,4 @@
 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

lzero/entry/train_muzero.py

@@ -15,9 +15,6 @@
 
 from lzero.entry.utils import log_buffer_memory_usage
 from lzero.policy import visit_count_temperature
-from lzero.policy.random_policy import LightZeroRandomPolicy
-from lzero.worker import MuZeroCollector as Collector
-from lzero.worker import MuZeroEvaluator as Evaluator
 from .utils import random_collect
 
 
@@ -47,8 +44,8 @@ def train_muzero(
     """
 
     cfg, create_cfg = input_cfg
-    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero'], \
-        "train_muzero entry now only support the following algo.: 'efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero'"
+    assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero', 'multi_agent_efficientzero', 'multi_agent_muzero'], \
+        "train_muzero entry now only support the following algo.: 'efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero', 'multi_agent_efficientzero', 'multi_agent_muzero'"
 
     if create_cfg.policy.type == 'muzero':
         from lzero.mcts import MuZeroGameBuffer as GameBuffer
@@ -58,6 +55,10 @@ def train_muzero(
         from lzero.mcts import SampledEfficientZeroGameBuffer as GameBuffer
     elif create_cfg.policy.type == 'gumbel_muzero':
         from lzero.mcts import GumbelMuZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'multi_agent_muzero':
+        from lzero.mcts import MultiAgentMuZeroGameBuffer as GameBuffer
+    elif create_cfg.policy.type == 'multi_agent_efficientzero':
+        from lzero.mcts import MultiAgentSampledEfficientZeroGameBuffer as GameBuffer
 
     if cfg.policy.cuda and torch.cuda.is_available():
         cfg.policy.device = 'cuda'
@@ -92,6 +93,14 @@ def train_muzero(
     batch_size = policy_config.batch_size
     # specific game buffer for MCTS+RL algorithms
     replay_buffer = GameBuffer(policy_config)
+
+    if policy_config.multi_agent:
+        from lzero.worker import MultiAgentMuZeroCollector as Collector
+        from lzero.worker import MuZeroEvaluator as Evaluator
+    else:
+        from lzero.worker import MuZeroCollector as Collector
+        from lzero.worker import MuZeroEvaluator as Evaluator
+
     collector = Collector(
         env=collector_env,
         policy=policy.collect_mode,
@@ -123,7 +132,11 @@ def train_muzero(
     # Exploration: The collection of random data aids the agent in exploring the environment and prevents premature convergence to a suboptimal policy.
     # Comparation: The agent's performance during random action-taking can be used as a reference point to evaluate the efficacy of reinforcement learning algorithms.
     if cfg.policy.random_collect_episode_num > 0:
-        random_collect(cfg.policy, policy, LightZeroRandomPolicy, collector, collector_env, replay_buffer)
+        if policy_config.multi_agent:
+            from lzero.policy.multi_agent_random_policy import MultiAgentLightZeroRandomPolicy as RandomPolicy
+        else:
+            from lzero.policy.random_policy import LightZeroRandomPolicy as RandomPolicy
+        random_collect(cfg.policy, policy, RandomPolicy, collector, collector_env, replay_buffer)
 
     while True:
         log_buffer_memory_usage(learner.train_iter, replay_buffer, tb_logger)

lzero/entry/utils.py

@@ -1,6 +1,8 @@
 import os
+from typing import Optional, Callable
 
 import psutil
+from easydict import EasyDict
 from pympler.asizeof import asizeof
 from tensorboardX import SummaryWriter
 from typing import Optional, Callable
@@ -39,7 +41,6 @@ def random_collect(
     # restore the policy
     collector.reset_policy(policy.collect_mode)
 
-
 def log_buffer_memory_usage(train_iter: int, buffer: "GameBuffer", writer: SummaryWriter) -> None:
     """
     Overview:

lzero/mcts/buffer/__init__.py

@@ -2,3 +2,5 @@
 from .game_buffer_efficientzero import EfficientZeroGameBuffer
 from .game_buffer_sampled_efficientzero import SampledEfficientZeroGameBuffer
 from .game_buffer_gumbel_muzero import GumbelMuZeroGameBuffer
+from .multi_agent_game_buffer_muzero import MultiAgentMuZeroGameBuffer
+from .multi_agent_game_buffer_efficientzero import MultiAgentSampledEfficientZeroGameBuffer

lzero/mcts/buffer/game_buffer_efficientzero.py

@@ -44,6 +44,8 @@ def __init__(self, cfg: dict):
         self.base_idx = 0
         self.clear_time = 0
 
+        self.tmp_obs = None # for value obs list [46 + 4(td_step)] not < 50(game_segment)
+
     def sample(self, batch_size: int, policy: Any) -> List[Any]:
         """
         Overview:
@@ -100,7 +102,6 @@ def _prepare_reward_value_context(
             - reward_value_context (:obj:`list`): value_obs_list, value_mask, pos_in_game_segment_list, rewards_list, game_segment_lens,
               td_steps_list, action_mask_segment, to_play_segment
         """
-        zero_obs = game_segment_list[0].zero_obs()
         value_obs_list = []
         # the value is valid or not (out of trajectory)
         value_mask = []
@@ -148,11 +149,12 @@ def _prepare_reward_value_context(
                     end_index = beg_index + self._cfg.model.frame_stack_num
                     # the stacked obs in time t
                     obs = game_obs[beg_index:end_index]
+                    self.tmp_obs = obs  # will be masked
                 else:
                     value_mask.append(0)
-                    obs = zero_obs
+                    obs = self.tmp_obs  # will be masked
 
-                value_obs_list.append(obs)
+                value_obs_list.append(obs.tolist())
 
         reward_value_context = [
             value_obs_list, value_mask, pos_in_game_segment_list, rewards_list, game_segment_lens, td_steps_list,

lzero/mcts/buffer/game_buffer_muzero.py

@@ -48,6 +48,8 @@ def __init__(self, cfg: dict):
         self.game_pos_priorities = []
         self.game_segment_game_pos_look_up = []
 
+        self.tmp_obs = None # a tmp value which records obs when value obs list [current_index + 4(td_step)] > 50(game_segment)
+
     def sample(
             self, batch_size: int, policy: Union["MuZeroPolicy", "EfficientZeroPolicy", "SampledEfficientZeroPolicy"]
     ) -> List[Any]:
@@ -198,7 +200,6 @@ def _prepare_reward_value_context(
             - reward_value_context (:obj:`list`): value_obs_list, value_mask, pos_in_game_segment_list, rewards_list, game_segment_lens,
               td_steps_list, action_mask_segment, to_play_segment
         """
-        zero_obs = game_segment_list[0].zero_obs()
         value_obs_list = []
         # the value is valid or not (out of game_segment)
         value_mask = []
@@ -238,11 +239,12 @@ def _prepare_reward_value_context(
                     end_index = beg_index + self._cfg.model.frame_stack_num
                     # the stacked obs in time t
                     obs = game_obs[beg_index:end_index]
+                    self.tmp_obs = obs  # will be masked
                 else:
                     value_mask.append(0)
-                    obs = zero_obs
+                    obs = self.tmp_obs  # will be masked
 
-                value_obs_list.append(obs)
+                value_obs_list.append(obs.tolist())
 
         reward_value_context = [
             value_obs_list, value_mask, pos_in_game_segment_list, rewards_list, game_segment_lens, td_steps_list,