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

lzero/entry/train_muzero.py

@@ -50,9 +50,9 @@ def train_muzero(
     assert create_cfg.policy.type in ['efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero', 'stochastic_muzero'], \
         "train_muzero entry now only support the following algo.: 'efficientzero', 'muzero', 'sampled_efficientzero', 'gumbel_muzero'"
 
-    if create_cfg.policy.type == 'muzero':
+    if create_cfg.policy.type == 'muzero' or create_cfg.policy.type == 'multi_agent_muzero':
         from lzero.mcts import MuZeroGameBuffer as GameBuffer
-    elif create_cfg.policy.type == 'efficientzero':
+    elif create_cfg.policy.type == 'efficientzero' or create_cfg.policy.type == 'multi_agent_efficientzero':
         from lzero.mcts import EfficientZeroGameBuffer as GameBuffer
     elif create_cfg.policy.type == 'sampled_efficientzero':
         from lzero.mcts import SampledEfficientZeroGameBuffer as GameBuffer
@@ -125,7 +125,11 @@ def train_muzero(
     # Exploration: Collecting random data helps the agent explore the environment and avoid getting stuck in a suboptimal policy prematurely.
     # Comparison: By observing the agent's performance during random action-taking, we can establish a baseline to evaluate the effectiveness 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)
@@ -192,4 +196,4 @@ def train_muzero(
 
     # Learner's after_run hook.
     learner.call_hook('after_run')
-    return policy
+    return policy

lzero/mcts/buffer/game_buffer_efficientzero.py

@@ -9,6 +9,8 @@
 from lzero.mcts.utils import prepare_observation
 from lzero.policy import to_detach_cpu_numpy, concat_output, concat_output_value, inverse_scalar_transform
 from .game_buffer_muzero import MuZeroGameBuffer
+from ding.torch_utils import to_device, to_tensor, to_ndarray
+from ding.utils.data import default_collate
 
 
 @BUFFER_REGISTRY.register('game_buffer_efficientzero')
@@ -44,6 +46,8 @@ def __init__(self, cfg: dict):
         self.base_idx = 0
         self.clear_time = 0
 
+        self.tmp_obs = None # since value obs list [46 + 4(td_step)] >= 50(game_segment), need pad
+
     def sample(self, batch_size: int, policy: Any) -> List[Any]:
         """
         Overview:
@@ -100,7 +104,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 +151,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,
@@ -196,7 +200,13 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A
                 beg_index = self._cfg.mini_infer_size * i
                 end_index = self._cfg.mini_infer_size * (i + 1)
 
-                m_obs = torch.from_numpy(value_obs_list[beg_index:end_index]).to(self._cfg.device).float()
+                if self._cfg.model.model_type and self._cfg.model.model_type in ['conv', 'mlp']:
+                    m_obs = torch.from_numpy(value_obs_list[beg_index:end_index]).to(self._cfg.device).float()
+                elif self._cfg.model.model_type and self._cfg.model.model_type == 'structure':
+                    m_obs = value_obs_list[beg_index:end_index]
+                    m_obs = sum(m_obs, [])
+                    m_obs = default_collate(m_obs)
+                    m_obs = to_device(m_obs, self._cfg.device)
 
                 # calculate the target value
                 m_output = model.initial_inference(m_obs)

lzero/mcts/buffer/game_buffer_muzero.py

@@ -9,6 +9,8 @@
 from lzero.mcts.utils import prepare_observation
 from lzero.policy import to_detach_cpu_numpy, concat_output, concat_output_value, inverse_scalar_transform
 from .game_buffer import GameBuffer
+from ding.torch_utils import to_device, to_tensor
+from ding.utils.data import default_collate
 
 if TYPE_CHECKING:
     from lzero.policy import MuZeroPolicy, EfficientZeroPolicy, SampledEfficientZeroPolicy
@@ -48,6 +50,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 +202,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 +241,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,
@@ -376,8 +380,14 @@ def _compute_target_reward_value(self, reward_value_context: List[Any], model: A
             for i in range(slices):
                 beg_index = self._cfg.mini_infer_size * i
                 end_index = self._cfg.mini_infer_size * (i + 1)
-
-                m_obs = torch.from_numpy(value_obs_list[beg_index:end_index]).to(self._cfg.device).float()
+
+                if self._cfg.model.model_type and self._cfg.model.model_type in ['conv', 'mlp']:
+                    m_obs = torch.from_numpy(value_obs_list[beg_index:end_index]).to(self._cfg.device).float()
+                elif self._cfg.model.model_type and self._cfg.model.model_type == 'structure':
+                    m_obs = value_obs_list[beg_index:end_index]
+                    m_obs = sum(m_obs, [])
+                    m_obs = default_collate(m_obs)
+                    m_obs = to_device(m_obs, self._cfg.device)
 
                 # calculate the target value
                 m_output = model.initial_inference(m_obs)

lzero/mcts/utils.py

@@ -92,7 +92,7 @@ def prepare_observation(observation_list, model_type='conv'):
         - observation_list (:obj:`List`): list of observations.
         - model_type (:obj:`str`): type of the model. (default is 'conv')
     """
-    assert model_type in ['conv', 'mlp']
+    assert model_type in ['conv', 'mlp', 'structure']
     observation_array = np.array(observation_list)
 
     if model_type == 'conv':
@@ -127,6 +127,9 @@ def prepare_observation(observation_list, model_type='conv'):
         observation_array = observation_array.reshape(observation_array.shape[0], -1)
         # print(observation_array.shape)
 
+    elif model_type == 'structure':
+        return observation_list
+
     return observation_array
 
 

lzero/model/efficientzero_model_mlp.py

@@ -4,6 +4,8 @@
 import torch.nn as nn
 from ding.torch_utils import MLP
 from ding.utils import MODEL_REGISTRY, SequenceType
+from ding.utils.default_helper import get_shape0
+
 from numpy import ndarray
 
 from .common import EZNetworkOutput, RepresentationNetworkMLP, PredictionNetworkMLP
@@ -36,6 +38,7 @@ def __init__(
         norm_type: Optional[str] = 'BN',
         discrete_action_encoding_type: str = 'one_hot',
         res_connection_in_dynamics: bool = False,
+        state_encoder=None,
         *args,
         **kwargs,
     ):
@@ -104,9 +107,12 @@ def __init__(
         self.state_norm = state_norm
         self.res_connection_in_dynamics = res_connection_in_dynamics
 
-        self.representation_network = RepresentationNetworkMLP(
-            observation_shape=observation_shape, hidden_channels=latent_state_dim, norm_type=norm_type
-        )
+        if state_encoder == None:
+            self.representation_network = RepresentationNetworkMLP(
+                observation_shape=observation_shape, hidden_channels=latent_state_dim, norm_type=norm_type
+            )
+        else:
+            self.representation_network = state_encoder
 
         self.dynamics_network = DynamicsNetworkMLP(
             action_encoding_dim=self.action_encoding_dim,
@@ -171,15 +177,16 @@ def initial_inference(self, obs: torch.Tensor) -> EZNetworkOutput:
             - latent_state (:obj:`torch.Tensor`): :math:`(B, H)`, where B is batch_size, H is the dimension of latent state.
             - reward_hidden_state (:obj:`Tuple[torch.Tensor]`): The shape of each element is :math:`(1, B, lstm_hidden_size)`, where B is batch_size.
         """
-        batch_size = obs.size(0)
+        batch_size = get_shape0(obs)
         latent_state = self._representation(obs)
+        device = latent_state.device
         policy_logits, value = self._prediction(latent_state)
         # zero initialization for reward hidden states
         # (hn, cn), each element shape is (layer_num=1, batch_size, lstm_hidden_size)
         reward_hidden_state = (
             torch.zeros(1, batch_size,
-                        self.lstm_hidden_size).to(obs.device), torch.zeros(1, batch_size,
-                                                                           self.lstm_hidden_size).to(obs.device)
+                        self.lstm_hidden_size).to(device), torch.zeros(1, batch_size,
+                                                                           self.lstm_hidden_size).to(device)
         )
         return EZNetworkOutput(value, [0. for _ in range(batch_size)], policy_logits, latent_state, reward_hidden_state)
 

lzero/model/muzero_model_mlp.py

@@ -4,6 +4,8 @@
 import torch.nn as nn
 from ding.torch_utils import MLP
 from ding.utils import MODEL_REGISTRY, SequenceType
+from ding.utils.default_helper import get_shape0
+
 
 from .common import MZNetworkOutput, RepresentationNetworkMLP, PredictionNetworkMLP
 from .utils import renormalize, get_params_mean, get_dynamic_mean, get_reward_mean
@@ -34,6 +36,7 @@ def __init__(
         discrete_action_encoding_type: str = 'one_hot',
         norm_type: Optional[str] = 'BN',
         res_connection_in_dynamics: bool = False,
+        state_encoder=None,
         *args,
         **kwargs
     ):
@@ -66,6 +69,7 @@ def __init__(
             - discrete_action_encoding_type (:obj:`str`): The encoding type of discrete action, which can be 'one_hot' or 'not_one_hot'.
             - norm_type (:obj:`str`): The type of normalization in networks. defaults to 'BN'.
             - res_connection_in_dynamics (:obj:`bool`): Whether to use residual connection for dynamics network, default set it to False.
+            - state_encoder (:obj:`Optional[nn.Module]`): The state encoder network, which is used to encode the raw observation to latent state.
         """
         super(MuZeroModelMLP, self).__init__()
         self.categorical_distribution = categorical_distribution
@@ -101,9 +105,12 @@ def __init__(
         self.state_norm = state_norm
         self.res_connection_in_dynamics = res_connection_in_dynamics
 
-        self.representation_network = RepresentationNetworkMLP(
-            observation_shape=observation_shape, hidden_channels=self.latent_state_dim, norm_type=norm_type
-        )
+        if state_encoder == None:
+            self.representation_network = RepresentationNetworkMLP(
+                observation_shape=observation_shape, hidden_channels=latent_state_dim, norm_type=norm_type
+            )
+        else:
+            self.representation_network = state_encoder
 
         self.dynamics_network = DynamicsNetwork(
             action_encoding_dim=self.action_encoding_dim,
@@ -166,7 +173,7 @@ def initial_inference(self, obs: torch.Tensor) -> MZNetworkOutput:
             - policy_logits (:obj:`torch.Tensor`): :math:`(B, action_dim)`, where B is batch_size.
             - latent_state (:obj:`torch.Tensor`): :math:`(B, H)`, where B is batch_size, H is the dimension of latent state.
         """
-        batch_size = obs.size(0)
+        batch_size = get_shape0(obs)
         latent_state = self._representation(obs)
         policy_logits, value = self._prediction(latent_state)
         return MZNetworkOutput(

lzero/policy/efficientzero.py

@@ -17,7 +17,13 @@
     DiscreteSupport, select_action, to_torch_float_tensor, ez_network_output_unpack, negative_cosine_similarity, \
     prepare_obs, \
     configure_optimizers
+<<<<<<< HEAD
+from collections import defaultdict
+from ding.torch_utils import to_device, to_tensor
+from ding.utils.data import default_collate
+=======
 from lzero.policy.muzero import MuZeroPolicy
+>>>>>>> origin
 
 
 @POLICY_REGISTRY.register('efficientzero')
@@ -191,6 +197,9 @@ class EfficientZeroPolicy(MuZeroPolicy):
             # (int) The decay steps from start to end eps.
             decay=int(1e5),
         ),
+
+        # (bool) Whether it is a multi-agent environment.
+        multi_agent=False,
     )
 
     def default_model(self) -> Tuple[str, List[str]]:
@@ -309,7 +318,7 @@ def _forward_learn(self, data: torch.Tensor) -> Dict[str, Union[float, int]]:
 
         target_value_prefix = target_value_prefix.view(self._cfg.batch_size, -1)
         target_value = target_value.view(self._cfg.batch_size, -1)
-        assert obs_batch.size(0) == self._cfg.batch_size == target_value_prefix.size(0)
+        assert self._cfg.batch_size == target_value_prefix.size(0)
 
         # ``scalar_transform`` to transform the original value to the scaled value,
         # i.e. h(.) function in paper https://arxiv.org/pdf/1805.11593.pdf.
@@ -395,9 +404,40 @@ def _forward_learn(self, data: torch.Tensor) -> Dict[str, Union[float, int]]:
             # calculate consistency loss for the next ``num_unroll_steps`` unroll steps.
             # ==============================================================
             if self._cfg.ssl_loss_weight > 0:
-                # obtain the oracle latent states from representation function.
-                beg_index, end_index = self._get_target_obs_index_in_step_k(step_k)
-                network_output = self._learn_model.initial_inference(obs_target_batch[:, beg_index:end_index])
+                # obtain the oracle hidden states from representation function.
+                if self._cfg.model.model_type == 'conv':
+                    beg_index = self._cfg.model.image_channel * step_i
+                    end_index = self._cfg.model.image_channel * (step_i + self._cfg.model.frame_stack_num)
+                    network_output = self._learn_model.initial_inference(obs_target_batch[:, beg_index:end_index, :, :])
+                elif self._cfg.model.model_type == 'mlp':
+                    beg_index = self._cfg.model.observation_shape * step_i
+                    end_index = self._cfg.model.observation_shape * (step_i + self._cfg.model.frame_stack_num)
+                    network_output = self._learn_model.initial_inference(obs_target_batch[:, beg_index:end_index])
+                elif self._cfg.model.model_type == 'structure':
+                    obs_target_batch_new = {}
+                    for k, v in obs_target_batch.items():
+                        if k == 'action_mask': 
+                            obs_target_batch_new[k] = v
+                            continue
+                        if isinstance(v, dict):
+                            obs_target_batch_new[k] = {}
+                            for k1, v1 in v.items():
+                                if len(v1.shape) == 1:
+                                    observation_shape = v1.shape[0]//self._cfg.num_unroll_steps
+                                    beg_index = observation_shape * step_i
+                                    end_index = observation_shape * (step_i + self._cfg.model.frame_stack_num)
+                                    obs_target_batch_new[k][k1] = v1[beg_index:end_index]
+                                else:
+                                    observation_shape = v1.shape[1]//self._cfg.num_unroll_steps
+                                    beg_index = observation_shape * step_i
+                                    end_index = observation_shape * (step_i + self._cfg.model.frame_stack_num)
+                                    obs_target_batch_new[k][k1] = v1[:, beg_index:end_index]
+                        else:
+                            observation_shape = v.shape[1]//self._cfg.num_unroll_steps
+                            beg_index = observation_shape * step_i
+                            end_index = observation_shape * (step_i + self._cfg.model.frame_stack_num)
+                            obs_target_batch_new[k] = v[:, beg_index:end_index]
+                    network_output = self._learn_model.initial_inference(obs_target_batch_new)
 
                 latent_state = to_tensor(latent_state)
                 representation_state = to_tensor(network_output.latent_state)
@@ -735,6 +775,7 @@ def _monitor_vars_learn(self) -> List[str]:
          """
         return [
             'collect_mcts_temperature',
+            'collect_epsilon',
             'cur_lr',
             'weighted_total_loss',
             'total_loss',