fix(pu): fix encoder-clip bug and num_channel/res bug

Author: tAnGjIa520

lzero/model/unizero_world_models/world_model.py

@@ -1885,6 +1885,10 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
         discounted_orig_policy_loss = (orig_policy_loss.view(-1, batch['actions'].shape[1]) * discounts).sum()/ batch['mask_padding'].sum()
         discounted_policy_entropy = (policy_entropy.view(-1, batch['actions'].shape[1]) * discounts).sum()/ batch['mask_padding'].sum()
 
+        # 为了让外部的训练循环能够获取encoder的输出，我们将其加入返回字典
+        # 使用 .detach() 是因为这个张量仅用于后续的clip操作，不应影响梯度计算
+        detached_obs_embeddings = obs_embeddings.detach()
+
         if self.continuous_action_space:
             return LossWithIntermediateLosses(
                 latent_recon_loss_weight=self.latent_recon_loss_weight,
@@ -1913,8 +1917,10 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 policy_mu=mu,
                 policy_sigma=sigma,
                 target_sampled_actions=target_sampled_actions,
+                
                 value_priority=value_priority,
                 intermediate_tensor_x=intermediate_tensor_x,
+                obs_embeddings=detached_obs_embeddings, # <-- 新增
             )
         else:
             return LossWithIntermediateLosses(
@@ -1941,8 +1947,10 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 e_rank_last_linear = e_rank_last_linear,
                 e_rank_sim_norm = e_rank_sim_norm,
                 latent_state_l2_norms=latent_state_l2_norms,
+
                 value_priority=value_priority,
                 intermediate_tensor_x=intermediate_tensor_x,
+                obs_embeddings=detached_obs_embeddings, # <-- 新增
             )
 
 

lzero/model/unizero_world_models/world_model_multitask.py

@@ -1899,6 +1899,10 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
         discounted_orig_policy_loss = (orig_policy_loss.view(-1, batch['actions'].shape[1]) * discounts).sum()/ batch['mask_padding'].sum()
         discounted_policy_entropy = (policy_entropy.view(-1, batch['actions'].shape[1]) * discounts).sum()/ batch['mask_padding'].sum()
 
+        # 为了让外部的训练循环能够获取encoder的输出，我们将其加入返回字典
+        # 使用 .detach() 是因为这个张量仅用于后续的clip操作，不应影响梯度计算
+        detached_obs_embeddings = obs_embeddings.detach()
+
         if self.continuous_action_space:
             return LossWithIntermediateLosses(
                 latent_recon_loss_weight=self.latent_recon_loss_weight,
@@ -1927,7 +1931,9 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 policy_mu=mu,
                 policy_sigma=sigma,
                 target_sampled_actions=target_sampled_actions,
+
                 value_priority=value_priority,
+                obs_embeddings=detached_obs_embeddings, # <-- 新增
 
             )
         else:
@@ -1955,7 +1961,10 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
                 e_rank_last_linear = e_rank_last_linear,
                 e_rank_sim_norm = e_rank_sim_norm,
                 latent_state_l2_norms=latent_state_l2_norms,
+                
                 value_priority=value_priority,
+                obs_embeddings=detached_obs_embeddings, # <-- 新增
+
 
             )
 

zoo/atari/config/atari_unizero_segment_config.py

@@ -73,6 +73,10 @@ def main(env_id, seed):
                 reward_support_range=(-300., 301., 1.),
                 value_support_range=(-300., 301., 1.),
                 norm_type=norm_type,
+                # num_res_blocks=1,
+                # num_channels=64,
+                num_res_blocks=2,
+                num_channels=128,
                 world_model_cfg=dict(
                     norm_type=norm_type,
                     final_norm_option_in_obs_head='LayerNorm',
@@ -138,10 +142,10 @@ def main(env_id, seed):
             # adaptive_entropy_alpha_lr=1e-3,
             target_entropy_start_ratio =0.98,
             # target_entropy_end_ratio =0.9,
-            target_entropy_end_ratio =0.7,
-            target_entropy_decay_steps = 100000, # 例如，在100k次迭代后达到最终值 需要与replay ratio协同调整
-            # target_entropy_end_ratio =0.5, # TODO=====
-            # target_entropy_decay_steps = 400000, # 例如，在100k次迭代后达到最终值 需要与replay ratio协同调整
+            # target_entropy_end_ratio =0.7,
+            # target_entropy_decay_steps = 100000, # 例如，在100k次迭代后达到最终值 需要与replay ratio协同调整
+            target_entropy_end_ratio =0.5, # TODO=====
+            target_entropy_decay_steps = 400000, # 例如，在100k次迭代后达到最终值 需要与replay ratio协同调整
 
 
             # ==================== START: Encoder-Clip Annealing Config ====================
@@ -217,7 +221,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_refactor1010/{env_id[:-14]}/{env_id[:-14]}_uz_targetentropy-alpha-100k-098-07-encoder-clip30-10-100k_label-smooth_resnet-encoder_priority_adamw-wd1e-2-encoder1-trans1-head1_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}'
+    main_config.exp_name = f'data_unizero_st_refactor1010/{env_id[:-14]}/{env_id[:-14]}_uz_ch128-res2_targetentropy-alpha-400k-098-05-encoder-clip30-10-100k-true_label-smooth_resnet-encoder_priority_adamw-wd1e-2-encoder1-trans1-head1_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)
 
 
@@ -240,21 +244,21 @@ def main(env_id, seed):
     # args.env = 'AlienNoFrameskip-v4'
 
     # 下面是atari8以外的2个代表环境
-    # args.env = 'QbertNoFrameskip-v4' # 记忆规划型环境 稀疏奖励
+    args.env = 'QbertNoFrameskip-v4' # 记忆规划型环境 稀疏奖励
     # args.env = 'SpaceInvadersNoFrameskip-v4' # 记忆规划型环境 稀疏奖励
 
     # 下面是已经表现不错的
     # args.env = 'BoxingNoFrameskip-v4' # 反应型环境 密集奖励
     # args.env = 'ChopperCommandNoFrameskip-v4'
-    args.env = 'RoadRunnerNoFrameskip-v4'
+    # args.env = 'RoadRunnerNoFrameskip-v4'
 
     main(args.env, args.seed)
 
     """
     tmux new -s uz-st-refactor-boxing
 
     conda activate /mnt/nfs/zhangjinouwen/puyuan/conda_envs/lz
-    export CUDA_VISIBLE_DEVICES=5
+    export CUDA_VISIBLE_DEVICES=1
     cd /mnt/nfs/zhangjinouwen/puyuan/LightZero
-    python /mnt/nfs/zhangjinouwen/puyuan/LightZero/zoo/atari/config/atari_unizero_segment_config.py 2>&1 | tee /mnt/nfs/zhangjinouwen/puyuan/LightZero/log/20251010_fix_uz_st_road.log
+    python /mnt/nfs/zhangjinouwen/puyuan/LightZero/zoo/atari/config/atari_unizero_segment_config.py 2>&1 | tee /mnt/nfs/zhangjinouwen/puyuan/LightZero/log/20251010_uz_st_ch128-res2_fix-encoder-clip_qbert.log
     """