feature(yzj): add ptz simple env

Author: jayyoung0802

lzero/mcts/buffer/game_buffer_muzero.py

@@ -201,12 +201,12 @@ def _prepare_reward_value_context(
               td_steps_list, action_mask_segment, to_play_segment
         """
         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),
-                     '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((3, 18), dtype=np.float32),
+        #              'global_state': np.zeros((30,), 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((8,), dtype=np.float32),
+                'global_state': np.zeros((1, 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 = []

lzero/model/muzero_model_mlp.py

@@ -186,7 +186,7 @@ def initial_inference(self, obs: torch.Tensor) -> MZNetworkOutput:
             value,
             [0. for _ in range(batch_size)],
             policy_logits,
-            latent_state[1],
+            latent_state,
         )
 
     def recurrent_inference(self, latent_state: torch.Tensor, action: torch.Tensor) -> MZNetworkOutput:
@@ -214,7 +214,7 @@ def recurrent_inference(self, latent_state: torch.Tensor, action: torch.Tensor)
         """
         next_latent_state, reward = self._dynamics(latent_state, action)
         policy_logits, value = self._prediction(next_latent_state)
-        return MZNetworkOutput(value, reward, policy_logits, next_latent_state[1])
+        return MZNetworkOutput(value, reward, policy_logits, next_latent_state)
 
     def _representation(self, observation: torch.Tensor) -> Tuple[torch.Tensor]:
         """

zoo/petting_zoo/config/ptz_simple_mz_config.py

@@ -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=6,
-            global_obs_shape=8,
+            global_obs_shape=14,
             discrete_action_encoding_type='one_hot',
             global_cooperation=True, # TODO: doesn't work now
             hidden_size_list=[256, 256],
@@ -97,7 +97,7 @@
         import_names=['zoo.petting_zoo.envs.petting_zoo_simple_spread_env'],
         type='petting_zoo',
     ),
-    env_manager=dict(type='base'),
+    env_manager=dict(type='subprocess'),
     policy=dict(
         type='muzero',
         import_names=['lzero.policy.muzero'],

zoo/petting_zoo/entry/train_muzero.py

@@ -79,7 +79,8 @@ def train_muzero(
     evaluator_env.seed(cfg.seed, dynamic_seed=False)
     set_pkg_seed(cfg.seed, use_cuda=cfg.policy.cuda)
 
-    model = Encoder(**cfg.policy.model, state_encoder=PettingZooEncoder(cfg), state_prediction=PettingZooPrediction(cfg), state_dynamics=PettingZooDynamics(cfg))
+    # model = Encoder(**cfg.policy.model, state_encoder=PettingZooEncoder(cfg), state_prediction=PettingZooPrediction(cfg), state_dynamics=PettingZooDynamics(cfg))
+    model = Encoder(**cfg.policy.model, state_encoder=PettingZooEncoder(cfg))
     policy = create_policy(cfg.policy, model=model, enable_field=['learn', 'collect', 'eval'])
 
     # load pretrained model

zoo/petting_zoo/model/model.py

@@ -23,7 +23,7 @@ def __init__(self, cfg):
                                                        norm_type='BN')
 
         self.global_encoder = RepresentationNetworkMLP(observation_shape=global_obs_shape, 
-                                                       hidden_channels=128,
+                                                       hidden_channels=256,
                                                        norm_type='BN')
 
         self.encoder = RepresentationNetworkMLP(observation_shape=128+128*self.agent_num, 
@@ -32,15 +32,16 @@ def __init__(self, cfg):
 
     def forward(self, x):
         # agent
-        batch_size, agent_num = x['agent_state'].shape[0], x['agent_state'].shape[1]
-        agent_state = x['agent_state'].reshape(batch_size*agent_num, -1)
-        agent_state = self.agent_encoder(agent_state)
-        agent_state_B = agent_state.reshape(batch_size, -1)
-        agent_state_B_A = agent_state.reshape(batch_size, agent_num, -1)
+        batch_size, agent_num = x['global_state'].shape[0], x['global_state'].shape[1]
+        latent_state = x['global_state'].reshape(batch_size*agent_num, -1)
+        latent_state = self.global_encoder(latent_state)
+        return latent_state
+        # agent_state_B = agent_state.reshape(batch_size, -1)
+        # agent_state_B_A = agent_state.reshape(batch_size, agent_num, -1)
         # global
-        global_state = self.global_encoder(x['global_state'])
-        global_state = self.encoder(torch.cat((agent_state_B, global_state),dim=1))
-        return (agent_state_B, global_state)
+        # global_state = self.global_encoder(x['global_state'])
+        # global_state = self.encoder(torch.cat((agent_state_B, global_state),dim=1))
+        # return (agent_state_B, global_state)
 
 
 class PettingZooPrediction(nn.Module):