polish(pu): polish weight decay and add latent_norm_loss

Author: puyuan1996

lzero/model/unizero_world_models/utils.py

@@ -257,24 +257,25 @@ def __init__(self, latent_recon_loss_weight=0, perceptual_loss_weight=0, continu
         if not kwargs:
             raise ValueError("At least one loss must be provided")
 
+
         # Get a reference device from one of the provided losses
         device = next(iter(kwargs.values())).device
 
         # NOTE: Define the weights for each loss type
         if not continuous_action_space:
             # orig, for atari and memory
-            # self.obs_loss_weight = 10
-            # self.value_loss_weight = 0.5
-            # self.reward_loss_weight = 1.
-            # self.policy_loss_weight = 1.
-            # self.ends_loss_weight = 0.
+            self.obs_loss_weight = 10
+            self.value_loss_weight = 0.5
+            self.reward_loss_weight = 1.
+            self.policy_loss_weight = 1.
+            self.ends_loss_weight = 0.
 
             # muzero loss weight
-            self.obs_loss_weight = 2
-            self.value_loss_weight = 0.25
-            self.reward_loss_weight = 1
-            self.policy_loss_weight = 1
-            self.ends_loss_weight = 0.
+            # self.obs_loss_weight = 2
+            # self.value_loss_weight = 0.25
+            # self.reward_loss_weight = 1
+            # self.policy_loss_weight = 1
+            # self.ends_loss_weight = 0.
 
             # EZV2, for atari and memory
             # self.obs_loss_weight = 5
@@ -297,6 +298,11 @@ def __init__(self, latent_recon_loss_weight=0, perceptual_loss_weight=0, continu
             # self.reward_loss_weight = 0.1
             # self.ends_loss_weight = 0.
 
+        # TODO(pu)
+        # self.latent_norm_loss_weight = 0.1
+        self.latent_norm_loss_weight = 0.01
+
+
         self.latent_recon_loss_weight = latent_recon_loss_weight
         self.perceptual_loss_weight = perceptual_loss_weight
 
@@ -317,6 +323,8 @@ def __init__(self, latent_recon_loss_weight=0, perceptual_loss_weight=0, continu
                 self.loss_total += self.latent_recon_loss_weight * v
             elif k == 'perceptual_loss':
                 self.loss_total += self.perceptual_loss_weight * v
+            elif k == 'latent_norm_loss':
+                self.loss_total += self.latent_norm_loss_weight * v
 
         self.intermediate_losses = {
             k: v if isinstance(v, dict) or isinstance(v, torch.Tensor) else (v if isinstance(v, float) else v.item())

lzero/model/unizero_world_models/world_model.py

@@ -1596,6 +1596,39 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
             # Calculate the L2 norm of the latent action
             latent_action_l2_norms = torch.norm(self.act_embedding_table(act_tokens), p=2, dim=2).mean()
 
+        if self.config.latent_norm_loss:
+            # ==================== L2惩罚损失计算（最终修复版 v2） ====================
+            # 1. 计算每个 latent_state 向量的L2范数的平方。
+            #    根据调试信息，obs_embeddings shape: (B*L, 1, E)
+            #    所以 latent_norm_sq shape: (B*L, 1)
+            latent_norm_sq = torch.norm(obs_embeddings, p=2, dim=-1).pow(2)
+            # 2. 获取源掩码。
+            #    根据调试信息，mask_source shape: (B, L)
+            mask_source = batch['mask_padding']
+            # 3. 将源掩码从 (B, L) reshape 为 (B*L, 1)，以匹配 latent_norm_sq 的形状。
+            #    这是解决维度不匹配错误的关键。
+            #    我们使用 view(-1, 1) 来实现这个变形。
+            correct_mask = mask_source.contiguous().view(-1, 1)
+            # 4. 检查变形后的形状是否匹配。
+            #    这是一个防御性编程，确保两个张量的第一个维度是相同的。
+            if latent_norm_sq.shape[0] != correct_mask.shape[0]:
+                # 如果形状不匹配，打印错误信息并抛出异常，这能帮助我们更快地定位未来可能出现的新问题。
+                raise RuntimeError(
+                    f"Shape mismatch for L2 norm loss calculation! "
+                    f"latent_norm_sq shape: {latent_norm_sq.shape}, "
+                    f"but correct_mask shape after reshape is: {correct_mask.shape}. "
+                    f"Original mask_source shape was: {mask_source.shape}"
+                )
+            # 5. 直接进行逐元素乘法。因为现在它们的形状都是 (B*L, 1)，所以可以安全相乘。
+            masked_latent_norm_sq = latent_norm_sq * correct_mask
+            # 6. 计算平均损失。分母是掩码中所有“1”的总和，代表有效的元素数量。
+            #    增加一个极小值 epsilon (1e-8) 防止分母为零。
+            latent_norm_loss = masked_latent_norm_sq.sum() / (correct_mask.sum() + 1e-8)
+            # =================================================================
+        else:
+            latent_norm_loss = torch.tensor(0.)
+
+
         # Forward pass to obtain predictions for observations, rewards, and policies
         outputs = self.forward({'obs_embeddings_and_act_tokens': (obs_embeddings, act_tokens)}, start_pos=start_pos)
 
@@ -1849,6 +1882,7 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 policy_mu=mu,
                 policy_sigma=sigma,
                 target_sampled_actions=target_sampled_actions,
+        latent_norm_loss=latent_norm_loss, # 新增
             )
         else:
             return LossWithIntermediateLosses(
@@ -1870,6 +1904,7 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 dormant_ratio_world_model=dormant_ratio_world_model,
                 latent_state_l2_norms=latent_state_l2_norms,
                 latent_action_l2_norms=latent_action_l2_norms,
+        latent_norm_loss=latent_norm_loss, # 新增
 
             )
 

zoo/atari/config/atari_muzero_segment_config.py

@@ -30,8 +30,10 @@ def main(env_id, seed):
 
     # Defines the frequency of reanalysis. E.g., 1 means reanalyze once per epoch, 2 means reanalyze once every two epochs.
     # buffer_reanalyze_freq = 1
+    buffer_reanalyze_freq = 1/2
+    # buffer_reanalyze_freq = 1/10
     # buffer_reanalyze_freq = 1/50
-    buffer_reanalyze_freq = 1/10000000000
+    # buffer_reanalyze_freq = 1/10000000000
     # Each reanalyze process will reanalyze <reanalyze_batch_size> sequences (<cfg.policy.num_unroll_steps> transitions per sequence)
     reanalyze_batch_size = 160
     # The partition of reanalyze. E.g., 1 means reanalyze_batch samples from the whole buffer, 0.5 means samples from the first half of the buffer.
@@ -165,16 +167,16 @@ def main(env_id, seed):
     parser.add_argument('--seed', type=int, help='The seed to use', default=0)
     args = parser.parse_args()
 
-    args.env = 'MsPacmanNoFrameskip-v4'
-    # args.env = 'QbertNoFrameskip-v4'
+    # args.env = 'MsPacmanNoFrameskip-v4'
+    args.env = 'QbertNoFrameskip-v4'
     # args.env = 'SeaquestNoFrameskip-v4'
     # args.env = 'BreakoutNoFrameskip-v4'
 
     args.seed = 0
     main(args.env, args.seed)
 
     """
-    export CUDA_VISIBLE_DEVICES=4
+    export CUDA_VISIBLE_DEVICES=3
     cd /fs-computility/niuyazhe/puyuan/code/LightZero
     python /fs-computility/niuyazhe/puyuan/code/LightZero/zoo/atari/config/atari_muzero_segment_config.py
     """

zoo/atari/config/atari_unizero_segment_config.py

@@ -9,13 +9,13 @@ def main(env_id, seed):
     # ==============================================================
     # begin of the most frequently changed config specified by the user
     # ==============================================================
-    collector_env_num = 8
-    num_segments = 8
-    evaluator_env_num = 3
+    # collector_env_num = 8
+    # num_segments = 8
+    # evaluator_env_num = 3
 
-    # collector_env_num = 1
-    # num_segments = 1
-    # evaluator_env_num = 1
+    collector_env_num = 1
+    num_segments = 1
+    evaluator_env_num = 1
 
     num_simulations = 50
     collect_num_simulations = 25
@@ -25,6 +25,8 @@ def main(env_id, seed):
     max_env_step = int(50e6)
     batch_size = 256
     # batch_size = 64 # debug
+    # batch_size = 4 # debug
+
     num_layers = 2
     # replay_ratio = 0.25
     replay_ratio = 0.1
@@ -33,6 +35,10 @@ def main(env_id, seed):
     num_unroll_steps = 10
     infer_context_length = 4
 
+    # game_segment_length = 40
+    # num_unroll_steps = 20
+    # infer_context_length = 8
+
     # game_segment_length = 200
     # num_unroll_steps = 16
     # infer_context_length = 8
@@ -93,6 +99,8 @@ def main(env_id, seed):
                     norm_type=norm_type,
                     num_res_blocks=2,
                     num_channels=128,
+                    # num_res_blocks=1, # TODO
+                    # num_channels=64,
                     support_size=601,
                     policy_entropy_weight=5e-3,
                     # policy_entropy_weight=5e-2, # TODO(pu)
@@ -125,6 +133,13 @@ def main(env_id, seed):
                     # final_norm_option_in_encoder="SimNorm",
                     # final_norm_option_in_obs_head="SimNorm",
                     # predict_latent_loss_type='group_kl',
+
+                    # weight_decay=1e-2,
+                    # latent_norm_loss=True,
+
+                    latent_norm_loss=False,
+                    weight_decay=1e-4, # TODO
+
                 ),
             ),
             # gradient_scale=True, #TODO
@@ -160,8 +175,8 @@ def main(env_id, seed):
             grad_clip_value=5,
             replay_buffer_size=int(1e6),
             # eval_freq=int(5e3),
-            # eval_freq=int(1e4),
-            eval_freq=int(2e4),
+            eval_freq=int(1e4), # TODO
+            # eval_freq=int(2e4),
             collector_env_num=collector_env_num,
             evaluator_env_num=evaluator_env_num,
             # ============= The key different params for reanalyze =============
@@ -193,8 +208,10 @@ 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_longrun_20250819/{env_id[:-14]}/{env_id[:-14]}_uz_fix-init-recur_clear20_mulossweight_spsi20_envnum{collector_env_num}_encoder-head-ln_soft-target-005_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}_c25_seed{seed}'
+    # main_config.exp_name = f'data_unizero_longrun_20250827/{env_id[:-14]}/{env_id[:-14]}_uz_wd1e-2_fix-init-recur_clear{game_segment_length}_originlossweight_spsi{game_segment_length}_envnum{collector_env_num}_encoder-head-ln_soft-target-005_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}_c25_seed{seed}'
+    main_config.exp_name = f'data_unizero_longrun_20250827/{env_id[:-14]}/{env_id[:-14]}_uz_lnlw001_fix-init-recur_clear{game_segment_length}_originlossweight_spsi{game_segment_length}_envnum{collector_env_num}_encoder-head-ln_soft-target-005_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}_c25_seed{seed}'
 
+    # main_config.exp_name = f'data_unizero_longrun_20250819/{env_id[:-14]}/{env_id[:-14]}_uz_fix-init-recur_clear{game_segment_length}_mulossweight_spsi{game_segment_length}_envnum{collector_env_num}_encoder-head-ln_soft-target-005_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}_c25_seed{seed}'
 
     # main_config.exp_name = f'data_unizero_longrun_20250819/{env_id[:-14]}/{env_id[:-14]}_uz_fix-init-recur_clear20_origlossweight_spsi20_envnum{collector_env_num}_encoder-head-l2norm_soft-target-005_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}_c25_seed{seed}'
 
@@ -256,7 +273,7 @@ def main(env_id, seed):
     main(args.env, args.seed)
 
     """
-    export CUDA_VISIBLE_DEVICES=0
+    export CUDA_VISIBLE_DEVICES=6
     cd /fs-computility/niuyazhe/puyuan/code/LightZero
     python /fs-computility/niuyazhe/puyuan/code/LightZero/zoo/atari/config/atari_unizero_segment_config.py
     """