fix(pu): fix task_id bug in balance pipeline, and polish benchmark_name option

Author: puyuan

lzero/entry/train_unizero_multitask_balance_segment_ddp.py

@@ -27,72 +27,6 @@
 import numpy as np                    # 计算均值
 from collections import defaultdict   # 保存所有任务最近一次评估分数
 
-# ====== UniZero-MT 需要用到的基准分数（与 26 个 Atari100k 任务 id 一一对应）======
-# 原始的 RANDOM_SCORES 和 HUMAN_SCORES
-
-
-global BENCHMARK_NAME
-BENCHMARK_NAME = "atari"
-# BENCHMARK_NAME = "dmc" # TODO
-if BENCHMARK_NAME == "atari":
-    RANDOM_SCORES = np.array([
-        227.8, 5.8, 222.4, 210.0, 14.2, 2360.0, 0.1, 1.7, 811.0, 10780.5,
-        152.1, 0.0, 65.2, 257.6, 1027.0, 29.0, 52.0, 1598.0, 258.5, 307.3,
-        -20.7, 24.9, 163.9, 11.5, 68.4, 533.4
-    ])
-    HUMAN_SCORES = np.array([
-        7127.7, 1719.5, 742.0, 8503.3, 753.1, 37187.5, 12.1, 30.5, 7387.8, 35829.4,
-        1971.0, 29.6, 4334.7, 2412.5, 30826.4, 302.8, 3035.0, 2665.5, 22736.3, 6951.6,
-        14.6, 69571.3, 13455.0, 7845.0, 42054.7, 11693.2
-    ])
-elif BENCHMARK_NAME == "dmc":
-    RANDOM_SCORES = np.array([0]*26)
-    HUMAN_SCORES = np.array([1000]*26)
-
-# 新顺序对应的原始索引列表
-# 新顺序： [Pong, MsPacman, Seaquest, Boxing, Alien, ChopperCommand, Hero, RoadRunner,
-#            Amidar, Assault, Asterix, BankHeist, BattleZone, CrazyClimber, DemonAttack,
-#            Freeway, Frostbite, Gopher, Jamesbond, Kangaroo, Krull, KungFuMaster,
-#            PrivateEye, UpNDown, Qbert, Breakout]
-# 映射为原始数组中的索引（注意：索引均从0开始）
-new_order = [
-    20,  # Pong
-    19,  # MsPacman
-    24,  # Seaquest
-    6,   # Boxing
-    0,   # Alien
-    8,   # ChopperCommand
-    14,  # Hero
-    23,  # RoadRunner
-    1,   # Amidar
-    2,   # Assault
-    3,   # Asterix
-    4,   # BankHeist
-    5,   # BattleZone
-    9,   # CrazyClimber
-    10,  # DemonAttack
-    11,  # Freeway
-    12,  # Frostbite
-    13,  # Gopher
-    15,  # Jamesbond
-    16,  # Kangaroo
-    17,  # Krull
-    18,  # KungFuMaster
-    21,  # PrivateEye
-    25,  # UpNDown
-    22,  # Qbert
-    7    # Breakout
-]
-
-# 根据 new_order 生成新的数组
-new_RANDOM_SCORES = RANDOM_SCORES[new_order]
-new_HUMAN_SCORES = HUMAN_SCORES[new_order]
-
-# 查看重排后的结果
-print("重排后的 RANDOM_SCORES:")
-print(new_RANDOM_SCORES)
-print("\n重排后的 HUMAN_SCORES:")
-print(new_HUMAN_SCORES)
 
 # 保存最近一次评估回报：{task_id: eval_episode_return_mean}
 from collections import defaultdict
@@ -354,6 +288,7 @@ def train_unizero_multitask_balance_segment_ddp(
         model_path: Optional[str] = None,
         max_train_iter: Optional[int] = int(1e10),
         max_env_step: Optional[int] = int(1e10),
+        benchmark_name: str = "atari"    
 ) -> 'Policy':
     """
     Overview:
@@ -378,6 +313,73 @@ def train_unizero_multitask_balance_segment_ddp(
     Returns:
         - policy (:obj:`Policy`): 收敛的策略。
     """
+
+    # ---------------------------------------------------------------
+    # ====== UniZero-MT 需要用到的基准分数（与 26 个 Atari100k 任务 id 一一对应）======
+    #   原始的 RANDOM_SCORES 和 HUMAN_SCORES
+    if benchmark_name == "atari":
+        RANDOM_SCORES = np.array([
+            227.8, 5.8, 222.4, 210.0, 14.2, 2360.0, 0.1, 1.7, 811.0, 10780.5,
+            152.1, 0.0, 65.2, 257.6, 1027.0, 29.0, 52.0, 1598.0, 258.5, 307.3,
+            -20.7, 24.9, 163.9, 11.5, 68.4, 533.4
+        ])
+        HUMAN_SCORES = np.array([
+            7127.7, 1719.5, 742.0, 8503.3, 753.1, 37187.5, 12.1, 30.5, 7387.8, 35829.4,
+            1971.0, 29.6, 4334.7, 2412.5, 30826.4, 302.8, 3035.0, 2665.5, 22736.3, 6951.6,
+            14.6, 69571.3, 13455.0, 7845.0, 42054.7, 11693.2
+        ])
+    elif benchmark_name == "dmc":
+        # RANDOM_SCORES = np.array([0]*26)
+        # HUMAN_SCORES = np.array([1000]*26)
+        RANDOM_SCORES = np.zeros(26)
+        HUMAN_SCORES  = np.ones(26) * 1000
+    else:
+        raise ValueError(f"Unsupported BENCHMARK_NAME: {BENCHMARK_NAME}")
+
+    # 新顺序对应的原始索引列表
+    # 新顺序： [Pong, MsPacman, Seaquest, Boxing, Alien, ChopperCommand, Hero, RoadRunner,
+    #            Amidar, Assault, Asterix, BankHeist, BattleZone, CrazyClimber, DemonAttack,
+    #            Freeway, Frostbite, Gopher, Jamesbond, Kangaroo, Krull, KungFuMaster,
+    #            PrivateEye, UpNDown, Qbert, Breakout]
+    # 映射为原始数组中的索引（注意：索引均从0开始）
+    new_order = [
+        20,  # Pong
+        19,  # MsPacman
+        24,  # Seaquest
+        6,   # Boxing
+        0,   # Alien
+        8,   # ChopperCommand
+        14,  # Hero
+        23,  # RoadRunner
+        1,   # Amidar
+        2,   # Assault
+        3,   # Asterix
+        4,   # BankHeist
+        5,   # BattleZone
+        9,   # CrazyClimber
+        10,  # DemonAttack
+        11,  # Freeway
+        12,  # Frostbite
+        13,  # Gopher
+        15,  # Jamesbond
+        16,  # Kangaroo
+        17,  # Krull
+        18,  # KungFuMaster
+        21,  # PrivateEye
+        25,  # UpNDown
+        22,  # Qbert
+        7    # Breakout
+    ]
+    # 根据 new_order 生成新的数组
+    new_RANDOM_SCORES = RANDOM_SCORES[new_order]
+    new_HUMAN_SCORES = HUMAN_SCORES[new_order]
+    # 查看重排后的结果
+    print("重排后的 RANDOM_SCORES:")
+    print(new_RANDOM_SCORES)
+    print("\n重排后的 HUMAN_SCORES:")
+    print(new_HUMAN_SCORES)
+    # ---------------------------------------------------------------
+
     # 初始化温度调度器
     initial_temperature = 10.0
     final_temperature = 1.0
@@ -552,7 +554,8 @@ def train_unizero_multitask_balance_segment_ddp(
             # TODO: ============
             # cfg.policy.target_return = 10
             #  ==================== 如果任务已解决，则不参与后续评估和采集 TODO: ddp ====================
-            if task_id in solved_task_pool:
+            # if task_id in solved_task_pool:
+            if cfg.policy.task_id in solved_task_pool:
                 continue
 
             # 记录缓冲区内存使用情况
@@ -601,8 +604,10 @@ def train_unizero_multitask_balance_segment_ddp(
 
                         # 如果达到目标奖励，将任务移入 solved_task_pool
                         if eval_mean_reward >= cfg.policy.target_return:
-                            print(f"任务 {task_id} 达到了目标奖励 {cfg.policy.target_return}, 移入 solved_task_pool.")
-                            solved_task_pool.add(task_id)
+                            cur_task_id = cfg.policy.task_id
+                            print(f"任务 {cur_task_id} 达到了目标奖励 {cfg.policy.target_return}, 移入 solved_task_pool.")
+                            solved_task_pool.add(cur_task_id)
+
 
                     except Exception as e:
                         print(f"提取评估奖励时发生错误: {e}")

lzero/entry/train_unizero_multitask_segment_ddp.py

@@ -28,8 +28,8 @@
 
 
 global BENCHMARK_NAME
-BENCHMARK_NAME = "atari"
-# BENCHMARK_NAME = "dmc" # TODO
+# BENCHMARK_NAME = "atari"
+BENCHMARK_NAME = "dmc" # TODO
 if BENCHMARK_NAME == "atari":
     RANDOM_SCORES = np.array([
         227.8, 5.8, 222.4, 210.0, 14.2, 2360.0, 0.1, 1.7, 811.0, 10780.5,

zoo/atari/config/atari_unizero_multitask_segment_ddp_balance_config.py

@@ -190,7 +190,7 @@ def create_config(env_id, action_space_size, collector_env_num, evaluator_env_nu
             n_episode=n_episode,
             replay_buffer_size=int(5e5),
             # eval_freq=int(1e4),
-            eval_freq=int(1.5e4),
+            eval_freq=int(1e4),
             # eval_freq=int(2),
             collector_env_num=collector_env_num,
             evaluator_env_num=evaluator_env_num,
@@ -208,7 +208,7 @@ def generate_configs(env_id_list, action_space_size, collector_env_num, n_episod
     # ===== only for debug =====
     # exp_name_prefix = f'data_lz/data_unizero_atari_mt_balance_20250509/atari_{len(env_id_list)}games_balance-total-stage{curriculum_stage_num}_vit-encoder-ps8_trans-nlayer8_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
     # exp_name_prefix = f'data_lz/data_unizero_atari_mt_balance_20250509/atari_{len(env_id_list)}games_balance-total-stage{curriculum_stage_num}_no-encoder-scale_cnn-encoder_moe8_trans-nlayer8_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
-    exp_name_prefix = f'data_lz/data_unizero_atari_mt_balance_20250509/atari_{len(env_id_list)}games_balance-total-stage{curriculum_stage_num}_vit-ln_moe8_trans-nlayer4_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
+    exp_name_prefix = f'data_lz/data_unizero_atari_mt_balance_20250514/atari_{len(env_id_list)}games_balance-total-stage{curriculum_stage_num}_vit-ln_moe8_trans-nlayer4_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
 
     for task_id, env_id in enumerate(env_id_list):
         config = create_config(
@@ -404,11 +404,10 @@ def get_atari_target_return_dict(ratio=1.0):
     ]
 
     global curriculum_stage_num
-
+    # TODO ==============
     curriculum_stage_num=3
-    # curriculum_stage_num=5
-    curriculum_stage_num=9
-
+    curriculum_stage_num=5
+    # curriculum_stage_num=9
 
     action_space_size = 18
     collector_env_num = 8
@@ -461,6 +460,6 @@ def get_atari_target_return_dict(ratio=1.0):
                                    num_segments, total_batch_size)
 
         with DDPContext():
-            train_unizero_multitask_balance_segment_ddp(configs, seed=seed, max_env_step=max_env_step)
+            train_unizero_multitask_balance_segment_ddp(configs, seed=seed, max_env_step=max_env_step, benchmark_name="atari")
             # ======== TODO: only for debug ========
             # train_unizero_multitask_segment_ddp(configs[:2], seed=seed, max_env_step=max_env_step) # train on the first four tasks

zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py

@@ -64,8 +64,8 @@ def create_config(env_id, action_space_size, collector_env_num, evaluator_env_nu
         policy=dict(
             multi_gpu=True,  # Very important for ddp
             only_use_moco_stats=False,
-            # use_moco=False,  # ==============TODO==============
-            use_moco=True,  # ==============TODO==============
+            use_moco=False,  # ==============TODO==============
+            # use_moco=True,  # ==============TODO==============
             learn=dict(learner=dict(hook=dict(save_ckpt_after_iter=200000))),
             grad_correct_params=dict(
                 MoCo_beta=0.5, MoCo_beta_sigma=0.5, MoCo_gamma=0.1, MoCo_gamma_sigma=0.5, MoCo_rho=0,
@@ -99,8 +99,8 @@ def create_config(env_id, action_space_size, collector_env_num, evaluator_env_nu
 
                     analysis_dormant_ratio_weight_rank=True, 
                     # analysis_dormant_ratio_weight_rank=False, # TODO
-                    analysis_dormant_ratio_interval=100,
-                    # analysis_dormant_ratio_interval=1000,
+                    # analysis_dormant_ratio_interval=100,
+                    analysis_dormant_ratio_interval=1000,
                     # analysis_dormant_ratio_interval=20,
 
                     continuous_action_space=False,
@@ -123,6 +123,7 @@ def create_config(env_id, action_space_size, collector_env_num, evaluator_env_nu
                     # num_heads=24,
 
                     num_layers=8,
+
                     # num_layers=12, # todo
                     num_heads=24,
 
@@ -134,8 +135,8 @@ def create_config(env_id, action_space_size, collector_env_num, evaluator_env_nu
                     obs_type='image',
                     env_num=8,
                     task_num=len(env_id_list),
-                    # encoder_type='vit',
-                    encoder_type='resnet',
+                    encoder_type='vit',
+                    # encoder_type='resnet',
 
                     use_normal_head=True,
                     use_softmoe_head=False,
@@ -197,7 +198,9 @@ def generate_configs(env_id_list, action_space_size, collector_env_num, n_episod
     configs = []
     # ===== only for debug =====
     # ========= TODO: global BENCHMARK_NAME =========
-    exp_name_prefix = f'data_lz/data_unizero_atari_mt_20250508/atari_{len(env_id_list)}games_orig-ln_moco_tran-nlayer8_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
+    exp_name_prefix = f'data_lz/data_unizero_atari_mt_20250508/atari_{len(env_id_list)}games_orig_vit-ln_tran-nlayer8_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
+
+    # exp_name_prefix = f'data_lz/data_unizero_atari_mt_20250508/atari_{len(env_id_list)}games_orig-ln_moco_tran-nlayer8_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
 
     # exp_name_prefix = f'data_lz/data_unizero_atari_mt_20250508/atari_{len(env_id_list)}games_orig_simnorm_tran-nlayer8_brf{buffer_reanalyze_freq}_not-share-head_seed{seed}/'
 
@@ -251,7 +254,7 @@ def create_env_manager():
     Overview:
         This script should be executed with <nproc_per_node> GPUs.
         Run the following command to launch the script:
-        python -m torch.distributed.launch --nproc_per_node=8 --master_port=29502 ./zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py 2>&1 | tee ./log/20250509/uz_mt_atari8_orig-ln_moco.log
+        python -m torch.distributed.launch --nproc_per_node=8 --master_port=29502 ./zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py 2>&1 | tee ./log/20250509/uz_mt_atari8_orig_vit-ln.log
         python -m torch.distributed.launch --nproc_per_node=8 --master_port=29502 ./zoo/atari/config/atari_unizero_multitask_segment_ddp_config.py 2>&1 | tee ./log/uz_mt_atari8_orig-simnorm.log
 
 

zoo/dmc2gym/config/dmc2gym_state_suz_multitask_ddp_balance_config.py

@@ -205,7 +205,7 @@ def generate_configs(env_id_list: List[str],
     configs = []
     # ========= TODO: global BENCHMARK_NAME =========
 
-    exp_name_prefix = f'data_lz/data_suz_dmc_mt_balance_20250509/dmc_{len(env_id_list)}tasks_frameskip8_balance-stage-total-{curriculum_stage_num}_moe8_nlayer4_not-share-head_brf{buffer_reanalyze_freq}_seed{seed}/'
+    exp_name_prefix = f'data_lz/data_suz_dmc_mt_balance_20250514/dmc_{len(env_id_list)}tasks_frameskip8_balance-stage-total-{curriculum_stage_num}_moe8_nlayer4_not-share-head_brf{buffer_reanalyze_freq}_seed{seed}/'
 
     # exp_name_prefix = f'data_lz/data_suz_dmc_mt_20250409_moco/dmc_{len(env_id_list)}tasks_notaskembed_nlayer8_not-share-head_final-ln_bs64_brf{buffer_reanalyze_freq}_seed{seed}/'
 
@@ -269,8 +269,6 @@ def create_env_manager():
     import os
     from zoo.dmc2gym.config.dmc_state_env_space_map import dmc_state_env_action_space_map, dmc_state_env_obs_space_map
 
-    global BENCHMARK_NAME
-    BENCHMARK_NAME='dmc'
 
     global curriculum_stage_num
 
@@ -303,9 +301,9 @@ def create_env_manager():
     ]
 
     target_return_dict = {
-        'acrobot-swingup': 500,
-        'cartpole-balance':950,
-        'cartpole-balance_sparse':950,
+        'acrobot-swingup': 500, # 0
+        'cartpole-balance':950, # 1
+        'cartpole-balance_sparse':950, # 2
         'cartpole-swingup': 800, # 3
         'cartpole-swingup_sparse': 750, # 4
         'cheetah-run': 650, # 5
@@ -405,6 +403,6 @@ def create_env_manager():
     )
 
     with DDPContext():
-        train_unizero_multitask_balance_segment_ddp(configs, seed=seed, max_env_step=max_env_step)
+        train_unizero_multitask_balance_segment_ddp(configs, seed=seed, max_env_step=max_env_step, benchmark_name="dmc")
         # 如果只想训练部分任务，可以修改 configs，例如:
         # train_unizero_multitask_segment_ddp(configs[:4], seed=seed, max_env_step=max_env_step)