fix(pu): fix tb log when gpu_num<task_num, fix total_loss += bug, polish

alpha_loss

Author: tAnGjIa520

lzero/entry/train_unizero_multitask_segment_ddp.py

@@ -437,12 +437,17 @@ def train_unizero_multitask_segment_ddp(
     tasks_per_rank = total_tasks // world_size
     remainder = total_tasks % world_size
 
+    # ==================== START: 关键修复 ====================
+    # 1. 精确计算当前Rank负责的任务数量
     if rank < remainder:
         start_idx = rank * (tasks_per_rank + 1)
         end_idx = start_idx + tasks_per_rank + 1
+        num_tasks_for_this_rank = tasks_per_rank + 1
     else:
         start_idx = rank * tasks_per_rank + remainder
         end_idx = start_idx + tasks_per_rank
+        num_tasks_for_this_rank = tasks_per_rank
+    # ==================== END: 关键修复 ====================
 
     tasks_for_this_rank = input_cfg_list[start_idx:end_idx]
 
@@ -465,8 +470,16 @@ def train_unizero_multitask_segment_ddp(
         # Use the config of the first task to create a shared policy.
         task_id, [cfg, create_cfg] = tasks_for_this_rank[0]
 
-        for config in tasks_for_this_rank:
-            config[1][0].policy.task_num = tasks_per_rank
+        # ==================== START: 关键修复 ====================
+        # 2. 将正确的任务数量设置到 *所有* 相关配置中
+        #    在创建Policy实例之前，必须确保配置是正确的
+        for config_tuple in tasks_for_this_rank:
+            # config_tuple is (task_id, [cfg_obj, create_cfg_obj])
+            config_tuple[1][0].policy.task_num = num_tasks_for_this_rank
+        
+        # 3. 确保用于创建Policy的那个cfg对象也拥有正确的task_num
+        cfg.policy.task_num = num_tasks_for_this_rank
+        # ==================== END: 关键修复 ====================
 
         # Ensure the specified policy type is supported.
         assert create_cfg.policy.type in ['unizero_multitask', 'sampled_unizero_multitask'], \
@@ -602,7 +615,9 @@ def train_unizero_multitask_segment_ddp(
                 collect_kwargs['epsilon'] = epsilon_greedy_fn(collector.envstep)
 
             # Check if it's time for evaluation.
-            if learner.train_iter > 10 and learner.train_iter % cfg.policy.eval_freq == 0:
+            # if learner.train_iter > 10 and learner.train_iter % cfg.policy.eval_freq == 0:
+            if learner.train_iter == 0 or learner.train_iter % cfg.policy.eval_freq == 0: # only for debug TODO
+            
                 print('=' * 20)
                 print(f'Rank {rank} 评估任务_id: {cfg.policy.task_id}...')
 

lzero/policy/unizero.py

@@ -68,14 +68,14 @@ def configure_optimizer_unizero(model, learning_rate, weight_decay, device_type,
             'params': list(tokenizer_params.values()),
             'lr': learning_rate,  # Tokenizer使用基础学习率，例如 1e-4
             # 'lr': learning_rate * 0.1,  # 为encoder设置一个较小的学习率，例如 1e-5
-            'weight_decay': weight_decay * 5.0  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
-
+            # 'weight_decay': weight_decay * 5.0  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
+            'weight_decay': weight_decay  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
         },
         {
             'params': list(head_params.values()),
             'lr': learning_rate,  # Heads也使用基础学习率率，例如 1e-4
-            'weight_decay': 0.0  # 通常Heads的权重不做衰减
-            # 'weight_decay': weight_decay
+            # 'weight_decay': 0.0  # 通常Heads的权重不做衰减
+            'weight_decay': weight_decay
 
         }
     ]

lzero/policy/unizero_multitask.py

@@ -71,6 +71,31 @@ def generate_task_loss_dict(multi_task_losses: List[Union[torch.Tensor, float]],
             task_loss_dict[task_name] = task_loss
     return task_loss_dict
 
+# # 修改后的函数:
+# def generate_task_loss_dict(
+#     multi_task_losses: List[Union[torch.Tensor, float]], 
+#     task_name_template: str, 
+#     global_task_ids: List[int]
+# ) -> Dict[str, float]:
+#     """
+#     Overview:
+#         Generates a dictionary for the losses of each task using their explicit global IDs.
+#     Arguments:
+#         - multi_task_losses (:obj:`List[Union[torch.Tensor, float]]`): A list containing the loss for each task.
+#         - task_name_template (:obj:`str`): The template for the task name, e.g., 'obs_loss_task{}'.
+#         - global_task_ids (:obj:`List[int]`): A list of global task IDs corresponding to each loss in multi_task_losses.
+#     Returns:
+#         - task_loss_dict (:obj:`Dict[str, float]`): A dictionary where keys are formatted task names and values are the corresponding losses.
+#     """
+#     task_loss_dict = {}
+#     # 使用 zip 将每个损失与其正确的全局ID配对
+#     for task_loss, global_id in zip(multi_task_losses, global_task_ids):
+#         task_name = task_name_template.format(global_id)
+#         try:
+#             task_loss_dict[task_name] = task_loss.item() if hasattr(task_loss, 'item') else task_loss
+#         except Exception as e:
+#             task_loss_dict[task_name] = task_loss
+#     return task_loss_dict
 
 
 class WrappedModel:
@@ -277,12 +302,15 @@ def configure_optimizer_unizero(model, learning_rate, weight_decay, device_type,
         {
             'params': tokenizer_params,
             'lr': learning_rate,  # Tokenizer使用基础学习率，例如 1e-4
-            'weight_decay': weight_decay * 5.0  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
+            # 'weight_decay': weight_decay * 5.0  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
+            'weight_decay': weight_decay  # <-- 为Encoder设置5倍的权重衰减！这是一个强力正则化
         },
         {
             'params': head_params,
             'lr': learning_rate,  # Heads也使用基础学习率率，例如 1e-4
-            'weight_decay': 0.0  # 通常Heads的权重不做衰减
+            # 'weight_decay': 0.0  # 通常Heads的权重不做衰减
+            'weight_decay': weight_decay 
+
         }
     ]
 
@@ -845,6 +873,7 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
         orig_policy_loss_multi_task = []
         policy_entropy_multi_task = []
         weighted_total_loss = 0.0  # Initialize to 0.0 to avoid in-place operations.
+        total_alpha_loss = 0.0
 
         latent_state_l2_norms_multi_task = []
         average_target_policy_entropy_multi_task = []
@@ -869,12 +898,27 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
         #     current_policy_label_eps = 0.0
         current_policy_label_eps = 0.01
 
+        # 新增一个列表来收集当前批次中所有任务的真实全局ID
+        global_task_ids_in_batch = [] 
+        alpha_loss = None
+
 
+        # 用于Alpha日志记录的新列表
+        alpha_loss_multi_task = []
+        target_entropy_multi_task = []
+
+        # 仅在自适应alpha启用时，预先获取当前alpha值，确保在单次迭代中对所有任务一致
+        current_alpha = self._cfg.model.world_model_cfg.policy_entropy_weight
+        if self.use_adaptive_entropy_weight:
+            current_alpha = self.log_alpha.exp().detach()
 
         losses_list = []  # Used to store the loss tensor for each task, required by gradient correction methods.
         for task_id, data_one_task in enumerate(data):
-            current_batch, target_batch, task_id = data_one_task
+            current_batch, target_batch, task_id = data_one_task # task_id 是真实的全局ID
 
+            # 将真实的全局ID添加到列表中
+            global_task_ids_in_batch.append(task_id)
+
             # TODO: Adapt RoPE for multitask settings (using timestep_batch).
             obs_batch_ori, action_batch, target_action_batch, mask_batch, indices, weights, make_time, timestep_batch  = current_batch
             target_reward, target_value, target_policy = target_batch
@@ -948,7 +992,7 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
             # )
 
             losses = self._learn_model.world_model.compute_loss(
-                batch_for_gpt, self._target_model.world_model.tokenizer, self.value_inverse_scalar_transform_handle, current_policy_label_eps=current_policy_label_eps,task_id=task_id
+                batch_for_gpt, self._target_model.world_model.tokenizer, self.value_inverse_scalar_transform_handle, current_policy_label_eps=current_policy_label_eps, task_id=task_id
             )
 
             # ==================== START MODIFICATION 2 ====================
@@ -960,7 +1004,7 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
 
 
             # TODO: Accumulate the weighted total loss. This assumes the loss from `compute_loss` is already weighted.
-            weighted_total_loss += losses.loss_total
+            weighted_total_loss += losses.loss_total # NOTE:+=
 
             # TODO: Add assertions to check for NaN or Inf values in the loss if needed for debugging.
             # assert not torch.isnan(losses.loss_total).any(), "Loss contains NaN values"
@@ -986,9 +1030,9 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
 
             # 从 losses 对象中提取策略熵
             # ==================== START: 目标熵正则化更新逻辑 ====================
-            alpha_loss = None
             current_alpha = self._cfg.model.world_model_cfg.policy_entropy_weight # 默认使用固定值
             if self.use_adaptive_entropy_weight:
+
                 # --- 动态计算目标熵 (这部分逻辑是正确的，予以保留) ---
                 progress = min(1.0, train_iter / self.target_entropy_decay_steps)
                 current_ratio = self.target_entropy_start_ratio * (1 - progress) + self.target_entropy_end_ratio * progress
@@ -999,12 +1043,19 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
                 # --- 计算 alpha_loss (已修正符号) ---
                 # 这是核心修正点：去掉了最前面的负号
                 # detach() 仍然是关键，确保 alpha_loss 的梯度只流向 log_alpha
-                alpha_loss = (self.log_alpha * (policy_entropy.detach() - current_target_entropy)).mean()
+                alpha_loss_task = (self.log_alpha * (policy_entropy.detach() - current_target_entropy)).mean() # NOTE:=
 
                 # # --- 更新 log_alpha ---
-                self.alpha_optimizer.zero_grad()
-                alpha_loss.backward()
-                self.alpha_optimizer.step()
+                # self.alpha_optimizer.zero_grad()
+                # alpha_loss.backward()
+                # self.alpha_optimizer.step()
+
+                # 累加alpha_loss
+                total_alpha_loss += alpha_loss_task
+                # 为日志记录收集每个任务的alpha_loss和目标熵
+                alpha_loss_multi_task.append(alpha_loss_task)
+                target_entropy_multi_task.append(current_target_entropy)
+
                 # --- [优化建议] 增加 log_alpha 裁剪作为安全措施 ---
                 with torch.no_grad():
                     # 将 alpha 限制在例如 [1e-4, 10.0] 的范围内
@@ -1030,7 +1081,7 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
                     self.obs_loss_weight  * obs_loss # 假设 ssl_loss_weight 是 obs_loss 的权重
                     # ... 如果还有其他损失项，也加进来 ...
                 )
-                weighted_total_loss = (weights * total_loss).mean()
+                weighted_total_loss += (weights * total_loss).mean() # NOTE:+=
             # ===================== END: 目标熵正则化更新逻辑 =====================
 
             # ============ For value-based priority calculation ============
@@ -1098,24 +1149,52 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
         # Core learn model update step.
         self._optimizer_world_model.zero_grad()
 
+        if self.use_adaptive_entropy_weight:
+            self.alpha_optimizer.zero_grad()
+        # 2. 计算最终的alpha loss (在累加后取平均)
+        final_alpha_loss = None
+        if self.use_adaptive_entropy_weight:
+            if len(data) > 0:
+                final_alpha_loss = total_alpha_loss / len(data)
+            else: # 防御性编程，避免除以0
+                final_alpha_loss = torch.tensor(0.0, device=self._cfg.device)
+
         # Assuming losses_list is a list of tensors with gradients, e.g., [loss1, loss2, ...].
         if self._cfg.use_moco:
             # Call MoCo's backward method, which handles gradient correction internally.
             if self._cfg.moco_version=="v0":
                 lambd, stats = self.grad_correct.backward(losses=losses_list, **self._cfg.grad_correct_params)
             elif self._cfg.moco_version=="v1":
                 lambd, stats = self.grad_correct.backward(losses_list)
+            
+            # 单独为alpha loss进行反向传播
+            if self.use_adaptive_entropy_weight:
+                final_alpha_loss.backward()
 
         elif self._cfg.only_use_moco_stats:
             # Only compute MoCo stats without applying gradient correction.
             lambd, stats = self.grad_correct.backward(losses=losses_list, **self._cfg.grad_correct_params)
+
             # Each rank performs its own backpropagation.
-            weighted_total_loss.backward()
+            # weighted_total_loss.backward()
+
+            # 如果启用自适应alpha，将alpha loss加到主损失上一起反向传播
+            if self.use_adaptive_entropy_weight:
+                (weighted_total_loss + final_alpha_loss).backward()
+            elif weighted_total_loss != 0.0: # 确保有损失可以反向传播
+                weighted_total_loss.backward()
+
         else:
             # If not using gradient correction, each rank performs standard backpropagation.
             lambd = torch.tensor([0. for _ in range(self.task_num_for_current_rank)], device=self._cfg.device)
-            weighted_total_loss.backward()
 
+            # weighted_total_loss.backward()
+
+            # 如果启用自适应alpha，将alpha loss加到主损失上一起反向传播
+            if self.use_adaptive_entropy_weight:
+                (weighted_total_loss + final_alpha_loss).backward()
+            elif weighted_total_loss != 0.0: # 确保有损失可以反向传播
+                weighted_total_loss.backward()
 
         # -----------------------------------------------------------------
         # 仍然在 torch.no_grad() 环境下执行
@@ -1150,9 +1229,6 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
                         scale_module_weights_vectorized(self._model.world_model.tokenizer.encoder, scale_factor)
 
 
-
-
-
         # For debugging purposes.
         # for name, param in self._learn_model.world_model.tokenizer.encoder.named_parameters():
         #     print('name, param.mean(), param.std():', name, param.mean(), param.std())
@@ -1179,6 +1255,13 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
 
         self._optimizer_world_model.step()
 
+        # 4. 更新Alpha优化器
+        if self.use_adaptive_entropy_weight:
+            self.alpha_optimizer.step()
+            # 裁剪log_alpha以保证稳定性
+            with torch.no_grad():
+                self.log_alpha.clamp_(np.log(1e-4), np.log(10.0))
+
         if self._cfg.cos_lr_scheduler or self._cfg.piecewise_decay_lr_scheduler:
             self.lr_scheduler.step()
 
@@ -1210,12 +1293,12 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
         if self.use_adaptive_entropy_weight:
             return_log_dict['adaptive_alpha'] = current_alpha.item()
             return_log_dict['adaptive_target_entropy_ratio'] = current_ratio
-            return_log_dict['alpha_loss'] = alpha_loss.item()
+            return_log_dict['final_alpha_loss'] = final_alpha_loss.item()
         # ==================== START: 添加新日志项 ====================
 
         # Generate task-related loss dictionaries and prefix each task-related loss with "noreduce_".
         multi_task_loss_dicts = {
-            **generate_task_loss_dict(obs_loss_multi_task, 'noreduce_obs_loss_task{}', task_id=self.task_id),
+            **generate_task_loss_dict(obs_loss_multi_task, 'noreduce_obs_loss_task{}', task_id=self.task_id), #global_task_ids=global_task_ids_in_batch), # task_id=self.task_id),
             **generate_task_loss_dict(latent_recon_loss_multi_task, 'noreduce_latent_recon_loss_task{}', task_id=self.task_id),
             **generate_task_loss_dict(perceptual_loss_multi_task, 'noreduce_perceptual_loss_task{}', task_id=self.task_id),
             **generate_task_loss_dict(latent_state_l2_norms_multi_task, 'noreduce_latent_state_l2_norms_task{}', task_id=self.task_id),
@@ -1230,6 +1313,10 @@ def _forward_learn(self, data: Tuple[torch.Tensor], task_weights=None, train_ite
             **generate_task_loss_dict(lambd, 'noreduce_lambd_task{}', task_id=self.task_id), 
             **generate_task_loss_dict(value_priority_multi_task, 'noreduce_value_priority_task{}', task_id=self.task_id),
             **generate_task_loss_dict(value_priority_mean_multi_task, 'noreduce_value_priority_mean_task{}', task_id=self.task_id),
+
+                    # 新增alpha相关日志
+            **generate_task_loss_dict(alpha_loss_multi_task, 'noreduce_alpha_loss_task{}', self.task_id),
+            **generate_task_loss_dict(target_entropy_multi_task, 'noreduce_target_entropy_task{}', self.task_id),
         }
         return_log_dict.update(multi_task_loss_dicts)
 
@@ -1319,7 +1406,7 @@ def _monitor_vars_learn(self, num_tasks: int = 2) -> List[str]:
             # 'value_priority',
             'adaptive_alpha',
             "adaptive_target_entropy_ratio",
-            'alpha_loss',
+            'final_alpha_loss',
         ]
 
 
@@ -1346,7 +1433,10 @@ def _monitor_vars_learn(self, num_tasks: int = 2) -> List[str]:
             'noreduce_avg_weight_mag_transformer',
             'noreduce_avg_weight_mag_head',
             'noreduce_e_rank_last_linear',
-            'noreduce_e_rank_sim_norm'
+            'noreduce_e_rank_sim_norm',
+            "noreduce_alpha_loss",
+            "noreduce_target_entropy",
+
         ]
 
         # Use self.task_num_for_current_rank as the number of tasks for the current rank.