fix(pu): fix init_recur kv share pool index

Author: puyuan1996

lzero/model/unizero_world_models/world_model.py

@@ -231,7 +231,7 @@ def _debug_check_for_stale_pointers(self, env_id: int, current_key: Any, index_t
 
                 # 打印详细的调试信息
                 print("="*60)
-                print(f"!!! BUG CONDITION DETECTED (Detection #{self.stale_pointer_detections}) !!!")
+                print(f"!!! INIT BUG CONDITION DETECTED (Detection #{self.stale_pointer_detections}) !!!")
                 print(f"    Environment ID: {env_id}")
                 print(f"    Pool Index to be overwritten: {index_to_be_written}")
                 print(f"    New state hash being written: '{current_key}'")
@@ -466,7 +466,18 @@ def _initialize_cache_structures(self) -> None:
         from collections import defaultdict
         # self.past_kv_cache_recurrent_infer = defaultdict(dict)
         # 使用 LRUCache 替换 defaultdict，并同步容量
-        self.past_kv_cache_recurrent_infer = LRUCache(self.shared_pool_size_recur)
+
+        # ========================= 核心修复与注释 (Recurrent Infer) =========================
+        # 问题: recurrent_infer 缓存同样存在 LRUCache 与环形缓冲区逻辑不匹配的问题。
+        #
+        # 修复方案:
+        # 1. 将 past_kv_cache_recurrent_infer 从 LRUCache 改为标准字典。
+        # 2. 引入辅助列表 pool_idx_to_key_map_recur_infer 来维护反向映射。
+        #    这确保了在覆写 recurrent 数据池中的条目时，可以同步删除旧的指针。
+        
+        self.past_kv_cache_recurrent_infer = {}
+        self.pool_idx_to_key_map_recur_infer = [None] * self.shared_pool_size_recur
+        # ========================== 修复结束 ==========================
 
         # self.past_kv_cache_init_infer_envs = [defaultdict(dict) for _ in range(self.env_num)]
 
@@ -490,7 +501,23 @@ def _initialize_cache_structures(self) -> None:
         #    完全同步。当数据池的索引0被新数据覆盖时，指向旧索引0的指针也已被自动清除。
         # 3. 杜绝污染: 从根本上解决了Episode内部的状态哈希碰撞问题。
 
-        self.past_kv_cache_init_infer_envs = [LRUCache(self.shared_pool_size_init) for _ in range(self.env_num)]
+        # self.past_kv_cache_init_infer_envs = [LRUCache(self.shared_pool_size_init-1) for _ in range(self.env_num)]
+        # ========================== 修复结束 ==========================
+
+        # ========================= 核心修复与注释 =========================
+        # 问题: LRUCache 的淘汰逻辑（基于访问顺序）与环形缓冲区的覆写逻辑（基于写入顺序）不匹配，导致指针过时。
+        #
+        # 修复方案:
+        # 1. 使用一个标准的字典 `past_kv_cache_init_infer_envs` 来存储 {state_hash -> pool_index}。
+        # 2. 引入一个辅助列表 `pool_idx_to_key_map_init_envs` 来维护反向映射 {pool_index -> state_hash}。
+        #
+        # 效果:
+        # 在向环形缓冲区的某个索引写入新数据之前，我们可以通过辅助列表立即找到即将被覆盖的旧 state_hash，
+        # 并从主字典中精确地删除这个过时的条目。这确保了字典和数据池的完全同步。
+        
+        self.past_kv_cache_init_infer_envs = [{} for _ in range(self.env_num)]
+        # 辅助数据结构，用于反向查找：pool_index -> key
+        self.pool_idx_to_key_map_init_envs = [[None] * self.shared_pool_size_init for _ in range(self.env_num)]
         # ========================== 修复结束 ==========================
 
         self.keys_values_wm_list = []
@@ -1365,34 +1392,68 @@ def update_cache_context(self, latent_state, is_init_infer=True, simulation_inde
 
             if is_init_infer:
                 # TODO
-                # ==================== DEBUG CODE INSERTION ====================
-                # 在写入之前，先获取将要写入的索引
+                # ==================== 主动淘汰修复逻辑 ====================
+                # 1. 获取即将被覆写的物理索引
                 index_to_write = self.shared_pool_index_init_envs[i]
+
+                # 2. 使用辅助列表查找该索引上存储的旧的 key
+                old_key_to_evict = self.pool_idx_to_key_map_init_envs[i][index_to_write]
+
+                # 3. 如果存在旧 key，就从主 cache map 中删除它
+                if old_key_to_evict is not None:
+                    # 确保要删除的键确实存在，避免意外错误
+                    if old_key_to_evict in self.past_kv_cache_init_infer_envs[i]:
+                        del self.past_kv_cache_init_infer_envs[i][old_key_to_evict]
+
+                # 现在可以安全地写入新数据了
+                cache_index = self.custom_copy_kv_cache_to_shared_init_envs(self.keys_values_wm_single_env, i)
+                
+                # 4. 在主 cache map 和辅助列表中同时更新新的映射关系
+                self.past_kv_cache_init_infer_envs[i][cache_key] = cache_index
+                self.pool_idx_to_key_map_init_envs[i][index_to_write] = cache_key
+
                 # 调用调试函数进行检查
                 self._debug_check_for_stale_pointers(env_id=i, current_key=cache_key, index_to_be_written=index_to_write)
                 # ============================================================
 
                 # Store the latest key-value cache for initial inference
-                cache_index = self.custom_copy_kv_cache_to_shared_init_envs(self.keys_values_wm_single_env, i)
-                self.past_kv_cache_init_infer_envs[i][cache_key] = cache_index
+                # cache_index = self.custom_copy_kv_cache_to_shared_init_envs(self.keys_values_wm_single_env, i)
+                # self.past_kv_cache_init_infer_envs[i][cache_key] = cache_index
             else:
                 # TODO 获取要存入的cache的某个唯一标识，例如tensor的和
                 # cache_to_store = self.keys_values_wm_single_env._keys_values[0]._k_cache._cache
                 # cache_sum = torch.sum(cache_to_store).item()
                 # cache_shape = cache_to_store.shape
                 # print(f"[CACHE WRITE] Storing for key={cache_key}, cache_shape={cache_shape}, cache_sum={cache_sum:.4f}")
 
-                # ==================== DEBUG CODE INSERTION ====================
-                # 在写入之前，获取将要写入的索引
+                # ==================== RECURRENT INFER FIX ====================
+                # 1. 获取即将被覆写的物理索引
                 index_to_write = self.shared_pool_index
+
+                # 2. 使用辅助列表查找该索引上存储的旧的 key
+                old_key_to_evict = self.pool_idx_to_key_map_recur_infer[index_to_write]
+
+                # 3. 如果存在旧 key，就从主 cache map 中删除它
+                if old_key_to_evict is not None:
+                    if old_key_to_evict in self.past_kv_cache_recurrent_infer:
+                        del self.past_kv_cache_recurrent_infer[old_key_to_evict]
+
+                # 4. 现在可以安全地写入新数据了
+                cache_index = self.custom_copy_kv_cache_to_shared_recur(self.keys_values_wm_single_env)
+
+                # 5. 在主 cache map 和辅助列表中同时更新新的映射关系
+                self.past_kv_cache_recurrent_infer[cache_key] = cache_index
+                self.pool_idx_to_key_map_recur_infer[index_to_write] = cache_key
+                # ============================================================
+
+                # ==================== DEBUG CODE INSERTION ====================
                 # 调用调试函数进行检查
                 self._debug_check_for_stale_pointers_recur(current_key=cache_key, index_to_be_written=index_to_write)
                 # ============================================================
 
-
                 # Store the latest key-value cache for recurrent inference
-                cache_index = self.custom_copy_kv_cache_to_shared_recur(self.keys_values_wm_single_env)
-                self.past_kv_cache_recurrent_infer[cache_key] = cache_index
+                # cache_index = self.custom_copy_kv_cache_to_shared_recur(self.keys_values_wm_single_env)
+                # self.past_kv_cache_recurrent_infer[cache_key] = cache_index
 
 
     def retrieve_or_generate_kvcache(self, latent_state: list, ready_env_num: int,

lzero/policy/unizero.py

@@ -955,7 +955,6 @@ def _reset_collect(self, env_id: int = None, current_steps: int = None, reset_in
 
         clear_interval = 2000 if getattr(self._cfg, 'sample_type', '') == 'episode' else self._cfg.game_segment_length
 
-        
 
         # Clear caches if the current steps are a multiple of the clear interval
         if current_steps is not None and current_steps % clear_interval == 0:
@@ -971,8 +970,7 @@ def _reset_collect(self, env_id: int = None, current_steps: int = None, reset_in
             # Free up GPU memory
             torch.cuda.empty_cache()
 
-            print('collector: collect_model clear()')
-            print(f'eps_steps_lst[{env_id}]: {current_steps}')
+            print(f'eps_steps_lst[{env_id}]: {current_steps}, collector: collect_model clear()')
 
 
     def _reset_eval(self, env_id: int = None, current_steps: int = None, reset_init_data: bool = True) -> None:

lzero/worker/muzero_segment_collector.py

@@ -603,9 +603,9 @@ def collect(self,
 
                     # ============ TODO(pu): only for UniZero now ============
                     if self._policy.get_attribute('cfg').type in ['unizero', 'sampled_unizero']:
-                        if eps_steps_lst[env_id]>=self.policy_config.game_segment_length:
+                        if eps_steps_lst[env_id]>self.policy_config.game_segment_length:
                             self._policy.reset(env_id=env_id, current_steps=eps_steps_lst[env_id], reset_init_data=False)
-                            print("eps_steps_lst[env_id]>=self.policy_config.game_segment_length")
+                            print(f"eps_steps_lst[env_id]>self.policy_config.game_segment_length:{eps_steps_lst[env_id]}>{self.policy_config.game_segment_length}")
 
                     # if self._policy.get_attribute('cfg').type in ['unizero', 'sampled_unizero']:
                     #     self._policy.reset(env_id=env_id, current_steps=eps_steps_lst[env_id], reset_init_data=False)

zoo/atari/config/atari_unizero_segment_config.py

@@ -26,8 +26,8 @@ def main(env_id, seed):
     # 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
@@ -192,7 +192,11 @@ 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_lrucache-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_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_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}'
+
+    # 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-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_debug/{env_id[:-14]}/{env_id[:-14]}_uz_lrucache_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_clear40_muzerolossweight_spsi20_envnum8_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}'
@@ -226,13 +230,13 @@ def main(env_id, seed):
 
     args.env = 'MsPacmanNoFrameskip-v4'
     # args.env = 'QbertNoFrameskip-v4'
+    # args.env = 'SeaquestNoFrameskip-v4' 
 
     # args.env = 'SpaceInvadersNoFrameskip-v4'
+
     # args.env = 'BeamRiderNoFrameskip-v4'
     # args.env = 'GravitarNoFrameskip-v4'
 
-
-    # args.env = 'SeaquestNoFrameskip-v4' 
     # args.env = 'BreakoutNoFrameskip-v4'
 
 
@@ -242,7 +246,7 @@ def main(env_id, seed):
     main(args.env, args.seed)
 
     """
-    export CUDA_VISIBLE_DEVICES=6
+    export CUDA_VISIBLE_DEVICES=4
     cd /fs-computility/niuyazhe/puyuan/code/LightZero
     python /fs-computility/niuyazhe/puyuan/code/LightZero/zoo/atari/config/atari_unizero_segment_config.py
     """