From 3038aa7ac709fb6d49292a2f2ee7b682d600e381 Mon Sep 17 00:00:00 2001 From: 8x54zj-m <8x54zj-m@motrixlab.local> Date: Wed, 1 Jul 2026 15:05:50 +0800 Subject: [PATCH] fix: critic called once (was 3x), adaptive KL schedule restored, upgrade 3m --- .../motrix_envs/locomotion/go1/dreamwaq.py | 2 +- .../go1/xmls/assets/flat_stairs.png | 3 + .../locomotion/go1/xmls/scene_stairs_box.xml | 36 ++++++++ .../rslrl/torch/train/dreamwaq_ppo.py | 27 ++++-- scripts/dreamwaq_sim2sim_mujoco.py | 4 +- scripts/export_dreamwaq_onnx_new.py | 86 +++++++++++++++++++ 6 files changed, 149 insertions(+), 9 deletions(-) create mode 100644 motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/flat_stairs.png create mode 100644 motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml create mode 100644 scripts/export_dreamwaq_onnx_new.py diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py b/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py index 1917c52..9e90458 100644 --- a/motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py +++ b/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 diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/flat_stairs.png b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/assets/flat_stairs.png new file mode 100644 index 0000000..0da3edb --- /dev/null +++ b/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 diff --git a/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml new file mode 100644 index 0000000..2bfc8d1 --- /dev/null +++ b/motrix_envs/src/motrix_envs/locomotion/go1/xmls/scene_stairs_box.xml @@ -0,0 +1,36 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py b/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py index c7d453d..2c53f7f 100644 --- a/motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py +++ b/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() diff --git a/scripts/dreamwaq_sim2sim_mujoco.py b/scripts/dreamwaq_sim2sim_mujoco.py index a6bd190..2db9726 100644 --- a/scripts/dreamwaq_sim2sim_mujoco.py +++ b/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) diff --git a/scripts/export_dreamwaq_onnx_new.py b/scripts/export_dreamwaq_onnx_new.py new file mode 100644 index 0000000..4250dcc --- /dev/null +++ b/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()