polish(xjy): delete unnecessary comments and translate CN comments into EN

Author: xiongjyu

lzero/entry/eval_muzero.py

@@ -60,13 +60,7 @@ def eval_muzero(
 
     # load pretrained model
     if model_path is not None:
-        logging.info(f"Loading pretrained model from {model_path}...")
         policy.learn_mode.load_state_dict(torch.load(model_path, map_location=cfg.policy.device))
-        logging.info("Pretrained model loaded successfully!")
-    else:
-        logging.warning("model_path is None!!!")
-
-    # print(policy._learn_model.representation_network.pretrained_model.encoder.layer[0].attention.output.LayerNorm.weight)
 
     # Create worker components: learner, collector, evaluator, replay buffer, commander.
     tb_logger = SummaryWriter(os.path.join('./{}/log/'.format(cfg.exp_name), 'serial'))

lzero/model/common.py

@@ -390,16 +390,6 @@ def __init__(self,
         # In distributed training, only the rank 0 process downloads the model, and other processes load from cache to speed up startup.
         if get_rank() == 0:
             self.pretrained_model = AutoModel.from_pretrained(model_path)
-        
-        # layernorm_all_ones = True
-        # for name, param in self.pretrained_model.named_parameters():
-        #     if 'LayerNorm.weight' in name:
-        #         # 检查所有值是否接近1。可以用torch.allclose进行判断：
-        #         if not torch.allclose(param, torch.ones_like(param), rtol=1e-05, atol=1e-08):
-        #             print(f"{name} is not all ones!")
-        #             layernorm_all_ones = False
-        # if layernorm_all_ones:
-        #     print("All LayerNorm.weight parameters are all ones.")
 
         if get_world_size() > 1:
             # Wait for rank 0 to finish loading the model.
@@ -422,8 +412,6 @@ def __init__(self,
         self.embedding_size = embedding_size
         self.embed_proj_head = nn.Linear(self.pretrained_model.config.hidden_size, self.embedding_size)
 
-        # self.sim_norm = SimNorm(simnorm_dim=group_size)
-
         # # Select the normalization method based on the final_norm_option_in_encoder parameter.
         if final_norm_option_in_encoder.lower() == "simnorm":
             self.norm = SimNorm(simnorm_dim=group_size)
@@ -484,7 +472,6 @@ def __init__(
             norm_type: str = 'BN',
             embedding_dim: int = 256,
             group_size: int = 8,
-            # final_norm_option_in_encoder: str = 'SimNorm',
             final_norm_option_in_encoder: str = 'LayerNorm', # TODO
     ) -> None:
         """
@@ -746,8 +733,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
             - output (:obj:`torch.Tensor`): :math:`(B, hidden_channels)`, where B is batch size.
         """
         x = self.fc_representation(x)
-        # TODO
-        # x = self.sim_norm(x)
         x = self.norm(x)
 
         return x

lzero/model/unizero_world_models/world_model.py

@@ -84,7 +84,7 @@ def __init__(self, config: TransformerConfig, tokenizer) -> None:
         # Head modules
         self.head_rewards = self._create_head(self.act_tokens_pattern, self.support_size)
         self.head_observations = self._create_head(self.all_but_last_latent_state_pattern, self.obs_per_embdding_dim, \
-                                                    self._get_final_norm(self.final_norm_option_in_obs_head)  # 使用指定的归一化方法
+                                                    self._get_final_norm(self.final_norm_option_in_obs_head)  # using the specified normalization method
                                                 #    self.sim_norm
                                                    )  # NOTE: we add a sim_norm to the head for observations
         if self.continuous_action_space:
@@ -158,7 +158,7 @@ def custom_init(module):
 
     def _get_final_norm(self, norm_option: str) -> nn.Module:
         """
-        根据指定的归一化选项返回相应的归一化模块。
+        Return the corresponding normalization module based on the specified normalization option.
         """
         if norm_option == 'LayerNorm':
             return nn.LayerNorm(self.config.embed_dim, eps=1e-5)
@@ -1299,6 +1299,15 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
         # Encode observations into latent state representations
         obs_embeddings = self.tokenizer.encode_to_obs_embeddings(batch['observations'])
 
+        # ========= for visual analysis =========
+        # Uncomment the lines below for visual analysis in Pong
+        # self.plot_latent_tsne_each_and_all_for_pong(obs_embeddings, suffix='pong_H10_H4_tsne')
+        # self.save_as_image_with_timestep(batch['observations'], suffix='pong_H10_H4_tsne')
+        # Uncomment the lines below for visual analysis in visual match
+        # self.plot_latent_tsne_each_and_all(obs_embeddings, suffix='visual_match_memlen1-60-15_tsne')
+        # self.save_as_image_with_timestep(batch['observations'], suffix='visual_match_memlen1-60-15_tsne')
+
+
         # ========= logging for analysis =========
         if self.analysis_dormant_ratio:
             # Calculate dormant ratio of the encoder