fix: critic called once (was 3x), adaptive KL schedule restored, upgrade 3m

motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py

@@ -84,29 +84,44 @@ class DreamWaQPPO(PPO):
             # NaN 检测 — 数据有 NaN 就跳过这个 batch
             if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any():
                 continue
-            # ── 标准 PPO 前向 ──
+            # ── 标准 PPO 前向（各调用一次，复用结果）──
             self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
-            self.critic(obs_batch, masks=masks_batch)
+            value_batch = self.critic(obs_batch, masks=masks_batch)
 
             # ── 动作对数概率 ──
             actions_log_prob_batch = self.actor.get_output_log_prob(actions_batch)
 
             # ── 价值损失 ──
-            value_batch = target_values_batch
             if self.use_clipped_value_loss:
-                value_clipped = target_values_batch + (self.critic(obs_batch).detach() - target_values_batch).clamp(
+                value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(
                     -self.clip_param, self.clip_param)
-                value_losses = torch.square(self.critic(obs_batch) - returns_batch)
+                value_losses = torch.square(value_batch - returns_batch)
                 value_losses_clipped = torch.square(value_clipped - returns_batch)
                 value_loss = torch.max(value_losses, value_losses_clipped).mean()
             else:
-                value_loss = torch.square(returns_batch - self.critic(obs_batch)).mean()
+                value_loss = torch.square(returns_batch - value_batch).mean()
 
             # ── 代理损失 ──
             surrogate_loss = self._compute_surrogate_loss(
                 actions_batch, actions_log_prob_batch,
                 old_actions_log_prob_batch, advantages_batch)
 
+            # ── KL 自适应 schedule ──
+            if self.desired_kl is not None and self.schedule == 'adaptive':
+                with torch.inference_mode():
+                    kl = torch.sum(
+                        torch.log(self.actor.output_std / old_sigma_batch + 1e-5)
+                        + (old_sigma_batch.pow(2) + (old_mu_batch - self.actor.output_mean).pow(2))
+                        / (2.0 * self.actor.output_std.pow(2))
+                        - 0.5, dim=-1)
+                    kl_mean = torch.mean(kl)
+                    if kl_mean > self.desired_kl * 2.0:
+                        self.learning_rate = max(1e-5, self.learning_rate / 1.5)
+                    elif kl_mean < self.desired_kl / 2.0 and kl_mean > 0.0:
+                        self.learning_rate = min(1e-2, self.learning_rate * 1.5)
+                    for param_group in self.optimizer.param_groups:
+                        param_group['lr'] = self.learning_rate
+
             # ── 熵 ──
             entropy_batch = self.actor.output_entropy
             entropy_loss = entropy_batch.mean()