fix(pu): fix recur kv pool index compatibility

Author: puyuan1996

lzero/model/unizero_world_models/world_model.py

@@ -167,6 +167,7 @@ def custom_init(module):
 
         self.num_simulations = getattr(self.config, 'num_simulations', 50)
 
+        # TODO: recur kv pool是否应该分成不同的环境有不同的pool呢
         self.shared_pool_size_recur = int(self.num_simulations*self.env_num)
 
         # self.shared_pool_size_init = int(50)  # NOTE: Will having too many cause incorrect retrieval of the kv cache?
@@ -1497,9 +1498,22 @@ def retrieve_or_generate_kvcache(self, latent_state: list, ready_env_num: int,
                     matched_value = None
 
                 # If not found, try to retrieve from past_kv_cache_recurrent_infer
+                # if matched_value is None:
+                #     matched_value = self.shared_pool_recur_infer[self.past_kv_cache_recurrent_infer.get(cache_key)]
+
+                # ==================== 核心修复 ====================
+                # 步骤 2: 仅当在 init_infer 中未找到时，才尝试从 recurrent_infer 缓存中查找
                 if matched_value is None:
-                    matched_value = self.shared_pool_recur_infer[self.past_kv_cache_recurrent_infer.get(cache_key)]
+                    # 2.1 安全地从字典中获取索引，它可能返回 None
+                    recur_cache_index = self.past_kv_cache_recurrent_infer.get(cache_key)
+                    # 2.2 只有在索引有效（不是 None）的情况下，才使用它来从物理池中检索值
+                    if recur_cache_index is not None:
+                        matched_value = self.shared_pool_recur_infer[recur_cache_index]
+                    
+                    if recur_cache_index is None:
+                        print(f"[CACHE MISS]  Not found for key={cache_key} in recurrent infer. Generating new cache.")
 
+                # =================================================
                     # # TODO
                     # retrieved_cache = matched_value._keys_values[0]._k_cache._cache
                     # retrieved_sum = torch.sum(retrieved_cache).item()

zoo/atari/config/atari_unizero_segment_config.py

@@ -9,25 +9,25 @@ 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
     # collect_num_simulations = 50
     eval_num_simulations = 50
     # max_env_step = int(5e5)
     max_env_step = int(50e6)
-    # batch_size = 256
-    batch_size = 64 # debug
+    batch_size = 256
+    # batch_size = 64 # debug
     num_layers = 2
-    replay_ratio = 0.25
-    # replay_ratio = 0.1
+    # replay_ratio = 0.25
+    replay_ratio = 0.1
 
     game_segment_length = 20
     num_unroll_steps = 10
@@ -114,14 +114,14 @@ def main(env_id, seed):
                     # final_norm_option_in_obs_head="LayerNorm",
                     # predict_latent_loss_type='mse',
 
-                    final_norm_option_in_encoder='L2Norm',
-                    final_norm_option_in_obs_head="L2Norm",
-                    predict_latent_loss_type='mse',
-
-                    # final_norm_option_in_encoder="LayerNorm",
-                    # final_norm_option_in_obs_head="LayerNorm",
+                    # final_norm_option_in_encoder='L2Norm',
+                    # final_norm_option_in_obs_head="L2Norm",
                     # predict_latent_loss_type='mse',
 
+                    final_norm_option_in_encoder="LayerNorm",
+                    final_norm_option_in_obs_head="LayerNorm",
+                    predict_latent_loss_type='mse',
+
                     # final_norm_option_in_encoder="SimNorm",
                     # final_norm_option_in_obs_head="SimNorm",
                     # predict_latent_loss_type='group_kl',
@@ -192,7 +192,17 @@ 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_muzerolossweight_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}'
+
+    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_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}'
+
+    # 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-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_muzerolossweight_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_20250819/{env_id[:-14]}/{env_id[:-14]}_uz_fix-init-recur_clear20_origlossweight_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_20250819/{env_id[:-14]}/{env_id[:-14]}_uz_lrucache-init-recur_clear20_muzerolossweight_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}'
 
@@ -246,7 +256,7 @@ def main(env_id, seed):
     main(args.env, args.seed)
 
     """
-    export CUDA_VISIBLE_DEVICES=4
+    export CUDA_VISIBLE_DEVICES=0
     cd /fs-computility/niuyazhe/puyuan/code/LightZero
     python /fs-computility/niuyazhe/puyuan/code/LightZero/zoo/atari/config/atari_unizero_segment_config.py
     """