feature(pu): add decode_loss for unizero atari

Author: puyuan1996

lzero/model/unizero_model.py

@@ -126,8 +126,9 @@ def __init__(
                     self.decoder_network_tokenizer = None
             else:
                 raise ValueError(f"Unsupported encoder option: {kwargs['encoder_option']}")     
+            
             self.tokenizer = Tokenizer(encoder=self.representation_network, decoder=self.decoder_network, decoder_network_tokenizer=self.decoder_network_tokenizer, 
-                                    with_lpips=False, projection=projection, encoder_option=kwargs['encoder_option'])     
+                                    with_lpips=False, projection=projection, encoder_option=kwargs['encoder_option']) 
             self.world_model = WorldModel(config=world_model_cfg, tokenizer=self.tokenizer)
             print(f'{sum(p.numel() for p in self.world_model.parameters())} parameters in agent.world_model')
             print('==' * 20)
@@ -168,8 +169,19 @@ def __init__(
             if world_model_cfg.analysis_sim_norm:
                 self.encoder_hook = FeatureAndGradientHook()
                 self.encoder_hook.setup_hooks(self.representation_network)
-                
-            self.tokenizer = Tokenizer(encoder=self.representation_network, decoder=None, with_lpips=False, obs_type=world_model_cfg.obs_type)
+            
+            if world_model_cfg.latent_recon_loss_weight==0:
+                self.tokenizer = Tokenizer(encoder=self.representation_network, decoder=None, with_lpips=False, obs_type=world_model_cfg.obs_type)
+            else:
+                # TODO =============
+                self.decoder_network = LatentDecoder(
+                    embedding_dim=world_model_cfg.embed_dim,
+                    output_shape=[3, 64, 64],
+                    num_channels = 64,
+                    activation=self.activation,
+                )
+                self.tokenizer = Tokenizer(encoder=self.representation_network, decoder=self.decoder_network, with_lpips=True, obs_type=world_model_cfg.obs_type)
+
             self.world_model = WorldModel(config=world_model_cfg, tokenizer=self.tokenizer)
             print(f'{sum(p.numel() for p in self.world_model.parameters())} parameters in agent.world_model')
             print('==' * 20)

lzero/model/unizero_world_models/lpips.py

@@ -13,6 +13,20 @@
 from torchvision import models
 from tqdm import tqdm
 
+# ==================================================================================
+# ================================ 核心修改部分 ====================================
+# ==================================================================================
+# 在导入 torch 和 torchvision 之后，但在实例化任何模型之前，设置 TORCH_HOME 环境变量。
+# 这会告诉 PyTorch 将所有通过 torch.hub 下载的模型（包括 torchvision.models 中的预训练模型）
+# 存放到您指定的目录下。
+# PyTorch 会自动在此目录下创建 'hub/checkpoints' 子文件夹。
+custom_torch_home = "/mnt/shared-storage-user/puyuan/code_20250828/LightZero/tokenizer_pretrained_vgg"
+os.environ['TORCH_HOME'] = custom_torch_home
+# 确保目录存在，虽然 torch.hub 也会尝试创建，但提前创建更稳妥
+os.makedirs(os.path.join(custom_torch_home, 'hub', 'checkpoints'), exist_ok=True)
+# ==================================================================================
+# ==================================================================================
+
 
 class LPIPS(nn.Module):
     # Learned perceptual metric
@@ -22,19 +36,23 @@ def __init__(self, use_dropout: bool = True):
         self.chns = [64, 128, 256, 512, 512]  # vg16 features
 
         # Comment out the following line if you don't need perceptual loss
-        # self.net = vgg16(pretrained=True, requires_grad=False)
-        # self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
-        # self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
-        # self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
-        # self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
-        # self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
-        # self.load_from_pretrained()
-        # for param in self.parameters():
-        #     param.requires_grad = False
+        # 现在，这一行将自动使用 TORCH_HOME 指定的路径
+        self.net = vgg16(pretrained=True, requires_grad=False)
+        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
+        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
+        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
+        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
+        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
+        self.load_from_pretrained()
+        for param in self.parameters():
+            param.requires_grad = False
 
     def load_from_pretrained(self) -> None:
-        ckpt = get_ckpt_path(name="vgg_lpips", root=Path.home() / ".cache/iris/tokenizer_pretrained_vgg")  # Download VGG if necessary
+        # 这一部分您已经修改正确，它用于加载 LPIPS 的线性层权重 (vgg.pth)
+        # 我们让它和 TORCH_HOME 使用相同的根目录，以保持一致性。
+        ckpt = get_ckpt_path(name="vgg_lpips", root=custom_torch_home)
         self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
+        print(f"Loaded LPIPS pretrained weights from: {ckpt}")
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
@@ -74,7 +92,10 @@ def __init__(self, chn_in: int, chn_out: int = 1, use_dropout: bool = False) ->
 class vgg16(torch.nn.Module):
     def __init__(self, requires_grad: bool = False, pretrained: bool = True) -> None:
         super(vgg16, self).__init__()
+        # 由于设置了 TORCH_HOME，这里的 pretrained=True 会在指定目录中查找或下载模型
+        print("Loading vgg16 backbone...")
         vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
+        print("vgg16 backbone loaded.")
         self.slice1 = torch.nn.Sequential()
         self.slice2 = torch.nn.Sequential()
         self.slice3 = torch.nn.Sequential()
@@ -160,10 +181,26 @@ def md5_hash(path: str) -> str:
 
 def get_ckpt_path(name: str, root: str, check: bool = False) -> str:
     assert name in URL_MAP
+    # 这个函数现在只为 vgg.pth 服务，路径是正确的
     path = os.path.join(root, CKPT_MAP[name])
     if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
         print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
         download(URL_MAP[name], path)
         md5 = md5_hash(path)
         assert md5 == MD5_MAP[name], md5
     return path
+
+# =======================
+# =====  运行示例  ======
+# =======================
+if __name__ == '__main__':
+    print(f"PyTorch Hub directory set to: {os.environ['TORCH_HOME']}")
+    
+    # 第一次运行时，你会看到两个下载过程：
+    # 1. 下载 vgg16-397923af.pth 到 /mnt/shared-storage-user/puyuan/code_20250828/LightZero/tokenizer_pretrained_vgg/hub/checkpoints/
+    # 2. 下载 vgg.pth 到 /mnt/shared-storage-user/puyuan/code_20250828/LightZero/tokenizer_pretrained_vgg/
+    # 之后再次运行，将不会有任何下载提示，直接从指定目录加载。
+    
+    print("\nInitializing LPIPS model...")
+    model = LPIPS()
+    print("\nLPIPS model initialized successfully.")

lzero/model/unizero_world_models/world_model.py

@@ -1759,7 +1759,7 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
 
         if self.obs_type == 'image':
             # Reconstruct observations from latent state representations
-            # reconstructed_images = self.tokenizer.decode_to_obs(obs_embeddings)
+            reconstructed_images = self.tokenizer.decode_to_obs(obs_embeddings)
 
             #  ========== for visualization ==========
             # Uncomment the lines below for visual analysis
@@ -1772,11 +1772,12 @@ def compute_loss(self, batch, target_tokenizer: Tokenizer = None, inverse_scalar
             #  ========== for visualization ==========
 
             # ========== Calculate reconstruction loss and perceptual loss ============
-            # latent_recon_loss = self.tokenizer.reconstruction_loss(batch['observations'].reshape(-1, 3, 64, 64), reconstructed_images) # NOTE: for stack=1
-            # perceptual_loss = self.tokenizer.perceptual_loss(batch['observations'].reshape(-1, 3, 64, 64), reconstructed_images) # NOTE: for stack=1
+            latent_recon_loss = self.tokenizer.reconstruction_loss(batch['observations'].reshape(-1, 3, 64, 64), reconstructed_images) # NOTE: for stack=1
+            perceptual_loss = self.tokenizer.perceptual_loss(batch['observations'].reshape(-1, 3, 64, 64), reconstructed_images) # NOTE: for stack=1
 
-            latent_recon_loss = self.latent_recon_loss
-            perceptual_loss = self.perceptual_loss
+            # TODO:
+            # latent_recon_loss = self.latent_recon_loss
+            # perceptual_loss = self.perceptual_loss
 
         elif self.obs_type == 'vector':
             perceptual_loss = torch.tensor(0., device=batch['observations'].device,

lzero/policy/unizero.py

@@ -840,9 +840,10 @@ def _forward_learn(self, data: Tuple[torch.Tensor]) -> Dict[str, Union[float, in
         # logits_policy_mean=self.intermediate_losses['logits_policy_mean']
         # logits_policy_max=self.intermediate_losses['logits_policy_max']
         # logits_policy_min=self.intermediate_losses['logits_policy_min']
-        # temperature_value=self.intermediate_losses['temperature_value']
-        # temperature_reward=self.intermediate_losses['temperature_reward']
-        # temperature_policy=self.intermediate_losses['temperature_policy']
+
+        temperature_value=self.intermediate_losses['temperature_value']
+        temperature_reward=self.intermediate_losses['temperature_reward']
+        temperature_policy=self.intermediate_losses['temperature_policy']
 
         assert not torch.isnan(losses.loss_total).any(), "Loss contains NaN values"
         assert not torch.isinf(losses.loss_total).any(), "Loss contains Inf values"
@@ -898,13 +899,30 @@ def _forward_learn(self, data: Tuple[torch.Tensor]) -> Dict[str, Union[float, in
             self.reward_loss_weight = 1.
             self.policy_loss_weight = 1.
             self.ends_loss_weight = 0.
-            total_loss = (
-                self.reward_loss_weight * reward_loss +
-                self.value_loss_weight * value_loss +
-                self.policy_loss_weight * weighted_policy_loss +
-                self.obs_loss_weight  * obs_loss # 假设 ssl_loss_weight 是 obs_loss 的权重
-                # ... 如果还有其他损失项，也加进来 ...
-            )
+
+            self.latent_recon_loss_weight = self._cfg.model.world_model_cfg.latent_recon_loss_weight # 默认使用固定值
+            self.perceptual_loss_weight = self._cfg.model.world_model_cfg.perceptual_loss_weight # 默认使用固定值
+
+            if self.latent_recon_loss_weight>0:
+                total_loss = (
+                    self.reward_loss_weight * reward_loss +
+                    self.value_loss_weight * value_loss +
+                    self.policy_loss_weight * weighted_policy_loss +
+                    self.obs_loss_weight  * obs_loss +   # 假设 ssl_loss_weight 是 obs_loss 的权重
+                    self.latent_recon_loss_weight * latent_recon_loss+
+                    self.perceptual_loss_weight*perceptual_loss
+                    # ... 如果还有其他损失项，也加进来 ...
+                )
+            else:
+
+                total_loss = (
+                    self.reward_loss_weight * reward_loss +
+                    self.value_loss_weight * value_loss +
+                    self.policy_loss_weight * weighted_policy_loss +
+                    self.obs_loss_weight  * obs_loss  # 假设 ssl_loss_weight 是 obs_loss 的权重
+
+                    # ... 如果还有其他损失项，也加进来 ...
+                )
             weighted_total_loss = (weights * total_loss).mean()
         # ===================== END: 目标熵正则化更新逻辑 =====================
 

zoo/atari/config/atari_unizero_segment_config.py

@@ -89,6 +89,9 @@ def main(env_id, seed):
                 num_res_blocks=2,
                 num_channels=128,
                 world_model_cfg=dict(
+                    latent_recon_loss_weight=1,
+                    perceptual_loss_weight=1,
+
                     # use_new_cache_manager=True,
                     use_new_cache_manager=False,
 
@@ -240,7 +243,9 @@ 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_refactor1024/{env_id[3:-3]}/{env_id[3:-3]}_uz_cossimloss_nokvcachemanager_ch128-res2_aug_targetentropy-alpha-100k-098-07-lr1e-3-encoder-clip30-10-100k_adamw-wd1e-2-encoder5-trans1-head0_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_refactor1024/{env_id[3:-3]}/{env_id[3:-3]}_uz_recon-perc-w1_cossimloss_nokvcachemanager_ch128-res2_aug_targetentropy-alpha-100k-098-07-lr1e-3-encoder-clip30-10-100k_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_refactor1024/{env_id[3:-3]}/{env_id[3:-3]}_uz_recon-perc-w1_cossimloss_nokvcachemanager_ch128-res2_aug_targetentropy-alpha-100k-098-07-lr1e-3-encoder-clip30-10-100k_adamw-wd1e-2-encoder5-trans1-head0_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_refactor1024/{env_id[3:-3]}/{env_id[3:-3]}_uz_cossimloss_nokvcachemanager_ch64-res1_targetentropy-alpha-100k-098-07-encoder-clip30-10-100k_adamw-wd1e-2-encoder5-trans1-head0_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}'
 
@@ -265,7 +270,7 @@ def main(env_id, seed):
 
     # 测试的atari8中的4个base环境
     # args.env = 'PongNoFrameskip-v4' # 反应型环境 密集奖励
-    args.env = 'MsPacmanNoFrameskip-v4' # 记忆规划型环境 稀疏奖励
+    # args.env = 'MsPacmanNoFrameskip-v4' # 记忆规划型环境 稀疏奖励
 
     # args.env = 'ALE/Pong-v5' # 记忆规划型环境 稀疏奖励
 
@@ -293,7 +298,7 @@ def main(env_id, seed):
     """
     tmux new -s uz-st-refactor-boxing
 
-    export CUDA_VISIBLE_DEVICES=1
+    export CUDA_VISIBLE_DEVICES=0
     cd /mnt/shared-storage-user/puyuan/code_20250828/LightZero/
     /mnt/shared-storage-user/puyuan/lz/bin/python /mnt/shared-storage-user/puyuan/code_20250828/LightZero/zoo/atari/config/atari_unizero_segment_config.py