polish(pu): polish minotor-log and adapt to ale/xxx-v5 style game

puyuan1996 · puyuan1996 · commit 3788eb7a5e70 · 2025-10-25T17:46:59.000+08:00
diff --git a/lzero/model/unizero_world_models/world_model.py b/lzero/model/unizero_world_models/world_model.py
@@ -1911,7 +1911,7 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
             # F.cosine_similarity 计算的是相似度，范围是 [-1, 1]。我们希望最大化它，
             # 所以最小化 1 - similarity。
             # reduction='none' 使得我们可以像原来一样处理mask
-            print("predict_latent_loss_type == 'cos_sim'")
+            # print("predict_latent_loss_type == 'cos_sim'")
             cosine_sim_loss = 1 - F.cosine_similarity(logits_observations, labels_observations, dim=-1)
             loss_obs = cosine_sim_loss
 
@@ -2034,6 +2034,16 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 value_priority=value_priority,
                 intermediate_tensor_x=intermediate_tensor_x,
                 obs_embeddings=detached_obs_embeddings, # <-- 新增
+
+                # logits_value_mean=outputs.logits_value.mean(),
+                # logits_value_max=outputs.logits_value.max(),
+                # logits_value_min=outputs.logits_value.min(),
+                # logits_policy_mean=outputs.logits_policy.mean(),
+                # logits_policy_max=outputs.logits_policy.max(),
+                # logits_policy_min=outputs.logits_policy.min(),
+                logits_value=outputs.logits_value.detach(),  # 使用detach()，因为它仅用于分析和裁剪，不参与梯度计算
+                logits_reward=outputs.logits_rewards.detach(),
+                logits_policy=outputs.logits_policy.detach(),
             )
         else:
             return LossWithIntermediateLosses(
@@ -2064,6 +2074,16 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 value_priority=value_priority,
                 intermediate_tensor_x=intermediate_tensor_x,
                 obs_embeddings=detached_obs_embeddings, # <-- 新增
+
+                # logits_value_mean=outputs.logits_value.mean(),
+                # logits_value_max=outputs.logits_value.max(),
+                # logits_value_min=outputs.logits_value.min(),
+                # logits_policy_mean=outputs.logits_policy.mean(),
+                # logits_policy_max=outputs.logits_policy.max(),
+                # logits_policy_min=outputs.logits_policy.min(),
+                logits_value=outputs.logits_value.detach(),  # 使用detach()，因为它仅用于分析和裁剪，不参与梯度计算
+                logits_reward=outputs.logits_rewards.detach(),
+                logits_policy=outputs.logits_policy.detach(),
             )
 
     
diff --git a/lzero/model/unizero_world_models/world_model_multitask.py b/lzero/model/unizero_world_models/world_model_multitask.py
@@ -1934,6 +1934,15 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
 
                 value_priority=value_priority,
                 obs_embeddings=detached_obs_embeddings, # <-- 新增
+                #                 logits_value_mean=outputs.logits_value.mean(),
+                # logits_value_max=outputs.logits_value.max(),
+                # logits_value_min=outputs.logits_value.min(),
+                # logits_policy_mean=outputs.logits_policy.mean(),
+                # logits_policy_max=outputs.logits_policy.max(),
+                # logits_policy_min=outputs.logits_policy.min(),
+                logits_value=outputs.logits_value.detach(),  # 使用detach()，因为它仅用于分析和裁剪，不参与梯度计算
+                logits_reward=outputs.logits_rewards.detach(),
+                logits_policy=outputs.logits_policy.detach(),
 
             )
         else:
@@ -1964,6 +1973,15 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 
                 value_priority=value_priority,
                 obs_embeddings=detached_obs_embeddings, # <-- 新增
+                # logits_value_mean=outputs.logits_value.mean(),
+                # logits_value_max=outputs.logits_value.max(),
+                # logits_value_min=outputs.logits_value.min(),
+                # logits_policy_mean=outputs.logits_policy.mean(),
+                # logits_policy_max=outputs.logits_policy.max(),
+                # logits_policy_min=outputs.logits_policy.min(),
+                logits_value=outputs.logits_value.detach(),  # 使用detach()，因为它仅用于分析和裁剪，不参与梯度计算
+                logits_reward=outputs.logits_rewards.detach(),
+                logits_policy=outputs.logits_policy.detach(),
 
 
             )
diff --git a/lzero/policy/unizero.py b/lzero/policy/unizero.py
@@ -833,15 +833,16 @@ def _forward_learn(self, data: Tuple[torch.Tensor]) -> Dict[str, Union[float, in
         latent_state_l2_norms = self.intermediate_losses['latent_state_l2_norms']
 
         latent_action_l2_norms = self.intermediate_losses['latent_action_l2_norms']
-        logits_value_mean=self.intermediate_losses['logits_value_mean']
-        logits_value_max=self.intermediate_losses['logits_value_max']
-        logits_value_min=self.intermediate_losses['logits_value_min']
-        logits_policy_mean=self.intermediate_losses['logits_policy_mean']
-        logits_policy_max=self.intermediate_losses['logits_policy_max']
-        logits_policy_min=self.intermediate_losses['logits_policy_min']
-        temperature_value=self.intermediate_losses['temperature_value']
-        temperature_reward=self.intermediate_losses['temperature_reward']
-        temperature_policy=self.intermediate_losses['temperature_policy']
+
+        # logits_value_mean=self.intermediate_losses['logits_value_mean']
+        # logits_value_max=self.intermediate_losses['logits_value_max']
+        # logits_value_min=self.intermediate_losses['logits_value_min']
+        # logits_policy_mean=self.intermediate_losses['logits_policy_mean']
+        # logits_policy_max=self.intermediate_losses['logits_policy_max']
+        # logits_policy_min=self.intermediate_losses['logits_policy_min']
+        # temperature_value=self.intermediate_losses['temperature_value']
+        # temperature_reward=self.intermediate_losses['temperature_reward']
+        # temperature_policy=self.intermediate_losses['temperature_policy']
 
         assert not torch.isnan(losses.loss_total).any(), "Loss contains NaN values"
         assert not torch.isinf(losses.loss_total).any(), "Loss contains Inf values"
@@ -875,7 +876,9 @@ def _forward_learn(self, data: Tuple[torch.Tensor]) -> Dict[str, Union[float, in
             # --- [优化建议] 增加 log_alpha 裁剪作为安全措施 ---
             with torch.no_grad():
                 # 将 alpha 限制在例如 [1e-4, 10.0] 的范围内
-                self.log_alpha.clamp_(np.log(1e-4), np.log(10.0))
+                # self.log_alpha.clamp_(np.log(1e-4), np.log(10.0))
+                self.log_alpha.clamp_(np.log(5e-3), np.log(10.0))
+
 
             # --- 使用当前更新后的 alpha (截断梯度流) ---
             current_alpha = self.log_alpha.exp().detach()
@@ -1047,12 +1050,13 @@ def _forward_learn(self, data: Tuple[torch.Tensor]) -> Dict[str, Union[float, in
             'analysis/l2_norm_after': self.l2_norm_after,
             'analysis/grad_norm_before': self.grad_norm_before,
             'analysis/grad_norm_after': self.grad_norm_after,
-                    "logits_value_mean":logits_value_mean,
-        "logits_value_max":logits_value_max,
-        "logits_value_min":logits_value_min,
-        "logits_policy_mean":logits_policy_mean,
-        "logits_policy_max":logits_policy_max,
-        "logits_policy_min":logits_policy_min,
+        
+        # "logits_value_mean":logits_value_mean,
+        # "logits_value_max":logits_value_max,
+        # "logits_value_min":logits_value_min,
+        # "logits_policy_mean":logits_policy_mean,
+        # "logits_policy_max":logits_policy_max,
+        # "logits_policy_min":logits_policy_min,
 
              "temperature_value":temperature_value,
         "temperature_reward":temperature_reward,
@@ -1621,12 +1625,12 @@ def _monitor_vars_learn(self) -> List[str]:
             'total_grad_norm_before_clip_wm',
 
             # ==================== Logits Statistics ====================
-            'logits_value_mean',
-            'logits_value_max',
-            'logits_value_min',
-            'logits_policy_mean',
-            'logits_policy_max',
-            'logits_policy_min',
+            # 'logits_value_mean',
+            # 'logits_value_max',
+            # 'logits_value_min',
+            # 'logits_policy_mean',
+            # 'logits_policy_max',
+            # 'logits_policy_min',
 
             # ==================== Temperature Parameters ====================
             'temperature_value',
diff --git a/lzero/policy/unizero_multitask.py b/lzero/policy/unizero_multitask.py
@@ -622,9 +622,9 @@ def _monitor_model_norms(self) -> Dict[str, float]:
         module_groups = {
             'encoder': world_model.tokenizer.encoder,
             'transformer': world_model.transformer,
-            'head_value': world_model.head_values,  # Note: multi-task uses head_values (plural)
-            'head_reward': world_model.head_rewards,
-            'head_policy': world_model.head_policies,  # Note: multi-task uses head_policies (plural)
+            'head_value': world_model.head_value_multi_task,  # Note: multi-task uses head_value (plural)
+            'head_reward': world_model.head_rewards_multi_task,
+            'head_policy': world_model.head_policy_multi_task,  # Note: multi-task uses head_policies (plural)
         }
 
         for group_name, group_module in module_groups.items():
@@ -669,9 +669,9 @@ def _monitor_gradient_norms(self) -> Dict[str, float]:
         module_groups = {
             'encoder': world_model.tokenizer.encoder,
             'transformer': world_model.transformer,
-            'head_value': world_model.head_values,
-            'head_reward': world_model.head_rewards,
-            'head_policy': world_model.head_policies,
+            'head_value': world_model.head_value_multi_task,
+            'head_reward': world_model.head_rewards_multi_task,
+            'head_policy': world_model.head_policy_multi_task,
         }
 
         for group_name, group_module in module_groups.items():
@@ -1169,7 +1169,9 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
                 # --- [优化建议] 增加 log_alpha 裁剪作为安全措施 ---
                 with torch.no_grad():
                     # 将 alpha 限制在例如 [1e-4, 10.0] 的范围内
-                    self.log_alpha.clamp_(np.log(1e-4), np.log(10.0))
+                    # self.log_alpha.clamp_(np.log(1e-4), np.log(10.0))
+                    self.log_alpha.clamp_(np.log(5e-3), np.log(10.0))
+
 
                 # --- 使用当前更新后的 alpha (截断梯度流) ---
                 current_alpha = self.log_alpha.exp().detach()
diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 DI-engine>=0.5.3
-gymnasium[atari]==0.28.0
+# gymnasium[atari]==0.28.0
 numpy>=1.24.1,<2
 pympler
 minigrid
diff --git a/zoo/atari/config/atari_env_action_space_map.py b/zoo/atari/config/atari_env_action_space_map.py
@@ -1,33 +1,33 @@
 from easydict import EasyDict
 
 atari_env_action_space_map = EasyDict({
-    'AlienNoFrameskip-v4': 18,
-    'AmidarNoFrameskip-v4': 10,
-    'AssaultNoFrameskip-v4': 7,
-    'AsterixNoFrameskip-v4': 9,
-    'BankHeistNoFrameskip-v4': 18,
-    'BattleZoneNoFrameskip-v4': 18,
-    'ChopperCommandNoFrameskip-v4': 18,
-    'CrazyClimberNoFrameskip-v4': 9,
-    'DemonAttackNoFrameskip-v4': 6,
-    'FreewayNoFrameskip-v4': 3,
-    'FrostbiteNoFrameskip-v4': 18,
-    'GopherNoFrameskip-v4': 8,
-    'HeroNoFrameskip-v4': 18,
-    'JamesbondNoFrameskip-v4': 18,
-    'KangarooNoFrameskip-v4': 18,
-    'KrullNoFrameskip-v4': 18,
-    'KungFuMasterNoFrameskip-v4': 14,
-    'PrivateEyeNoFrameskip-v4': 18,
-    'RoadRunnerNoFrameskip-v4': 18,
-    'UpNDownNoFrameskip-v4': 6,
-    'PongNoFrameskip-v4': 6,
-    'MsPacmanNoFrameskip-v4': 9,
-    'QbertNoFrameskip-v4': 6,
-    'SeaquestNoFrameskip-v4': 18,
-    'BoxingNoFrameskip-v4': 18,
-    'BreakoutNoFrameskip-v4': 4,
-    'SpaceInvadersNoFrameskip-v4': 6,
-    'BeamRiderNoFrameskip-v4': 9,
-    'GravitarNoFrameskip-v4': 18,
+    'ALE/Alien-v5': 18,
+    'ALE/Amidar-v5': 10,
+    'ALE/Assault-v5': 7,
+    'ALE/Asterix-v5': 9,
+    'ALE/BankHeist-v5': 18,
+    'ALE/BattleZone-v5': 18,
+    'ALE/ChopperCommand-v5': 18,
+    'ALE/CrazyClimber-v5': 9,
+    'ALE/DemonAttack-v5': 6,
+    'ALE/Freeway-v5': 3,
+    'ALE/Frostbite-v5': 18,
+    'ALE/Gopher-v5': 8,
+    'ALE/Hero-v5': 18,
+    'ALE/Jamesbond-v5': 18,
+    'ALE/Kangaroo-v5': 18,
+    'ALE/Krull-v5': 18,
+    'ALE/KungFuMaster-v5': 14,
+    'ALE/PrivateEye-v5': 18,
+    'ALE/RoadRunner-v5': 18,
+    'ALE/UpNDown-v5': 6,
+    'ALE/Pong-v5': 6,
+    'ALE/MsPacman-v5': 9,
+    'ALE/Qbert-v5': 6,
+    'ALE/Seaquest-v5': 18,
+    'ALE/Boxing-v5': 18,
+    'ALE/Breakout-v5': 4,
+    'ALE/SpaceInvaders-v5': 6,
+    'ALE/BeamRider-v5': 9,
+    'ALE/Gravitar-v5': 18,
 })
diff --git a/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py b/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py
@@ -93,6 +93,7 @@ def create_config(
     """
     return EasyDict(dict(
         env=dict(
+            frame_skip=1, # TODO
             stop_value=int(1e6),
             env_id=env_id,
             observation_shape=(3, 64, 64),
@@ -162,8 +163,8 @@ def create_config(
                     # use_priority=False, # TODO=====
                     priority_prob_alpha=1,
                     priority_prob_beta=1,
-                    # encoder_type='vit',
-                    encoder_type='resnet',
+                    encoder_type='vit',
+                    # encoder_type='resnet',
                     use_normal_head=True,
                     use_softmoe_head=False,
                     use_moe_head=False,
@@ -195,7 +196,8 @@ def create_config(
             # use_adaptive_entropy_weight=False,
 
             # (float) 自适应alpha优化器的学习率
-            adaptive_entropy_alpha_lr=1e-4,
+            # adaptive_entropy_alpha_lr=1e-4,
+            adaptive_entropy_alpha_lr=1e-3,
             target_entropy_start_ratio =0.98,
             # target_entropy_end_ratio =0.9, # TODO=====
             # target_entropy_end_ratio =0.7,
@@ -289,15 +291,18 @@ def generate_configs(
     # --- Experiment Name Template ---
     # Replace placeholders like [BENCHMARK_TAG] and [MODEL_TAG] to define the experiment name.
     # benchmark_tag = "data_unizero_mt_refactor1010_debug"  # e.g., unizero_atari_mt_20250612
-    benchmark_tag = "data_unizero_mt_refactor1012"  # e.g., unizero_atari_mt_20250612
+    benchmark_tag = "data_unizero_mt_refactor1024"  # e.g., unizero_atari_mt_20250612
 
     # model_tag = f"vit-small_moe8_tbs512_tran-nlayer{num_layers}_brf{buffer_reanalyze_freq}_not-share-head"
     # model_tag = f"resnet_noprior_noalpha_nomoe_head-inner-ln_adamw-wd1e-2_tbs512_tran-nlayer{num_layers}_brf{buffer_reanalyze_freq}"
     
     # model_tag = f"vit_prior_alpha-100k-098-07_encoder-100k-30-10_moe8_head-inner-ln_adamw-wd1e-2_tbs512_tran-nlayer{num_layers}_brf{buffer_reanalyze_freq}"
 
     # model_tag = f"resnet_encoder-100k-30-10-true_label-smooth_prior_alpha-100k-098-07_moe8_head-inner-ln_adamw-wd1e-2-all_tbs512_tran-nlayer{num_layers}_brf{buffer_reanalyze_freq}"
-    model_tag = f"resnet_tran-nlayer{num_layers}_moe8_encoder-100k-30-10-true_alpha-100k-098-05_prior_adamw-wd1e-2-all_tbs512_brf{buffer_reanalyze_freq}_label-smooth_head-inner-ln"
+    model_tag = f"vit_tran-nlayer{num_layers}_moe8_encoder-100k-30-10-true_alpha-100k-098-05_prior_adamw-wd1e-2-all_tbs512_brf{buffer_reanalyze_freq}_label-smooth_head-inner-ln"
+    
+    # model_tag = f"resnet_tran-nlayer{num_layers}_moe8_encoder-100k-30-10-true_alpha-100k-098-05_prior_adamw-wd1e-2-all_tbs512_brf{buffer_reanalyze_freq}_label-smooth_head-inner-ln"
+    
     # model_tag = f"resnet_encoder-100k-30-10-true_label-smooth_prior_alpha-150k-098-05_moe8_head-inner-ln_adamw-wd1e-2-all_tbs512_tran-nlayer{num_layers}_brf{buffer_reanalyze_freq}"
 
     exp_name_prefix = f'{benchmark_tag}/atari_{len(env_id_list)}games_{model_tag}_seed{seed}/'
@@ -309,7 +314,10 @@ def generate_configs(
             buffer_reanalyze_freq, reanalyze_batch_size, reanalyze_partition, num_segments, total_batch_size, num_layers
         )
         config.policy.task_id = task_id
-        config.exp_name = exp_name_prefix + f"{env_id.split('NoFrameskip')[0]}_seed{seed}"
+        # --- MODIFIED LINE ---
+        # Correctly extract the game name from 'ALE/GameName-v5' format.
+        game_name = env_id.split('/')[1].split('-')[0]
+        config.exp_name = exp_name_prefix + f"{game_name}_seed{seed}"
         configs.append([task_id, [config, create_env_manager()]])
     return configs
 
@@ -348,6 +356,8 @@ def create_env_manager() -> EasyDict:
         export CUDA_VISIBLE_DEVICES=4,5,6,7
 
         cd /path/to/your/project/
+        /mnt/shared-storage-user/puyuan/lz/bin/python -m torch.distributed.launch --nproc_per_node=4 --master_port=29502 /mnt/shared-storage-user/puyuan/code_20250828/LightZero/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py  2>&1 | tee /mnt/shared-storage-user/puyuan/code_20250828/LightZero/log/20251024_vit_nlayer4_alpha-100k-098-05.log
+
         python -m torch.distributed.launch --nproc_per_node=6 --master_port=29502 /mnt/nfs/zhangjinouwen/puyuan/LightZero/zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py 2>&1 | tee /mnt/nfs/zhangjinouwen/puyuan/LightZero/log/20251012_resnet_nlayer4_alpha-100k-098-05.log
             /path/to/this/script.py 2>&1 | tee /path/to/your/log/file.log
     """
@@ -370,22 +380,23 @@ def create_env_manager() -> EasyDict:
     max_env_step = int(5e6) # TODO
     reanalyze_ratio = 0.0
 
+    # --- MODIFIED SECTION: Standardized env_id_list formats ---
     if num_games == 3:
-        env_id_list = ['PongNoFrameskip-v4', 'MsPacmanNoFrameskip-v4', 'SeaquestNoFrameskip-v4']
+        env_id_list = ['ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5']
     elif num_games == 8:
         env_id_list = [
-            'PongNoFrameskip-v4', 'MsPacmanNoFrameskip-v4', 'SeaquestNoFrameskip-v4', 'BoxingNoFrameskip-v4',
-            'AlienNoFrameskip-v4', 'ChopperCommandNoFrameskip-v4', 'HeroNoFrameskip-v4', 'RoadRunnerNoFrameskip-v4',
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
         ]
     elif num_games == 26:
         env_id_list = [
-            'PongNoFrameskip-v4', 'MsPacmanNoFrameskip-v4', 'SeaquestNoFrameskip-v4', 'BoxingNoFrameskip-v4',
-            'AlienNoFrameskip-v4', 'ChopperCommandNoFrameskip-v4', 'HeroNoFrameskip-v4', 'RoadRunnerNoFrameskip-v4',
-            'AmidarNoFrameskip-v4', 'AssaultNoFrameskip-v4', 'AsterixNoFrameskip-v4', 'BankHeistNoFrameskip-v4',
-            'BattleZoneNoFrameskip-v4', 'CrazyClimberNoFrameskip-v4', 'DemonAttackNoFrameskip-v4', 'FreewayNoFrameskip-v4',
-            'FrostbiteNoFrameskip-v4', 'GopherNoFrameskip-v4', 'JamesbondNoFrameskip-v4', 'KangarooNoFrameskip-v4',
-            'KrullNoFrameskip-v4', 'KungFuMasterNoFrameskip-v4', 'PrivateEyeNoFrameskip-v4', 'UpNDownNoFrameskip-v4',
-            'QbertNoFrameskip-v4', 'BreakoutNoFrameskip-v4',
+            'ALE/Pong-v5', 'ALE/MsPacman-v5', 'ALE/Seaquest-v5', 'ALE/Boxing-v5',
+            'ALE/Alien-v5', 'ALE/ChopperCommand-v5', 'ALE/Hero-v5', 'ALE/RoadRunner-v5',
+            'ALE/Amidar-v5', 'ALE/Assault-v5', 'ALE/Asterix-v5', 'ALE/BankHeist-v5',
+            'ALE/BattleZone-v5', 'ALE/CrazyClimber-v5', 'ALE/DemonAttack-v5', 'ALE/Freeway-v5',
+            'ALE/Frostbite-v5', 'ALE/Gopher-v5', 'ALE/Jamesbond-v5', 'ALE/Kangaroo-v5',
+            'ALE/Krull-v5', 'ALE/KungFuMaster-v5', 'ALE/PrivateEye-v5', 'ALE/UpNDown-v5',
+            'ALE/Qbert-v5', 'ALE/Breakout-v5',
         ]
     else:
         raise ValueError(f"Unsupported number of environments: {num_games}")
diff --git a/zoo/atari/config/atari_unizero_segment_config.py b/zoo/atari/config/atari_unizero_segment_config.py
diff --git a/zoo/atari/envs/atari_lightzero_env.py b/zoo/atari/envs/atari_lightzero_env.py
diff --git a/zoo/atari/envs/atari_wrappers.py b/zoo/atari/envs/atari_wrappers.py
