feature(yzj): ptz simple mz cfg is ready and add ptz simple ez cfg

jayyoung0802 · jayyoung0802 · commit 0e6dfd32a518 · 2023-10-23T15:12:25.000+08:00
diff --git a/lzero/mcts/buffer/game_buffer_efficientzero.py b/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
+from ding.utils.data import default_collate
 
 
 @BUFFER_REGISTRY.register('game_buffer_efficientzero')
@@ -100,7 +102,15 @@ 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()
+        # zero_obs = game_segment_list[0].zero_obs()
+        # zero_obs = np.array([{'agent_state': np.zeros((3, 18), dtype=np.float32),
+        #         'global_state': np.zeros((84,), dtype=np.float32),
+        #         'agent_alone_state': np.zeros((3, 14), dtype=np.float32),
+        #         'agent_alone_padding_state': np.zeros((3, 18), dtype=np.float32),}])
+        zero_obs = np.array([{'agent_state': np.zeros((1, 6), dtype=np.float32),
+                'global_state': np.zeros((14, ), dtype=np.float32),
+                'agent_alone_state': np.zeros((1, 12), dtype=np.float32),
+                'agent_alone_padding_state': np.zeros((1, 12), dtype=np.float32),}])
         value_obs_list = []
         # the value is valid or not (out of trajectory)
         value_mask = []
@@ -152,7 +162,7 @@ def _prepare_reward_value_context(
                     value_mask.append(0)
                     obs = zero_obs
 
-                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 +206,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)
diff --git a/lzero/mcts/buffer/game_buffer_muzero.py b/lzero/mcts/buffer/game_buffer_muzero.py
@@ -202,11 +202,11 @@ def _prepare_reward_value_context(
         """
         zero_obs = game_segment_list[0].zero_obs()
         # zero_obs = np.array([{'agent_state': np.zeros((3, 18), dtype=np.float32),
-        #              'global_state': np.zeros((30,), dtype=np.float32),
+        #              'global_state': np.zeros((84,), dtype=np.float32),
         #              'agent_alone_state': np.zeros((3, 14), dtype=np.float32),
         #              'agent_alone_padding_state': np.zeros((3, 18), dtype=np.float32),}])
         zero_obs = np.array([{'agent_state': np.zeros((1, 6), dtype=np.float32),
-                'global_state': np.zeros((1, 14), dtype=np.float32),
+                'global_state': np.zeros((14, ), dtype=np.float32),
                 'agent_alone_state': np.zeros((1, 12), dtype=np.float32),
                 'agent_alone_padding_state': np.zeros((1, 12), dtype=np.float32),}])
         value_obs_list = []
diff --git a/lzero/model/efficientzero_model_mlp.py b/lzero/model/efficientzero_model_mlp.py
@@ -8,6 +8,7 @@
 
 from .common import EZNetworkOutput, RepresentationNetworkMLP, PredictionNetworkMLP
 from .utils import renormalize, get_params_mean, get_dynamic_mean, get_reward_mean
+from ding.utils.default_helper import get_shape0
 
 
 @MODEL_REGISTRY.register('EfficientZeroModelMLP')
@@ -36,6 +37,9 @@ def __init__(
         norm_type: Optional[str] = 'BN',
         discrete_action_encoding_type: str = 'one_hot',
         res_connection_in_dynamics: bool = False,
+        state_encoder=None,
+        state_prediction=None,
+        state_dynamics=None,
         *args,
         **kwargs,
     ):
@@ -104,31 +108,40 @@ 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
-        )
-
-        self.dynamics_network = DynamicsNetworkMLP(
-            action_encoding_dim=self.action_encoding_dim,
-            num_channels=latent_state_dim + self.action_encoding_dim,
-            common_layer_num=2,
-            lstm_hidden_size=lstm_hidden_size,
-            fc_reward_layers=fc_reward_layers,
-            output_support_size=self.reward_support_size,
-            last_linear_layer_init_zero=self.last_linear_layer_init_zero,
-            norm_type=norm_type,
-            res_connection_in_dynamics=self.res_connection_in_dynamics,
-        )
-
-        self.prediction_network = PredictionNetworkMLP(
-            action_space_size=action_space_size,
-            num_channels=latent_state_dim,
-            fc_value_layers=fc_value_layers,
-            fc_policy_layers=fc_policy_layers,
-            output_support_size=self.value_support_size,
-            last_linear_layer_init_zero=self.last_linear_layer_init_zero,
-            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
+
+        if state_dynamics == None:
+            self.dynamics_network = DynamicsNetworkMLP(
+                action_encoding_dim=self.action_encoding_dim,
+                num_channels=latent_state_dim + self.action_encoding_dim,
+                common_layer_num=2,
+                lstm_hidden_size=lstm_hidden_size,
+                fc_reward_layers=fc_reward_layers,
+                output_support_size=self.reward_support_size,
+                last_linear_layer_init_zero=self.last_linear_layer_init_zero,
+                norm_type=norm_type,
+                res_connection_in_dynamics=self.res_connection_in_dynamics,
+            )
+        else:
+            self.dynamics_network = state_dynamics
+
+        if state_prediction == None:
+            self.prediction_network = PredictionNetworkMLP(
+                action_space_size=action_space_size,
+                num_channels=latent_state_dim,
+                fc_value_layers=fc_value_layers,
+                fc_policy_layers=fc_policy_layers,
+                output_support_size=self.value_support_size,
+                last_linear_layer_init_zero=self.last_linear_layer_init_zero,
+                norm_type=norm_type
+            )
+        else:
+            self.prediction_network = state_prediction
 
         if self.self_supervised_learning_loss:
             # self_supervised_learning_loss related network proposed in EfficientZero
@@ -171,15 +184,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)
 
diff --git a/lzero/policy/efficientzero.py b/lzero/policy/efficientzero.py
@@ -18,6 +18,8 @@
     prepare_obs, \
     configure_optimizers
 from lzero.policy.muzero import MuZeroPolicy
+from ding.utils.data import default_collate
+from ding.torch_utils import to_device, to_tensor
 
 
 @POLICY_REGISTRY.register('efficientzero')
@@ -307,7 +309,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 obs_batch.size(0) == 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.
@@ -562,7 +564,13 @@ def _forward_collect(
         self._collect_model.eval()
         self._collect_mcts_temperature = temperature
         self.collect_epsilon = epsilon
-        active_collect_env_num = data.shape[0]
+        active_collect_env_num = len(data)
+        # 
+        data = sum(data, [])
+        data = default_collate(data)
+        data = to_device(data, self._device)
+        to_play = np.array(to_play).reshape(-1).tolist()
+
         with torch.no_grad():
             # data shape [B, S x C, W, H], e.g. {Tensor:(B, 12, 96, 96)}
             network_output = self._collect_model.initial_inference(data)
@@ -667,7 +675,12 @@ def _forward_eval(self, data: torch.Tensor, action_mask: list, to_play: -1, read
                  ``visit_count_distribution_entropy``, ``value``, ``pred_value``, ``policy_logits``.
          """
         self._eval_model.eval()
-        active_eval_env_num = data.shape[0]
+        active_eval_env_num = len(data)
+        #
+        data = sum(data, [])
+        data = default_collate(data)
+        data = to_device(data, self._device)
+        to_play = np.array(to_play).reshape(-1).tolist()
         with torch.no_grad():
             # data shape [B, S x C, W, H], e.g. {Tensor:(B, 12, 96, 96)}
             network_output = self._eval_model.initial_inference(data)
diff --git a/zoo/petting_zoo/config/ptz_simple_ez_config.py b/zoo/petting_zoo/config/ptz_simple_ez_config.py
@@ -0,0 +1,115 @@
+from easydict import EasyDict
+
+env_name = 'ptz_simple'
+multi_agent = True
+
+# ==============================================================
+# begin of the most frequently changed config specified by the user
+# ==============================================================
+seed = 0
+n_agent = 1
+n_landmark = n_agent
+collector_env_num = 8
+evaluator_env_num = 8
+n_episode = 8
+batch_size = 256
+num_simulations = 50
+update_per_collect = 50
+reanalyze_ratio = 0.
+action_space_size = 5
+eps_greedy_exploration_in_collect = True
+# ==============================================================
+# end of the most frequently changed config specified by the user
+# ==============================================================
+
+main_config = dict(
+    exp_name=
+    f'data_ez_ctree/{env_name}_efficientzero_ns{num_simulations}_upc{update_per_collect}_rr{reanalyze_ratio}_seed{seed}',
+    env=dict(
+        env_family='mpe',
+        env_id='simple_v2',
+        n_agent=n_agent,
+        n_landmark=n_landmark,
+        max_cycles=25,
+        agent_obs_only=False,
+        agent_specific_global_state=True,
+        continuous_actions=False,
+        stop_value=0,
+        collector_env_num=collector_env_num,
+        evaluator_env_num=evaluator_env_num,
+        n_evaluator_episode=evaluator_env_num,
+        manager=dict(shared_memory=False, ),
+    ),
+    policy=dict(
+        multi_agent=multi_agent,
+        ignore_done=False,
+        model=dict(
+            model_type='structure',
+            latent_state_dim=256,
+            action_space='discrete',
+            action_space_size=action_space_size,
+            agent_num=n_agent,
+            self_supervised_learning_loss=False,  # default is False
+            agent_obs_shape=6,
+            global_obs_shape=14,
+            discrete_action_encoding_type='one_hot',
+            global_cooperation=True, # TODO: doesn't work now
+            hidden_size_list=[256, 256],
+            norm_type='BN',
+        ),
+        cuda=True,
+        mcts_ctree=True,
+        gumbel_algo=False,
+        env_type='not_board_games',
+        game_segment_length=30,
+        random_collect_episode_num=0,
+        eps=dict(
+            eps_greedy_exploration_in_collect=eps_greedy_exploration_in_collect,
+            type='linear',
+            start=1.,
+            end=0.05,
+            decay=int(2e5),
+        ),
+        use_augmentation=False,
+        update_per_collect=update_per_collect,
+        batch_size=batch_size,
+        optim_type='SGD',
+        lr_piecewise_constant_decay=True,
+        learning_rate=0.2,
+        ssl_loss_weight=0,  # NOTE: default is 0.
+        num_simulations=num_simulations,
+        reanalyze_ratio=reanalyze_ratio,
+        n_episode=n_episode,
+        eval_freq=int(2e3),
+        replay_buffer_size=int(1e6),  # the size/capacity of replay_buffer, in the terms of transitions.
+        collector_env_num=collector_env_num,
+        evaluator_env_num=evaluator_env_num,
+    ),
+    learn=dict(learner=dict(
+        log_policy=True,
+        hook=dict(log_show_after_iter=10, ),
+    ), ),
+)
+main_config = EasyDict(main_config)
+create_config = dict(
+    env=dict(
+        import_names=['zoo.petting_zoo.envs.petting_zoo_simple_spread_env'],
+        type='petting_zoo',
+    ),
+    env_manager=dict(type='subprocess'),
+    policy=dict(
+        type='efficientzero',
+        import_names=['lzero.policy.efficientzero'],
+    ),
+    collector=dict(
+        type='episode_muzero',
+        import_names=['lzero.worker.muzero_collector'],
+    )
+)
+create_config = EasyDict(create_config)
+ptz_simple_spread_efficientzero_config = main_config
+ptz_simple_spread_efficientzero_create_config = create_config
+
+if __name__ == "__main__":
+    from zoo.petting_zoo.entry import train_muzero
+    train_muzero([main_config, create_config], seed=seed)
diff --git a/zoo/petting_zoo/config/ptz_simple_mz_config.py b/zoo/petting_zoo/config/ptz_simple_mz_config.py
@@ -46,7 +46,6 @@
         model=dict(
             model_type='structure',
             latent_state_dim=256,
-            frame_stack_num=1,
             action_space='discrete',
             action_space_size=action_space_size,
             agent_num=n_agent,
@@ -69,7 +68,7 @@
             type='linear',
             start=1.,
             end=0.05,
-            decay=int(1e5),
+            decay=int(2e5),
         ),
         use_augmentation=False,
         update_per_collect=update_per_collect,
@@ -111,6 +110,6 @@
 ptz_simple_spread_muzero_config = main_config
 ptz_simple_spread_muzero_create_config = create_config
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     from zoo.petting_zoo.entry import train_muzero
     train_muzero([main_config, create_config], seed=seed)
diff --git a/zoo/petting_zoo/config/ptz_simple_spread_mz_config.py b/zoo/petting_zoo/config/ptz_simple_spread_mz_config.py
@@ -13,7 +13,7 @@
 evaluator_env_num = 8
 n_episode = 8
 batch_size = 256
-num_simulations = 50
+num_simulations = 200
 update_per_collect = 50
 reanalyze_ratio = 0.
 action_space_size = 5*5*5
@@ -32,7 +32,7 @@
         n_landmark=n_landmark,
         max_cycles=25,
         agent_obs_only=False,
-        agent_specific_global_state=False,
+        agent_specific_global_state=True,
         continuous_actions=False,
         stop_value=0,
         collector_env_num=collector_env_num,
@@ -52,7 +52,7 @@
             agent_num=n_agent,
             self_supervised_learning_loss=False,  # default is False
             agent_obs_shape=18,
-            global_obs_shape=30,
+            global_obs_shape=18*n_agent+30, # 84
             discrete_action_encoding_type='one_hot',
             global_cooperation=True, # TODO: doesn't work now
             hidden_size_list=[256, 256],
@@ -69,7 +69,7 @@
             type='linear',
             start=1.,
             end=0.05,
-            decay=int(1e5),
+            decay=int(2e5),
         ),
         use_augmentation=False,
         update_per_collect=update_per_collect,
diff --git a/zoo/petting_zoo/entry/train_muzero.py b/zoo/petting_zoo/entry/train_muzero.py
@@ -56,6 +56,7 @@ def train_muzero(
         from lzero.model.muzero_model_mlp import MuZeroModelMLP as Encoder
     elif create_cfg.policy.type == 'efficientzero':
         from lzero.mcts import EfficientZeroGameBuffer as GameBuffer
+        from lzero.model.efficientzero_model_mlp import EfficientZeroModelMLP as Encoder
     elif create_cfg.policy.type == 'sampled_efficientzero':
         from lzero.mcts import SampledEfficientZeroGameBuffer as GameBuffer
     elif create_cfg.policy.type == 'gumbel_muzero':
diff --git a/zoo/petting_zoo/envs/petting_zoo_simple_spread_env.py b/zoo/petting_zoo/envs/petting_zoo_simple_spread_env.py
diff --git a/zoo/petting_zoo/model/model.py b/zoo/petting_zoo/model/model.py
