fix(pu): fix configure_optimizer_unizero in unizero_mt

Author: tAnGjIa520

lzero/model/unizero_world_models/world_model.py

@@ -334,7 +334,9 @@ def _initialize_patterns(self) -> None:
     def _create_head(self, block_mask: torch.Tensor, output_dim: int, norm_layer=None) -> Head:
         """Create head modules for the transformer."""
         modules = [
+            nn.LayerNorm(self.config.embed_dim),  # <-- 核心优化！ # TODO
             nn.Linear(self.config.embed_dim, self.config.embed_dim),
+            nn.LayerNorm(self.config.embed_dim),      # 2. <-- 新增！稳定内部激活
             nn.GELU(approximate='tanh'),
             nn.Linear(self.config.embed_dim, output_dim)
         ]

lzero/policy/unizero_multitask.py

@@ -194,50 +194,103 @@ def zero_grad(self, set_to_none: bool = False) -> None:
         self.act_embedding_table.zero_grad(set_to_none=set_to_none)
 
 
+# def configure_optimizer_unizero(model, learning_rate, weight_decay, device_type, betas):
+#     """
+#     为UniZero模型配置带有差异化学习率的优化器。
+#     """
+#     # 1. 定义需要特殊处理的参数
+#     param_dict = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
+
+#     # 2. 将参数分为三组：Transformer主干、Tokenizer、Heads
+#     transformer_params = {pn: p for pn, p in param_dict.items() if 'transformer' in pn}
+#     tokenizer_params = {pn: p for pn, p in param_dict.items() if 'tokenizer' in pn}
+
+#     # Heads的参数是那些既不属于transformer也不属于tokenizer的
+#     head_params = {
+#         pn: p for pn, p in param_dict.items() 
+#         if 'transformer' not in pn and 'tokenizer' not in pn
+#     }
+
+#     # 3. 为每组设置不同的优化器参数（特别是学习率）
+#     #    这里我们仍然使用AdamW，但学习率设置更合理
+#     optim_groups = [
+#         {
+#             'params': list(transformer_params.values()),
+#             'lr': learning_rate,  # 1e-4
+#             # 'lr': learning_rate * 0.2,  # 为Transformer主干设置一个较小的学习率，例如 1e-5
+#             'weight_decay': weight_decay
+#             # 'weight_decay': weight_decay * 5.0 
+#         },
+#         {
+#             'params': list(tokenizer_params.values()),
+#             'lr': learning_rate,  # Tokenizer使用基础学习率，例如 1e-4
+#             # 'lr': learning_rate * 0.1,  # 为encoder设置一个较小的学习率，例如 1e-5
+#             'weight_decay': weight_decay * 5.0  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
+
+#         },
+#         {
+#             'params': list(head_params.values()),
+#             'lr': learning_rate,  # Heads也使用基础学习率率，例如 1e-4
+#             'weight_decay': 0.0  # 通常Heads的权重不做衰减
+#             # 'weight_decay': weight_decay
+
+#         }
+#     ]
+
+#     print("--- Optimizer Groups ---")
+#     print(f"Transformer LR: {learning_rate}")
+#     print(f"Tokenizer/Heads LR: {learning_rate}")
+
+#     optimizer = torch.optim.AdamW(optim_groups, betas=betas)
+#     return optimizer
+
 def configure_optimizer_unizero(model, learning_rate, weight_decay, device_type, betas):
     """
     为UniZero模型配置带有差异化学习率的优化器。
+    (修正版，确保参数组互斥)
     """
-    # 1. 定义需要特殊处理的参数
-    param_dict = {pn: p for pn, p in model.named_parameters() if p.requires_grad}
-
-    # 2. 将参数分为三组：Transformer主干、Tokenizer、Heads
-    transformer_params = {pn: p for pn, p in param_dict.items() if 'transformer' in pn}
-    tokenizer_params = {pn: p for pn, p in param_dict.items() if 'tokenizer' in pn}
-
-    # Heads的参数是那些既不属于transformer也不属于tokenizer的
-    head_params = {
-        pn: p for pn, p in param_dict.items() 
-        if 'transformer' not in pn and 'tokenizer' not in pn
-    }
-
-    # 3. 为每组设置不同的优化器参数（特别是学习率）
+    # 1. 创建空的参数列表用于分组
+    transformer_params = []
+    tokenizer_params = []
+    head_params = []
+
+    # 2. 遍历所有可训练参数，并使用 if/elif/else 结构确保每个参数只被分配到一个组
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+
+        if 'transformer' in name:
+            transformer_params.append(param)
+        elif 'tokenizer' in name:
+            tokenizer_params.append(param)
+        else:
+            head_params.append(param)
+            
+    # 3. 为每组设置不同的优化器参数
     #    这里我们仍然使用AdamW，但学习率设置更合理
     optim_groups = [
         {
-            'params': list(transformer_params.values()),
+            'params': transformer_params,
             'lr': learning_rate,  # 1e-4
-            # 'lr': learning_rate * 0.2,  # 为Transformer主干设置一个较小的学习率，例如 1e-5
             'weight_decay': weight_decay
-            # 'weight_decay': weight_decay * 5.0 
         },
         {
-            'params': list(tokenizer_params.values()),
+            'params': tokenizer_params,
             'lr': learning_rate,  # Tokenizer使用基础学习率，例如 1e-4
-            # 'lr': learning_rate * 0.1,  # 为encoder设置一个较小的学习率，例如 1e-5
             'weight_decay': weight_decay * 5.0  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
-
         },
         {
-            'params': list(head_params.values()),
+            'params': head_params,
             'lr': learning_rate,  # Heads也使用基础学习率率，例如 1e-4
             'weight_decay': 0.0  # 通常Heads的权重不做衰减
-            # 'weight_decay': weight_decay
-
         }
     ]
 
     print("--- Optimizer Groups ---")
+    # 打印每个组的参数数量以供调试
+    print(f"Transformer params: {len(transformer_params)}")
+    print(f"Tokenizer params: {len(tokenizer_params)}")
+    print(f"Head params: {len(head_params)}")
     print(f"Transformer LR: {learning_rate}")
     print(f"Tokenizer/Heads LR: {learning_rate}")
 

zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py

@@ -155,20 +155,20 @@ def create_config(
                     task_num=len(env_id_list),
                     # game_segment_length=game_segment_length,
                     game_segment_length=20, # TODO
-                    # use_priority=True,
-                    use_priority=False, # TODO=====
+                    use_priority=True,
+                    # 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,
                     num_experts_in_moe_head=4,
                     moe_in_transformer=False,
 
-                    # multiplication_moe_in_transformer=True,
-                    multiplication_moe_in_transformer=False, # TODO=====
+                    multiplication_moe_in_transformer=True,
+                    # multiplication_moe_in_transformer=False, # TODO=====
 
                     n_shared_experts=1,
                     num_experts_per_tok=1,
@@ -188,8 +188,8 @@ def create_config(
             learning_rate=0.0001,
 
             # (bool) 是否启用自适应策略熵权重 (alpha)
-            # use_adaptive_entropy_weight=True,
-            use_adaptive_entropy_weight=False,
+            use_adaptive_entropy_weight=True,
+            # use_adaptive_entropy_weight=False,
 
             # (float) 自适应alpha优化器的学习率
             adaptive_entropy_alpha_lr=1e-4,
@@ -216,8 +216,8 @@ def create_config(
             total_batch_size=total_batch_size,
             allocated_batch_sizes=False,
             train_start_after_envsteps=int(0),
-            use_priority=False, # TODO=====
-            # use_priority=True,
+            # use_priority=False, # TODO=====
+            use_priority=True,
             priority_prob_alpha=1,
             priority_prob_beta=1,
             print_task_priority_logs=False,
@@ -271,7 +271,10 @@ def generate_configs(
     # Replace placeholders like [BENCHMARK_TAG] and [MODEL_TAG] to define the experiment name.
     benchmark_tag = "data_unizero_mt_refactor0929"  # 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"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}"
+
     exp_name_prefix = f'{benchmark_tag}/atari_{len(env_id_list)}games_{model_tag}_seed{seed}/'
 
     for task_id, env_id in enumerate(env_id_list):

zoo/atari/config/atari_unizero_segment_config.py

@@ -14,11 +14,13 @@ def main(env_id, seed):
     evaluator_env_num = 3
     num_simulations = 50
     # max_env_step = int(4e5)
-    max_env_step = int(10e6) # TODO
+    max_env_step = int(5e6) # TODO
 
-    batch_size = 64
+    # batch_size = 64
+    batch_size = 256
     num_layers = 2
-    replay_ratio = 0.25
+    replay_ratio = 0.1
+    # replay_ratio = 0.25
     num_unroll_steps = 10
     infer_context_length = 4
 
@@ -131,6 +133,7 @@ def main(env_id, seed):
             use_adaptive_entropy_weight=True,
             # (float) 自适应alpha优化器的学习率
             adaptive_entropy_alpha_lr=1e-4,
+            # adaptive_entropy_alpha_lr=1e-3,
             target_entropy_start_ratio =0.98,
             # target_entropy_end_ratio =0.9,
             target_entropy_end_ratio =0.7,
@@ -200,7 +203,7 @@ def main(env_id, seed):
 
     # ============ use muzero_segment_collector instead of muzero_collector =============
     from lzero.entry import train_unizero_segment
-    main_config.exp_name = f'data_unizero_st_refactor0929/{env_id[:-14]}/{env_id[:-14]}_uz_resnet-encoder_priority_adamw-wd1e-2_brf{buffer_reanalyze_freq}-rbs{reanalyze_batch_size}-rp{reanalyze_partition}_nlayer{num_layers}_numsegments-{num_segments}_gsl{game_segment_length}_rr{replay_ratio}_Htrain{num_unroll_steps}-Hinfer{infer_context_length}_bs{batch_size}_seed{seed}'
+    main_config.exp_name = f'data_unizero_st_refactor0929/{env_id[:-14]}/{env_id[:-14]}_uz_resnet-encoder_priority_adamw-wd1e-2_ln-inner-ln_brf{buffer_reanalyze_freq}-rbs{reanalyze_batch_size}-rp{reanalyze_partition}_nlayer{num_layers}_numsegments-{num_segments}_gsl{game_segment_length}_rr{replay_ratio}_Htrain{num_unroll_steps}-Hinfer{infer_context_length}_bs{batch_size}_seed{seed}'
     train_unizero_segment([main_config, create_config], seed=seed, model_path=main_config.policy.model_path, max_env_step=max_env_step)