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

motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py

@@ -436,7 +436,7 @@ class DreamWaQTask(Go1WalkTask):
             base_pose = self._body.get_pose(state.data)
             base_pos = base_pose[done, :2]
             distance = np.linalg.norm(base_pos - old_origins, axis=1)
-            move_up = distance > (self._cell_size / 2.0)
+            move_up = distance > 3.0  # 走过去 3m 就升级（上游是 4m）
             cmd_speed = np.linalg.norm(old_commands[:, :2], axis=1)
             required_dist = cmd_speed * (self.cfg.max_episode_steps * self.cfg.ctrl_dt) * 0.5
             move_down = (distance < required_dist) & ~move_up

motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/flat_stairs.png

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72618dd9378e12b591ef20eeb089fb930db05012e1dac9bf42439a097740bd58
+size 83456

motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml

@@ -0,0 +1,36 @@
+<mujoco model="go1 box stairs scene">
+  <include file="go1_motor_actuator.xml" />
+  <include file="materials.xml" />
+  <statistic center="0 0 0.3" extent="2" meansize="0.04" />
+
+  <visual>
+    <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0" />
+    <rgba haze="0.15 0.25 0.35 1" />
+    <global azimuth="120" elevation="-20" />
+    <map force="0.01" />
+    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2" />
+    <quality shadowsize="8192" />
+  </visual>
+
+  <worldbody>
+    <light pos="0 0 4" dir="0 0 -1" directional="true" />
+
+    <geom name="floor" pos="0 0 -0.001" size="0 0 0.001" type="plane"
+          material="motphys-ground" contype="1" conaffinity="0" priority="0" friction="0.6" />
+
+    <geom name="step0" type="box" size="0.15 10 0.03" pos="0.0 0 0.03" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step1" type="box" size="0.15 10 0.03" pos="0.3 0 0.09" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step2" type="box" size="0.15 10 0.03" pos="0.6 0 0.15" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step3" type="box" size="0.15 10 0.03" pos="0.8999999999999999 0 0.21" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step4" type="box" size="0.15 10 0.03" pos="1.2 0 0.27" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step5" type="box" size="0.15 10 0.03" pos="1.5 0 0.32999999999999996" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step6" type="box" size="0.15 10 0.03" pos="1.7999999999999998 0 0.39" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step7" type="box" size="0.15 10 0.03" pos="2.1 0 0.44999999999999996" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step8" type="box" size="0.15 10 0.03" pos="2.4 0 0.51" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="step9" type="box" size="0.15 10 0.03" pos="2.6999999999999997 0 0.5700000000000001" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+    <geom name="platform" type="box" size="0.5 10 0.03" pos="3.5 0 0.63" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>
+
+    <!-- Fill under stairs -->
+    <geom name="fill" type="box" size="1.5 10 0.3" pos="1.5 0 -0.3" rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
+  </worldbody>
+</mujoco>

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()

scripts/dreamwaq_sim2sim_mujoco.py

@@ -235,8 +235,8 @@ def main():
 
             # ONNX inference
             outputs = session.run(None, {
-                'observations': obs.reshape(1, -1).astype(np.float32),
-                'obs_history': history.astype(np.float32),
+                'obs': obs.reshape(1, -1).astype(np.float32),
+                'obs_history': history.reshape(1, -1).astype(np.float32),
             })
             action = outputs[0][0]
             action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)

scripts/export_dreamwaq_onnx_new.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Export DreamWaQ CENetActorModel to ONNX for MuJoCo sim2sim.
+
+输入: obs(45) + obs_history(225) → 输出: action(12)（确定性，mean CENet code）
+
+用法:
+    uv run scripts/export_dreamwaq_onnx_new.py
+    uv run scripts/export_dreamwaq_onnx_new.py --checkpoint runs/.../model_2850.pt
+"""
+import argparse, os, sys, glob
+import torch
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from motrix_rl.rslrl.torch.models.cenet_actor import CENetActorModel
+from tensordict import TensorDict
+
+PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
+
+class DwaqInfer(torch.nn.Module):
+    """确定性推理：VAE mean code + Actor MLP。"""
+    def __init__(self, model: CENetActorModel):
+        super().__init__()
+        self.vae = model.vae
+        self.actor_mlp = model.mlp
+
+    def forward(self, obs, obs_history):
+        """obs: (N,45), obs_history: (N,225) → action: (N,12)"""
+        code = self.vae.deterministic_code(obs_history)  # (N,19)
+        latent = torch.cat([code, obs], dim=-1)  # (N,64)
+        return self.actor_mlp(latent)
+
+
+def main():
+    p = argparse.ArgumentParser()
+    p.add_argument("--checkpoint", default=None)
+    p.add_argument("--output", default=None)
+    args = p.parse_args()
+
+    if args.checkpoint:
+        ckpt = args.checkpoint
+    else:
+        models = sorted(glob.glob(os.path.join(PROJECT, "runs/go1-dreamwaq-walk/rslrl/*/model_*.pt")))
+        if not models:
+            print("No checkpoints found"); sys.exit(1)
+        ckpt = models[-1]
+    print(f"[Export] checkpoint: {ckpt}")
+
+    # 加载模型
+    dummy = TensorDict({
+        "policy": torch.zeros(1, 45),
+        "obs_history": torch.zeros(1, 225),
+        "privileged_obs": torch.zeros(1, 247),
+    }, batch_size=[1])
+    model = CENetActorModel(
+        dummy, {"actor": ["policy", "obs_history"]}, "actor", 12,
+        hidden_dims=[512, 256, 128], activation="elu", stochastic=True,
+        init_noise_std=1.0, cenet_in_dim=225, cenet_out_dim=19,
+    )
+    data = torch.load(ckpt, map_location="cpu")
+    if "actor_state_dict" in data:
+        model.load_state_dict(data["actor_state_dict"])
+    else:
+        model.load_state_dict(data)
+    model.eval()
+
+    # 包装为推理模型
+    infer = DwaqInfer(model)
+    infer.eval()
+
+    # 导出 ONNX
+    out_path = args.output or os.path.join(os.path.dirname(ckpt), "policy.onnx")
+    dummy_obs = torch.zeros(1, 45)
+    dummy_hist = torch.zeros(1, 225)
+    torch.onnx.export(
+        infer, (dummy_obs, dummy_hist), out_path,
+        input_names=["obs", "obs_history"],
+        output_names=["actions"],
+        opset_version=11,
+        dynamic_axes={"obs": {0: "batch"}, "obs_history": {0: "batch"}, "actions": {0: "batch"}},
+    )
+    print(f"[Export] saved: {out_path}")
+
+
+if __name__ == "__main__":
+    main()