cyy/Motrixlab
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: clamp std before distribution, lower init_noise to 0.5, NaN guard

Browse Source
This commit is contained in:
8x54zj-m
2026-06-30 15:35:58 +08:00
parent f2a8e0e2ff
commit b0f1da4596
61 changed files with 6522 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
cmd.md Normal file
View File

@@ -0,0 +1,16 @@
uv run tensorboard --logdir runs/go1-stairs-terrain-walk-no-linevel
uv run scripts/train.py --env go1-stairs-terrain-walk-no-linevel \
--checkpoint runs/go1-stairs-terrain-walk-no-linevel/skrl/26-06-23_01-37-34-357547_PPO/checkpoints/best_agent.pickle
uv run scripts/play.py --env go1-stairs-terrain-walk-no-linevel --num-envs 1
# 只训练楼梯地形，线速度奖励 3 倍，继承 best_agent 权重
uv run scripts/train.py --env go1-stairs-terrain-walk-no-linevel --force-phase 0 --tracking-linvel-scale 1.0
uv run scripts/play.py --env go1-stairs-terrain-walk-no-linevel --num-envs 1 --force-phase 0 --num-envs 1024
DREAMWAQ_TERRAIN=flat_stairs uv run scripts/train_dreamwaq_rsl.py \
--num-envs 2048 --init-noise-std 2.0 \
--resume runs/go1-dreamwaq-walk/rsl_dwaq/06-29_16-01-56/model_3000.pt

236
deploy_orig/go1_sim2sim.py Normal file
View File

@@ -0,0 +1,236 @@
#!/usr/bin/env python3
"""Go1 sim2sim MuJoCo viewer — Original 30k flat training.
Usage: python go1_sim2sim.py
Requires: mujoco, onnxruntime, pynput
Install: pip install mujoco onnxruntime pynput
Controls:
W/S: forward/back Q/E: strafe left/right
A/D: rotate Space: stop R: reset Esc: quit
"""
import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue
from mujoco import viewer
from pynput import keyboard
HERE = os.path.dirname(os.path.abspath(__file__))
ONNX = os.path.join(HERE, "policy.onnx")
# ── Parameters (original MotrixLab Go1 config) ──
NUM_OBS = 48
KP, KD = 80.0, 0.5 # KD=0.5 + MuJoCo joint_damping(0.5) = 1.0 = training kd
ACTION_SCALE = 0.05
CLIP = 23.7
DEFAULT_ANGLES = np.array([
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
], dtype=np.float32)
# ── Keyboard ──
class KB:
def __init__(s):
s._q = queue.Queue(); s.running = True; s.held = set()
def _n(s, k):
try:
if hasattr(k, 'char') and k.char: return k.char.lower()
except: pass
return str(k).lower()
def _w(s):
while s.running:
try:
et, k = s._q.get(timeout=0.05)
n = s._n(k)
if et == 'press': s.held.add(n)
elif et == 'release': s.held.discard(n)
except queue.Empty: pass
def init(s):
s._l = keyboard.Listener(
on_press=lambda k: s._q.put(('press', k)),
on_release=lambda k: s._q.put(('release', k)))
s._l.start()
s._t = threading.Thread(target=s._w, daemon=True); s._t.start()
def keys(s): return s.held.copy()
def stop(s): s.running = False; s._l.stop()
# ── Main ──
def main():
# The model XML is embedded below
xml = '''<mujoco model="go1 scene">
<compiler angle="radian" autolimits="true"/>
<option timestep="0.005" integrator="Euler" iterations="60">
<flag eulerdamp="disable"/>
</option>
<custom>
<numeric data="30" name="max_contact_points"/>
<numeric data="12" name="max_geom_pairs"/>
</custom>
<default>
<default class="go1">
<geom condim="1"/>
<joint axis="0 1 0" armature="0.005" damping="0.5"/>
<default class="abduction">
<joint axis="1 0 0" range="-0.863 0.863" frictionloss="0.3"/>
</default>
<default class="hip">
<joint range="-0.686 4.501" frictionloss="0.3"/>
</default>
<default class="knee">
<joint range="-2.818 -0.888" frictionloss="1.0"/>
</default>
</default>
</default>
<asset>
<texture name="skybox" type="skybox" builtin="gradient" rgb1="0.4 0.4 0.4" rgb2="0 0 0" width="512" height="512"/>
<texture name="ground" type="2d" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="ground" texture="ground" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
</asset>
<worldbody>
<light pos="0 0 1.5" dir="0 0 -1" directional="true"/>
<geom name="floor" size="0 0 0.01" type="plane" material="ground" contype="1" conaffinity="0" priority="1" friction="0.6" condim="3"/>
<body name="trunk" pos="0 0 0.4" childclass="go1">
<freejoint/>
<inertial pos="0.0223 0.002 -0.0005" quat="-0.00342088 0.705204 0.000106698 0.708996" mass="5.204" diaginertia="0.0716565 0.0630105 0.0168101"/>
<geom name="trunk_geom" contype="0" conaffinity="0" group="2" type="box" size="0.35 0.12 0.08" rgba="0.4 0.4 0.4 1"/>
<geom name="trunk_col" contype="1" conaffinity="1" group="3" pos="0.24 0 0" size="0.05 0.05 0.05" type="box"/>
<site name="imu" pos="-0.01592 -0.06659 -0.00617" group="5"/>
<!-- FR leg -->
<body name="FR_hip" pos="0.1881 -0.04675 0">
<joint class="abduction" name="FR_hip_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
<body name="FR_thigh" pos="0 -0.08 0">
<joint class="hip" name="FR_thigh_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<body name="FR_calf" pos="0 0 -0.213">
<joint class="knee" name="FR_calf_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<geom name="FR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
</body>
</body>
</body>
<!-- FL leg -->
<body name="FL_hip" pos="0.1881 0.04675 0">
<joint class="abduction" name="FL_hip_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
<body name="FL_thigh" pos="0 0.08 0">
<joint class="hip" name="FL_thigh_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<body name="FL_calf" pos="0 0 -0.213">
<joint class="knee" name="FL_calf_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<geom name="FL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
</body>
</body>
</body>
<!-- RR leg -->
<body name="RR_hip" pos="-0.1881 -0.04675 0">
<joint class="abduction" name="RR_hip_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
<body name="RR_thigh" pos="0 -0.08 0">
<joint class="hip" name="RR_thigh_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<body name="RR_calf" pos="0 0 -0.213">
<joint class="knee" name="RR_calf_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<geom name="RR_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
</body>
</body>
</body>
<!-- RL leg -->
<body name="RL_hip" pos="-0.1881 0.04675 0">
<joint class="abduction" name="RL_hip_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.08" rgba="0.5 0.5 0.5 1"/>
<body name="RL_thigh" pos="0 0.08 0">
<joint class="hip" name="RL_thigh_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.05 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<body name="RL_calf" pos="0 0 -0.213">
<joint class="knee" name="RL_calf_joint"/>
<geom contype="0" conaffinity="0" group="2" type="capsule" size="0.04 0.11" pos="0 0 -0.1" rgba="0.5 0.5 0.5 1"/>
<geom name="RL_foot" contype="1" conaffinity="1" group="3" type="sphere" size="0.023" pos="0 0 -0.213" priority="10" condim="3"/>
</body>
</body>
</body>
</body>
</worldbody>
<actuator>
<motor class="abduction" name="FR_hip" joint="FR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="FR_thigh" joint="FR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="FR_calf" joint="FR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="abduction" name="FL_hip" joint="FL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="FL_thigh" joint="FL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="FL_calf" joint="FL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="abduction" name="RR_hip" joint="RR_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="RR_thigh" joint="RR_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="RR_calf" joint="RR_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="abduction" name="RL_hip" joint="RL_hip_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="hip" name="RL_thigh" joint="RL_thigh_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
<motor class="knee" name="RL_calf" joint="RL_calf_joint" ctrllimited="true" ctrlrange="-23.7 23.7"/>
</actuator>
<sensor>
<gyro site="imu" name="gyro"/>
<velocimeter site="imu" name="local_linvel"/>
</sensor>
</mujoco>'''
model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)
data.qpos[0:3] = [0, 0, 0.42]
data.qpos[3:7] = [1, 0, 0, 0]
data.qpos[7:19] = DEFAULT_ANGLES
mujoco.mj_forward(model, data)
session = ort.InferenceSession(ONNX, providers=['CPUExecutionProvider'])
print(f"[Go1 sim2sim] ONNX={ONNX}")
print(f"[Go1 sim2sim] PD kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} obs={NUM_OBS}-dim")
print(f"[Go1 sim2sim] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
kb = KB(); kb.init()
view = viewer.launch_passive(model, data)
step, vx, vy, wz = 0, 0.0, 0.0, 0.0
last_a = np.zeros(12, dtype=np.float32)
while view.is_running():
keys = kb.keys()
if 'escape' in keys: break
if 'r' in keys:
data.qpos[0:3] = [0, 0, 0.42]
data.qpos[3:7] = [1, 0, 0, 0]
data.qpos[7:19] = DEFAULT_ANGLES
data.qvel[:] = 0; last_a[:] = 0
mujoco.mj_forward(model, data)
vx = 1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.0)
vy = 1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.0)
wz = 1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.0)
if ' ' in keys: vx = vy = wz = 0.0
if step % 2 == 0: # 100Hz control (MuJoCo dt=0.005)
obs = np.zeros(NUM_OBS, dtype=np.float32)
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro")
adr = model.sensor_adr[sid]
obs[0:3] = data.sensordata[adr:adr+3] * 0.25
R = data.xmat[1].reshape(3, 3)
obs[6:9] = (R.T @ np.array([0., 0., -1.])).astype(np.float32)
obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
obs[21:33] = data.qvel[6:18] * 0.05
obs[33:45] = last_a
obs[45:48] = np.array([vx, vy, wz]) * np.array([2., 2., 0.25])
obs = np.clip(obs, -100., 100.)
action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
action = np.clip(action, -CLIP, CLIP)
last_a = action.copy()
target = DEFAULT_ANGLES + action * ACTION_SCALE
torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
data.ctrl[:] = np.clip(torques, -CLIP, CLIP)
mujoco.mj_step(model, data)
view.sync()
step += 1
time.sleep(0.001)
kb.stop(); view.close()
if __name__ == "__main__":
main()

BIN
deploy_orig/policy.onnx Normal file
View File

Binary file not shown.

189
docs/dreamwaq_alignment.md Normal file
View File

@@ -0,0 +1,189 @@
# DreamWaQ 复现与对齐记录
> 目标：让本地 `go1-dreamwaq-walk` 完全对齐并复现 [Manaro-Alpha/DreamWaQ](https://github.com/Manaro-Alpha/DreamWaQ.git)（基于 IsaacGym + rsl_rl 的实现），迁移到 MotrixLab（MotrixSim + SKRL-JAX）。
最后更新：2026-06-25
---
## 1. 架构总览
```
CENet (VAE): history(5×45=225) → [128,64] → latent(16) + vel_est(3) = code(19)
Decoder: code(19) → [64,128] → obs_recon(45)
Actor: code(19) + obs(45) = 64 → [512,256,128] → action(12)
Critic: privileged_obs(235) → [512,256,128] → value(1)
观测(45): ang_vel(3) + gravity(3) + commands(3) + dof_pos(12) + dof_vel(12) + actions(12)
特权观测(235): obs(45) + base_vel(3) + heights(187, 17×11 网格)
增广观测(254, 训练用): [code(19) | obs(45) | base_vel(3) | heights(187)]
```
---
## 2. 🔴 真正让机器人站不起来的 Bug（最关键）
排查发现并非训练不充分，而是两个环境 bug 让**约一半的训练数据是垃圾**：
### Bug A — 机器人生成在地形地下
- **现象**：reset 时机器人固定生成在绝对高度 z=0.40，但 DreamWaQ 地形高度为 0–1.6m。在高地形格子上，机器人生成在**地表以下最多 1.2m**。
- **影响**：31/64 的环境在生成瞬间就卡进地形 → 爆炸/立即终止/疯狂抖动 → dof_acc 高达 150 万 → 奖励被 `only_positive_rewards` 裁剪为 0。一半训练样本无效。
- **修复**：新增 `_sample_terrain_height(xy)`，按生成点的局部地形高度设置 z = `terrain_z + 0.40`。修复后 0/64 在地下。
- **位置**：[dreamwaq.py](../motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py) `reset()` + `_sample_terrain_height()`
### Bug B — feet_air_time 奖励恒为 0
- **现象**：腾空时间在 `update_feet_air_time` 中被 `*= ~contacts` 清零，**早于** reward 函数读取 `first_contact = (air_time > 0) * contacts`。落地瞬间 air_time 已是 0 → first_contact 永远为 0 → 步态奖励完全失效。
- **影响**：机器人没有迈步激励，倾向于站着不动/拖行，学不会正常步态。
- **修复**：按上游顺序重构 —— 先用 `contact_filt = contact OR last_contact` 算出 `first_contact` 和落地时的 air_time，**再**清零；reward 函数改为读取预计算值（无状态，对 DreamWaQ 的 `_get_reward` 二次调用安全）。
- **位置**：[walk_np.py](../motrix_envs/src/motrix_envs/locomotion/go1/walk_np.py) `update_feet_air_time()` + `_reward_feet_air_time()`
### Bug C — 初始地形难度过高
- **现象**：首次 reset 在 0-9 级随机生成（含最难地形）。上游用 `max_init_terrain_level=5`。
- **修复**：首次生成限制在 0-5 级，由课程学习逐步提升。
- **位置**：[dreamwaq.py](../motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py) `reset()`
---
## 3. 对齐改动（与上游一致）
### 控制 / 物理（cfg.py + dreamwaq.py）
| 参数 | 改动前 | 改动后（上游） |
|---|---|---|
| stiffness (Kp) | 80 | **28.0** |
| damping (Kd) | 1.0 | **0.7** |
| action_scale | 0.05 | **0.25** |
| sim_dt / ctrl_dt | 0.01 / 0.01 | **0.005 / 0.02**（decimation=4，200Hz 物理 / 50Hz 控制） |
### 默认关节角度（cfg.py，非对称）
```
hips: ±0.1 (FL/RL: +0.1, FR/RR: -0.1)
thighs: 前 0.8 / 后 1.0
calves: -1.5 (全部)
```
### 观测顺序（walk_np.py）
改为上游顺序 `[ang_vel, gravity, commands, dof_pos, dof_vel, actions]`，并加入观测噪声（`add_noise=True, noise_level=1.0`）。
### 命令范围（dreamwaq.py）
`lin_vel_x/y/yaw` 全部 `[-1, 1]`（之前 vel_x 上限是 2.0）。
### 奖励函数对齐（dreamwaq.py）
- `smoothness`：基于 `joint_pos_target`（= action×scale + default）的二阶差分，而非原始 action
- `base_height`：`square(base_z - mean(measured_heights) - 0.30)`，扣除局部地形高度
- `power_distribution`：`var(torque×vel)` 跨全部 12 关节，而非按腿分组
- 12 项奖励尺度与上游 Go1 base config 精确一致
### 域随机化（dreamwaq.py）
motor_strength / Kp_factor / Kd_factor ∈ [0.9, 1.1]，每 4 秒重采样（`rand_interval_s=4`），在 `_compute_torques` 中应用。
### 训练（dreamwaq_ppo.py + go1_dreamwaq.py）
- **关键 bug 修复**：SKRL 字段名 `entropy_coef→entropy_loss_scale`、`clip_param→ratio_clip`、`gamma→discount_factor`。改前 entropy 被静默设为 0.0（字段名不存在）。
- CENet 解码目标：重构**当前观测**（历史最后一帧），而非预测下一帧（匹配上游 `decode_target = obs_batch`）
- 正交权重初始化 `orthogonal(sqrt(2))`，应用于 CENet/Actor/Critic 所有 Dense 层
- PPO 超参：rollouts=24, lr=1e-3, entropy=0.01, gamma=0.99, lam=0.95, clip=0.2, epochs=5, mini_batches=4
### 地形（gen_dreamwaq_terrain.py）
降低高度尺度（最大 2.73m → 1.66m），更适合 0.33m 高的机器人。生成后更新 XML 的 hfield `size`。
---
## 3.5 🔴🔴 部署缺少观测归一化（"站不起来"的真正根因）
这是导致导出/play/sim2sim 中机器人**完全站不起来**的根本原因，独立于训练质量。
- **现象**：训练奖励正常（mean 35 / max 924），但导出 ONNX / `play_dreamwaq.py` / MuJoCo sim2sim 中机器人塌成一团，actor 输出动作幅度高达 12（关节目标偏移 ~3 rad，直接打到限位）。
- **根因**：SKRL PPO 通过 `_add_runtime_config` 启用了 **`RunningStandardScaler` 状态预处理器**——训练时观测被 `clip((x - running_mean) / (sqrt(var) + 1e-8), -5, 5)` 归一化（running_mean 幅度达 ±3.2，std 0.14–3.2）。但部署管线直接喂**原始观测**给 actor，输入尺度完全错误 → 策略输出垃圾动作 → 倒地。
- **验证**：加上归一化后，站立从 **0/16 → 16/16 upright**。
- **修复**：
- [export_dreamwaq_onnx.py](../scripts/export_dreamwaq_onnx.py)：从 checkpoint 提取 `state_preprocessor` 的 mean/var（前 64 维 = code+obs），烘焙进 ONNX 模型（actor 前做归一化+clip）
- [play_dreamwaq.py](../scripts/play_dreamwaq.py)：`policy_forward` 中应用同样的归一化
- sim2sim 用 ONNX，重新导出后自动修复
- **注意**：`RunningStandardScaler` 是 SKRL 的附加项，上游 rsl_rl 没有（上游只用固定 obs_scales）。任何脱离 SKRL agent 的部署都**必须**复现这个归一化，否则策略失效。
---
## 3.6 🔴 训练不收敛：初始探索过猛（IsaacGym→MotrixSim reality gap）
修复部署归一化后发现策略**根本没学会站立**——在 MotrixSim 训练环境里 0/64 站立/行走，比"零动作"还差。
- **诊断**：零动作（纯 PD 保持默认姿态）奖励 **1.27/step**，27/64 站立；训练后的策略奖励 **0.035/step**，0/64 站立。**训练把策略训得比啥都不做还差**。
- **根因**：`initial_log_std=0.0`（std=1）× `action_scale=0.25` = 每步 ±0.25 rad 随机抖动 @50Hz。实测：
| 探索 std | 抖动 | 站立率 | 平均奖励 |
|---|---|---|---|
| 1.0（上游） | 0.25rad | 27/64 | **0.025** |
| 0.37 | 0.09rad | 29/64 | 0.512 |
| 0.2 | 0.05rad | 23/64 | 0.730 |
机器人在所有 std 下都能站（站立率相近），但 std=1 时剧烈抖动产生的 dof_acc/action_rate 惩罚把站立的正奖励**完全淹没**（净奖励 0.025 ≈ 噪声）→ PPO 没有指向站立的梯度 → 收敛到塌倒局部最优。
- **为何上游能用**：DreamWaQ 的 std=1 是为 IsaacGym 调的；MotrixSim 接触/动力学不同（你能用的 go1 平地任务用的是 5× 更温和的 `action_scale=0.05`）。
- **修复**：[go1_dreamwaq.py](../motrix_rl/src/motrix_rl/tasks/go1_dreamwaq.py) 设 `initial_log_std=-1.0`（std=0.37，抖动 0.09rad），让机器人早期保持直立、获得清晰的站立奖励信号后再 bootstrap 行走。这是 reality-gap 的必要偏离（上游用 0.0）。
---
## 4. 🛠 原生 play 渲染崩溃修复
- **现象**：`play_dreamwaq.py` 启动 MotrixSim 渲染器时崩溃 `Couldn't get swap chain texture`，而其他 go1 任务的 play 正常。
- **根因**：JAX 默认预分配 **75% 显存**（12GB 中 ~9GB），饿死了 Vulkan 渲染器。标准 play.py 因渲染器创建与 JAX 初始化间隔充足而不受影响。
- **修复**：在 `import jax` 前设置 `os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"`，并优雅处理 `RenderClosedError`。
- **注意**：训练占用 GPU 时无法同时跑原生渲染；请在训练结束后查看。
---
## 5. 修改的文件
| 文件 | 改动 |
|---|---|
| [motrix_envs/.../go1/cfg.py](../motrix_envs/src/motrix_envs/locomotion/go1/cfg.py) | PD 增益、关节角度、噪声尺度 |
| [motrix_envs/.../go1/dreamwaq.py](../motrix_envs/src/motrix_envs/locomotion/go1/dreamwaq.py) | **地形生成修复**、ctrl 频率、命令范围、奖励、域随机化、观测噪声 |
| [motrix_envs/.../go1/walk_np.py](../motrix_envs/src/motrix_envs/locomotion/go1/walk_np.py) | 观测顺序、**feet_air_time 修复** |
| [motrix_rl/.../skrl/jax/train/dreamwaq_ppo.py](../motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py) | CENet 解码目标、VAE 缓冲、正交初始化 |
| [motrix_rl/.../tasks/go1_dreamwaq.py](../motrix_rl/src/motrix_rl/tasks/go1_dreamwaq.py) | **SKRL 字段名修复**、超参 |
| [scripts/gen_dreamwaq_terrain.py](../scripts/gen_dreamwaq_terrain.py) | 地形高度尺度 |
| [scripts/play_dreamwaq.py](../scripts/play_dreamwaq.py) | **GPU 预分配修复** |
| [scripts/dreamwaq_sim2sim_mujoco.py](../scripts/dreamwaq_sim2sim_mujoco.py) | PD/角度/decimation 与训练对齐 |
| [scripts/export_dreamwaq_onnx.py](../scripts/export_dreamwaq_onnx.py) | CENet 编码器 + Actor → ONNX |
---
## 6. 使用方法
### 训练
```bash
uv run scripts/train_dreamwaq.py --num-envs 2048 --timesteps 100M --seed 42
```
> 用 2048 envs 避免 4096 在 ~70% 处的 OOM（exit 144）。检查点存于 `runs/go1-dreamwaq-walk/skrl/`。
### 原生查看（MotrixSim，GPU 空闲时）
```bash
uv run scripts/play_dreamwaq.py --num-envs 9
```
### 导出 ONNX
```bash
uv run scripts/export_dreamwaq_onnx.py \
--checkpoint runs/go1-dreamwaq-walk/skrl/<run>/checkpoints/best_agent.pickle \
--vae runs/go1-dreamwaq-walk/skrl/vae_<N>.pkl \
--output exports_go1_dreamwaq/policy.onnx
```
### MuJoCo sim2sim
```bash
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain dreamwaq
# 控件: W/S 前后 Q/E 左右 A/D 旋转 Space 停 R 重置 Esc 退出
```
### TensorBoard
```bash
uv run tensorboard --logdir runs/go1-dreamwaq-walk
```
---
## 7. 待办 / 注意事项
- **检查点兼容性**：改动前的检查点（含旧 PD/观测顺序/地形 bug）不兼容，应丢弃重训。
- **未对齐项（有意保留）**：
- 奖励未乘 `dt`（上游 `_prepare_reward_function` 会 `scale *= dt`）。这是全局均匀缩放，PPO 优势归一化会抵消，不影响行为；为避免干扰已调好的 value 学习而保留。
- VAE 训练为 PPO 更新后独立一步，而非上游的联合优化（功能等价，更易维护）。
- 地形几何为自定义生成器，非 IsaacGym `terrain_utils`（概念相似，具体形状不同）。
- **OOM**：4096 envs 曾在 70% 处被杀（exit 144），改用 2048。

168
docs/dreamwaq_rsl_handoff.md Normal file
View File

@@ -0,0 +1,168 @@
# DreamWaQ on MotrixSim — rsl_rl 复现交接文档
> 目标:在 MotrixLab(MotrixSim + numpy 物理)上**忠实复现** [Manaro-Alpha/DreamWaQ](https://github.com/Manaro-Alpha/DreamWaQ)(原版基于 IsaacGym + 定制 rsl_rl-1.0.2),让 Go1 在地形上行走。
>
> 最后更新:2026-06-26。当前状态:rsl_rl 路线训练中,趋势健康(站立解决、行走学习中)。
---
## 0. 一句话总览
原版 DreamWaQ = **ActorCritic_DWAQ(CENet+Actor+Critic 一体网络) + PPO 联合训练 VAE + 金字塔地形课程**。我把上游定制的 `rsl_rl-1.0.2` 核心**原样移植**到 MotrixLab,写了个 numpy↔torch 环境适配器,并针对 MotrixSim 的两个 reality gap(hfield 大 z_scale 接触不稳、IsaacGym 超参偏激进)做了适配。
---
## 1. 两条路线(背景)
复现走过两条路线,**rsl_rl 是当前主线**:
| | SKRL-JAX 路线(旧,参考) | **rsl_rl 路线(当前主线)** |
|---|---|---|
| 代码 | `motrix_rl/skrl/jax/train/dreamwaq_ppo.py` | `motrix_rl/dwaq_rsl/`(新模块) |
| 网络 | CENet 拼在 wrapper,Actor/Critic 在 SKRL | 上游 `ActorCritic_DWAQ` 一体网络 |
| VAE 训练 | **单独一步**(CENet 不参与 PPO 梯度) | **PPO 联合优化**(和上游一致) |
| 结果 | 站立解决,但**学不会走**(机器人趴/站着不动) | 趋势健康,行走学习中 |
**为什么切到 rsl_rl**:SKRL 路线机器人学不会走,根因是 **VAE 单独训练** —— CENet 的 code(latent+vel)早期是垃圾,actor 用不上。上游把 VAE 和 PPO 联合优化(一个 optimizer step),code 早期就有意义。切到 rsl_rl 后同样的 `std=1.0` 下 reward 从卡死的 0.025 变成快速上涨,这是关键转折。
> SKRL 路线的详细 bug 修复记录见 [dreamwaq_alignment.md](dreamwaq_alignment.md)(地下生成、feet_air_time、部署归一化等,这些 env 层 bug 两条路线都受益)。
---
## 2. 新建模块:`motrix_rl/src/motrix_rl/dwaq_rsl/`
把上游 `rsl_rl-1.0.2` 的核心**原样移植**(纯 torch,去掉 isaacgym 依赖,改相对 import):
| 文件 | 来源 | 作用 |
|---|---|---|
| `actor_critic_dwaq.py` | 上游 `modules/actor_critic_DWAQ.py` **原样** | CENet(encoder 225→128→64 + latent/vel 头)+ decoder + Actor[512,256,128] + Critic[512,256,128]。`act(obs, obs_history)` 把 CENet code(19)和 obs(45)拼成 64 维喂 actor |
| `ppo.py` | 上游 `algorithms/ppo.py` **原样** | PPO + **VAE 联合 loss**:`loss = surrogate + value + autoenc`,其中 `autoenc = MSE(vel, prev_priv[:,45:48]) + MSE(decode, obs) + beta*KL` |
| `rollout_storage.py` | 上游 `storage/rollout_storage.py` **原样** | 经验存储,含 `prev_critic_obs`(VAE 的 vel_target 来源)、`observation_history` |
| `utils.py` | 上游 `utils/utils.py` | `split_and_pad_trajectories` |
| `runner.py` | 上游 `on_policy_runner.py` **适配** | 去掉 rsl_rl 包 import + isaacgym 噪声,改用本地模块。learn() 主循环、save/load、tensorboard 日志 |
| `env_wrapper.py` | **新写** | `DwaqVecEnv`:MotrixLab `DreamWaQTask`(numpy)→ 上游 VecEnv 接口(torch)。见下 |
### `DwaqVecEnv`(env_wrapper.py)关键点
上游 runner 要求 `env.step()` 返回 **7 元组**:`(obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, extras)`。MotrixLab 的 `DreamWaQTask` 已提供 obs(45)、`info["privileged_obs"]`(235)、`info["obs_history"]`(N,5,45)。wrapper 做:
- numpy↔torch 转换(MotrixSim 物理在 CPU,网络在 GPU,每步传输)
- `obs_history` flatten:(N,5,45) → (N,225)
- 维护 `prev_privileged_obs`(step 前的 privileged,VAE 用其 [45:48]=base_vel 作 vel_target)
- `extras["time_outs"]` = truncated
- **rewards/dones 是 (N,) 不是 (N,1)**(上游约定,storage 内部 view 成 (N,1)) — 这是个踩过的坑
### 维度速查
```
obs=45, privileged=235(=obs45+base_vel3+heights187), obs_hist=225(=5*45)
actor_in=64(=code19+obs45), critic_in=235, cenet_in=225, cenet_out=19(=vel3+latent16)
```
---
## 3. 环境层改动:`motrix_envs/.../go1/dreamwaq.py`
(这些改动 rsl_rl 路线和 SKRL 路线共用同一个 `DreamWaQTask`)
1. **去掉命令课程**(关键):原来我加过一个命令课程(cmd_scale 0.3→1.0),但它让机器人**趴着刷小命令分**(小命令下趴着也能 tracking)。改回**上游 full range [-1,1]**(`commands.curriculum=False`),机器人被要求走,趴着不划算。见 `resample_commands`。
2. **level 0 起步**:`reset()` 里 init terrain level 全设 0(平地),靠地形难度课程逐级升。原来 0-5 随机会让机器人 spawn 在难地形秒摔。
3. **中心 spawn**:`_make_origins` 坐标中心化(hfield 中心在世界原点 (0,0)),spawn 在 cell 的 3m 平台中心 ±1m。
4. **地形高度采样**:`_sample_terrain_height` / `_get_heights` 从 hfield 采样(`hf.height_matrix` + `hf.bound`),用于 spawn 高度 + privileged obs 的 heights。
5. 奖励对齐上游(smoothness 用 joint_pos_target + step mask、power_distribution 用 `var(abs(power))`、base_height 扣地形高度等)。
---
## 4. 地形:`scripts/gen_dreamwaq_terrain.py` + ⚠️ MotrixSim hfield 坑
### 生成器
移植了上游 IsaacGym `terrain_utils` 的金字塔地形(纯 numpy):`pyramid_sloped_terrain` / `pyramid_stairs_terrain` / `discrete_obstacles_terrain` / `random_uniform_terrain`,10 行难度 × 20 列类型网格,每格 3m 中心平台,输出 16-bit PNG hfield + XML `<size>` 行。
### ⚠️⚠️ 最大的坑:MotrixSim hfield 在大 z_scale 下接触不稳
**实测结论**(零动作站立 20 步摔倒数):
| hfield z_scale | 机器人 above-terrain | 20步摔倒 |
|---|---|---|
| 3.78(上游全高金字塔) | 0.12(**趴**) | 1267 |
| 1.66 | 0.14(趴) | 1049 |
| **0.54(矮金字塔)** | **0.35(正常站立)** | **120** |
| plane geom(对比) | 0.33 | **0** |
- MotrixSim 用 **hfield(高度场)**,陡台阶垂直面只能近似成陡坡,机器人脚(sphere geom)打滑;**z_scale 越大越不稳**。
- 上游用 **trimesh(三角网格)**,垂直面精确。**调研确认 MotrixSim 不支持非凸 trimesh 碰撞**(mesh 是凸包),接触参数(condim/solref/solimp)调优也无效。
- **解决**:把金字塔楼梯/坡的高度**缩小 ~4 倍**(`make_terrain` 里 `step_height=0.012+0.020*diff`、`slope=0.10*diff`),让 z_scale≈0.54。保留金字塔结构 + 难度课程,但高度适配 MotrixSim hfield。**这是必要的 reality-gap 妥协**(不完全忠实上游高度,但能训)。
- `scene_dreamwaq_terrain.xml` 的 hfield `size="85 45 0.54 0.01"`:**第1/2 是半径**(全宽 170×90m),第3 是 z_scale,第4 是底部厚度。
**重新生成地形后必须同步更新 XML 的 z_scale**(gen 脚本会打印 `size="..."` 那行的值)。
---
## 5. 脚本(训练 / 导出 / 可视化)
| 脚本 | 作用 |
|---|---|
| `scripts/train_dreamwaq_rsl.py` | rsl_rl 训练。超参 dict 内嵌(对齐上游 Go1RoughCfgPPO)。`--num-envs 2048 --iterations 3000 --init-noise-std 1.0` |
| `scripts/export_dreamwaq_rsl_onnx.py` | rsl 的 `.pt` checkpoint → ONNX。**rsl 无 state 归一化**(不像 SKRL),所以 ONNX 简单:`obs(1,45)+obs_history(1,5,45)→action(1,12)`,用 mean CENet code |
| `scripts/gen_dreamwaq_terrain.py` | 生成矮金字塔地形 PNG |
| `scripts/dreamwaq_sim2sim_mujoco.py` | MuJoCo sim2sim 可视化(用 ONNX)。PD/角度/decimation 已对齐(KP=28,KD=0.7,action_scale=0.25,decimation=4) |
### 命令
```bash
# 训练 (后台跑, print 会被 buffer, 用 tensorboard 看)
PYTHONUNBUFFERED=1 uv run scripts/train_dreamwaq_rsl.py --num-envs 2048 --iterations 3000 --seed 1 --init-noise-std 1.0
# 监控 (print 被 buffer, 看 tensorboard)
uv run tensorboard --logdir runs/go1-dreamwaq-walk/rsl_dwaq
# 导出最新 checkpoint
JAX_PLATFORMS=cpu uv run scripts/export_dreamwaq_rsl_onnx.py # 自动找最新 model_*.pt
# MuJoCo 可视化 (本地有显示的机器)
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain flat # 平地最干净
```
---
## 6. 关键决策 + 为什么(reality gap 适配)
复现不是 1:1 照搬,MotrixSim vs IsaacGym 有真实差异,做了几处**有意的适配**(都在代码注释里标了原因):
1. **VAE 联合训练**(用 rsl_rl 而非 SKRL):核心。SKRL 路线 VAE 单独训练→学不会走。
2. **去命令课程**:命令课程让机器人趴着刷小命令分;上游本来就 full range 无课程。
3. **矮金字塔地形**(z_scale 0.54):MotrixSim hfield 大 z_scale 接触不稳的妥协。
4. **level 0 平地起步**:机器人先在平地学走,靠地形课程升级。
5. 其余超参/网络/奖励**严格对齐上游**(PD 28/0.7、action_scale 0.25、decimation 4、init_noise_std 1.0、rollouts 24、lr 1e-3、entropy 0.01 等)。
---
## 7. 当前状态(2026-06-26)
训练中(`runs/go1-dreamwaq-walk/rsl_dwaq/06-26_00-10-27`),趋势健康:
| 指标 | it59 → it182 |
|---|---|
| mean_reward | 0.02 → 2.6 → **18.6** |
| episode 长度 | 12 → 101 → **143** |
| noise_std | 1.0 → 0.84 → **0.57**(在学) |
| base_height | 稳定(**不趴**) |
机器人能站(above ~0.35)+ episode 越来越长 + std 持续降 = 在学走。**需等训练跑完(3000 iter,~2.5h)再导出验证实际行走**(站立率/前进速度跟踪/摔倒率,别只看奖励数字 — 之前吃过"奖励高但实际趴着"的亏)。
---
## 8. 已知问题 / 坑(接手必读)
1. **训练 print 被 buffer**:重定向到文件时 Python print 不实时 flush。**用 tensorboard 监控**,别等 print。(已加 `PYTHONUNBUFFERED=1` 缓解)
2. **`ps` 的 TIME 是 CPU 累计时间**(跨核),不是 wall time。看 wall 用 `etime` 或 tfevents 时间戳。
3. **VAE checkpoint 文件混在共享目录**(SKRL 路线遗留):`runs/.../skrl/vae_*.pkl`。rsl 路线 checkpoint 是 `rsl_dwaq/<run>/model_*.pt`,别搞混。
4. **物理在 CPU**(MotrixSim numpy backend),网络在 GPU。每步 numpy↔torch 传输有开销,~16 it/min(2048 envs)。物理是瓶颈。
5. **MotrixSim hfield 不支持陡台阶**(根本限制):如果要更忠实的高地形,得调研 MotrixSim 是否有 mesh 凸分解地形方案,或接受矮地形。
6. **验证要看行为不看奖励**:`only_positive_rewards` + 一堆"别动"惩罚,机器人容易学"趴着/站着不动"刷分。验证时务必测**实际前进速度跟踪 + 摔倒率**,不能只看 mean_reward。
---
## 9. 下一步
1. 等当前训练跑完,导出 + MuJoCo 验证实际行走(前进/转向跟踪)。
2. 若行走 OK:逐步提高地形难度(z_scale),找 MotrixSim hfield 能稳定接触的上限。
3. 若行走不够好:检查地形课程是否真在升级(机器人有没有从 level 0 走到 level 1+)、调 init_noise_std。
4. 长期:调研 MotrixSim mesh 地形(凸分解)以更忠实复现陡台阶。

BIN
exports_go1_cts/normalizer.npz Normal file
View File

Binary file not shown.

BIN
exports_go1_cts/policy.onnx Normal file
View File

Binary file not shown.

BIN
exports_go1_cts/policy_full.onnx Normal file
View File

Binary file not shown.

BIN
exports_go1_dreamwaq/policy.onnx Normal file
View File

Binary file not shown.

21
exports_go1_flat/metadata.txt Normal file
View File

@@ -0,0 +1,21 @@
# Go1 Flat Terrain Walk - ONNX Policy Metadata
obs_dim: 45
action_dim: 12
hidden_dims: [256, 128, 64]
observation_layout: gyro(3) + gravity(3) + joint_angle(12) + joint_vel(12) + last_action(12) + command(3)
- NO linear velocity in observation
# Joint order: ['FR_hip', 'FR_thigh', 'FR_calf', 'FL_hip', 'FL_thigh', 'FL_calf', 'RR_hip', 'RR_thigh', 'RR_calf', 'RL_hip', 'RL_thigh', 'RL_calf']
default_angles: [-0.0, 0.8999999761581421, -1.7999999523162842, 0.0, 0.8999999761581421, -1.7999999523162842, -0.0, 0.8999999761581421, -1.7999999523162842, 0.0, 0.8999999761581421, -1.7999999523162842]
action_scale: 0.05
kp: 80.0
kd: 1.0
clip_actions: 23.7
clip_observations: 100.0
# Observation scales (applied BEFORE ONNX normalization):
lin_vel: 2.0
ang_vel: 0.25
dof_pos: 1.0
dof_vel: 0.05
command_scale: [2.0, 2.0, 0.25] # for [vx, vy, wz]

BIN
exports_go1_flat/normalizer.npz Normal file
View File

Binary file not shown.

BIN
exports_go1_flat/policy.onnx Normal file
View File

Binary file not shown.

BIN
exports_go1_go2style/normalizer.npz Normal file
View File

Binary file not shown.

BIN
exports_go1_go2style/policy.onnx Normal file
View File

Binary file not shown.

BIN
exports_go1_no_linevel/better1_policy.onnx Normal file
View File

Binary file not shown.

BIN
exports_go1_no_linevel/better2.onnx Normal file
View File

Binary file not shown.

BIN
exports_go1_no_linevel/better3.onnx Normal file
View File

Binary file not shown.

29
exports_go1_no_linevel/metadata.txt Normal file
View File

@@ -0,0 +1,29 @@
# Go1 No-Linevel Terrain Walk - ONNX Policy Metadata
env: go1-stairs-terrain-walk-no-linevel
obs_dim: 57
action_dim: 12
hidden_dims: [512, 256, 128]
# Observation layout (57 dims, NO linear velocity):
[0:3] gyro * ang_vel_scale
[3:6] gravity (body frame)
[6:18] joint_angle_deviation * dof_pos_scale
[18:30] joint_vel * dof_vel_scale
[30:42] last_actions (raw)
[42:45] commands [vx*2.0, vy*2.0, wz*0.25]
[45:57] foot_contact_forces (body frame, raw)
# Joint order: ['FR_hip', 'FR_thigh', 'FR_calf', 'FL_hip', 'FL_thigh', 'FL_calf', 'RR_hip', 'RR_thigh', 'RR_calf', 'RL_hip', 'RL_thigh', 'RL_calf']
default_angles: [-0.0, 0.8999999761581421, -1.7999999523162842, 0.0, 0.8999999761581421, -1.7999999523162842, -0.0, 0.8999999761581421, -1.7999999523162842, 0.0, 0.8999999761581421, -1.7999999523162842]
action_scale: 0.05
kp: 80.0
kd: 1.0
clip_actions: 23.7
clip_observations: 100.0
# Observation scales:
ang_vel: 0.25
dof_pos: 1.0
dof_vel: 0.05
contact_force: 1.0
command_scale: [2.0, 2.0, 0.25]

BIN
exports_go1_no_linevel/normalizer.npz Normal file
View File

Binary file not shown.

BIN
exports_go1_orig/policy.onnx Normal file
View File

Binary file not shown.

BIN
exports_orig_comb/normalizer.npz Normal file
View File

Binary file not shown.

BIN
exports_orig_comb/policy.onnx Normal file
View File

Binary file not shown.

13
motrix_rl/src/motrix_rl/dwaq_rsl/__init__.py Normal file
View File

@@ -0,0 +1,13 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""Faithful port of Manaro-Alpha/DreamWaQ's custom rsl_rl-1.0.2 for MotrixLab.
Provides ActorCritic_DWAQ (CENet+Actor+Critic), PPO (with joint VAE training),
RolloutStorage, OnPolicyRunner, and a MotrixLab NpEnv -> VecEnv adapter.
"""
from .actor_critic_dwaq import ActorCritic_DWAQ
from .ppo import PPO
from .rollout_storage import RolloutStorage
from .runner import OnPolicyRunner
from .env_wrapper import DwaqVecEnv
__all__ = ["ActorCritic_DWAQ", "PPO", "RolloutStorage", "OnPolicyRunner", "DwaqVecEnv"]

169
motrix_rl/src/motrix_rl/dwaq_rsl/actor_critic_dwaq.py Normal file
View File

@@ -0,0 +1,169 @@
from __future__ import annotations
import torch
import torch.nn as nn
from torch.distributions import Normal
class ActorCritic_DWAQ(nn.Module):
def __init__(self, num_actor_obs, num_critic_obs, num_actions, cenet_in_dim, cenet_out_dim, activation="elu", init_noise_std=1.0,):
super().__init__()
self.activation = get_activation(activation)
actor_input_dim = num_actor_obs
critic_input_dim = num_critic_obs
self.actor = nn.Sequential(
nn.Linear(actor_input_dim,512),
self.activation,
nn.Linear(512,256),
self.activation,
nn.Linear(256,128),
self.activation,
nn.Linear(128,num_actions)
)
self.critic = nn.Sequential(
nn.Linear(critic_input_dim,512),
self.activation,
nn.Linear(512,256),
self.activation,
nn.Linear(256,128),
self.activation,
nn.Linear(128,1)
)
self.encoder = nn.Sequential(
nn.Linear(cenet_in_dim,128),
self.activation,
nn.Linear(128,64),
self.activation,
)
self.encode_mean_latent = nn.Linear(64,cenet_out_dim-3)
self.encode_logvar_latent = nn.Linear(64,cenet_out_dim-3)
self.encode_mean_vel = nn.Linear(64,3)
self.encode_logvar_vel = nn.Linear(64,3)
self.decoder = nn.Sequential(
nn.Linear(cenet_out_dim,64),
self.activation,
nn.Linear(64,128),
self.activation,
nn.Linear(128,45)
)
self.std = nn.Parameter(init_noise_std * torch.ones(num_actions))
self.distribution = None
# disable args validation for speedup
Normal.set_default_validate_args = False
# seems that we get better performance without init
# self.init_memory_weights(self.memory_a, 0.001, 0.)
# self.init_memory_weights(self.memory_c, 0.001, 0.)
@staticmethod
# not used at the moment
def init_weights(sequential, scales):
[
torch.nn.init.orthogonal_(module.weight, gain=scales[idx])
for idx, module in enumerate(mod for mod in sequential if isinstance(mod, nn.Linear))
]
def reset(self, dones=None):
pass
def forward(self):
raise NotImplementedError
def reparameterise(self,mean,logvar):
var = torch.exp(logvar*0.5)
code_temp = torch.randn_like(var)
code = mean + var*code_temp
return code
def cenet_forward(self,obs_history):
distribution = self.encoder(obs_history)
mean_latent = self.encode_mean_latent(distribution)
logvar_latent = self.encode_logvar_latent(distribution)
# var = torch.exp(logvar_latent*0.5)
# code_temp = torch.randn_like(var)
# code = mean_latent + var*code_temp
# print("latent : ",code[0])
mean_vel = self.encode_mean_vel(distribution)
logvar_vel = self.encode_mean_vel(distribution)
code_latent = self.reparameterise(mean_latent,logvar_latent)
code_vel = self.reparameterise(mean_vel,logvar_vel)
code = torch.cat((code_vel,code_latent),dim=-1)
decode = self.decoder(code)
return code,code_vel,decode,mean_vel,logvar_vel,mean_latent,logvar_latent
@property
def action_mean(self):
return self.distribution.mean
@property
def action_std(self):
return self.distribution.stddev
@property
def entropy(self):
return self.distribution.entropy().sum(dim=-1)
def update_distribution(self, observations):
mean = self.actor(observations)
# Clamp raw actor output before building the Normal distribution:
# on unstable terrain (hfield) the network can produce extreme values
# that overflow → NaN → crash the full PPO pipeline.
mean = torch.clamp(mean, -10.0, 10.0)
self.distribution = Normal(mean, mean * 0.0 + self.std)
def act(self, observations, obs_history, **kwargs):
code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
observations = torch.cat((code,observations),dim=-1)
self.update_distribution(observations)
return self.distribution.sample()
def get_actions_log_prob(self, actions):
return self.distribution.log_prob(actions).sum(dim=-1)
def act_inference(self, observations,obs_history):
code,_,decode,_,_,_,_ = self.cenet_forward(obs_history)
observations = torch.cat((code,observations),dim=-1)
actions_mean = self.actor(observations)
return actions_mean
def evaluate(self, critic_observations, **kwargs):
value = self.critic(critic_observations)
return value
def get_activation(act_name):
if act_name == "elu":
return nn.ELU()
elif act_name == "selu":
return nn.SELU()
elif act_name == "relu":
return nn.ReLU()
elif act_name == "crelu":
return nn.CReLU()
elif act_name == "lrelu":
return nn.LeakyReLU()
elif act_name == "tanh":
return nn.Tanh()
elif act_name == "sigmoid":
return nn.Sigmoid()
else:
print("invalid activation function!")
return None

106
motrix_rl/src/motrix_rl/dwaq_rsl/env_wrapper.py Normal file
View File

@@ -0,0 +1,106 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""Adapter: MotrixLab DreamWaQ NpEnv -> upstream rsl_rl-1.0.2 VecEnv interface.
Upstream OnPolicyRunner expects env.step() to return a 7-tuple:
(obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, extras)
and env.reset() to return (obs, privileged_obs, prev_privileged_obs, obs_hist).
MotrixLab's DreamWaQTask (numpy, CPU physics) already provides:
state.obs -> obs (N, 45)
info["privileged_obs"] -> privileged (N, 235) = obs(45)+base_vel(3)+heights(187)
info["obs_history"] -> history (N, 5, 45) (flattened to 225)
state.terminated/truncated -> dones / time_outs
This wrapper does numpy<->torch transfer and maintains prev_privileged_obs.
"""
import numpy as np
import torch
class DwaqVecEnv:
"""Wraps a MotrixLab DreamWaQ NpEnv for the upstream DreamWaQ rsl_rl runner."""
def __init__(self, env, device, num_obs=45, num_privileged_obs=235,
num_obs_hist=5, num_actions=12, clip_actions=100.0):
self._env = env
self.device = device
self.num_envs = env.num_envs
self.num_obs = num_obs
self.num_privileged_obs = num_privileged_obs
self.num_obs_hist = num_obs_hist
self.num_actions = num_actions
self._clip_actions = clip_actions
self.max_episode_length = int(getattr(env.cfg, "max_episode_steps", None) or 1000)
n = self.num_envs
self.obs_buf = torch.zeros(n, num_obs, device=device)
self.obs_hist_buf = torch.zeros(n, num_obs_hist * num_obs, device=device)
self.privileged_obs_buf = torch.zeros(n, num_privileged_obs, device=device)
self.prev_privileged_obs_buf = torch.zeros(n, num_privileged_obs, device=device)
self.rew_buf = torch.zeros(n, device=device)
self.reset_buf = torch.ones(n, device=device, dtype=torch.long)
self.episode_length_buf = torch.zeros(n, device=device, dtype=torch.long)
self.extras = {}
self._env.init_state()
self._sync_from_state(self._env._state)
# ── helpers ──
def _t(self, arr):
return torch.from_numpy(np.ascontiguousarray(arr, dtype=np.float32)).to(self.device)
def _sync_from_state(self, state):
self.obs_buf = self._t(state.obs)
info = state.info
priv = info.get("privileged_obs")
if priv is None:
priv = np.zeros((self.num_envs, self.num_privileged_obs), dtype=np.float32)
self.privileged_obs_buf = self._t(priv)
hist = info.get("obs_history") # (N, 5, 45)
if hist is None:
hist = np.zeros((self.num_envs, self.num_obs_hist, self.num_obs), dtype=np.float32)
self.obs_hist_buf = self._t(hist).reshape(self.num_envs, -1)
self.episode_length_buf = torch.from_numpy(
np.asarray(info.get("steps", np.zeros(self.num_envs)), dtype=np.int64)).to(self.device)
# ── VecEnv interface ──
def step(self, actions):
actions = torch.clip(actions, -self._clip_actions, self._clip_actions)
actions_np = actions.detach().cpu().numpy().astype(np.float32)
# prev_privileged = privileged BEFORE this step (matches upstream)
self.prev_privileged_obs_buf = self.privileged_obs_buf.clone()
state = self._env.step(actions_np)
self._sync_from_state(state)
done = np.logical_or(state.terminated, state.truncated)
self.rew_buf = self._t(state.reward).reshape(-1)
self.reset_buf = torch.from_numpy(done.astype(np.int64)).to(self.device)
self.extras = {
"time_outs": torch.from_numpy(np.asarray(state.truncated, dtype=np.float32)).to(self.device),
}
# pass per-term reward logging through if present
terms = state.info.get("reward_terms")
if terms:
self.extras["episode"] = {f"rew_{k}": float(v) for k, v in terms.items()}
# Log terrain_level so tensorboard tracks curriculum progress.
tl = state.info.get("terrain_level")
if tl is not None:
self.extras["episode"] = self.extras.get("episode", {})
self.extras["episode"]["terrain_level"] = float(np.asarray(tl).mean())
# rewards/dones are (N,) like upstream (storage views them to (N,1) internally)
return (self.obs_buf, self.privileged_obs_buf, self.prev_privileged_obs_buf,
self.obs_hist_buf, self.rew_buf, self.reset_buf, self.extras)
def reset(self, env_ids=None):
self._env.init_state()
self._sync_from_state(self._env._state)
self.prev_privileged_obs_buf = torch.zeros_like(self.privileged_obs_buf)
return (self.obs_buf, self.privileged_obs_buf,
self.prev_privileged_obs_buf, self.obs_hist_buf)
def get_observations(self):
return self.obs_buf, self.obs_hist_buf
def get_privileged_observations(self):
return self.privileged_obs_buf, self.prev_privileged_obs_buf

203
motrix_rl/src/motrix_rl/dwaq_rsl/ppo.py Normal file
View File

@@ -0,0 +1,203 @@
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
import torch
import torch.nn as nn
import torch.optim as optim
from .actor_critic_dwaq import ActorCritic_DWAQ
from .rollout_storage import RolloutStorage
class PPO:
actor_critic: ActorCritic_DWAQ
def __init__(self,
actor_critic,
num_learning_epochs=1,
num_mini_batches=1,
clip_param=0.2,
gamma=0.99,
lam=0.95,
value_loss_coef=1.0,
entropy_coef=0.0,
learning_rate=1e-3,
max_grad_norm=1.0,
use_clipped_value_loss=True,
schedule="fixed",
desired_kl=0.01,
device='cpu',
):
self.device = device
self.desired_kl = desired_kl
self.schedule = schedule
self.learning_rate = learning_rate
# PPO components
self.actor_critic = actor_critic
self.actor_critic.to(self.device)
self.storage = None # initialized later
self.optimizer = optim.Adam(self.actor_critic.parameters(), lr=learning_rate)
self.transition = RolloutStorage.Transition()
# PPO parameters
self.clip_param = clip_param
self.num_learning_epochs = num_learning_epochs
self.num_mini_batches = num_mini_batches
self.value_loss_coef = value_loss_coef
self.entropy_coef = entropy_coef
self.gamma = gamma
self.lam = lam
self.max_grad_norm = max_grad_norm
self.use_clipped_value_loss = use_clipped_value_loss
def init_storage(self, num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, obs_hist_shape, action_shape):
self.storage = RolloutStorage(num_envs, num_transitions_per_env, actor_obs_shape, critic_obs_shape, obs_hist_shape, action_shape, self.device)
def test_mode(self):
self.actor_critic.test()
def train_mode(self):
self.actor_critic.train()
def act(self, obs, critic_obs, prev_critic_obs, obs_history):
# if self.actor_critic.is_recurrent:
# self.transition.hidden_states = self.actor_critic.get_hidden_states()
# Compute the actions and values
self.transition.actions = self.actor_critic.act(obs,obs_history).detach()
self.transition.values = self.actor_critic.evaluate(critic_obs).detach()
self.transition.actions_log_prob = self.actor_critic.get_actions_log_prob(self.transition.actions).detach()
self.transition.action_mean = self.actor_critic.action_mean.detach()
self.transition.action_sigma = self.actor_critic.action_std.detach()
# need to record obs and critic_obs before env.step()
self.transition.observations = obs
self.transition.observation_history = obs_history
self.transition.critic_observations = critic_obs
self.transition.prev_critic_obs = prev_critic_obs
return self.transition.actions
def process_env_step(self, rewards, dones, infos):
self.transition.rewards = rewards.clone()
self.transition.dones = dones
# Bootstrapping on time outs
if 'time_outs' in infos:
self.transition.rewards += self.gamma * torch.squeeze(self.transition.values * infos['time_outs'].unsqueeze(1).to(self.device), 1)
# Record the transition
self.storage.add_transitions(self.transition)
self.transition.clear()
self.actor_critic.reset(dones)
def compute_returns(self, last_critic_obs):
last_values= self.actor_critic.evaluate(last_critic_obs).detach()
self.storage.compute_returns(last_values, self.gamma, self.lam)
def update(self,beta=1):
mean_value_loss = 0
mean_surrogate_loss = 0
mean_autoenc_loss = 0
# if self.actor_critic.is_recurrent:
# generator = self.storage.reccurent_mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
# else:
generator = self.storage.mini_batch_generator(self.num_mini_batches, self.num_learning_epochs)
for obs_batch, critic_obs_batch, prev_critic_obs_batch, obs_hist_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, old_actions_log_prob_batch, \
old_mu_batch, old_sigma_batch, hid_states_batch, masks_batch in generator:
self.actor_critic.act(obs_batch, obs_hist_batch, masks=masks_batch, hidden_states=hid_states_batch[0])
actions_log_prob_batch = self.actor_critic.get_actions_log_prob(actions_batch)
value_batch = self.actor_critic.evaluate(critic_obs_batch, masks=masks_batch, hidden_states=hid_states_batch[1])
mu_batch = self.actor_critic.action_mean
sigma_batch = self.actor_critic.action_std
entropy_batch = self.actor_critic.entropy
# KL
if self.desired_kl != None and self.schedule == 'adaptive':
with torch.inference_mode():
kl = torch.sum(
torch.log(sigma_batch / old_sigma_batch + 1.e-5) + (torch.square(old_sigma_batch) + torch.square(old_mu_batch - mu_batch)) / (2.0 * torch.square(sigma_batch)) - 0.5, axis=-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
#Beta VAE loss
code,code_vel,decode,mean_vel,logvar_vel,mean_latent,logvar_latent = self.actor_critic.cenet_forward(obs_hist_batch)
vel_target = prev_critic_obs_batch[:,45:48]
decode_target = obs_batch
vel_target.requires_grad = False
decode_target.requires_grad = False
autoenc_loss = (nn.MSELoss()(code_vel,vel_target) + nn.MSELoss()(decode,decode_target) + beta*(-0.5 * torch.sum(1 + logvar_latent - mean_latent.pow(2) - logvar_latent.exp())))/self.num_mini_batches
# estimation_loss = (code[:,0:3] - prev_critic_obs_batch[:,45:48]).pow(2).mean()
# reconst_loss = (decode - obs_batch).pow(2).mean()
# latent_loss = beta*(-0.5 * torch.sum(1 + logvar - mean.pow(2) - logvar.exp()))/mean.shape[0]
# autoenc_loss = estimation_loss + reconst_loss + latent_loss
# Surrogate loss
ratio = torch.exp(actions_log_prob_batch - torch.squeeze(old_actions_log_prob_batch))
surrogate = -torch.squeeze(advantages_batch) * ratio
surrogate_clipped = -torch.squeeze(advantages_batch) * torch.clamp(ratio, 1.0 - self.clip_param,
1.0 + self.clip_param)
surrogate_loss = torch.max(surrogate, surrogate_clipped).mean()
# Value function loss
if self.use_clipped_value_loss:
value_clipped = target_values_batch + (value_batch - target_values_batch).clamp(-self.clip_param,
self.clip_param)
value_losses = (value_batch - returns_batch).pow(2)
value_losses_clipped = (value_clipped - returns_batch).pow(2)
value_loss = torch.max(value_losses, value_losses_clipped).mean()
else:
value_loss = (returns_batch - value_batch).pow(2).mean()
loss = surrogate_loss + self.value_loss_coef * value_loss - self.entropy_coef * entropy_batch.mean() + autoenc_loss
# Gradient step
self.optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(self.actor_critic.parameters(), self.max_grad_norm)
self.optimizer.step()
mean_value_loss += value_loss.item()
mean_surrogate_loss += surrogate_loss.item()
mean_autoenc_loss += autoenc_loss.item()
num_updates = self.num_learning_epochs * self.num_mini_batches
mean_value_loss /= num_updates
mean_surrogate_loss /= num_updates
self.storage.clear()
return mean_value_loss, mean_surrogate_loss, mean_autoenc_loss

246
motrix_rl/src/motrix_rl/dwaq_rsl/rollout_storage.py Normal file
View File

@@ -0,0 +1,246 @@
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
import torch
import numpy as np
from .utils import split_and_pad_trajectories
class RolloutStorage:
class Transition:
def __init__(self):
self.observations = None
self.critic_observations = None
self.prev_critic_obs = None
self.observation_history = None
self.actions = None
self.rewards = None
self.dones = None
self.values = None
self.actions_log_prob = None
self.action_mean = None
self.action_sigma = None
self.hidden_states = None
def clear(self):
self.__init__()
def __init__(self, num_envs, num_transitions_per_env, obs_shape, privileged_obs_shape, obs_hist_shape, actions_shape, device="cpu"):
self.device = device
self.obs_shape = obs_shape
self.privileged_obs_shape = privileged_obs_shape
self.actions_shape = actions_shape
# Core
self.observations = torch.zeros(num_transitions_per_env, num_envs, *obs_shape, device=self.device)
if privileged_obs_shape[0] is not None:
self.privileged_observations = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
else:
self.privileged_observations = None
self.prev_critic_obs = torch.zeros(num_transitions_per_env, num_envs, *privileged_obs_shape, device=self.device)
self.observation_history = torch.zeros(num_transitions_per_env, num_envs, *obs_hist_shape, device=self.device)
self.rewards = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
self.actions = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
self.dones = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device).byte()
# For PPO
self.actions_log_prob = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
self.values = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
self.returns = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
self.advantages = torch.zeros(num_transitions_per_env, num_envs, 1, device=self.device)
self.mu = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
self.sigma = torch.zeros(num_transitions_per_env, num_envs, *actions_shape, device=self.device)
self.num_transitions_per_env = num_transitions_per_env
self.num_envs = num_envs
# rnn
self.saved_hidden_states_a = None
self.saved_hidden_states_c = None
self.step = 0
def add_transitions(self, transition: Transition):
if self.step >= self.num_transitions_per_env:
raise AssertionError("Rollout buffer overflow")
self.observations[self.step].copy_(transition.observations)
if self.privileged_observations is not None: self.privileged_observations[self.step].copy_(transition.critic_observations)
self.prev_critic_obs[self.step].copy_(transition.prev_critic_obs)
self.observation_history[self.step].copy_(transition.observation_history)
self.actions[self.step].copy_(transition.actions)
self.rewards[self.step].copy_(transition.rewards.view(-1, 1))
self.dones[self.step].copy_(transition.dones.view(-1, 1))
self.values[self.step].copy_(transition.values)
self.actions_log_prob[self.step].copy_(transition.actions_log_prob.view(-1, 1))
self.mu[self.step].copy_(transition.action_mean)
self.sigma[self.step].copy_(transition.action_sigma)
self._save_hidden_states(transition.hidden_states)
self.step += 1
def _save_hidden_states(self, hidden_states):
if hidden_states is None or hidden_states==(None, None):
return
# make a tuple out of GRU hidden state sto match the LSTM format
hid_a = hidden_states[0] if isinstance(hidden_states[0], tuple) else (hidden_states[0],)
hid_c = hidden_states[1] if isinstance(hidden_states[1], tuple) else (hidden_states[1],)
# initialize if needed
if self.saved_hidden_states_a is None:
self.saved_hidden_states_a = [torch.zeros(self.observations.shape[0], *hid_a[i].shape, device=self.device) for i in range(len(hid_a))]
self.saved_hidden_states_c = [torch.zeros(self.observations.shape[0], *hid_c[i].shape, device=self.device) for i in range(len(hid_c))]
# copy the states
for i in range(len(hid_a)):
self.saved_hidden_states_a[i][self.step].copy_(hid_a[i])
self.saved_hidden_states_c[i][self.step].copy_(hid_c[i])
def clear(self):
self.step = 0
def compute_returns(self, last_values, gamma, lam):
advantage = 0
for step in reversed(range(self.num_transitions_per_env)):
if step == self.num_transitions_per_env - 1:
next_values = last_values
else:
next_values = self.values[step + 1]
next_is_not_terminal = 1.0 - self.dones[step].float()
delta = self.rewards[step] + next_is_not_terminal * gamma * next_values - self.values[step]
advantage = delta + next_is_not_terminal * gamma * lam * advantage
self.returns[step] = advantage + self.values[step]
# Compute and normalize the advantages
self.advantages = self.returns - self.values
self.advantages = (self.advantages - self.advantages.mean()) / (self.advantages.std() + 1e-8)
def get_statistics(self):
done = self.dones
done[-1] = 1
flat_dones = done.permute(1, 0, 2).reshape(-1, 1)
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero(as_tuple=False)[:, 0]))
trajectory_lengths = (done_indices[1:] - done_indices[:-1])
return trajectory_lengths.float().mean(), self.rewards.mean()
def mini_batch_generator(self, num_mini_batches, num_epochs=8):
batch_size = self.num_envs * self.num_transitions_per_env
mini_batch_size = batch_size // num_mini_batches
indices = torch.randperm(num_mini_batches*mini_batch_size, requires_grad=False, device=self.device)
observations = self.observations.flatten(0, 1)
if self.privileged_observations is not None:
critic_observations = self.privileged_observations.flatten(0, 1)
else:
critic_observations = observations
prev_critic_obs = self.prev_critic_obs.flatten(0,1)
obs_history = self.observation_history.flatten(0,1)
actions = self.actions.flatten(0, 1)
values = self.values.flatten(0, 1)
returns = self.returns.flatten(0, 1)
old_actions_log_prob = self.actions_log_prob.flatten(0, 1)
advantages = self.advantages.flatten(0, 1)
old_mu = self.mu.flatten(0, 1)
old_sigma = self.sigma.flatten(0, 1)
for epoch in range(num_epochs):
for i in range(num_mini_batches):
start = i*mini_batch_size
end = (i+1)*mini_batch_size
batch_idx = indices[start:end]
obs_batch = observations[batch_idx]
critic_observations_batch = critic_observations[batch_idx]
prev_critic_obs_batch = prev_critic_obs[batch_idx]
obs_hist_batch = obs_history[batch_idx]
actions_batch = actions[batch_idx]
target_values_batch = values[batch_idx]
returns_batch = returns[batch_idx]
old_actions_log_prob_batch = old_actions_log_prob[batch_idx]
advantages_batch = advantages[batch_idx]
old_mu_batch = old_mu[batch_idx]
old_sigma_batch = old_sigma[batch_idx]
yield obs_batch, critic_observations_batch,prev_critic_obs_batch,obs_hist_batch, actions_batch, target_values_batch, advantages_batch, returns_batch, \
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (None, None), None
# for RNNs only
def reccurent_mini_batch_generator(self, num_mini_batches, num_epochs=8):
padded_obs_trajectories, trajectory_masks = split_and_pad_trajectories(self.observations, self.dones)
if self.privileged_observations is not None:
padded_critic_obs_trajectories, _ = split_and_pad_trajectories(self.privileged_observations, self.dones)
else:
padded_critic_obs_trajectories = padded_obs_trajectories
mini_batch_size = self.num_envs // num_mini_batches
for ep in range(num_epochs):
first_traj = 0
for i in range(num_mini_batches):
start = i*mini_batch_size
stop = (i+1)*mini_batch_size
dones = self.dones.squeeze(-1)
last_was_done = torch.zeros_like(dones, dtype=torch.bool)
last_was_done[1:] = dones[:-1]
last_was_done[0] = True
trajectories_batch_size = torch.sum(last_was_done[:, start:stop])
last_traj = first_traj + trajectories_batch_size
masks_batch = trajectory_masks[:, first_traj:last_traj]
obs_batch = padded_obs_trajectories[:, first_traj:last_traj]
critic_obs_batch = padded_critic_obs_trajectories[:, first_traj:last_traj]
actions_batch = self.actions[:, start:stop]
old_mu_batch = self.mu[:, start:stop]
old_sigma_batch = self.sigma[:, start:stop]
returns_batch = self.returns[:, start:stop]
advantages_batch = self.advantages[:, start:stop]
values_batch = self.values[:, start:stop]
old_actions_log_prob_batch = self.actions_log_prob[:, start:stop]
# reshape to [num_envs, time, num layers, hidden dim] (original shape: [time, num_layers, num_envs, hidden_dim])
# then take only time steps after dones (flattens num envs and time dimensions),
# take a batch of trajectories and finally reshape back to [num_layers, batch, hidden_dim]
last_was_done = last_was_done.permute(1, 0)
hid_a_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj].transpose(1, 0).contiguous()
for saved_hidden_states in self.saved_hidden_states_a ]
hid_c_batch = [ saved_hidden_states.permute(2, 0, 1, 3)[last_was_done][first_traj:last_traj].transpose(1, 0).contiguous()
for saved_hidden_states in self.saved_hidden_states_c ]
# remove the tuple for GRU
hid_a_batch = hid_a_batch[0] if len(hid_a_batch)==1 else hid_a_batch
hid_c_batch = hid_c_batch[0] if len(hid_c_batch)==1 else hid_a_batch
yield obs_batch, critic_obs_batch, actions_batch, values_batch, advantages_batch, returns_batch, \
old_actions_log_prob_batch, old_mu_batch, old_sigma_batch, (hid_a_batch, hid_c_batch), masks_batch
first_traj = last_traj

180
motrix_rl/src/motrix_rl/dwaq_rsl/runner.py Normal file
View File

@@ -0,0 +1,180 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ OnPolicyRunner — adapted from upstream rsl_rl-1.0.2.
Faithful port of Manaro-Alpha/DreamWaQ's on_policy_runner.py, with the
rsl_rl-package imports replaced by local (relative) imports and the isaacgym
noise removed. Drives the PPO + CENet (VAE) joint training loop.
"""
import time
import os
from collections import deque
import statistics
from torch.utils.tensorboard import SummaryWriter
import torch
from .ppo import PPO
from .actor_critic_dwaq import ActorCritic_DWAQ
class OnPolicyRunner:
def __init__(self, env, train_cfg, log_dir=None, device="cpu"):
self.cfg = train_cfg["runner"]
self.alg_cfg = train_cfg["algorithm"]
self.policy_cfg = train_cfg["policy"]
self.device = device
self.env = env
num_critic_obs = env.num_privileged_obs if env.num_privileged_obs is not None else env.num_obs
cenet_in_dim = env.num_obs_hist * env.num_obs # 5*45 = 225
cenet_out_dim = 19 # vel(3) + latent(16)
actor_critic = ActorCritic_DWAQ(
env.num_obs + cenet_out_dim, # actor input: obs(45) + code(19) = 64
num_critic_obs, # critic input: privileged_obs
env.num_actions, # 12
cenet_in_dim, # 225
cenet_out_dim, # 19
**self.policy_cfg,
).to(self.device)
self.alg = PPO(actor_critic, device=self.device, **self.alg_cfg)
self.num_steps_per_env = self.cfg["num_steps_per_env"]
self.save_interval = self.cfg["save_interval"]
self.alg.init_storage(
env.num_envs, self.num_steps_per_env,
[env.num_obs], [env.num_privileged_obs],
[env.num_obs_hist * env.num_obs], [env.num_actions],
)
self.log_dir = log_dir
self.writer = None
self.tot_timesteps = 0
self.tot_time = 0
self.current_learning_iteration = 0
_, _, _, _ = self.env.reset()
def learn(self, num_learning_iterations, init_at_random_ep_len=False):
if self.log_dir is not None and self.writer is None:
self.writer = SummaryWriter(log_dir=self.log_dir, flush_secs=10)
if init_at_random_ep_len:
self.env.episode_length_buf = torch.randint_like(
self.env.episode_length_buf, high=int(self.env.max_episode_length))
obs, obs_hist = self.env.get_observations()
privileged_obs, prev_critic_obs = self.env.get_privileged_observations()
critic_obs = privileged_obs if privileged_obs is not None else obs
obs, critic_obs, prev_critic_obs, obs_hist = (
obs.to(self.device), critic_obs.to(self.device),
prev_critic_obs.to(self.device), obs_hist.to(self.device))
self.alg.actor_critic.train()
ep_infos = []
rewbuffer = deque(maxlen=100)
lenbuffer = deque(maxlen=100)
cur_reward_sum = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
cur_episode_length = torch.zeros(self.env.num_envs, dtype=torch.float, device=self.device)
tot_iter = self.current_learning_iteration + num_learning_iterations
for it in range(self.current_learning_iteration, tot_iter):
start = time.time()
with torch.inference_mode():
for i in range(self.num_steps_per_env):
actions = self.alg.act(obs, critic_obs, prev_critic_obs, obs_hist)
obs, privileged_obs, prev_privileged_obs, obs_hist, rewards, dones, infos = self.env.step(actions)
critic_obs = privileged_obs if privileged_obs is not None else obs
prev_critic_obs = prev_privileged_obs
obs, critic_obs, prev_critic_obs, obs_hist, rewards, dones = (
obs.to(self.device), critic_obs.to(self.device),
prev_critic_obs.to(self.device), obs_hist.to(self.device),
rewards.to(self.device), dones.to(self.device))
self.alg.process_env_step(rewards, dones, infos)
if self.log_dir is not None:
if "episode" in infos:
ep_infos.append(infos["episode"])
cur_reward_sum += rewards
cur_episode_length += 1
new_ids = (dones > 0).nonzero(as_tuple=False)
rewbuffer.extend(cur_reward_sum[new_ids][:, 0].cpu().numpy().tolist())
lenbuffer.extend(cur_episode_length[new_ids][:, 0].cpu().numpy().tolist())
cur_reward_sum[new_ids] = 0
cur_episode_length[new_ids] = 0
stop = time.time()
collection_time = stop - start
start = stop
self.alg.compute_returns(critic_obs)
mean_value_loss, mean_surrogate_loss, mean_autoenc_loss = self.alg.update()
stop = time.time()
learn_time = stop - start
if self.log_dir is not None:
self.log(locals())
if it % self.save_interval == 0:
self.save(os.path.join(self.log_dir, "model_{}.pt".format(it)))
ep_infos.clear()
self.current_learning_iteration += num_learning_iterations
self.save(os.path.join(self.log_dir, "model_{}.pt".format(self.current_learning_iteration)))
def log(self, locs, width=80, pad=35):
self.tot_timesteps += self.num_steps_per_env * self.env.num_envs
self.tot_time += locs["collection_time"] + locs["learn_time"]
iteration_time = locs["collection_time"] + locs["learn_time"]
if locs["ep_infos"]:
for key in locs["ep_infos"][0]:
infotensor = torch.tensor([], device=self.device)
for ep_info in locs["ep_infos"]:
v = ep_info[key]
if not isinstance(v, torch.Tensor):
v = torch.Tensor([v])
if len(v.shape) == 0:
v = v.unsqueeze(0)
infotensor = torch.cat((infotensor, v.to(self.device)))
self.writer.add_scalar("Episode/" + key, torch.mean(infotensor), locs["it"])
mean_std = self.alg.actor_critic.std.mean()
fps = int(self.num_steps_per_env * self.env.num_envs / (locs["collection_time"] + locs["learn_time"]))
self.writer.add_scalar("Loss/value_function", locs["mean_value_loss"], locs["it"])
self.writer.add_scalar("Loss/surrogate", locs["mean_surrogate_loss"], locs["it"])
self.writer.add_scalar("Loss/autoenc_function", locs["mean_autoenc_loss"], locs["it"])
self.writer.add_scalar("Loss/learning_rate", self.alg.learning_rate, locs["it"])
self.writer.add_scalar("Policy/mean_noise_std", mean_std.item(), locs["it"])
self.writer.add_scalar("Perf/total_fps", fps, locs["it"])
if len(locs["rewbuffer"]) > 0:
self.writer.add_scalar("Train/mean_reward", statistics.mean(locs["rewbuffer"]), locs["it"])
self.writer.add_scalar("Train/mean_episode_length", statistics.mean(locs["lenbuffer"]), locs["it"])
if it_is_log := (locs["it"] % 10 == 0):
rew = statistics.mean(locs["rewbuffer"]) if len(locs["rewbuffer"]) > 0 else 0.0
elen = statistics.mean(locs["lenbuffer"]) if len(locs["lenbuffer"]) > 0 else 0.0
print(f"[it {locs['it']}/{self.current_learning_iteration + locs['num_learning_iterations']}] "
f"fps={fps} val_loss={locs['mean_value_loss']:.3f} surr={locs['mean_surrogate_loss']:.4f} "
f"vae={locs['mean_autoenc_loss']:.3f} std={mean_std.item():.2f} "
f"rew={rew:.2f} eplen={elen:.0f}")
def save(self, path, infos=None):
torch.save({
"model_state_dict": self.alg.actor_critic.state_dict(),
"optimizer_state_dict": self.alg.optimizer.state_dict(),
"iter": self.current_learning_iteration,
"infos": infos,
}, path)
def load(self, path, load_optimizer=True):
loaded_dict = torch.load(path, map_location=self.device)
self.alg.actor_critic.load_state_dict(loaded_dict["model_state_dict"])
if load_optimizer:
self.alg.optimizer.load_state_dict(loaded_dict["optimizer_state_dict"])
self.current_learning_iteration = loaded_dict["iter"]
return loaded_dict["infos"]
def get_inference_policy(self, device=None):
self.alg.actor_critic.eval()
if device is not None:
self.alg.actor_critic.to(device)
return self.alg.actor_critic.act_inference

71
motrix_rl/src/motrix_rl/dwaq_rsl/utils.py Normal file
View File

@@ -0,0 +1,71 @@
# SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
# Copyright (c) 2021 ETH Zurich, Nikita Rudin
import torch
def split_and_pad_trajectories(tensor, dones):
""" Splits trajectories at done indices. Then concatenates them and padds with zeros up to the length og the longest trajectory.
Returns masks corresponding to valid parts of the trajectories
Example:
Input: [ [a1, a2, a3, a4 | a5, a6],
[b1, b2 | b3, b4, b5 | b6]
]
Output:[ [a1, a2, a3, a4], | [ [True, True, True, True],
[a5, a6, 0, 0], | [True, True, False, False],
[b1, b2, 0, 0], | [True, True, False, False],
[b3, b4, b5, 0], | [True, True, True, False],
[b6, 0, 0, 0] | [True, False, False, False],
] | ]
Assumes that the inputy has the following dimension order: [time, number of envs, aditional dimensions]
"""
dones = dones.clone()
dones[-1] = 1
# Permute the buffers to have order (num_envs, num_transitions_per_env, ...), for correct reshaping
flat_dones = dones.transpose(1, 0).reshape(-1, 1)
# Get length of trajectory by counting the number of successive not done elements
done_indices = torch.cat((flat_dones.new_tensor([-1], dtype=torch.int64), flat_dones.nonzero()[:, 0]))
trajectory_lengths = done_indices[1:] - done_indices[:-1]
trajectory_lengths_list = trajectory_lengths.tolist()
# Extract the individual trajectories
trajectories = torch.split(tensor.transpose(1, 0).flatten(0, 1),trajectory_lengths_list)
padded_trajectories = torch.nn.utils.rnn.pad_sequence(trajectories)
trajectory_masks = trajectory_lengths > torch.arange(0, tensor.shape[0], device=tensor.device).unsqueeze(1)
return padded_trajectories, trajectory_masks
def unpad_trajectories(trajectories, masks):
""" Does the inverse operation of split_and_pad_trajectories()
"""
# Need to transpose before and after the masking to have proper reshaping
return trajectories.transpose(1, 0)[masks.transpose(1, 0)].view(-1, trajectories.shape[0], trajectories.shape[-1]).transpose(1, 0)

6
motrix_rl/src/motrix_rl/rslrl/torch/models/cenet_actor.py
View File

@@ -116,15 +116,15 @@ class CENetActorModel(MLPModel):
self._last_cenet_output = None
def _update_distribution(self, obs: torch.Tensor) -> None:
"""覆盖父类 — 确保 std 始终为正，防止 NaN。"""
super()._update_distribution(obs)
# 如果 std 因数值问题变负，clamp 到最小值
"""覆盖父类 — 强制 std > 0 再创建 Normal 分布（防止 NaN）。"""
# 先 clamp std，再调父类创建分布
if self.stochastic and not self.state_dependent_std:
with torch.no_grad():
if self.noise_std_type == "scalar":
self.std.clamp_(min=1e-6)
elif self.noise_std_type == "log":
self.log_std.clamp_(min=-20.0, max=10.0)
super()._update_distribution(obs)
def _get_latent_dim(self) -> int:
"""Actor 实际输入：code(19) + policy(45) = 64。"""

3
motrix_rl/src/motrix_rl/rslrl/torch/train/dreamwaq_ppo.py
View File

@@ -81,6 +81,9 @@ class DreamWaQPPO(PPO):
returns_batch, old_actions_log_prob_batch, old_mu_batch,
old_sigma_batch, hid_states_batch, masks_batch,
) in generator:
# NaN 检测 — 数据有 NaN 就跳过这个 batch
if torch.isnan(obs_batch["policy"]).any() or torch.isnan(obs_batch["obs_history"]).any():
continue
# ── 标准 PPO 前向 ──
self.actor(obs_batch, masks=masks_batch, stochastic_output=True)
self.critic(obs_batch, masks=masks_batch)

306
motrix_rl/src/motrix_rl/skrl/jax/train/cts_ppo.py Normal file
View File

@@ -0,0 +1,306 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""CTS PPO v2 — Teacher/Student encoders with distillation loss.
Matches go2_rl_gym CTS training flow:
Phase 1: Standard PPO update with teacher-computed latent
Phase 2: Distillation: MSE(teacher_latent, student_latent) on student envs
Architecture:
Teacher: privileged_obs(76) → MLP(512,256) → L2Norm → latent(32)
Student: history_obs(225) → MLP(512,256,256) → L2Norm → latent(32)
Actor: student_obs(45) + latent(32) = 77 → [512,256,128] → action(12)
Critic: student_obs(45) + latent(32) = 77 → [512,256,128] → value(1)
The wrapper computes latent online (teacher for 75% envs, student for 25%)
and augments observations to 77-dim for standard PPO. After each PPO update,
a distillation step trains the student encoder to match the teacher's latent.
"""
import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import optax
from skrl.agents.jax.ppo import PPO as BasePPO
from skrl.memories.jax import RandomMemory
from skrl.trainers.jax import SequentialTrainer
from skrl.utils import set_seed
from motrix_envs import registry as env_registry
from motrix_rl import registry, utils
from motrix_rl.skrl import get_log_dir
# ═══════════════════════════════════════════════════════════════════════════
# Encoder Modules (Flax)
# ═══════════════════════════════════════════════════════════════════════════
class TeacherEncoder(nn.Module):
latent_dim: int = 32
hidden_dims: tuple = (512, 256)
@nn.compact
def __call__(self, x):
for d in self.hidden_dims:
x = nn.elu(nn.Dense(d)(x))
return nn.Dense(self.latent_dim)(x)
class StudentEncoder(nn.Module):
latent_dim: int = 32
hidden_dims: tuple = (512, 256, 256)
@nn.compact
def __call__(self, x):
for d in self.hidden_dims:
x = nn.elu(nn.Dense(d)(x))
return nn.Dense(self.latent_dim)(x)
# ═══════════════════════════════════════════════════════════════════════════
# CTS Env Wrapper
# ═══════════════════════════════════════════════════════════════════════════
class CTSEnvWrapper:
"""Wraps CTS env. Computes latent online, returns augmented obs (77-dim)."""
def __init__(self, env, teacher_ratio=0.75, latent_dim=32,
teacher_enc=None, teacher_params=None,
student_enc=None, student_params=None):
self.__env = env
self._teacher_ratio = teacher_ratio
self._latent_dim = latent_dim
self._num_envs = env.num_envs
self._t_enc = teacher_enc
self._t_params = teacher_params
self._s_enc = student_enc
self._s_params = student_params
from gymnasium import spaces
self._obs_space = spaces.Box(-np.inf, np.inf,
(45 + latent_dim,), dtype=np.float32)
@property
def num_envs(self): return self._num_envs
@property
def num_teacher(self): return int(self._num_envs * self._teacher_ratio)
@property
def teacher_mask(self):
m = np.zeros(self._num_envs, dtype=np.bool)
m[:self.num_teacher] = True
return m
@property
def student_mask(self): return ~self.teacher_mask
@property
def privileged_obs(self):
return self.__env._state.info.get("privileged_obs",
np.zeros((self._num_envs, 76), dtype=np.float32))
@property
def history_obs(self): return self.__env.history_obs
@property
def observation_space(self): return self._obs_space
@property
def action_space(self): return self.__env.action_space
@property
def device(self): return None
@property
def num_agents(self): return 1
@property
def state_preprocessor(self): return None
@property
def value_preprocessor(self): return None
@property
def world_info_idx(self): return None
@property
def random_timestep(self): return False
@property
def running(self): return True
@property
def _env(self): return self.__env
def _augment(self, states):
n = states.shape[0]
latent = np.zeros((n, self._latent_dim), dtype=np.float32)
t_mask = self.teacher_mask
if t_mask.any() and self._t_params is not None:
priv = jnp.array(self.privileged_obs[t_mask])
latent[t_mask] = np.array(self._t_enc.apply(self._t_params, priv))
s_mask = self.student_mask
if s_mask.any() and self._s_params is not None:
hist = jnp.array(self.history_obs[s_mask])
latent[s_mask] = np.array(self._s_enc.apply(self._s_params, hist))
return np.concatenate([states, latent], axis=-1)
def reset(self):
state = self.__env.init_state()
return self._augment(state.obs), state.info
def step(self, actions):
state = self.__env.step(np.array(actions))
return (self._augment(state.obs), state.reward.reshape(-1, 1),
state.terminated.reshape(-1, 1),
state.truncated.reshape(-1, 1), state.info)
def render(self, *a, **kw): pass
def close(self): pass
# ═══════════════════════════════════════════════════════════════════════════
# CTS Agent — adds distillation step after PPO update
# ═══════════════════════════════════════════════════════════════════════════
class CTSAgent(BasePPO):
"""PPO agent with post-update distillation on student encoder."""
def __init__(self, models, memory, cfg, observation_space, action_space, device,
teacher_enc=None, student_enc=None, cts_wrapper=None,
teacher_params=None, student_params=None,
distill_coef=0.1, student_lr=1e-3):
super().__init__(models=models, memory=memory, cfg=cfg,
observation_space=observation_space,
action_space=action_space, device=device)
self._t_enc = teacher_enc
self._s_enc = student_enc
self._t_params = teacher_params
self._s_params = student_params
self._wrapper = cts_wrapper
self._distill_coef = distill_coef
self._latent_dim = getattr(teacher_enc, 'latent_dim', 32) if teacher_enc else 32
# Separate optimizer for student encoder
if student_params is not None:
self._s_optimizer = optax.adam(student_lr)
self._s_opt_state = self._s_optimizer.init(student_params)
else:
self._s_optimizer = None
def update(self):
"""PPO update + distillation phase."""
result = super().update()
# Phase 2: Distillation — MSE(teacher_latent, student_latent) on student envs
if self._s_params is not None and self._s_optimizer is not None:
self._distillation_step()
return result
def _distillation_step(self):
"""Train student encoder to match teacher latent.
go2_rl_gym Phase 2: only on student env samples.
Loss = || teacher_enc(privileged_obs) - student_enc(history_obs) ||²
"""
wrapper = self._wrapper
if wrapper is None:
return
s_mask = wrapper.student_mask
if not s_mask.any():
return
# Get current student env observations
priv = jnp.array(wrapper.privileged_obs[s_mask]) # [S, 76]
hist = jnp.array(wrapper.history_obs[s_mask]) # [S, 225]
# L2 normalize both before computing MSE (match go2_rl_gym)
def l2_norm(x):
return x / (jnp.linalg.norm(x, axis=-1, keepdims=True) + 1e-8)
def distill_loss(s_params):
t_latent = l2_norm(jax.lax.stop_gradient(
self._t_enc.apply(self._t_params, priv)))
s_latent = l2_norm(self._s_enc.apply(s_params, hist))
return jnp.mean(jnp.square(t_latent - s_latent))
loss, grads = jax.value_and_grad(distill_loss)(self._s_params)
updates, self._s_opt_state = self._s_optimizer.update(
grads, self._s_opt_state, self._s_params)
self._s_params = optax.apply_updates(self._s_params, updates)
# Track
self.tracking_data["Loss / Distillation"].append(loss)
# ═══════════════════════════════════════════════════════════════════════════
# Trainer
# ═══════════════════════════════════════════════════════════════════════════
class CTSTrainer:
def __init__(self, env_name="go1-cts-flat-walk-go2style",
sim_backend=None, enable_render=False, cfg_override=None):
self._env_name = env_name
self._sim_backend = sim_backend
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
if cfg_override:
rlcfg = utils.cfg_override(rlcfg, cfg_override)
self._rlcfg = rlcfg
def train(self):
cfg = self._rlcfg
set_seed(cfg.runner.seed)
env = env_registry.make(self._env_name, num_envs=cfg.num_envs)
env_cfg = env._cfg
latent_dim = getattr(env_cfg, 'latent_dim', 32)
priv_dim = getattr(env_cfg, 'num_privileged_obs', 76)
hist_len = getattr(env_cfg, 'history_length', 5)
student_dim = getattr(env_cfg, 'num_student_obs', 45)
teacher_ratio = getattr(env_cfg, 'teacher_env_ratio', 0.75)
# Init Flax encoders
rng = jax.random.PRNGKey(cfg.runner.seed)
rng, t_key, s_key = jax.random.split(rng, 3)
teacher_enc = TeacherEncoder(latent_dim=latent_dim)
student_enc = StudentEncoder(latent_dim=latent_dim)
t_params = teacher_enc.init(t_key, jnp.ones((1, priv_dim)))
s_params = student_enc.init(s_key, jnp.ones((1, hist_len * student_dim)))
# CTS wrapper
wrapper = CTSEnvWrapper(env, teacher_ratio=teacher_ratio, latent_dim=latent_dim,
teacher_enc=teacher_enc, teacher_params=t_params,
student_enc=student_enc, student_params=s_params)
env.init_state()
from gymnasium import spaces
aug_obs = spaces.Box(-np.inf, np.inf, (student_dim + latent_dim,), dtype=np.float32)
wrapper._obs_space = aug_obs
nT = int(cfg.num_envs * teacher_ratio)
nS = cfg.num_envs - nT
print(f"[CTS v2] {cfg.num_envs} envs ({nT}T/{nS}S), latent={latent_dim}, "
f"priv={priv_dim}, hist={hist_len}×{student_dim}, distill_coef=0.1")
# Standard model creation
from motrix_rl.skrl.jax.train.ppo import Trainer as StdTrainer, _add_runtime_config
std = StdTrainer(self._env_name, sim_backend=None, enable_render=False)
std._rlcfg = cfg
models = std._make_model(wrapper, cfg)
ppo_cfg = cfg.runner.agent.to_dict()
_add_runtime_config(ppo_cfg, wrapper,
log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="CTS_PPO"))
mem_size = cfg.runner.memory.memory_size
if mem_size == -1:
mem_size = ppo_cfg.get("rollouts", 24)
memory = RandomMemory(memory_size=mem_size, num_envs=cfg.num_envs, device=None)
agent = CTSAgent(
models=models, memory=memory, cfg=ppo_cfg,
observation_space=aug_obs, action_space=env.action_space,
device=None,
teacher_enc=teacher_enc, student_enc=student_enc,
cts_wrapper=wrapper,
teacher_params=t_params, student_params=s_params,
distill_coef=0.1, student_lr=1e-3,
)
trainer_cfg = {"timesteps": cfg.runner.trainer.timesteps, "headless": True}
SequentialTrainer(cfg=trainer_cfg, env=wrapper, agents=agent).train()
# Save encoder params for deployment
import pickle as _pk
_log = get_log_dir(self._env_name, rllib="skrl", agent_name="CTS_PPO")
_pk.dump(t_params, open(f"{_log}/teacher_params.pkl", "wb"))
_pk.dump(s_params, open(f"{_log}/student_params.pkl", "wb"))
print(f"[CTS] Encoder weights saved → {_log}")

540
motrix_rl/src/motrix_rl/skrl/jax/train/dreamwaq_ppo.py Normal file
View File

@@ -0,0 +1,540 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
"""DreamWaQ PPO — aligned with Manaro-Alpha/DreamWaQ.
CENet (VAE): history(225) → [128,64] → latent(16) + vel_est(3) = code(19)
Decoder: code(19) → [64,128] → next_obs(45)
Actor: code(19) + obs(45) = 64 → [512,256,128] → action(12)
Critic: privileged_obs(235) → [512,256,128] → value(1)
VAE Loss: reconstruction_MSE + velocity_MSE + beta * KL
"""
import flax.linen as nn
import jax
import jax.numpy as jnp
import numpy as np
import optax
import pickle, os
from skrl.agents.jax.ppo import PPO as BasePPO
from skrl.memories.jax import RandomMemory
from skrl.models.jax import DeterministicMixin, GaussianMixin, Model
from skrl.trainers.jax import SequentialTrainer
from skrl.utils import set_seed
from motrix_envs import registry as env_registry
from motrix_envs.locomotion.go1.dreamwaq import DreamWaQTask
from motrix_rl import registry, utils
from motrix_rl.skrl import get_log_dir
# Ensure RL config registered
import motrix_rl.tasks.go1_dreamwaq # noqa: F401
# ═══════════════════════════════════════════════════════════════════════
# CENet (VAE) — exact Manaro-Alpha architecture
# ═══════════════════════════════════════════════════════════════════════
class CENet(nn.Module):
"""VAE for terrain imagination + velocity estimation.
Encoder: 225 → Linear(128) → ELU → Linear(64) → ELU
4 heads on 64-dim bottleneck:
- latent_mu, latent_logvar: 64 → 16 each
- vel_mu, vel_logvar: 64 → 3 each
Decoder: (latent+vel=19) → Linear(64) → ELU → Linear(128) → ELU → Linear(45)
"""
num_obs: int = 45
num_history: int = 5
num_latent: int = 16 # latent z dim
cenet_out_dim: int = 19 # code = vel(3) + latent(16)
def setup(self):
# Orthogonal weight init matching upstream torch.nn.init.orthogonal_
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
# Encoder (matches Manaro-Alpha: 225 → 128 → 64)
self.enc_fc1 = nn.Dense(128, kernel_init=kernel_init)
self.enc_fc2 = nn.Dense(64, kernel_init=kernel_init)
# 4 output heads
self.latent_mu = nn.Dense(self.num_latent, kernel_init=kernel_init)
self.latent_logvar = nn.Dense(self.num_latent, kernel_init=kernel_init)
self.vel_mu = nn.Dense(3, kernel_init=kernel_init)
self.vel_logvar = nn.Dense(3, kernel_init=kernel_init)
# Decoder (matches Manaro-Alpha: 19 → 64 → 128 → 45)
self.dec_fc1 = nn.Dense(64, kernel_init=kernel_init)
self.dec_fc2 = nn.Dense(128, kernel_init=kernel_init)
self.dec_out = nn.Dense(self.num_obs, kernel_init=kernel_init)
def encode(self, obs_history):
"""obs_history: (B, T, 45) or (B, T*45) → (latent_mu, latent_logvar, vel_mu, vel_logvar)."""
x = obs_history.reshape(obs_history.shape[0], -1)
x = nn.elu(self.enc_fc1(x))
x = nn.elu(self.enc_fc2(x))
return (self.latent_mu(x), self.latent_logvar(x),
self.vel_mu(x), self.vel_logvar(x))
def reparameterize(self, mu, logvar, rng):
std = jnp.exp(0.5 * logvar)
return mu + std * jax.random.normal(rng, mu.shape)
def decode(self, z, vel):
x = jnp.concatenate([z, vel], axis=-1)
x = nn.elu(self.dec_fc1(x))
x = nn.elu(self.dec_fc2(x))
return self.dec_out(x)
def __call__(self, obs_history, rng):
lm, ll, vm, vl = self.encode(obs_history)
rng_z, rng_v = jax.random.split(rng)
z = self.reparameterize(lm, ll, rng_z)
vel = self.reparameterize(vm, vl, rng_v)
return {
'z': z, 'vel_est': vel,
'latent_mu': lm, 'latent_logvar': ll,
'vel_mu': vm, 'vel_logvar': vl,
'obs_pred': self.decode(z, vel),
}
def sample(self, obs_history, rng):
"""Sample (z, vel) for rollout (adds noise)."""
lm, ll, vm, vl = self.encode(obs_history)
rng_z, rng_v = jax.random.split(rng)
z = self.reparameterize(lm, ll, rng_z)
vel = self.reparameterize(vm, vl, rng_v)
return z, vel
def inference(self, obs_history):
"""Mean (z, vel) for deployment (no noise)."""
lm, ll, vm, vl = self.encode(obs_history)
return lm, vm
def loss_fn(self, obs_history, current_obs, base_vel, rng, kl_weight=1.0):
"""VAE loss: reconstructs current obs (last history frame), NOT next obs.
Upstream: decode_target = obs_batch (the same timestep's observation,
not a future prediction). This is an autoencoder over current observation.
"""
lm, ll, vm, vl = self.encode(obs_history)
rng_z, rng_v = jax.random.split(rng)
z = self.reparameterize(lm, ll, rng_z)
vel = self.reparameterize(vm, vl, rng_v)
pred = self.decode(z, vel)
recons = jnp.mean(jnp.square(pred - current_obs), axis=-1)
vel_loss = jnp.mean(jnp.square(vel - base_vel), axis=-1)
kld = -0.5 * jnp.sum(1 + ll - jnp.square(lm) - jnp.exp(ll), axis=-1)
return {
'loss': jnp.mean(recons + vel_loss + kl_weight * kld),
'recons_loss': jnp.mean(recons),
'vel_loss': jnp.mean(vel_loss),
'kld_loss': jnp.mean(kld),
}
# ═══════════════════════════════════════════════════════════════════════
# Environment Wrapper
# ═══════════════════════════════════════════════════════════════════════
class DreamWaQWrapper:
"""Wraps DreamWaQTask, computes CENet code online.
Augmented observation layout (254-dim):
[code_vel(3) | code_latent(16) | obs(45) | base_vel(3) | heights(187)]
Actor uses: [0:64] = code(19) + obs(45)
Critic uses: [19:254] = obs(45) + base_vel(3) + heights(187)
"""
def __init__(self, env, cenet, cenet_params, rng=None, vae_buf_size=24):
self._env = env
self._cenet = cenet
self._cenet_params = cenet_params
self._num_envs = env.num_envs
self._rng = rng if rng is not None else jax.random.PRNGKey(0)
self._code_dim = 19 # vel(3) + latent(16)
self._aug_dim = self._code_dim + 45 + 3 + 187 # 254
from gymnasium import spaces
self._obs_space = spaces.Box(-np.inf, np.inf, (self._aug_dim,), dtype=np.float32)
# VAE training buffer
self._vae_buf = []
self._vae_buf_size = vae_buf_size
@property
def num_envs(self): return self._num_envs
@property
def observation_space(self): return self._obs_space
@property
def action_space(self): return self._env.action_space
@property
def device(self): return None
@property
def num_agents(self): return 1
@property
def state_preprocessor(self): return None
@property
def value_preprocessor(self): return None
@property
def world_info_idx(self): return None
@property
def random_timestep(self): return False
@property
def running(self): return True
@property
def privileged_obs(self):
return self._env._state.info.get("privileged_obs",
np.zeros((self._num_envs, 235), dtype=np.float32))
@property
def obs_history(self):
return self._env._state.info.get("obs_history",
np.zeros((self._num_envs, 5, 45), dtype=np.float32))
@property
def base_vel(self):
return self._env._state.info.get("base_vel",
np.zeros((self._num_envs, 3), dtype=np.float32))
@property
def vae_training_data(self):
return self._vae_buf
def _augment(self, obs):
"""Compute CENet code and build augmented obs."""
hist = jnp.array(self.obs_history)
self._rng, rng_key = jax.random.split(self._rng)
z, vel = self._cenet.apply(self._cenet_params, hist, rng_key, method=self._cenet.sample)
code = np.concatenate([np.array(vel), np.array(z)], axis=-1) # (N, 19)
priv = self.privileged_obs
heights = priv[:, 45+3:] # (N, 187)
base_vel_norm = self.base_vel * np.array([2.0, 2.0, 1.0], dtype=np.float32)
return np.concatenate([code, obs, base_vel_norm, heights], axis=-1)
def reset(self):
state = self._env.init_state()
self._vae_buf = []
return self._augment(state.obs), state.info
def step(self, actions):
old_hist = self.obs_history.copy()
state = self._env.step(np.array(actions, dtype=np.float32))
# Store for VAE training: history → current_obs (last frame of history)
# Upstream: decode target = obs_batch (current observation, not next)
self._vae_buf.append((old_hist, old_hist[:, -1, :].copy(), self.base_vel.copy()))
if len(self._vae_buf) > self._vae_buf_size:
self._vae_buf.pop(0)
return (self._augment(state.obs),
state.reward.reshape(-1, 1).astype(np.float32),
state.terminated.reshape(-1, 1).astype(np.float32),
state.truncated.reshape(-1, 1).astype(np.float32),
state.info)
def render(self, *a, **kw): pass
def close(self): pass
# ═══════════════════════════════════════════════════════════════════════
# DreamWaQ Agent — PPO + CENet training
# ═══════════════════════════════════════════════════════════════════════
class DreamWaQAgent(BasePPO):
"""PPO agent with asymmetric critic + CENet (VAE) training."""
def __init__(self, models, memory, cfg, observation_space, action_space, device,
cenet=None, cenet_params=None, wrapper=None,
kl_weight=1.0, vae_lr=1e-3, save_dir=None):
super().__init__(models=models, memory=memory, cfg=cfg,
observation_space=observation_space,
action_space=action_space, device=device)
self._cenet = cenet
self._cenet_params = cenet_params
self._wrapper = wrapper
self._kl_weight = kl_weight
self._save_dir = save_dir
self._vae_count = 0
self._rng = jax.random.PRNGKey(42)
if cenet_params is not None:
self._vae_opt = optax.adam(vae_lr)
self._vae_opt_state = self._vae_opt.init(cenet_params)
else:
self._vae_opt = None
# Override SKRL's internal _update (NOT 'update')
def _update(self, timestep: int, timesteps: int):
super()._update(timestep, timesteps)
# Track detailed reward terms
wrapper = self._wrapper
if wrapper is not None:
try:
terms = wrapper._env._state.info.get("reward_terms", {})
for name, val in terms.items():
self.tracking_data[f"Reward / {name}"].append(float(val))
except Exception:
pass
# CENet training after each PPO update
if self._cenet_params is not None and self._vae_opt is not None:
self._vae_step()
def _vae_step(self):
"""Train CENet: reconstruct current obs (last history frame) + velocity supervision + KL."""
wrapper = self._wrapper
if wrapper is None or len(wrapper.vae_training_data) == 0:
return
hist_np, current_obs_np, base_vel_np = wrapper.vae_training_data[-1]
hist = jnp.array(hist_np)
current_obs = jnp.array(current_obs_np)
base_vel = jnp.array(base_vel_np)
n = wrapper.num_envs
valid = jnp.linalg.norm(hist.reshape(n, -1), axis=-1) > 1e-6
if not valid.any():
return
self._rng, vae_rng = jax.random.split(self._rng)
def loss_fn(p):
return self._cenet.apply(
p, hist, current_obs, base_vel, vae_rng, self._kl_weight,
method=self._cenet.loss_fn)['loss']
loss, grads = jax.value_and_grad(loss_fn)(self._cenet_params)
updates, self._vae_opt_state = self._vae_opt.update(
grads, self._vae_opt_state, self._cenet_params)
self._cenet_params = optax.apply_updates(self._cenet_params, updates)
# Sync wrapper
wrapper._cenet_params = self._cenet_params
self._vae_count += 1
if self._save_dir and self._vae_count % 200 == 0:
pickle.dump(self._cenet_params,
open(f"{self._save_dir}/vae_{self._vae_count}.pkl", "wb"))
self.tracking_data["Loss / VAE"].append(loss)
# ═══════════════════════════════════════════════════════════════════════
# Trainer
# ═══════════════════════════════════════════════════════════════════════
class DreamWaQTrainer:
def __init__(self, env_name="go1-dreamwaq-walk", cfg_override=None):
self._env_name = env_name
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
if cfg_override:
rlcfg = utils.cfg_override(rlcfg, cfg_override)
self._rlcfg = rlcfg
def train(self):
cfg = self._rlcfg
set_seed(cfg.runner.seed)
env = env_registry.make(self._env_name, num_envs=cfg.num_envs)
env_cfg = env._cfg
# Init CENet
rng = jax.random.PRNGKey(cfg.runner.seed)
rng, crng = jax.random.split(rng)
cenet = CENet(num_obs=45, num_history=env_cfg.num_history,
num_latent=env_cfg.num_latent,
cenet_out_dim=env_cfg.cenet_out_dim)
cenet_params = cenet.init(crng, jnp.ones((1, env_cfg.num_history, 45)), jax.random.PRNGKey(0))
# Wrapper
wrapper = DreamWaQWrapper(env, cenet, cenet_params,
rng=jax.random.PRNGKey(cfg.runner.seed + 1))
env.init_state()
print(f"[DreamWaQ] {cfg.num_envs} envs, CENet[128,64], latent=16+3=19, "
f"aug_dim={wrapper.observation_space.shape[0]}")
# Build models
from motrix_rl.skrl.jax.train.ppo import Trainer as StdTrainer, _add_runtime_config
policy_cfg = cfg.runner.models.policy
value_cfg = cfg.runner.models.value
class Policy(GaussianMixin, Model):
def __init__(self, observation_space, action_space, device=None, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
GaussianMixin.__init__(self, policy_cfg.clip_actions,
policy_cfg.clip_log_std,
policy_cfg.min_log_std, policy_cfg.max_log_std,
policy_cfg.reduction)
@nn.compact
def __call__(self, inputs, role):
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
x = inputs["states"]
# Actor: code(19) + obs(45) = first 64 dims
x = x[:, :64]
for d in policy_cfg.hiddens:
x = nn.elu(nn.Dense(d, kernel_init=kernel_init)(x))
x = nn.Dense(self.num_actions, kernel_init=kernel_init)(x)
log_std = self.param("log_std",
lambda _: jnp.full(self.num_actions, float(policy_cfg.initial_log_std)))
return x, log_std, {}
class Value(DeterministicMixin, Model):
def __init__(self, observation_space, action_space, device=None, **kwargs):
Model.__init__(self, observation_space, action_space, device, **kwargs)
DeterministicMixin.__init__(self, value_cfg.clip_actions)
@nn.compact
def __call__(self, inputs, role):
kernel_init = nn.initializers.orthogonal(jnp.sqrt(2))
x = inputs["states"]
# Critic: obs(45) + base_vel(3) + heights(187) = 235
# Layout: [code(19) | obs(45) | base_vel(3) | heights(187)]
x_c = jnp.concatenate([x[:, 19:64], x[:, 64:254]], axis=-1)
for d in value_cfg.hiddens:
x_c = nn.elu(nn.Dense(d, kernel_init=kernel_init)(x_c))
return nn.Dense(1, kernel_init=kernel_init)(x_c), {}
models = {
"policy": Policy(wrapper.observation_space, wrapper.action_space, None),
"value": Value(wrapper.observation_space, wrapper.action_space, None),
}
for role, m in models.items():
m.init_state_dict(role)
# Agent
ppo_cfg = cfg.runner.agent.to_dict()
log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
_add_runtime_config(ppo_cfg, wrapper, log_dir=log_dir)
mem_size = cfg.runner.memory.memory_size
if mem_size == -1:
mem_size = ppo_cfg.get("rollouts", 24)
memory = RandomMemory(memory_size=mem_size, num_envs=cfg.num_envs, device=None)
agent = DreamWaQAgent(
models=models, memory=memory, cfg=ppo_cfg,
observation_space=wrapper.observation_space,
action_space=wrapper.action_space, device=None,
cenet=cenet, cenet_params=cenet_params, wrapper=wrapper,
kl_weight=1.0, vae_lr=1e-3, save_dir=log_dir,
)
trainer_cfg = {"timesteps": cfg.runner.trainer.timesteps, "headless": True}
SequentialTrainer(cfg=trainer_cfg, env=wrapper, agents=agent).train()
# Save final VAE params
pickle.dump(agent._cenet_params, open(f"{log_dir}/cenet_params.pkl", "wb"))
print(f"[DreamWaQ] CENet weights saved → {log_dir}")
def play(self, policy_path=None):
"""Render env with trained policy (uses wrapper + CENet)."""
from motrix_rl.skrl.jax.wrap_np import SkrlNpWrapper as RenderWrapper
cfg = self._rlcfg
raw_env = env_registry.make(self._env_name, num_envs=min(64, cfg.play_num_envs))
env = RenderWrapper(raw_env) # adds render() method
env_cfg = raw_env._cfg
import msgpack
rng = jax.random.PRNGKey(cfg.runner.seed)
rng, crng = jax.random.split(rng)
cenet = CENet()
log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
# Load CENet params
if policy_path and os.path.exists(policy_path):
# Load full agent checkpoint and extract
with open(policy_path, 'rb') as f:
ckpt = pickle.load(f)
# Load policy params
policy_raw = msgpack.unpackb(ckpt['policy'])
policy_params = {}
for name, val in policy_raw['params'].items():
if isinstance(val, dict):
policy_params[name] = {
k: _decode_msgpack_array(v) for k, v in val.items()
}
else:
policy_params[name] = _decode_msgpack_array(val)
# Try loading VAE from base skrl dir
log_dir = get_log_dir(self._env_name, rllib="skrl", agent_name="DreamWaQ_PPO")
vae_path = os.path.join(log_dir, "cenet_params.pkl")
if not os.path.exists(vae_path):
# Find latest vae_N.pkl
vae_files = sorted([f for f in os.listdir(log_dir) if f.startswith("vae_")])
if vae_files:
vae_path = os.path.join(log_dir, vae_files[-1])
if os.path.exists(vae_path):
with open(vae_path, 'rb') as f:
cenet_params = pickle.load(f)
print(f"[Play] VAE loaded: {vae_path}")
else:
print("[Play] WARNING: VAE params not found, using random init")
cenet_params = cenet.init(crng, jnp.ones((1, 5, 45)), jax.random.PRNGKey(0))
else:
print("[Play] No policy, random actions")
policy_params = None
cenet_params = cenet.init(crng, jnp.ones((1, 5, 45)), jax.random.PRNGKey(0))
wrapper = DreamWaQWrapper(raw_env, cenet, cenet_params,
rng=jax.random.PRNGKey(cfg.runner.seed + 1))
raw_env.init_state()
import time
n = raw_env._num_envs
print(f"[Play] {n} envs, press Ctrl+C to stop")
try:
while True:
if policy_params is not None:
aug_obs = wrapper._augment(raw_env._state.obs)
actions = np.array(_policy_forward(jnp.array(aug_obs), policy_params))
else:
actions = np.zeros((n, 12), dtype=np.float32)
wrapper.step(actions)
env.render()
time.sleep(0.01)
except KeyboardInterrupt:
env.close()
print("[Play] Stopped")
def _decode_msgpack_array(ext) -> np.ndarray | None:
"""Decode flax-serialized msgpack ExtType."""
import msgpack
if not hasattr(ext, "code"): return None
parts = msgpack.unpackb(ext.data, raw=False)
if not isinstance(parts, list) or len(parts) < 3: return None
shape = []
def _flatten(s):
if isinstance(s, list):
for x in s: _flatten(x)
elif isinstance(s, int): shape.append(s)
_flatten(parts[0])
return np.frombuffer(parts[2], dtype=np.dtype(parts[1])).reshape(shape)
def _policy_forward(x, params):
"""Forward pass through policy (first 64 dims only, apply elu)."""
import jax.numpy as jnp
import jax
w0, b0 = jnp.array(params['Dense_0']['kernel']), jnp.array(params['Dense_0']['bias'])
w1, b1 = jnp.array(params['Dense_1']['kernel']), jnp.array(params['Dense_1']['bias'])
w2, b2 = jnp.array(params['Dense_2']['kernel']), jnp.array(params['Dense_2']['bias'])
w3, b3 = jnp.array(params['Dense_3']['kernel']), jnp.array(params['Dense_3']['bias'])
x = x[:, :64] # code(19) + obs(45)
x = jax.nn.elu(x @ w0 + b0)
x = jax.nn.elu(x @ w1 + b1)
x = jax.nn.elu(x @ w2 + b2)
return x @ w3 + b3

28
motrix_rl/src/motrix_rl/skrl/jax/train/ppo.py
View File

@@ -185,6 +185,7 @@ class Trainer:
sim_backend: str = None,
enable_render: bool = False,
cfg_override: dict = None,
env_cfg_override: dict = None,
) -> None:
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
if cfg_override is not None:
@@ -193,13 +194,18 @@ class Trainer:
self._env_name = env_name
self._sim_backend = sim_backend
self._enable_render = enable_render
self._env_cfg_override = env_cfg_override
def train(self) -> None:
def train(self, checkpoint: str = None) -> None:
"""
Start training the agent.
Args:
checkpoint: Optional path to a checkpoint to resume from.
"""
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
env = env_registry.make(self._env_name, sim_backend=self._sim_backend,
num_envs=rlcfg.num_envs, env_cfg_override=self._env_cfg_override)
set_seed(rlcfg.runner.seed)
skrl_env = wrap_env(env, self._enable_render)
@@ -209,6 +215,12 @@ class Trainer:
# Add runtime-specific configuration
_add_runtime_config(ppo_cfg, skrl_env, log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="PPO"))
agent = self._make_agent(models, skrl_env, ppo_cfg, rlcfg.runner.memory)
if checkpoint:
agent.load(checkpoint)
logger = __import__("logging").getLogger(__name__)
logger.info(f"Resumed from checkpoint: {checkpoint}")
cfg_trainer = {
"timesteps": rlcfg.runner.trainer.timesteps,
"headless": not self._enable_render,
@@ -220,7 +232,17 @@ class Trainer:
import time
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
# Enable play_mode + merge constructor-level overrides (e.g. force_phase)
play_override = {"play_mode": True}
if self._env_cfg_override:
play_override.update(self._env_cfg_override)
try:
env = env_registry.make(
self._env_name, sim_backend=self._sim_backend,
num_envs=rlcfg.play_num_envs, env_cfg_override=play_override,
)
except ValueError:
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
set_seed(rlcfg.runner.seed)
env = wrap_env(env, self._enable_render)

28
motrix_rl/src/motrix_rl/skrl/torch/train/ppo.py
View File

@@ -185,6 +185,7 @@ class Trainer:
sim_backend: str = None,
enable_render: bool = False,
cfg_override: dict = None,
env_cfg_override: dict = None,
) -> None:
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="torch")
if cfg_override is not None:
@@ -193,13 +194,18 @@ class Trainer:
self._env_name = env_name
self._sim_backend = sim_backend
self._enable_render = enable_render
self._env_cfg_override = env_cfg_override
def train(self) -> None:
def train(self, checkpoint: str = None) -> None:
"""
Start training the agent.
Args:
checkpoint: Optional path to a checkpoint to resume from.
"""
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.num_envs)
env = env_registry.make(self._env_name, sim_backend=self._sim_backend,
num_envs=rlcfg.num_envs, env_cfg_override=self._env_cfg_override)
set_seed(rlcfg.runner.seed)
skrl_env = wrap_env(env, self._enable_render)
models = self._make_model(skrl_env, rlcfg)
@@ -208,6 +214,12 @@ class Trainer:
# Add runtime-specific configuration
_add_runtime_config(ppo_cfg, skrl_env, log_dir=get_log_dir(self._env_name, rllib="skrl", agent_name="PPO"))
agent = self._make_agent(models, skrl_env, ppo_cfg, rlcfg.runner.memory)
if checkpoint:
agent.load(checkpoint)
logger = __import__("logging").getLogger(__name__)
logger.info(f"Resumed from checkpoint: {checkpoint}")
cfg_trainer = {
"timesteps": rlcfg.runner.trainer.timesteps,
"headless": not self._enable_render,
@@ -219,7 +231,17 @@ class Trainer:
import time
rlcfg = self._rlcfg
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
# Enable play_mode if the env config supports it (for multi-terrain random spawn)
play_override = {"play_mode": True}
if self._env_cfg_override:
play_override.update(self._env_cfg_override)
try:
env = env_registry.make(
self._env_name, sim_backend=self._sim_backend,
num_envs=rlcfg.play_num_envs, env_cfg_override=play_override,
)
except ValueError:
env = env_registry.make(self._env_name, sim_backend=self._sim_backend, num_envs=rlcfg.play_num_envs)
set_seed(rlcfg.runner.seed)
env = wrap_env(env, self._enable_render)
models = self._make_model(env, rlcfg)

2
motrix_rl/src/motrix_rl/tasks/go1.py
View File

@@ -135,7 +135,7 @@ class rslrl:
runner.actor.class_name = (
"motrix_rl.rslrl.torch.models.cenet_actor:CENetActorModel")
runner.actor.hidden_dims = [512, 256, 128]
runner.actor.init_noise_std = 1.0
runner.actor.init_noise_std = 0.5 # 降噪声先学站稳，再探索行走
# Critic：标准 MLPModel，输入 privileged_obs
runner.critic.class_name = "MLPModel"

42
motrix_rl/src/motrix_rl/tasks/go1_dreamwaq.py Normal file
View File

@@ -0,0 +1,42 @@
"""DreamWaQ training config — Manaro-Alpha aligned.
Actor [512,256,128], Critic [512,256,128], CENet [128,64].
lr=1e-3, vae_lr=1e-3, KL_weight=1.0, entropy=0.01.
"""
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401
from motrix_rl.registry import rlcfg
from motrix_rl.skrl.config import SkrlCfg
def _apply(cfg: SkrlCfg):
# PPO hyperparams (exact DreamWaQ match)
cfg.runner.agent.rollouts = 24 # num_steps_per_env
cfg.runner.agent.learning_epochs = 5 # num_learning_epochs
cfg.runner.agent.mini_batches = 4 # num_mini_batches
cfg.runner.agent.learning_rate = 1e-3
cfg.runner.agent.entropy_loss_scale = 0.01
cfg.runner.agent.ratio_clip = 0.2
cfg.runner.agent.discount_factor = 0.99
cfg.runner.agent.lam = 0.95
# Network (DreamWaQ: [512,256,128])
cfg.runner.models.policy.hiddens = [512, 256, 128]
cfg.runner.models.policy.hidden_activation = ["elu"]
# initial_log_std: upstream uses 0.0 (std=1) in IsaacGym, but in MotrixSim with
# action_scale=0.25 that gives ±0.25 rad jitter/step → robot falls immediately and
# never learns to stand. -1.0 (std=0.37, ±0.09 rad) lets it stay upright early and
# bootstrap walking. (MotrixSim<->IsaacGym reality-gap adjustment.)
cfg.runner.models.policy.initial_log_std = -1.0
cfg.runner.models.value.hiddens = [512, 256, 128]
cfg.runner.models.value.hidden_activation = ["elu"]
cfg.runner.memory.memory_size = 24
cfg.runner.agent.experiment.checkpoint_interval = 200
cfg.runner.agent.experiment.write_interval = 200
@rlcfg("go1-dreamwaq-walk")
class DreamWaQ(SkrlCfg):
def __init__(self, **kwargs):
super().__init__(**kwargs)
_apply(self)

64
motrix_rl/src/motrix_rl/tasks/go1_go2style.py Normal file
View File

@@ -0,0 +1,64 @@
"""Training configs for Go1 go2style environments.
Importing triggers registration of: flat, stairs, combined.
"""
from dataclasses import dataclass
import motrix_envs.locomotion.go1.go2style_flat # noqa: F401
import motrix_envs.locomotion.go1.go2style_stairs # noqa: F401
import motrix_envs.locomotion.go1.go2style_combined # noqa: F401
import motrix_envs.locomotion.go1.go2style_cts # noqa: F401
import motrix_envs.locomotion.go1.orig_combined # noqa: F401
from motrix_rl.registry import rlcfg
from motrix_rl.skrl.config import SkrlCfg
def _apply_go2style_ppo(cfg: SkrlCfg):
"""Apply go2_rl_gym-style PPO params to config."""
cfg.runner.models.policy.hiddens = [512, 256, 128]
cfg.runner.models.value.hiddens = [512, 256, 128]
cfg.runner.agent.rollouts = 24
cfg.runner.agent.learning_epochs = 5
cfg.runner.agent.mini_batches = 4
cfg.runner.agent.learning_rate = 1e-3
cfg.runner.trainer.timesteps = 60000
def _apply_orig_ppo(cfg: SkrlCfg):
"""Match original Go1 stairs config: [512,256,128], 30k steps."""
cfg.runner.models.policy.hiddens = [512, 256, 128]
cfg.runner.models.value.hiddens = [512, 256, 128]
cfg.runner.agent.rollouts = 24
cfg.runner.agent.learning_epochs = 5
cfg.runner.agent.mini_batches = 4
cfg.runner.agent.learning_rate = 1e-3
cfg.runner.trainer.timesteps = 30000 # same as original stairs
class skrl:
@rlcfg("go1-flat-terrain-walk-go2style")
@dataclass
class Flat(SkrlCfg):
def __post_init__(self): _apply_go2style_ppo(self)
@rlcfg("go1-stairs-terrain-walk-go2style")
@dataclass
class Stairs(SkrlCfg):
def __post_init__(self): _apply_go2style_ppo(self)
@rlcfg("go1-combined-terrain-walk-go2style")
@dataclass
class Combined(SkrlCfg):
def __post_init__(self): _apply_go2style_ppo(self)
@rlcfg("go1-cts-flat-walk-go2style")
@dataclass
class CTS(SkrlCfg):
def __post_init__(self): _apply_go2style_ppo(self)
@rlcfg("go1-orig-combined-walk")
@dataclass
class OrigCombined(SkrlCfg):
"""Original Go1 (PD 80, 45-dim, [256,128,64]) on flat+stairs."""
def __post_init__(self): _apply_orig_ppo(self)

264
scripts/dreamwaq_sim2sim_mujoco.py Normal file
View File

@@ -0,0 +1,264 @@
#!/usr/bin/env python3
"""DreamWaQ MuJoCo sim2sim — VAE encoder + Actor, 5-frame history buffer.
Usage:
uv run scripts/dreamwaq_sim2sim_mujoco.py # flat
uv run scripts/dreamwaq_sim2sim_mujoco.py --terrain rough
uv run scripts/dreamwaq_sim2sim_mujoco.py --onnx path/to/policy.onnx
Controls:
W/S: forward/back Q/E: left/right A/D: rotate
Space: stop R: reset Esc: quit
"""
import numpy as np
import mujoco
from mujoco import viewer
import onnxruntime as ort
import os, sys, threading, queue, argparse, time, signal
g_exit_requested = False
signal.signal(signal.SIGINT, lambda *a: globals().update(g_exit_requested=True))
# ═══════════════════════════════════════════════════════════════════════
_PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
XML_DIR = os.path.join(_PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
DEFAULT_ONNX = os.path.join(_PROJECT, "exports_go1_dreamwaq", "policy.onnx")
# ── DreamWaQ params (matching training: PD 28/0.7, action_scale 0.25, ctrl_dt=0.02) ──
NUM_OBS = 45
NUM_ACTIONS = 12
HISTORY_LEN = 5
ACTION_SCALE = 0.25
KP = 28.0
KD = 0.7
CLIP_ACTIONS = 23.7
CLIP_OBS = 100.0
MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
# DreamWaQ default joint angles — MUST match MuJoCo XML joint order:
# qpos[7:19] = FR_hip,FR_thigh,FR_calf, FL_hip,FL_thigh,FL_calf, RR_hip,RR_thigh,RR_calf, RL_hip,RL_thigh,RL_calf
DEFAULT_ANGLES = np.array([
-0.1, 0.8, -1.5, # FR
0.1, 0.8, -1.5, # FL
-0.1, 1.0, -1.5, # RR
0.1, 1.0, -1.5, # RL
], dtype=np.float32)
# ═══════════════════════════════════════════════════════════════════════
# Keyboard
# ═══════════════════════════════════════════════════════════════════════
from pynput import keyboard
class KB:
def __init__(self):
self._q = queue.Queue(); self.running = True
self.held = set(); self._t = None; self._l = None
def _n(self, k):
try:
if hasattr(k, 'char') and k.char: return k.char.lower()
except: pass
return str(k).lower()
def _w(self):
while self.running:
try:
et, k = self._q.get(timeout=0.05)
n = self._n(k)
if et == 'press': self.held.add(n)
elif et == 'release': self.held.discard(n)
except queue.Empty: pass
def init(self):
self._l = keyboard.Listener(on_press=lambda k: self._q.put(('press', k)),
on_release=lambda k: self._q.put(('release', k)))
self._l.start()
self._t = threading.Thread(target=self._w, daemon=True); self._t.start()
def stop(self): self.running = False; self._l.stop()
# ═══════════════════════════════════════════════════════════════════════
# Sensor
# ═══════════════════════════════════════════════════════════════════════
def get_sensor(m, d, name):
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, name)
if sid < 0: return None
adr = m.sensor_adr[sid]; dim = m.sensor_dim[sid]
return d.sensordata[adr:adr+dim].copy()
def compute_obs(model, data, commands, last_action):
"""DreamWaQ observation (Manaro-Alpha order):
ang_vel(3) + gravity(3) + commands(3) + joint_pos(12) + joint_vel(12) + actions(12) = 45
"""
obs = np.zeros(NUM_OBS, dtype=np.float32)
# ang_vel [0:3]
g = get_sensor(model, data, "gyro")
obs[0:3] = (g if g is not None else data.qvel[3:6]) * 0.25
# gravity [3:6] (read from MuJoCo model, matching training)
grav_world = model.opt.gravity.copy()
grav_world = grav_world / np.linalg.norm(grav_world) # normalize
R = data.xmat[1].reshape(3, 3)
obs[3:6] = (R.T @ grav_world).astype(np.float32)
# commands [6:9]
obs[6:9] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32)
# joint_pos [9:21]
obs[9:21] = (data.qpos[7:19] - DEFAULT_ANGLES) * 1.0
# joint_vel [21:33]
obs[21:33] = data.qvel[6:18] * 0.05
# last_action [33:45]
obs[33:45] = last_action
return np.clip(obs, -CLIP_OBS, CLIP_OBS)
# ═══════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════
def main():
p = argparse.ArgumentParser()
p.add_argument("--onnx", default=DEFAULT_ONNX)
p.add_argument("--terrain", default="flat", choices=["flat", "rough", "stairs", "dreamwaq", "stairs_test", "stairs_box", "flat_stairs"])
p.add_argument("--level", type=int, default=0, help="terrain difficulty level 0-9 (0=flat, 9=hardest)")
args = p.parse_args()
# Select XML scene
terrain_map = {
"flat": "scene_motor_actuator.xml",
"rough": "scene_rough_terrain.xml",
"stairs": "scene_stairs_terrain.xml",
"dreamwaq": "scene_dreamwaq_terrain.xml",
"stairs_test": "scene_stairs_test.xml",
"stairs_box": "scene_stairs_box.xml",
"flat_stairs": "scene_flat_stairs.xml",
}
xml_file = os.path.join(XML_DIR, terrain_map[args.terrain])
if not os.path.exists(args.onnx):
print(f"[ERROR] ONNX not found: {args.onnx}")
print("Run: uv run scripts/export_dreamwaq_onnx.py (after training completes)")
sys.exit(1)
os.chdir(XML_DIR)
with open(xml_file) as f:
model = mujoco.MjModel.from_xml_string(f.read())
data = mujoco.MjData(model)
# Spawn pose. Hfield heights: MuJoCo z = gp[2] + sbase + (hd * ztop).
# The stairs_test terrain has sbase=0, flat platform z=0; just lift by clearance.
if args.terrain == "flat_stairs":
lvl = max(0, min(1, args.level))
col = np.random.randint(0, 4)
spawn_y = 4.0 - lvl * 8.0 # level 0 flat at y=+4, level 1 stairs at y=-4
spawn_x = -12.0 + col * 8.0 # platform center (cell center x)
elif args.terrain == "stairs_test":
spawn_x, spawn_y = -7.5, -4.0 # flat approach before first step (1m zone)
elif args.terrain == "stairs_box":
spawn_x, spawn_y = -2.0, 0.0 # flat ground before stairs
elif args.terrain == "dreamwaq":
lvl = max(0, min(9, args.level))
col = np.random.randint(0, 4) # NUM_COLS=4
spawn_y = 36.0 - lvl * 8.0 # level 0 flat at y=+36, level 9 stairs at y=-36
spawn_x = -12.0 + col * 8.0 + 4.0 # centre of cell
print(f"[Level {lvl}] type={col} spawn=({spawn_x:.1f}, {spawn_y:.1f})")
else:
spawn_x, spawn_y = (0.0, 0.0)
# When DISPLAY is a virtual framebuffer (Xvfb), MuJoCo headless rendering is
# handled transparently; on a real display this opens a normal GUI window.
# No explicit headless flag needed — MuJoCo glfw detects the display type.
def hfield_z(mx, my):
gid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_GEOM, "floor")
if gid < 0 or model.geom_type[gid] != mujoco.mjtGeom.mjGEOM_HFIELD:
return 0.0
hf = model.geom_dataid[gid]
nrow, ncol = int(model.hfield_nrow[hf]), int(model.hfield_ncol[hf])
sx, sy, ztop, sbase = model.hfield_size[hf]
adr = model.hfield_adr[hf]
hd = model.hfield_data[adr:adr + nrow * ncol].reshape(nrow, ncol)
gp = model.geom_pos[gid]
col = int(np.clip(((mx - gp[0]) / sx * 0.5 + 0.5) * (ncol - 1), 0, ncol - 1))
row = int(np.clip(((my - gp[1]) / sy * 0.5 + 0.5) * (nrow - 1), 0, nrow - 1))
return float(gp[2] + sbase + hd[row, col] * ztop)
spawn_z = hfield_z(spawn_x, spawn_y) + 0.45 # standing clearance above terrain
def reset_state():
data.qpos[:] = 0
data.qpos[0:3] = [spawn_x, spawn_y, spawn_z]; data.qpos[3:7] = [1, 0, 0, 0]
data.qpos[7:19] = DEFAULT_ANGLES; data.qvel[:] = 0
mujoco.mj_forward(model, data)
reset_state()
# ONNX (2 inputs: observations + obs_history)
session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
print(f"[DreamWaQ] {args.onnx}")
print(f"[Terrain] {args.terrain}")
print(f"[CTRL] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
kb = KB(); kb.init()
view = viewer.launch_passive(model, data)
# Camera tracking: follow the trunk body
trunk_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk")
view.cam.lookat = data.body(trunk_id).xpos.copy()
view.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
view.cam.trackbodyid = trunk_id
step = 0
vx, vy, wz = 0.0, 0.0, 0.0
last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
history = np.zeros((1, HISTORY_LEN, NUM_OBS), dtype=np.float32)
decimation = 4 # MuJoCo dt=0.005, policy dt=0.02 (DreamWaQ aligned)
loop_t0 = time.time()
while view.is_running() and not g_exit_requested:
keys = kb.held
if 'escape' in keys: break
if 'r' in keys:
reset_state()
last_action[:] = 0; history[:] = 0
print("[R] Reset")
if ' ' in keys: vx = vy = wz = 0.0
vx = MAX_VX if 'w' in keys else (-MAX_VX if 's' in keys else 0.0)
vy = MAX_VY if 'q' in keys else (-MAX_VY if 'e' in keys else 0.0)
wz = MAX_WZ if 'a' in keys else (-MAX_WZ if 'd' in keys else 0.0)
if step % decimation == 0:
cmd = np.array([vx, vy, wz], dtype=np.float32)
obs = compute_obs(model, data, cmd, last_action)
# Shift history + add new obs
history = np.concatenate([history[:, 1:, :], obs.reshape(1, 1, -1)], axis=1)
# ONNX inference
outputs = session.run(None, {
'observations': obs.reshape(1, -1).astype(np.float32),
'obs_history': history.astype(np.float32),
})
action = outputs[0][0]
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
last_action = action.copy()
# PD control
target = DEFAULT_ANGLES + action * ACTION_SCALE
torques = KP * (target - data.qpos[7:19]) - KD * data.qvel[6:18]
data.ctrl[:] = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
mujoco.mj_step(model, data)
view.sync()
# Time sync (policy at 50Hz = 0.02s per step)
expected = step * 0.02
elapsed = time.time() - loop_t0
if elapsed < expected:
time.sleep(expected - elapsed)
step += 1
kb.stop(); view.close()
if __name__ == "__main__":
main()

272
scripts/eval_go1_commands.py Normal file
View File

@@ -0,0 +1,272 @@
# Copyright (C) 2020-2025 Motphys Technology Co., Ltd. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Evaluate trained Go1 policy with specific velocity commands.
Tests three command patterns:
1. Forward-Backward (前后往返): vx oscillates +1.0 <-> -1.0
2. Left-Right (左右往返): vy oscillates +1.0 <-> -1.0
3. Rotation (旋转): wz oscillates +1.0 <-> -1.0
For each pattern, the command toggles direction every N seconds.
Metrics (tracking errors, etc.) are logged to CSV files for analysis.
"""
import logging
import time
from pathlib import Path
import numpy as np
from absl import app, flags
from motrix_envs import registry as env_registry
from motrix_rl import registry, utils
from motrix_rl.skrl.jax import wrap_env
from motrix_rl.skrl.jax.train.ppo import Trainer as SkrlJaxTrainer
logger = logging.getLogger(__name__)
_ENV = flags.DEFINE_string("env", "go1-flat-terrain-walk", "The env to evaluate")
_POLICY = flags.DEFINE_string(
"policy",
None,
"Path to policy checkpoint. Auto-discovers the latest best_agent.pickle if not specified.",
)
_DURATION = flags.DEFINE_float("duration", 4.0, "Seconds per command direction before toggling")
_OUTDIR = flags.DEFINE_string("outdir", None, "Output directory for CSV logs")
_RENDER = flags.DEFINE_bool("render", False, "Enable rendering (may crash on headless)")
# ── Helpers ─────────────────────────────────────────────────────────────
def _find_best_policy(env_name: str) -> Path:
"""Auto-discover the best SKRL policy for the given env."""
base = Path(f"runs/{env_name}/skrl")
if not base.exists():
raise FileNotFoundError(f"No training runs found at {base}")
runs = sorted([d for d in base.iterdir() if d.is_dir()], key=lambda d: d.stat().st_mtime, reverse=True)
if not runs:
raise FileNotFoundError(f"No training runs found at {base}")
ckpt_dir = runs[0] / "checkpoints"
best = list(ckpt_dir.glob("best_agent.*"))
if best:
return best[0]
# Fallback: highest timestep
ckpts = list(ckpt_dir.glob("agent_*.pickle"))
if not ckpts:
raise FileNotFoundError(f"No checkpoints found in {ckpt_dir}")
def _ts(p):
try:
return int(p.stem.split("_")[1])
except (IndexError, ValueError):
return 0
return max(ckpts, key=_ts)
def _save_log(outdir: Path, label: str, records: list[dict]) -> None:
"""Save records to a CSV file."""
outdir.mkdir(parents=True, exist_ok=True)
import csv
path = outdir / f"{label}.csv"
with open(path, "w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=records[0].keys())
writer.writeheader()
writer.writerows(records)
logger.info(f"Saved {len(records)} records → {path}")
# ── Command generators ──────────────────────────────────────────────────
def _command_forward_backward(t: float, period: float):
"""vx oscillates +1.0 / -1.0, vy=0, wz=0."""
half = period / 2.0
phase = (t % period) / half # 0..1 forward, 1..2 backward
vx = 1.0 if phase < 1.0 else -1.0
return np.array([vx, 0.0, 0.0], dtype=np.float32)
def _command_left_right(t: float, period: float):
"""vy oscillates +1.0 / -1.0, vx=0, wz=0."""
half = period / 2.0
phase = (t % period) / half
vy = 1.0 if phase < 1.0 else -1.0
return np.array([0.0, vy, 0.0], dtype=np.float32)
def _command_rotation(t: float, period: float):
"""wz oscillates +1.0 / -1.0, vx=0, vy=0."""
half = period / 2.0
phase = (t % period) / half
wz = 1.0 if phase < 1.0 else -1.0
return np.array([0.0, 0.0, wz], dtype=np.float32)
# ── Main evaluation logic ───────────────────────────────────────────────
def _run_pattern(
trainer: SkrlJaxTrainer,
policy_path: str,
pattern_name: str,
command_fn,
period: float,
total_steps: int,
ctrl_dt: float,
) -> list[dict]:
"""Run one command pattern and return tracking records."""
logger.info(f"--- {pattern_name} ---")
rlcfg = trainer._rlcfg
env = env_registry.make(trainer._env_name, sim_backend=trainer._sim_backend, num_envs=1)
env = wrap_env(env, enable_render=False)
# Build fresh agent for this run
models = trainer._make_model(env, rlcfg)
ppo_cfg = rlcfg.runner.agent.to_dict()
from motrix_rl.skrl.jax.train.ppo import _add_runtime_config
_add_runtime_config(ppo_cfg, env)
agent = trainer._make_agent(models, env, ppo_cfg, rlcfg.runner.memory)
agent.load(policy_path)
obs, info = env.reset()
state = env._env.state
records = []
for step in range(total_steps):
t = step * ctrl_dt
# Override command in env state
cmd = command_fn(t, period)
state.info["commands"] = cmd.reshape(1, -1)
# Recompute observation with the new command
new_obs = env._env._get_obs(state.data, state.info)
state = state.replace(obs=new_obs)
obs = new_obs
# Agent inference
outputs = agent.act(obs, timestep=0, timesteps=0)
actions = outputs[-1].get("mean_actions", outputs[0])
obs, reward, terminated, truncated, info = env.step(actions)
state = env._env.state
# Collect tracking data
lin_vel = env._env.get_local_linvel(state.data)[0] # [vx, vy, vz] (body frame)
gyro = env._env.get_gyro(state.data)[0]
tracking_err_xy = np.linalg.norm(cmd[:2] - lin_vel[:2])
tracking_err_yaw = abs(cmd[2] - gyro[2])
records.append({
"step": step,
"time": round(t, 3),
"cmd_vx": round(float(cmd[0]), 4),
"cmd_vy": round(float(cmd[1]), 4),
"cmd_wz": round(float(cmd[2]), 4),
"actual_vx": round(float(lin_vel[0]), 4),
"actual_vy": round(float(lin_vel[1]), 4),
"actual_vz": round(float(lin_vel[2]), 4),
"actual_wz": round(float(gyro[2]), 4),
"tracking_err_xy": round(float(tracking_err_xy), 6),
"tracking_err_yaw": round(float(tracking_err_yaw), 6),
"reward": round(float(reward[0][0]), 6),
})
if step % 100 == 0:
logger.info(
f" [{pattern_name}] step {step:5d}/{total_steps} "
f"cmd=[{cmd[0]:+.1f},{cmd[1]:+.1f},{cmd[2]:+.1f}] "
f"actual_v=[{lin_vel[0]:+.3f},{lin_vel[1]:+.3f},{gyro[2]:+.3f}] "
f"track_err_xy={tracking_err_xy:.4f} "
f"reward={float(reward[0][0]):.4f}"
)
env.close()
return records
def main(argv):
env_name = _ENV.value
ctrl_dt = 0.01 # matches Go1WalkNpEnvCfg.ctrl_dt
period = _DURATION.value * 2.0 # full cycle: forward + backward
# Resolve policy path
if _POLICY.present:
policy_path = _POLICY.value
else:
policy_path = str(_find_best_policy(env_name))
logger.info(f"Policy: {policy_path}")
# Build trainer (used to construct models & agent)
rlcfg = registry.default_rl_cfg(env_name, "skrl", backend="jax")
trainer = SkrlJaxTrainer(env_name, sim_backend=None, enable_render=False)
trainer._rlcfg = rlcfg
# Total steps per pattern = enough full cycles
cycles = 3
total_steps = int(cycles * period / ctrl_dt)
# Output directory
if _OUTDIR.present:
outdir = Path(_OUTDIR.value)
else:
ts = time.strftime("%y-%m-%d_%H-%M-%S")
outdir = Path(f"runs/{env_name}/eval_{ts}")
logger.info(f"Output directory: {outdir}")
patterns = [
("forward_backward", _command_forward_backward),
("left_right", _command_left_right),
("rotation", _command_rotation),
]
summary = {}
for label, cmd_fn in patterns:
records = _run_pattern(
trainer, policy_path, label, cmd_fn,
period=period, total_steps=total_steps, ctrl_dt=ctrl_dt,
)
_save_log(outdir, label, records)
# Summary stats (steady-state, skip first 2s for stabilization)
warmup = int(2.0 / ctrl_dt)
steady = records[warmup:]
if steady:
avg_err = np.mean([r["tracking_err_xy"] for r in steady])
avg_reward = np.mean([r["reward"] for r in steady])
else:
avg_err, avg_reward = float("nan"), float("nan")
summary[label] = {"avg_tracking_err_xy": avg_err, "avg_reward": avg_reward}
logger.info(f" [{label}] steady-state avg tracking_err_xy = {avg_err:.4f}, avg_reward = {avg_reward:.4f}")
# Print summary
print("\n" + "=" * 70)
print("EVALUATION SUMMARY")
print("=" * 70)
for label, stats in summary.items():
print(f" {label:25s} track_err_xy={stats['avg_tracking_err_xy']:.4f} avg_reward={stats['avg_reward']:.4f}")
print(f"\nDetailed CSV logs saved to: {outdir}")
print("=" * 70)
if __name__ == "__main__":
app.run(main)

217
scripts/export_dreamwaq_onnx.py Normal file
View File

@@ -0,0 +1,217 @@
#!/usr/bin/env python3
"""DreamWaQ ONNX export: CENet encoder + Actor → ONNX for MuJoCo deployment.
Model: 2 inputs, 1 output
- observations: (1, 45)
- obs_history: (1, 5, 45)
→ actions: (1, 12)
Usage:
uv run scripts/export_dreamwaq_onnx.py
uv run scripts/export_dreamwaq_onnx.py --checkpoint PATH --vae PATH --output PATH
"""
import argparse, os, pickle, sys
import msgpack
import numpy as np
import jax, jax.numpy as jnp
import flax.linen as nn
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# ═══════════════════════════════════════════════════════════════════════
# Export helpers
# ═══════════════════════════════════════════════════════════════════════
def _decode_flax_array(ext) -> np.ndarray | None:
"""Decode flax-serialized msgpack ExtType to numpy array."""
if not hasattr(ext, "code"): return None
parts = msgpack.unpackb(ext.data, raw=False)
if not isinstance(parts, list) or len(parts) < 3: return None
shape = []
def _flatten(s):
if isinstance(s, list):
for x in s: _flatten(x)
elif isinstance(s, int):
shape.append(s)
_flatten(parts[0])
dtype_str = parts[1]
raw_bytes = parts[2]
return np.frombuffer(raw_bytes, dtype=np.dtype(dtype_str)).reshape(shape)
def load_skrl_policy(path):
"""Extract actor weights from SKRL checkpoint. Slice first layer: 254→64."""
with open(path, 'rb') as f:
ckpt = pickle.load(f)
raw = msgpack.unpackb(ckpt['policy'])['params']
params = {}
for name, val in raw.items():
if isinstance(val, dict):
params[name] = {k: _decode_flax_array(v) for k, v in val.items()}
else:
params[name] = _decode_flax_array(val)
# Slice first Dense layer: (254, 512) → (64, 512)
params['Dense_0'] = {
'kernel': params['Dense_0']['kernel'][:64, :],
'bias': params['Dense_0']['bias'],
}
return params
def load_state_preprocessor(path):
"""Load RunningStandardScaler stats (running_mean, running_variance) for the
first 64 dims = [code(19), obs(45)] that feed the actor.
CRITICAL: the policy was trained on NORMALIZED observations. Deployment must
apply: clip((x - mean) / (sqrt(var) + 1e-8), -5, 5) before the actor.
"""
with open(path, 'rb') as f:
ckpt = pickle.load(f)
if 'state_preprocessor' not in ckpt:
print("[WARN] No state_preprocessor in checkpoint — skipping normalization")
return None, None
sp = msgpack.unpackb(ckpt['state_preprocessor'], raw=False)
mean = _decode_flax_array(sp['running_mean'])[:64].astype(np.float32)
var = _decode_flax_array(sp['running_variance'])[:64].astype(np.float32)
return mean, var
def export_onnx(actor_params, vae_params, output_path, obs_mean=None, obs_var=None):
"""Build PyTorch model from Flax params, export to ONNX.
If obs_mean/obs_var given, bakes in the state-preprocessor normalization
(applied to [code(19), obs(45)] before the actor) — REQUIRED for the policy
to behave correctly, since it was trained on normalized observations.
"""
import torch, torch.nn as tnn
class DreamWaQTorch(tnn.Module):
def __init__(self):
super().__init__()
# CENet encoder (Manaro-Alpha: 225→128→64)
self.enc1 = tnn.Linear(225, 128)
self.enc2 = tnn.Linear(128, 64)
self.latent_mu = tnn.Linear(64, 16)
self.vel_mu = tnn.Linear(64, 3)
# Actor (64→512→256→128→12)
self.act1 = tnn.Linear(64, 512)
self.act2 = tnn.Linear(512, 256)
self.act3 = tnn.Linear(256, 128)
self.act_out = tnn.Linear(128, 12)
# State-preprocessor normalization buffers (for [code(19), obs(45)] = 64)
self.register_buffer("obs_mean", torch.zeros(64))
self.register_buffer("obs_std", torch.ones(64))
self.normalize = False
def forward(self, obs, history):
h = history.reshape(history.shape[0], -1)
h = tnn.functional.elu(self.enc1(h))
h = tnn.functional.elu(self.enc2(h))
z = self.latent_mu(h)
vel = self.vel_mu(h)
x = torch.cat([vel, z, obs], dim=-1)
# Apply state-preprocessor normalization (clip((x-mean)/(std+eps), -5, 5))
if self.normalize:
x = torch.clamp((x - self.obs_mean) / (self.obs_std + 1e-8), -5.0, 5.0)
x = tnn.functional.elu(self.act1(x))
x = tnn.functional.elu(self.act2(x))
x = tnn.functional.elu(self.act3(x))
return self.act_out(x)
model = DreamWaQTorch()
if obs_mean is not None and obs_var is not None:
model.obs_mean.data = torch.from_numpy(obs_mean.copy())
model.obs_std.data = torch.from_numpy(np.sqrt(obs_var).copy())
model.normalize = True
print("[ONNX] State-preprocessor normalization baked in")
# Transfer CENet encoder weights (from Flax frozen dict)
vp = vae_params['params']
model.enc1.weight.data = torch.from_numpy(np.array(vp['enc_fc1']['kernel']).T.copy())
model.enc1.bias.data = torch.from_numpy(np.array(vp['enc_fc1']['bias']).copy())
model.enc2.weight.data = torch.from_numpy(np.array(vp['enc_fc2']['kernel']).T.copy())
model.enc2.bias.data = torch.from_numpy(np.array(vp['enc_fc2']['bias']).copy())
model.latent_mu.weight.data = torch.from_numpy(np.array(vp['latent_mu']['kernel']).T.copy())
model.latent_mu.bias.data = torch.from_numpy(np.array(vp['latent_mu']['bias']).copy())
model.vel_mu.weight.data = torch.from_numpy(np.array(vp['vel_mu']['kernel']).T.copy())
model.vel_mu.bias.data = torch.from_numpy(np.array(vp['vel_mu']['bias']).copy())
# Transfer Actor weights (64-dim, already sliced, from msgpack decoded)
ap = actor_params
model.act1.weight.data = torch.from_numpy(ap['Dense_0']['kernel'].T.copy())
model.act1.bias.data = torch.from_numpy(ap['Dense_0']['bias'].copy())
model.act2.weight.data = torch.from_numpy(ap['Dense_1']['kernel'].T.copy())
model.act2.bias.data = torch.from_numpy(ap['Dense_1']['bias'].copy())
model.act3.weight.data = torch.from_numpy(ap['Dense_2']['kernel'].T.copy())
model.act3.bias.data = torch.from_numpy(ap['Dense_2']['bias'].copy())
model.act_out.weight.data = torch.from_numpy(ap['Dense_3']['kernel'].T.copy())
model.act_out.bias.data = torch.from_numpy(ap['Dense_3']['bias'].copy())
model.eval()
os.makedirs(os.path.dirname(output_path), exist_ok=True)
torch.onnx.export(
model,
(torch.randn(1, 45), torch.randn(1, 5, 45)),
output_path,
input_names=['observations', 'obs_history'],
output_names=['actions'],
opset_version=11,
dynamic_axes={'observations': {0: 'batch'}, 'obs_history': {0: 'batch'}, 'actions': {0: 'batch'}},
)
print(f"[ONNX] Exported → {output_path}")
# ═══════════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════════
def main():
p = argparse.ArgumentParser()
p.add_argument("--checkpoint", default=None, help="SKRL agent checkpoint")
p.add_argument("--vae", default=None, help="CENet params .pkl")
p.add_argument("--output", default=os.path.join(PROJECT, "exports_go1_dreamwaq", "policy.onnx"))
args = p.parse_args()
run_dir = os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "skrl")
# Auto-find checkpoint
if not args.checkpoint:
runs = sorted([d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)) and d.startswith("26-")])
if runs:
ckpt_dir = os.path.join(run_dir, runs[-1], "checkpoints")
args.checkpoint = os.path.join(ckpt_dir, "best_agent.pickle")
# Auto-find VAE params
if not args.vae:
vae_files = sorted([f for f in os.listdir(run_dir) if f.startswith("vae_") and f.endswith(".pkl")],
key=lambda x: int(x.split("_")[1].split(".")[0]))
if vae_files:
args.vae = os.path.join(run_dir, vae_files[-1])
else:
# Try cenet_params.pkl (saved at end of training)
cpath = os.path.join(run_dir, "cenet_params.pkl")
if os.path.exists(cpath):
args.vae = cpath
if not args.checkpoint or not os.path.exists(args.checkpoint):
print(f"[ERROR] Checkpoint not found: {args.checkpoint}")
sys.exit(1)
if not args.vae or not os.path.exists(args.vae):
print(f"[ERROR] VAE params not found: {args.vae}")
sys.exit(1)
print(f"Policy: {args.checkpoint}")
print(f"VAE: {args.vae}")
with open(args.vae, 'rb') as f:
vae_params = pickle.load(f)
actor_params = load_skrl_policy(args.checkpoint)
obs_mean, obs_var = load_state_preprocessor(args.checkpoint)
export_onnx(actor_params, vae_params, args.output, obs_mean, obs_var)
if __name__ == "__main__":
main()

88
scripts/export_dreamwaq_rsl_onnx.py Normal file
View File

@@ -0,0 +1,88 @@
#!/usr/bin/env python3
"""Export a DreamWaQ rsl_rl checkpoint (ActorCritic_DWAQ) to ONNX for MuJoCo.
rsl_rl has NO state preprocessor, so no normalization is needed (unlike SKRL).
The exported model uses the MEAN CENet code (deterministic deploy).
ONNX: inputs observations(1,45) + obs_history(1,5,45) -> actions(1,12)
(matches scripts/dreamwaq_sim2sim_mujoco.py interface)
Usage:
uv run scripts/export_dreamwaq_rsl_onnx.py # auto-find latest
uv run scripts/export_dreamwaq_rsl_onnx.py --checkpoint runs/.../model_700.pt
"""
import argparse, glob, os, sys
import torch
import torch.nn as tnn
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from motrix_rl.dwaq_rsl.actor_critic_dwaq import ActorCritic_DWAQ
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
class DwaqInfer(tnn.Module):
"""Deterministic inference: CENet mean code + actor."""
def __init__(self, ac: ActorCritic_DWAQ):
super().__init__()
self.encoder = ac.encoder
self.encode_mean_vel = ac.encode_mean_vel
self.encode_mean_latent = ac.encode_mean_latent
self.actor = ac.actor
def forward(self, obs, obs_history):
h = self.encoder(obs_history.reshape(obs_history.shape[0], -1)) # (B,225)->(B,64)
vel = self.encode_mean_vel(h) # (B,3) mean velocity estimate
latent = self.encode_mean_latent(h) # (B,16) mean latent
code = torch.cat([vel, latent], dim=-1) # (B,19) = [vel, latent]
x = torch.cat([code, obs], dim=-1) # (B,64) = [code, obs]
return self.actor(x) # (B,12)
def main():
p = argparse.ArgumentParser()
p.add_argument("--checkpoint", default=None)
p.add_argument("--output", default=os.path.join(PROJECT, "exports_go1_dreamwaq", "policy.onnx"))
p.add_argument("--num-obs", type=int, default=45)
p.add_argument("--num-priv", type=int, default=235)
p.add_argument("--num-hist", type=int, default=5)
p.add_argument("--num-act", type=int, default=12)
args = p.parse_args()
if args.checkpoint is None:
runs = sorted(glob.glob(os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "rsl_dwaq", "*")),
key=os.path.getmtime)
if not runs:
print("[ERROR] no rsl_dwaq runs found"); sys.exit(1)
models = glob.glob(os.path.join(runs[-1], "model_*.pt"))
args.checkpoint = max(models, key=os.path.getmtime)
print(f"[rsl-ONNX] checkpoint: {args.checkpoint}")
cenet_out = 19
ac = ActorCritic_DWAQ(
args.num_obs + cenet_out, # actor in = 64
args.num_priv, # critic in = 235
args.num_act, # 12
args.num_hist * args.num_obs, # cenet in = 225
cenet_out, # 19
)
ckpt = torch.load(args.checkpoint, map_location="cpu")
ac.load_state_dict(ckpt["model_state_dict"])
ac.eval()
model = DwaqInfer(ac).eval()
os.makedirs(os.path.dirname(args.output), exist_ok=True)
torch.onnx.export(
model,
(torch.zeros(1, args.num_obs), torch.zeros(1, args.num_hist, args.num_obs)),
args.output,
input_names=["observations", "obs_history"],
output_names=["actions"],
opset_version=11,
dynamic_axes={"observations": {0: "batch"}, "obs_history": {0: "batch"}, "actions": {0: "batch"}},
)
print(f"[rsl-ONNX] exported -> {args.output} (no normalization; mean CENet code)")
if __name__ == "__main__":
main()

289
scripts/export_go1_no_linevel_onnx.py Normal file
View File

@@ -0,0 +1,289 @@
#!/usr/bin/env python3
"""Export JAX/Flax-trained SKRL Go1 (no-linevel, 57-dim) policy to ONNX.
Converts Flax weights → PyTorch → ONNX, baking in the RunningStandardScaler
normalization so the ONNX model accepts raw (scaled) observations directly.
Usage:
uv run scripts/export_go1_no_linevel_onnx.py
uv run scripts/export_go1_no_linevel_onnx.py --output ./my_exports
"""
import argparse
import os
import pickle
from pathlib import Path
import msgpack
import numpy as np
def _decode_flax_array(ext) -> np.ndarray | None:
if not hasattr(ext, "code"):
return None
parts = msgpack.unpackb(ext.data, raw=False)
if not isinstance(parts, list) or len(parts) < 3:
return None
def _flatten(s):
if isinstance(s, list):
out = []
for item in s:
out.extend(_flatten(item))
return out
return [s]
shape = tuple(_flatten(parts[0]))
dtype_str = parts[1]
if isinstance(dtype_str, bytes):
dtype_str = dtype_str.decode("utf-8")
raw = parts[2]
return np.frombuffer(raw, dtype=np.dtype(dtype_str)).reshape(shape)
def load_jax_checkpoint(ckpt_path: str) -> dict:
with open(ckpt_path, "rb") as f:
data = pickle.load(f)
policy_raw = msgpack.unpackb(data["policy"])
flax_params = {}
for name, val in policy_raw["params"].items():
if isinstance(val, dict):
flax_params[name] = {k: _decode_flax_array(v) for k, v in val.items()}
else:
flax_params[name] = _decode_flax_array(val)
prep = msgpack.unpackb(data["state_preprocessor"])
running_mean = _decode_flax_array(prep["running_mean"])
running_var = _decode_flax_array(prep["running_variance"])
count_arr = _decode_flax_array(prep["current_count"])
count = int(count_arr.flat[0]) if count_arr is not None else 0
return {
"flax_params": flax_params,
"running_mean": running_mean,
"running_var": running_var,
"count": count,
}
# -- PyTorch model -----------------------------------------------------------
import torch
import torch.nn as nn
class PolicyTorch(nn.Module):
def __init__(self, obs_dim: int, action_dim: int, hidden_dims: list[int]):
super().__init__()
self.obs_dim = obs_dim
self.action_dim = action_dim
self.hidden_dims = hidden_dims
layers = []
in_dim = obs_dim
for h in hidden_dims:
layers.extend([nn.Linear(in_dim, h), nn.ELU()])
in_dim = h
self.net = nn.Sequential(*layers)
self.mean_layer = nn.Linear(in_dim, action_dim)
def forward(self, x):
return self.mean_layer(self.net(x))
class ONNXExporter(nn.Module):
def __init__(self, policy: PolicyTorch, mean: np.ndarray, std: np.ndarray):
super().__init__()
self.policy = policy
self.register_buffer("mean", torch.from_numpy(mean).float())
self.register_buffer("std", torch.from_numpy(std).float())
self.clip_threshold = 5.0
def forward(self, x):
x = (x - self.mean) / (self.std + 1e-8)
x = torch.clamp(x, min=-self.clip_threshold, max=self.clip_threshold)
return self.policy(x)
def flax_to_torch_weights(flax_params: dict, obs_dim: int, hidden_dims: list[int], action_dim: int) -> dict:
state_dict = {}
layer_names = sorted([k for k in flax_params if k.startswith("Dense_")])
hidden_dense = layer_names[:-1]
layer_idx = 0
for name in hidden_dense:
layer_params = flax_params[name]
kernel = layer_params["kernel"]
bias = layer_params["bias"]
state_dict[f"net.{layer_idx}.weight"] = torch.from_numpy(kernel.T.copy()).float()
state_dict[f"net.{layer_idx}.bias"] = torch.from_numpy(bias.copy()).float()
layer_idx += 2
last_name = layer_names[-1]
last_params = flax_params[last_name]
state_dict["mean_layer.weight"] = torch.from_numpy(last_params["kernel"].T.copy()).float()
state_dict["mean_layer.bias"] = torch.from_numpy(last_params["bias"].copy()).float()
return state_dict
# -- Config ------------------------------------------------------------------
GO1_JOINT_NAMES = [
"FR_hip", "FR_thigh", "FR_calf",
"FL_hip", "FL_thigh", "FL_calf",
"RR_hip", "RR_thigh", "RR_calf",
"RL_hip", "RL_thigh", "RL_calf",
]
GO1_DEFAULT_ANGLES = np.array([
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
], dtype=np.float32)
# 57-dim observation layout (NO linear velocity):
# [0:3] gyro (scaled *0.25)
# [3:6] gravity vector (body frame)
# [6:18] joint angle deviation (scaled *1.0)
# [18:30] joint velocity (scaled *0.05)
# [30:42] last action (raw)
# [42:45] command [vx*2.0, vy*2.0, wz*0.25]
# [45:57] foot contact forces (body frame, raw)
OBS_SCALES = {
"ang_vel": 0.25,
"dof_pos": 1.0,
"dof_vel": 0.05,
"contact_force": 1.0, # raw, no scaling
}
ACTION_SCALE = 0.05
KP, KD = 80.0, 1.0
CLIP_ACTIONS = 23.7
CLIP_OBS = 100.0
def auto_discover_checkpoint(env_name: str) -> str:
"""Find the latest best_agent checkpoint for the given env."""
base_dir = Path(f"runs/{env_name}/skrl")
if not base_dir.exists():
raise FileNotFoundError(f"No training results found: {base_dir}")
runs = sorted([d for d in base_dir.iterdir() if d.is_dir()], key=lambda d: d.stat().st_mtime, reverse=True)
for run_dir in runs:
ckpt = run_dir / "checkpoints" / "best_agent.pickle"
if ckpt.exists():
return str(ckpt)
raise FileNotFoundError(f"No best_agent.pickle found in {base_dir}")
def export(checkpoint_path: str, output_dir: str):
ckpt = load_jax_checkpoint(checkpoint_path)
flax_params = ckpt["flax_params"]
running_mean = ckpt["running_mean"]
running_var = ckpt["running_var"]
running_std = np.sqrt(running_var)
dense_keys = sorted([k for k in flax_params if k.startswith("Dense_")])
hidden_dims = [flax_params[k]["bias"].shape[0] for k in dense_keys[:-1]]
obs_dim = flax_params[dense_keys[0]]["kernel"].shape[0]
action_dim = flax_params[dense_keys[-1]]["bias"].shape[0]
print(f"Architecture: obs={obs_dim}, hidden={hidden_dims}, action={action_dim}")
print(f"Normalizer mean range: [{running_mean.min():.4f}, {running_mean.max():.4f}]")
print(f"Normalizer std range: [{running_std.min():.6f}, {running_std.max():.6f}]")
policy = PolicyTorch(obs_dim, action_dim, hidden_dims)
torch_weights = flax_to_torch_weights(flax_params, obs_dim, hidden_dims, action_dim)
policy.load_state_dict(torch_weights, strict=True)
policy.eval()
rng = np.random.RandomState(42)
test_obs = rng.randn(1, obs_dim).astype(np.float32)
with torch.no_grad():
torch_out = policy(torch.from_numpy(test_obs)).numpy()
print(f"Test forward pass: input shape={test_obs.shape}, output shape={torch_out.shape}")
print(f" output sample: {np.array2string(torch_out[0, :4], precision=4, suppress_small=True)} ...")
os.makedirs(output_dir, exist_ok=True)
onnx_path = os.path.join(output_dir, "policy.onnx")
exporter = ONNXExporter(policy, running_mean, running_std)
exporter.eval()
dummy = torch.zeros(1, obs_dim, dtype=torch.float32)
torch.onnx.export(
exporter, dummy, onnx_path,
export_params=True, opset_version=11,
input_names=["observations"], output_names=["actions"],
dynamic_axes={},
)
print(f"ONNX exported to: {onnx_path}")
# Normalizer stats
npz_path = os.path.join(output_dir, "normalizer.npz")
np.savez(npz_path, mean=running_mean, std=running_std)
print(f"Normalizer saved to: {npz_path}")
# Metadata
meta_path = os.path.join(output_dir, "metadata.txt")
with open(meta_path, "w") as f:
f.write("# Go1 No-Linevel Terrain Walk - ONNX Policy Metadata\n")
f.write(f"env: go1-stairs-terrain-walk-no-linevel\n")
f.write(f"obs_dim: {obs_dim}\n")
f.write(f"action_dim: {action_dim}\n")
f.write(f"hidden_dims: {hidden_dims}\n")
f.write("\n# Observation layout (57 dims, NO linear velocity):\n")
f.write(" [0:3] gyro * ang_vel_scale\n")
f.write(" [3:6] gravity (body frame)\n")
f.write(" [6:18] joint_angle_deviation * dof_pos_scale\n")
f.write(" [18:30] joint_vel * dof_vel_scale\n")
f.write(" [30:42] last_actions (raw)\n")
f.write(" [42:45] commands [vx*2.0, vy*2.0, wz*0.25]\n")
f.write(" [45:57] foot_contact_forces (body frame, raw)\n")
f.write(f"\n# Joint order: {GO1_JOINT_NAMES}\n")
f.write(f"default_angles: {GO1_DEFAULT_ANGLES.tolist()}\n")
f.write(f"action_scale: {ACTION_SCALE}\n")
f.write(f"kp: {KP}\n")
f.write(f"kd: {KD}\n")
f.write(f"clip_actions: {CLIP_ACTIONS}\n")
f.write(f"clip_observations: {CLIP_OBS}\n")
f.write("\n# Observation scales:\n")
for k, v in OBS_SCALES.items():
f.write(f" {k}: {v}\n")
f.write(" command_scale: [2.0, 2.0, 0.25]\n")
print(f"Metadata saved to: {meta_path}")
return onnx_path
def main():
parser = argparse.ArgumentParser(description="Export JAX-trained Go1 no-linevel policy to ONNX")
parser.add_argument("--checkpoint", type=str, default=None,
help="Path to SKRL JAX checkpoint (auto-discovered if not set)")
parser.add_argument("--output", type=str, default="exports_go1_no_linevel",
help="Output directory")
args = parser.parse_args()
if args.checkpoint:
ckpt_path = args.checkpoint
else:
ckpt_path = auto_discover_checkpoint("go1-stairs-terrain-walk-no-linevel")
if not os.path.exists(ckpt_path):
print(f"Error: checkpoint not found: {ckpt_path}")
return 1
print(f"Loading checkpoint: {ckpt_path}")
onnx_path = export(ckpt_path, args.output)
print(f"\nDone! ONNX model ready for sim2sim:")
print(f" {onnx_path}")
return 0
if __name__ == "__main__":
exit(main())

314
scripts/export_go1_onnx.py Normal file
View File

@@ -0,0 +1,314 @@
#!/usr/bin/env python3
"""Export JAX/Flax-trained SKRL Go1 policy to ONNX for sim2sim deployment.
Converts Flax weights → PyTorch → ONNX, baking in the RunningStandardScaler
normalization so the ONNX model accepts raw (scaled) observations directly.
Usage:
uv run scripts/export_go1_onnx.py
uv run scripts/export_go1_onnx.py --output ./my_exports
Output files (in output_dir):
policy.onnx - ONNX model with normalization baked in
normalizer.npz - Normalizer stats (for reference/debugging)
metadata.txt - Policy metadata (obs dim, joint order, scales, etc.)
"""
import argparse
import os
import pickle
from pathlib import Path
import msgpack
import numpy as np
# ── Flax weight decoder ──────────────────────────────────────────────────
def _decode_flax_array(ext) -> np.ndarray | None:
"""Decode a flax-serialized msgpack ExtType to a numpy array."""
if not hasattr(ext, "code"):
return None
# The ExtType data is a msgpack array: [shape_list, dtype_str, raw_bytes]
parts = msgpack.unpackb(ext.data, raw=False)
if not isinstance(parts, list) or len(parts) < 3:
return None
# parts[0]: nested shape list, e.g. [[12]] or [[256, 45]]
# parts[1]: dtype string, e.g. "float32"
# parts[2]: raw bytes of array data
def _flatten(s):
if isinstance(s, list):
out = []
for item in s:
out.extend(_flatten(item))
return out
return [s]
shape = tuple(_flatten(parts[0]))
dtype_str = parts[1]
if isinstance(dtype_str, bytes):
dtype_str = dtype_str.decode("utf-8")
raw = parts[2]
return np.frombuffer(raw, dtype=np.dtype(dtype_str)).reshape(shape)
def load_jax_checkpoint(ckpt_path: str) -> dict:
"""Load a SKRL JAX checkpoint and extract all arrays.
Returns dict with keys:
flax_params: {layer_name: {kernel, bias} | array} – Flax-format weights
running_mean: np.ndarray
running_var: np.ndarray
count: int
"""
with open(ckpt_path, "rb") as f:
data = pickle.load(f)
# Decode policy params
policy_raw = msgpack.unpackb(data["policy"])
flax_params = {}
for name, val in policy_raw["params"].items():
if isinstance(val, dict):
flax_params[name] = {
k: _decode_flax_array(v) for k, v in val.items()
}
else:
flax_params[name] = _decode_flax_array(val)
# Decode state preprocessor
prep = msgpack.unpackb(data["state_preprocessor"])
running_mean = _decode_flax_array(prep["running_mean"])
running_var = _decode_flax_array(prep["running_variance"])
count_arr = _decode_flax_array(prep["current_count"])
count = int(count_arr.flat[0]) if count_arr is not None else 0
return {
"flax_params": flax_params,
"running_mean": running_mean,
"running_var": running_var,
"count": count,
}
# ── PyTorch model (for ONNX export) ─────────────────────────────────────
import torch
import torch.nn as nn
class PolicyTorch(nn.Module):
"""PyTorch MLP matching the SKRL Flax policy architecture."""
def __init__(self, obs_dim: int, action_dim: int, hidden_dims: list[int]):
super().__init__()
self.obs_dim = obs_dim
self.action_dim = action_dim
self.hidden_dims = hidden_dims
layers = []
in_dim = obs_dim
for h in hidden_dims:
layers.extend([nn.Linear(in_dim, h), nn.ELU()])
in_dim = h
self.net = nn.Sequential(*layers)
self.mean_layer = nn.Linear(in_dim, action_dim)
def forward(self, x):
return self.mean_layer(self.net(x))
class ONNXExporter(nn.Module):
"""Wraps policy with RunningStandardScaler normalization baked in."""
def __init__(self, policy: PolicyTorch, mean: np.ndarray, std: np.ndarray):
super().__init__()
self.policy = policy
self.register_buffer("mean", torch.from_numpy(mean).float())
self.register_buffer("std", torch.from_numpy(std).float())
self.clip_threshold = 5.0
def forward(self, x):
x = (x - self.mean) / (self.std + 1e-8)
x = torch.clamp(x, min=-self.clip_threshold, max=self.clip_threshold)
return self.policy(x)
# ── Flax → PyTorch weight conversion ────────────────────────────────────
def flax_to_torch_weights(flax_params: dict, obs_dim: int, hidden_dims: list[int], action_dim: int) -> dict:
"""Convert Flax-format params to PyTorch state_dict.
Flax Dense kernel: shape [in_dim, out_dim]
PyTorch Linear weight: shape [out_dim, in_dim] → needs transpose
Architecture: Dense_0..Dense_{N-1} → net hidden layers (Linear+ELU pairs)
Dense_N → mean_layer (Linear, no activation)
"""
state_dict = {}
layer_names = sorted([k for k in flax_params if k.startswith("Dense_")])
# Hidden layers: all Dense except the last
hidden_dense = layer_names[:-1]
layer_idx = 0
for name in hidden_dense:
layer_params = flax_params[name]
kernel = layer_params["kernel"] # Flax: [in_dim, out_dim]
bias = layer_params["bias"] # [out_dim]
state_dict[f"net.{layer_idx}.weight"] = torch.from_numpy(kernel.T.copy()).float()
state_dict[f"net.{layer_idx}.bias"] = torch.from_numpy(bias.copy()).float()
layer_idx += 2 # skip ELU activation (no params)
# Output layer (mean_layer)
last_name = layer_names[-1]
last_params = flax_params[last_name]
state_dict["mean_layer.weight"] = torch.from_numpy(last_params["kernel"].T.copy()).float()
state_dict["mean_layer.bias"] = torch.from_numpy(last_params["bias"].copy()).float()
return state_dict
# ── Main export ─────────────────────────────────────────────────────────
GO1_JOINT_NAMES = [
"FR_hip", "FR_thigh", "FR_calf",
"FL_hip", "FL_thigh", "FL_calf",
"RR_hip", "RR_thigh", "RR_calf",
"RL_hip", "RL_thigh", "RL_calf",
]
GO1_DEFAULT_ANGLES = np.array([
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
], dtype=np.float32)
# Observation layout for our 45-dim policy (NO linear velocity):
# [0:3] gyro (scaled *0.25)
# [3:6] gravity vector (body frame, no scale)
# [6:18] joint angle deviation from default (scaled *1.0)
# [18:30] joint velocity (scaled *0.05)
# [30:42] last action (raw)
# [42:45] command [vx, vy, wz] (scaled *[2.0, 2.0, 0.25])
OBS_SCALES = {
"lin_vel": 2.0, # NOT used in 45-dim obs (kept for reference)
"ang_vel": 0.25,
"dof_pos": 1.0,
"dof_vel": 0.05,
}
ACTION_SCALE = 0.05
KP, KD = 80.0, 1.0
CLIP_ACTIONS = 23.7
CLIP_OBS = 100.0
def export(checkpoint_path: str, output_dir: str):
"""Main export pipeline."""
ckpt = load_jax_checkpoint(checkpoint_path)
flax_params = ckpt["flax_params"]
running_mean = ckpt["running_mean"]
running_var = ckpt["running_var"]
running_std = np.sqrt(running_var)
# Infer architecture from Flax params
dense_keys = sorted([k for k in flax_params if k.startswith("Dense_")])
hidden_dims = [flax_params[k]["bias"].shape[0] for k in dense_keys[:-1]]
obs_dim = flax_params[dense_keys[0]]["kernel"].shape[0]
action_dim = flax_params[dense_keys[-1]]["bias"].shape[0]
print(f"Architecture: obs={obs_dim}, hidden={hidden_dims}, action={action_dim}")
print(f"Normalizer mean range: [{running_mean.min():.4f}, {running_mean.max():.4f}]")
print(f"Normalizer std range: [{running_std.min():.6f}, {running_std.max():.6f}]")
# Build PyTorch model and load weights
policy = PolicyTorch(obs_dim, action_dim, hidden_dims)
torch_weights = flax_to_torch_weights(flax_params, obs_dim, hidden_dims, action_dim)
policy.load_state_dict(torch_weights, strict=True)
policy.eval()
# Verify conversion with a random input
rng = np.random.RandomState(42)
test_obs = rng.randn(1, obs_dim).astype(np.float32)
with torch.no_grad():
torch_out = policy(torch.from_numpy(test_obs)).numpy()
print(f"Test forward pass: input shape={test_obs.shape}, output shape={torch_out.shape}")
print(f" output sample: {np.array2string(torch_out[0, :4], precision=4, suppress_small=True)} ...")
# Export ONNX
os.makedirs(output_dir, exist_ok=True)
onnx_path = os.path.join(output_dir, "policy.onnx")
exporter = ONNXExporter(policy, running_mean, running_std)
exporter.eval()
dummy = torch.zeros(1, obs_dim, dtype=torch.float32)
torch.onnx.export(
exporter,
dummy,
onnx_path,
export_params=True,
opset_version=11,
input_names=["observations"],
output_names=["actions"],
dynamic_axes={},
)
print(f"✓ ONNX exported to: {onnx_path}")
# Save normalizer stats for reference
npz_path = os.path.join(output_dir, "normalizer.npz")
np.savez(npz_path, mean=running_mean, std=running_std)
print(f"✓ Normalizer saved to: {npz_path}")
# Save metadata
meta_path = os.path.join(output_dir, "metadata.txt")
with open(meta_path, "w") as f:
f.write(f"# Go1 Flat Terrain Walk - ONNX Policy Metadata\n")
f.write(f"obs_dim: {obs_dim}\n")
f.write(f"action_dim: {action_dim}\n")
f.write(f"hidden_dims: {hidden_dims}\n")
f.write(f"observation_layout: gyro(3) + gravity(3) + joint_angle(12) + joint_vel(12) + last_action(12) + command(3)\n")
f.write(f" - NO linear velocity in observation\n")
f.write(f"\n# Joint order: {GO1_JOINT_NAMES}\n")
f.write(f"default_angles: {GO1_DEFAULT_ANGLES.tolist()}\n")
f.write(f"action_scale: {ACTION_SCALE}\n")
f.write(f"kp: {KP}\n")
f.write(f"kd: {KD}\n")
f.write(f"clip_actions: {CLIP_ACTIONS}\n")
f.write(f"clip_observations: {CLIP_OBS}\n")
f.write(f"\n# Observation scales (applied BEFORE ONNX normalization):\n")
for k, v in OBS_SCALES.items():
f.write(f" {k}: {v}\n")
f.write(f" command_scale: [2.0, 2.0, 0.25] # for [vx, vy, wz]\n")
print(f"✓ Metadata saved to: {meta_path}")
return onnx_path
def main():
parser = argparse.ArgumentParser(description="Export JAX-trained Go1 policy to ONNX")
parser.add_argument("--checkpoint", type=str,
default="runs/go1-flat-terrain-walk/skrl/26-06-19_15-05-34-538657_PPO/checkpoints/best_agent.pickle",
help="Path to SKRL JAX checkpoint (.pickle)")
parser.add_argument("--output", type=str, default="exports_go1_flat",
help="Output directory for ONNX model and artifacts")
args = parser.parse_args()
if not os.path.exists(args.checkpoint):
print(f"Error: checkpoint not found: {args.checkpoint}")
print("Train first: uv run scripts/train.py --env go1-flat-terrain-walk")
return 1
print(f"Loading checkpoint: {args.checkpoint}")
onnx_path = export(args.checkpoint, args.output)
print(f"\nDone! ONNX model ready for sim2sim deployment:")
print(f" {onnx_path}")
return 0
if __name__ == "__main__":
exit(main())

119
scripts/gen_flat_stairs.py Normal file
View File

@@ -0,0 +1,119 @@
#!/usr/bin/env python3
"""Generate 2-level terrain: level 0=flat, level 1=pyramid stairs.
1cm = 1px. OpenCV draws concentric filled rectangles (outside→in, 10 steps).
Higher values overwrite lower = convex pyramid (stairs up toward center).
Lower values overwrite higher = concave pyramid (stairs down from center).
Usage:
uv run scripts/gen_flat_stairs.py # convex (stairs UP)
uv run scripts/gen_flat_stairs.py --concave # concave (stairs DOWN)
uv run scripts/gen_flat_stairs.py --step-h 0.10 --num-steps 5
"""
import cv2, numpy as np, os, argparse
# ═══ fixed params ═══
HS = 0.01 # 1cm/px
VS = 0.005 # height unit = 0.5cm
CELL_M = 8.0 # 8m cell
BORDER_M = 5.0 # 5m border
NUM_ROWS = 2 # flat + stairs
NUM_COLS = 4 # columns
CELL_PX = int(CELL_M / HS) # 800
BORDER_PX = int(BORDER_M / HS) # 500
PLATFORM_PX = int(1.0 / HS) # 1m platform = 100px
TOT_ROWS = NUM_ROWS * CELL_PX + 2 * BORDER_PX
TOT_COLS = NUM_COLS * CELL_PX + 2 * BORDER_PX
# ═══ stairs params (override via CLI) ═══
NUM_STEPS = 10 # 10 steps
STEP_H_CM = 20 # 20cm rise per step
STEP_D_CM = 20 # 20cm tread per step
STEP_H_VS = int(STEP_H_CM / 100.0 / VS) # 0.20 / 0.005 = 40
STEP_D_PX = int(STEP_D_CM / 100.0 / HS) # 0.20 / 0.01 = 20
def draw_pyramid(canvas, x0, y0, num_steps, step_d_px, step_h_vs, concave=False):
"""Draw concentric rectangles from outside→in.
Convex: edge=0 → platform=max (stairs up toward center)
Concave: raise whole cell to max, then draw pit: edge=max → platform=0
"""
cx, cy = x0 + CELL_PX // 2, y0 + CELL_PX // 2
p2 = PLATFORM_PX // 2
h_max = step_h_vs * num_steps
# Fill cell to cell boundary with reference-plane height
half_max = CELL_PX // 2 # extend to cell edge
cv2.rectangle(canvas, (cx - half_max, cy - half_max),
(cx + half_max, cy + half_max), int(h_max), -1)
if concave:
# Pit: rings going DOWN from reference plane
for i in range(num_steps + 1):
half = p2 + (num_steps - i) * step_d_px
x1, y1 = cx - half, cy - half
x2, y2 = cx + half, cy + half
h = h_max - step_h_vs * i
cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
else:
# Mound: rings going UP from reference plane
for i in range(num_steps + 1):
half = p2 + (num_steps - i) * step_d_px
x1, y1 = cx - half, cy - half
x2, y2 = cx + half, cy + half
h = h_max + step_h_vs * i
cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
def main():
p = argparse.ArgumentParser()
p.add_argument("--concave", action="store_true", help="concave pyramid (stairs down from center)")
p.add_argument("--step-h", type=float, default=0.20, help="step rise (m)")
p.add_argument("--step-d", type=float, default=0.20, help="step tread (m)")
p.add_argument("--num-steps", type=int, default=10, help="number of steps")
args = p.parse_args()
step_h_vs = int(args.step_h / VS)
step_d_px = int(args.step_d / HS)
total_h = step_h_vs * args.num_steps * VS
print(f"Building {TOT_COLS}×{TOT_ROWS}px ({TOT_COLS*HS:.0f}×{TOT_ROWS*HS:.0f}m)")
print(f" type={'concave' if args.concave else 'convex'} "
f"steps={args.num_steps} rise={step_h_vs*VS*100:.0f}cm "
f"tread={step_d_px*HS*100:.0f}cm total_h={total_h*100:.0f}cm")
canvas = np.zeros((TOT_ROWS, TOT_COLS), dtype=np.uint16)
for row in range(NUM_ROWS):
for col in range(NUM_COLS):
x0 = BORDER_PX + col * CELL_PX
y0 = BORDER_PX + row * CELL_PX
if row == 1:
# alternate convex/concave across cols
concave_cell = (col % 2 == 1)
draw_pyramid(canvas, x0, y0, args.num_steps,
step_d_px, step_h_vs, concave_cell)
hf_m = canvas.astype(np.float32) * VS
z_min, z_max = float(hf_m.min()), float(hf_m.max())
z_range = max(z_max - z_min, 0.001)
png = ((hf_m - z_min) / z_range * 65535).astype(np.uint16)
out_dir = os.path.join(os.path.dirname(__file__), "..",
"motrix_envs", "src", "motrix_envs", "locomotion",
"go1", "xmls", "assets")
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, "flat_stairs.png")
cv2.imwrite(out_path, png)
print(f" saved: {out_path}")
print(f" XML: size=\"{TOT_COLS*HS/2:.1f} {TOT_ROWS*HS/2:.1f} "
f"{z_range:.3f} {max(z_min,0.001):.3f}\"")
if __name__ == "__main__":
main()

95
scripts/gen_stairs_box.py Normal file
View File

@@ -0,0 +1,95 @@
#!/usr/bin/env python3
"""Generate box-geom stairs XML for MuJoCo sim2sim.
Each step is a separate box with vertical rises — much steeper than hfield.
Usage:
uv run scripts/gen_stairs_box.py # default: 10 steps × 6cm = 60cm
uv run scripts/gen_stairs_box.py --step-height 0.04 --num-steps 5
uv run scripts/gen_stairs_box.py --step-height 0.10 --num-steps 8 --step-depth 0.4
"""
import argparse, os
TPL = '''<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" />
{steps}
<!-- Fill under stairs -->
<geom name="fill" type="box" size="{fill_sx} 10 {fill_sz}" pos="{fill_x} 0 {fill_z}" rgba="0.5 0.4 0.3 1" friction="0.8 0.3 0.3"/>
</worldbody>
</mujoco>
'''
STEP_TPL = ' <geom name="step{n}" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
PLAT_TPL = ' <geom name="platform" type="box" size="{sx} 10 {sz}" pos="{x} 0 {z}" rgba="0.6 0.5 0.4 1" friction="0.8 0.3 0.3"/>\n'
def main():
p = argparse.ArgumentParser()
p.add_argument("--step-height", type=float, default=0.06, help="rise per step [m]")
p.add_argument("--step-depth", type=float, default=0.30, help="tread depth per step [m]")
p.add_argument("--num-steps", type=int, default=10, help="number of steps")
p.add_argument("--box-thickness", type=float, default=0.03, help="box half-height [m]")
args = p.parse_args()
h = args.step_height
d = args.step_depth
n = args.num_steps
sz = args.box_thickness # half-height of each box
steps_xml = ""
for i in range(n):
x = i * d
z = i * h + sz # center of box = step top surface - sz
steps_xml += STEP_TPL.format(n=i, sx=d/2, sz=sz, x=x, z=z)
# Platform at top
plat_x = n * d + 0.5
plat_z = n * h + sz
steps_xml += PLAT_TPL.format(sx=0.5, sz=sz, x=plat_x, z=plat_z)
# Fill box under stairs
total_depth = n * d
total_height = n * h
fill_sx = total_depth / 2
fill_sz = total_height / 2
fill_x = total_depth / 2
fill_z = -fill_sz
out = TPL.format(steps=steps_xml.rstrip(),
fill_sx=fill_sx, fill_sz=fill_sz,
fill_x=fill_x, fill_z=fill_z)
out_dir = os.path.join(os.path.dirname(__file__), "..",
"motrix_envs", "src", "motrix_envs", "locomotion",
"go1", "xmls")
out_path = os.path.join(out_dir, "scene_stairs_box.xml")
with open(out_path, "w") as f:
f.write(out)
max_h = n * h
print(f"Generated {n} steps × {h*100:.0f}cm = {max_h*100:.0f}cm total")
print(f" step depth: {d*100:.0f}cm box thickness: {sz*200:.0f}cm")
print(f" saved: {out_path}")
if __name__ == "__main__":
main()

108
scripts/gen_stairs_test.py Normal file
View File

@@ -0,0 +1,108 @@
#!/usr/bin/env python3
"""Generate standard linear stairs for MuJoCo sim2sim testing.
Each cell: flat 2m approach -> N linear steps -> flat 2m platform -> N steps down -> flat edge.
Treads are horizontal, rises are vertical (1px = 0.1m wide, acceptable for hfield).
Usage:
uv run scripts/gen_stairs_test.py --step-height 0.07 --step-depth 0.31
"""
import numpy as np, os, argparse
from PIL import Image
HS = 0.1 # horizontal scale [m/px]
VS = 0.005 # vertical scale [m/unit]
CELL_M = 8.0
PLATFORM_M = 4.0 # bigger flat platform -> fewer steps
CELL_PX = int(CELL_M / HS) # 80
PLATFORM_PX = int(PLATFORM_M / HS) # 40
BORDER_M = 2.0
BORDER_PX = int(BORDER_M / HS) # 20
NUM_CELLS = 2
TOT_PX = NUM_CELLS * CELL_PX + 2 * BORDER_PX
TOTAL_M = TOT_PX * HS
np.random.seed(42)
def make_linear_stairs(step_height_m, step_depth_m=0.31):
"""Linear stairs: flat approach -> N steps up -> flat platform -> edge.
Each step has a flat horizontal tread and (essentially) vertical rise."""
t = np.zeros((CELL_PX, CELL_PX), dtype=np.int16)
sd = int(step_depth_m / HS) # tread depth in px
sh = int(step_height_m / VS) # rise height in pixel units
# How many steps fit on each side of the platform?
avail = (CELL_PX - PLATFORM_PX) // 2
n_steps = avail // max(sd, 1)
if n_steps < 1:
n_steps = 1
edge = (CELL_PX - PLATFORM_PX - n_steps * sd) // 2 # remaining flat on each side
# Draw steps going UP from left (in +x direction)
# Each step: flat tread at current height, then rise to next height
x = edge
h = 0
for i in range(n_steps):
x_next = x + sd
t[:, x:x_next] = h # tread at current height
x = x_next
h += sh
# Platform (flat at max height)
plat_start = x
plat_end = plat_start + PLATFORM_PX
t[:, plat_start:plat_end] = h
# Continue stairs going DOWN on the right (optional: mirror)
x = plat_end
for i in range(n_steps):
h -= sh
x_next = x + sd
t[:, x:x_next] = h
x = x_next
return t
def main():
p = argparse.ArgumentParser()
p.add_argument("--step-height", type=float, default=0.15,
help="step rise height in metres (default 0.15)")
p.add_argument("--step-depth", type=float, default=0.50,
help="step tread depth in metres (default 0.50)")
args = p.parse_args()
print(f"Linear stairs: step_h={args.step_height:.2f}m tread={args.step_depth:.2f}m platform={PLATFORM_M:.0f}m")
hf_raw = np.zeros((TOT_PX, TOT_PX), dtype=np.int16)
for i in range(NUM_CELLS):
for j in range(NUM_CELLS):
cell = make_linear_stairs(args.step_height, args.step_depth)
y0 = BORDER_PX + i * CELL_PX
x0 = BORDER_PX + j * CELL_PX
hf_raw[y0:y0 + CELL_PX, x0:x0 + CELL_PX] = cell
hf_m = hf_raw.astype(np.float32) * VS
z_min = float(hf_m.min())
z_max = float(hf_m.max())
z_range = max(z_max - z_min, 0.001)
print(f" height: [{z_min:.3f}, {z_max:.3f}]m z_scale={z_range:.3f} max={z_max*100:.0f}cm")
png = ((hf_m - z_min) / z_range * 65535.0).astype(np.uint16)
out_dir = os.path.join(os.path.dirname(__file__), "..",
"motrix_envs", "src", "motrix_envs", "locomotion",
"go1", "xmls", "assets")
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, "stairs_test.png")
Image.fromarray(png).save(out_path)
print(f" saved: {out_path}")
sbase = max(z_min, 0.001)
print(f" XML: size=\"{TOTAL_M/2:.1f} {TOTAL_M/2:.1f} {z_range:.3f} {sbase:.3f}\"")
if __name__ == "__main__":
main()

428
scripts/go1_no_linevel_sim2sim_mujoco.py Normal file
View File

@@ -0,0 +1,428 @@
#!/usr/bin/env python3
"""MuJoCo sim2sim for go1-stairs-terrain-walk-no-linevel (57-dim obs, no linvel).
Loads the ONNX policy exported by export_go1_no_linevel_onnx.py and runs
inference in MuJoCo with PD control.
Usage:
# Default: combined flat+rough+stairs terrain, random spawn
uv run scripts/go1_no_linevel_sim2sim_mujoco.py
# Specific terrain
uv run scripts/go1_no_linevel_sim2sim_mujoco.py --terrain flat
uv run scripts/go1_no_linevel_sim2sim_mujoco.py --terrain rough
uv run scripts/go1_no_linevel_sim2sim_mujoco.py --terrain stairs
# Custom ONNX path
uv run scripts/go1_no_linevel_sim2sim_mujoco.py --onnx ./exports_go1_no_linevel/policy.onnx
Keyboard controls:
W/S - forward/backward
A/D - turn left/right
Q/E - strafe left/right
Space - stop
R - reset robot
1/2/3 - switch terrain (flat/rough/stairs)
Esc - quit
"""
import numpy as np
import mujoco
from mujoco import viewer
import os
import threading
import signal
import queue
import argparse
import time
g_exit_requested = False
def signal_handler(signum, frame):
global g_exit_requested
g_exit_requested = True
signal.signal(signal.SIGINT, signal_handler)
# ============================================================
# Paths
# ============================================================
_PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_ONNX_PATH = os.path.join(_PROJECT_DIR, "exports_go1_no_linevel", "policy.onnx")
MOTRIX_XML_DIR = os.path.join(_PROJECT_DIR, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
# ============================================================
# MotrixLab parameters (matching cfg.py + walk_stairs_terrain_no_linevel.py)
# ============================================================
NUM_OBS = 57 # 57-dim: NO linear velocity, WITH contact forces
NUM_ACTIONS = 12
OBS_SCALES = {"ang_vel": 0.25, "dof_pos": 1.0, "dof_vel": 0.05}
ACTION_SCALE = 0.05
KP, KD = 80.0, 1.0
CLIP_ACTIONS = 23.7
CLIP_OBSERVATIONS = 100.0
MAX_LIN_VEL_X = 1.0
MAX_LIN_VEL_Y = 1.0
MAX_ANG_VEL = 1.0
# ============================================================
# Joint names and order
# ============================================================
POLICY_JOINT_NAMES = [
"FR_hip", "FR_thigh", "FR_calf",
"FL_hip", "FL_thigh", "FL_calf",
"RR_hip", "RR_thigh", "RR_calf",
"RL_hip", "RL_thigh", "RL_calf",
]
DEFAULT_JOINT_ANGLES = np.array([
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
-0.0, 0.9, -1.8,
0.0, 0.9, -1.8,
], dtype=np.float32)
FEET = ["FR", "FL", "RR", "RL"]
# Terrain spawn positions (world Y)
TERRAIN_SPAWN = {
"flat": np.array([0.0, 54.0, 0.42], dtype=np.float64),
"rough": np.array([0.0, 32.0, 0.42], dtype=np.float64),
"stairs": np.array([0.0, 0.0, 0.42], dtype=np.float64),
}
# ============================================================
# Keyboard input
# ============================================================
from pynput import keyboard
class KeyboardReader:
def __init__(self):
self._event_queue = queue.Queue()
self.running = True
self.shared_keys_held = set()
self.shared_one_shot = set()
self._reader_thread = None
self._listener = None
def _normalize_key(self, key):
try:
if hasattr(key, "char") and key.char is not None:
return key.char.lower()
except Exception:
pass
key_str = str(key)
if key_str == "Key.esc":
return "escape"
elif key_str == "Key.space":
return "space"
elif key_str.startswith("Key."):
return key_str.lower()
return key_str.lower()
def _reader_worker(self):
while self.running:
try:
event_type, key = self._event_queue.get(timeout=0.05)
k = self._normalize_key(key)
if event_type == "press":
self.shared_keys_held.add(k)
self.shared_one_shot.discard(k)
elif event_type == "release":
self.shared_keys_held.discard(k)
self.shared_one_shot.discard(k)
except queue.Empty:
pass
def init(self):
def on_press(key):
self._event_queue.put(("press", key))
def on_release(key):
self._event_queue.put(("release", key))
try:
self._listener = keyboard.Listener(on_press=on_press, on_release=on_release)
self._listener.start()
self._reader_thread = threading.Thread(target=self._reader_worker, daemon=True)
self._reader_thread.start()
print("[INFO] Keyboard listener started")
except Exception as e:
print(f"[WARN] Cannot init keyboard: {e}")
def is_key_pressed(self, key):
k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
if k not in self.shared_keys_held or k in self.shared_one_shot:
return False
self.shared_one_shot.add(k)
return True
def is_key_held(self, key):
k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
return k in self.shared_keys_held
def restore(self):
self.running = False
if self._listener:
self._listener.stop()
# ============================================================
# Sensor reading
# ============================================================
def get_sensor(model, data, name):
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
if sid < 0:
return None
adr = model.sensor_adr[sid]
dim = model.sensor_dim[sid]
return data.sensordata[adr : adr + dim].copy()
def read_contact_forces(model, data, base_rot):
"""Read foot contact forces (12-dim, body frame) from MuJoCo contact sensors.
Tries _stairs, _rough, _flat suffixes for each foot, picking the first
sensor that returns non-zero data. In MuJoCo, `data="force"` returns a
scalar (normal force). We construct a 3D force vector by projecting onto
the body-frame Z axis as an approximation.
"""
forces = np.zeros(12, dtype=np.float32)
for i, foot in enumerate(FEET):
f_scalar = 0.0
for suffix in ["_stairs", "_rough", "_flat"]:
name = f"{foot}_foot_contact{suffix}"
v = get_sensor(model, data, name)
if v is not None and np.abs(v[0]) > 1e-6:
f_scalar = v[0]
break
# Assume contact force is approximately vertical (world Z),
# rotate into body frame
force_world = np.array([0.0, 0.0, f_scalar], dtype=np.float64)
force_body = base_rot.T @ force_world
forces[i * 3 : i * 3 + 3] = force_body.astype(np.float32)
return forces
def compute_observations(model, data, commands, last_actions, base_rot):
"""Compute 57-dim observation matching go1-stairs-terrain-walk-no-linevel.
Layout (57 dims, NO linear velocity):
[0:3] gyro (ang_vel * 0.25)
[3:6] gravity vector (body frame)
[6:18] joint angle deviation (dof_pos * 1.0)
[18:30] joint velocity (dof_vel * 0.05)
[30:42] last actions (raw)
[42:45] commands [vx*2.0, vy*2.0, wz*0.25]
[45:57] foot contact forces (body frame, raw)
"""
obs = np.zeros(NUM_OBS, dtype=np.float32)
# Gyro
gyro = get_sensor(model, data, "gyro")
if gyro is not None:
obs[0:3] = gyro * OBS_SCALES["ang_vel"]
else:
obs[0:3] = data.qvel[3:6] * OBS_SCALES["ang_vel"]
# Gravity vector (body frame)
gravity_world = np.array([0.0, 0.0, -1.0], dtype=np.float64)
local_gravity = base_rot.T @ gravity_world
obs[3:6] = local_gravity.astype(np.float32)
# Joint position deviation
joint_pos = data.qpos[7:19]
dof_pos_rel = (joint_pos - DEFAULT_JOINT_ANGLES) * OBS_SCALES["dof_pos"]
obs[6:18] = dof_pos_rel
# Joint velocity
joint_vel = data.qvel[6:18]
obs[18:30] = joint_vel * OBS_SCALES["dof_vel"]
# Last actions
obs[30:42] = last_actions
# Commands (scale matching MotrixLab: [2.0, 2.0, 0.25])
obs[42] = commands[0] * 2.0
obs[43] = commands[1] * 2.0
obs[44] = commands[2] * 0.25
# Contact forces
# obs[45:57] = read_contact_forces(model, data, base_rot) # disabled: test with zeros
obs[45:57] = np.zeros(12, dtype=np.float32)
obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
return obs
# ============================================================
# Main
# ============================================================
def main():
import onnxruntime as ort
parser = argparse.ArgumentParser(description="MotrixLab Go1 No-Linevel Policy Inference in MuJoCo")
parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX_PATH)
parser.add_argument(
"--terrain", type=str, default="combined",
choices=["flat", "rough", "stairs", "combined"],
help="Terrain type (combined = flat+rough+stairs in one scene)",
)
args = parser.parse_args()
os.chdir(MOTRIX_XML_DIR)
# Select XML
if args.terrain == "combined":
xml_file = f"{MOTRIX_XML_DIR}/scene_combined_flat_rough_stairs.xml"
elif args.terrain == "rough":
xml_file = f"{MOTRIX_XML_DIR}/scene_rough_terrain.xml"
elif args.terrain == "stairs":
xml_file = f"{MOTRIX_XML_DIR}/scene_stairs_terrain.xml"
else:
xml_file = f"{MOTRIX_XML_DIR}/scene_motor_actuator.xml"
with open(xml_file, "r") as f:
xml_content = f.read()
model = mujoco.MjModel.from_xml_string(xml_content)
data = mujoco.MjData(model)
print(f"[INFO] Terrain: {args.terrain}")
print(f"[INFO] Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
print(f"[INFO] Timestep: {model.opt.timestep}")
# Initial spawn (default to stairs at origin - visible to default camera)
current_terrain = args.terrain if args.terrain != "combined" else "stairs"
spawn_xyz = TERRAIN_SPAWN.get(current_terrain, TERRAIN_SPAWN["stairs"]).copy()
# Start higher so robot drops onto terrain safely
spawn_xyz[2] = 1.0
data.qpos[0:3] = spawn_xyz
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
data.qpos[7:19] = DEFAULT_JOINT_ANGLES
data.qvel[:] = 0.0
data.ctrl[:] = 0.0
mujoco.mj_forward(model, data)
# Load ONNX
session = ort.InferenceSession(args.onnx, providers=["CPUExecutionProvider"])
print(f"[INFO] ONNX loaded: {args.onnx}")
# Main loop
ctrl_dt = 0.01
num_steps_per_inference = int(ctrl_dt / model.opt.timestep)
print(f"[INFO] Inference every {num_steps_per_inference} sim steps")
step_count = 0
inference_step = 0
commands = np.zeros(3, dtype=np.float32)
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
keyboard_reader = KeyboardReader()
keyboard_reader.init()
viewer_handle = viewer.launch_passive(model, data)
print("[INFO] Viewer launched!")
print("[KEYS] WASD=move, QE=strafe, Space=stop, R=reset, 1/2/3=terrain, Esc=quit")
loop_start_time = time.time()
terrain_changed = False
while viewer_handle.is_running() and not g_exit_requested:
# --- Keyboard input ---
x_vel, y_vel, yaw_vel = 0.0, 0.0, 0.0
if keyboard_reader.is_key_held("w"):
x_vel = MAX_LIN_VEL_X
elif keyboard_reader.is_key_held("s"):
x_vel = -MAX_LIN_VEL_X
if keyboard_reader.is_key_held("q"):
y_vel = MAX_LIN_VEL_Y
elif keyboard_reader.is_key_held("e"):
y_vel = -MAX_LIN_VEL_Y
if keyboard_reader.is_key_held("a"):
yaw_vel = MAX_ANG_VEL
elif keyboard_reader.is_key_held("d"):
yaw_vel = -MAX_ANG_VEL
if keyboard_reader.is_key_pressed("space"):
x_vel = y_vel = yaw_vel = 0.0
# Terrain switching
for key, terrain_name in [("1", "flat"), ("2", "rough"), ("3", "stairs")]:
if keyboard_reader.is_key_pressed(key):
current_terrain = terrain_name
terrain_changed = True
print(f"[TERRAIN] Switch to: {current_terrain}")
# Reset
if keyboard_reader.is_key_pressed("r") or terrain_changed:
if terrain_changed and args.terrain == "combined":
spawn_xyz = TERRAIN_SPAWN[current_terrain].copy()
elif keyboard_reader.is_key_pressed("r"):
spawn_xyz = TERRAIN_SPAWN.get(current_terrain, TERRAIN_SPAWN["flat"]).copy()
data.qpos[0:3] = spawn_xyz
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
data.qpos[7:19] = DEFAULT_JOINT_ANGLES
data.qvel[:] = 0.0
data.ctrl[:] = 0.0
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
mujoco.mj_forward(model, data)
terrain_changed = False
print(f"[RESET] Terrain={current_terrain}, pos={spawn_xyz}")
if keyboard_reader.is_key_pressed("escape"):
break
# --- Inference ---
if inference_step == 0:
commands[0] = x_vel
commands[1] = y_vel
commands[2] = yaw_vel
base_rot = data.xmat[1].reshape(3, 3)
obs = compute_observations(model, data, commands, last_actions, base_rot)
action = session.run(None, {"observations": obs.reshape(1, -1).astype(np.float32)})[0][0]
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
last_actions = action.copy()
# --- PD control ---
joint_targets = DEFAULT_JOINT_ANGLES + action * ACTION_SCALE
current_pos = data.qpos[7:19]
current_vel = data.qvel[6:18]
torques = KP * (joint_targets - current_pos) - KD * current_vel
torques = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
data.ctrl[:] = torques
mujoco.mj_step(model, data)
viewer_handle.sync()
expected_time = step_count * ctrl_dt
elapsed = time.time() - loop_start_time
sleep_time = expected_time - elapsed
if sleep_time > 0:
time.sleep(sleep_time)
step_count += 1
inference_step = (inference_step + 1) % num_steps_per_inference
if step_count % 500 == 0:
trunk_z = data.qpos[2]
print(f"[{step_count}] cmd=({x_vel:.1f},{y_vel:.1f},{yaw_vel:.1f}) "
f"z={trunk_z:.3f}m terrain={current_terrain}")
keyboard_reader.restore()
viewer_handle.close()
if __name__ == "__main__":
main()

371
scripts/go1_sim2sim_mujoco.py Normal file
View File

@@ -0,0 +1,371 @@
#!/usr/bin/env python3
"""
MOTRIXLAB_UNTRIEE_GO1_SIM2SIM
source /opt/mujoco/venv/bin/activate
cd /opt/unitree_mujoco
python demo/go1_sim2sim_mujoco.py
"""
import numpy as np
import mujoco
from mujoco import viewer
import os
import threading
import signal
import queue
import argparse
import time
g_exit_requested = False
def signal_handler(signum, frame):
global g_exit_requested
g_exit_requested = True
signal.signal(signal.SIGINT, signal_handler)
# ============================================================
# 配置
# ============================================================
_PROJECT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
DEFAULT_ONNX_PATH = os.path.join(_PROJECT_DIR, "exports_go1_flat", "policy.onnx")
MOTRIX_XML_DIR = os.path.join(_PROJECT_DIR, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
XML_PATH = f"{MOTRIX_XML_DIR}/go1_motor_actuator.xml"
TERRAIN = "none"
# ============================================================
# MotrixLab 参数 (来自 cfg.py)
# ============================================================
NUM_OBS = 45 # 去掉线速度观测 (原来是48)
NUM_ACTIONS = 12
OBS_SCALES = {'lin_vel': 2.0, 'ang_vel': 0.25, 'dof_pos': 1.0, 'dof_vel': 0.05}
ACTION_SCALE = 0.05
KP = 80.0
KD = 1.0
CLIP_ACTIONS = 23.7
CLIP_OBSERVATIONS = 100.0
MAX_LIN_VEL_X = 1.0
MAX_LIN_VEL_Y = 1.0
MAX_ANG_VEL = 1.0 # 匹配训练时的角速度命令范围 [-1.0, 1.0]
# ============================================================
# 关节名称和顺序
# ============================================================
POLICY_JOINT_NAMES = [
"FR_hip", "FR_thigh", "FR_calf",
"FL_hip", "FL_thigh", "FL_calf",
"RR_hip", "RR_thigh", "RR_calf",
"RL_hip", "RL_thigh", "RL_calf",
]
DEFAULT_JOINT_ANGLES = np.array([
-0.0, 0.9, -1.8, # FR_hip, FR_thigh, FR_calf
0.0, 0.9, -1.8, # FL_hip, FL_thigh, FL_calf
-0.0, 0.9, -1.8, # RR_hip, RR_thigh, RR_calf
0.0, 0.9, -1.8, # RL_hip, RL_thigh, RL_calf
], dtype=np.float32)
MUJOCO_TO_POLICY = np.arange(12, dtype=np.int64)
POLICY_TO_MUJOCO = np.arange(12, dtype=np.int64)
# ============================================================
# 键盘输入
# ============================================================
from pynput import keyboard
class KeyboardReader:
def __init__(self):
self._event_queue = queue.Queue()
self.running = True
self.shared_keys_held = set()
self.shared_one_shot = set()
self._reader_thread = None
self._listener = None
def _normalize_key(self, key):
try:
if hasattr(key, 'char') and key.char is not None:
return key.char.lower()
except:
pass
key_str = str(key)
if key_str == 'Key.esc':
return 'escape'
elif key_str == 'Key.space':
return 'space'
elif key_str.startswith('Key.'):
return key_str.lower()
return key_str.lower()
def _reader_worker(self):
while self.running:
try:
event_type, key = self._event_queue.get(timeout=0.05)
k = self._normalize_key(key)
if event_type == 'press':
self.shared_keys_held.add(k)
self.shared_one_shot.discard(k)
elif event_type == 'release':
self.shared_keys_held.discard(k)
self.shared_one_shot.discard(k)
except queue.Empty:
pass
def init(self):
def on_press(key):
self._event_queue.put(('press', key))
def on_release(key):
self._event_queue.put(('release', key))
try:
self._listener = keyboard.Listener(on_press=on_press, on_release=on_release)
self._listener.start()
self._reader_thread = threading.Thread(target=self._reader_worker, daemon=True)
self._reader_thread.start()
print("[INFO] 键盘监听已启动")
except Exception as e:
print(f"[WARN] 无法初始化键盘监听: {e}")
def is_key_pressed(self, key):
k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
if k not in self.shared_keys_held or k in self.shared_one_shot:
return False
self.shared_one_shot.add(k)
return True
def is_key_held(self, key):
k = self._normalize_key(key) if isinstance(key, str) else self._normalize_key(key)
return k in self.shared_keys_held
def restore(self):
self.running = False
if self._listener:
self._listener.stop()
# ============================================================
# Sensor 读取
# ============================================================
def get_sensor(model, data, name):
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, name)
if sid < 0:
return None
adr = model.sensor_adr[sid]
dim = model.sensor_dim[sid]
return data.sensordata[adr:adr + dim].copy()
def compute_observations_motrix(model, data, commands, last_actions):
"""计算 45 维观测 (去掉局部线速度，策略仅靠命令+关节信息+陀螺仪来推理)
布局 (45 dims):
[0:3] 陀螺仪 (ang_vel * 0.25)
[3:6] 重力向量 (躯干坐标系，无缩放)
[6:18] 关节位置偏差 (dof_pos * 1.0)
[18:30] 关节速度 (dof_vel * 0.05)
[30:42] 上一步动作 (原始值)
[42:45] 命令 [vx*2.0, vy*2.0, wz*0.25]
"""
obs = np.zeros(NUM_OBS, dtype=np.float32)
# 陀螺仪
gyro = get_sensor(model, data, "gyro")
if gyro is not None:
obs[0:3] = gyro * OBS_SCALES['ang_vel']
else:
obs[0:3] = data.qvel[3:6] * OBS_SCALES['ang_vel']
# 重力向量 (躯干坐标系)
base_rot = data.xmat[1].reshape(3, 3)
gravity_world = np.array([0., 0., -1.], dtype=np.float64)
local_gravity = base_rot.T @ gravity_world
obs[3:6] = local_gravity.astype(np.float32)
# 关节位置偏差
joint_pos = data.qpos[7:19]
dof_pos_rel = (joint_pos - DEFAULT_JOINT_ANGLES) * OBS_SCALES['dof_pos']
obs[6:18] = dof_pos_rel
# 关节速度
joint_vel = data.qvel[6:18]
obs[18:30] = joint_vel * OBS_SCALES['dof_vel']
# 上一步动作
obs[30:42] = last_actions
# 命令
obs[42:45] = commands * np.array([OBS_SCALES['lin_vel'], OBS_SCALES['lin_vel'], OBS_SCALES['ang_vel']], dtype=np.float32)
# 限幅
obs = np.clip(obs, -CLIP_OBSERVATIONS, CLIP_OBSERVATIONS)
return obs
def main():
import re
import onnxruntime as ort
parser = argparse.ArgumentParser(description="MotrixLab Go1 Policy Inference in MuJoCo")
parser.add_argument("--onnx", type=str, default=DEFAULT_ONNX_PATH)
parser.add_argument("--terrain", type=str, default=TERRAIN, choices=["none", "rough", "stairs"])
args = parser.parse_args()
os.chdir(MOTRIX_XML_DIR)
if args.terrain == "rough":
xml_file = f"{MOTRIX_XML_DIR}/scene_rough_terrain.xml"
elif args.terrain == "stairs":
xml_file = f"{MOTRIX_XML_DIR}/scene_stairs_terrain.xml"
else:
xml_file = f"{MOTRIX_XML_DIR}/scene_motor_actuator.xml" # flat floor
with open(xml_file, 'r') as f:
xml_content = f.read()
model = mujoco.MjModel.from_xml_string(xml_content)
data = mujoco.MjData(model)
print(f"[INFO] Model: {model.nbody} bodies, {model.nq} DoF, {model.nu} actuators")
print(f"[INFO] MuJoCo timestep: {model.opt.timestep}")
print(f"[INFO] 关节顺序 (qpos[7:19]): {[mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, i) for i in range(1, 13)]}")
for sensor_name in ["gyro", "local_linvel"]:
sid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, sensor_name)
print(f"[SENSOR] {sensor_name}: {'存在' if sid >= 0 else '不存在'}")
# 初始化
data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0]) # 默认四元数
data.qpos[7:19] = DEFAULT_JOINT_ANGLES
data.qvel[:] = 0.0
data.ctrl[:] = 0.0
mujoco.mj_forward(model, data)
print(f"[INIT] qpos[2]={data.qpos[2]:.3f}")
print(f"[INIT] qpos[7:19]={data.qpos[7:19]}")
print(f"[INIT] DEFAULT_JOINT_ANGLES={DEFAULT_JOINT_ANGLES}")
# 加载onnx
session = ort.InferenceSession(args.onnx, providers=['CPUExecutionProvider'])
print(f"[INFO] loaded")
# 主循环
ctrl_dt = 0.01 # 100Hz
num_steps_per_inference = int(ctrl_dt / model.opt.timestep)
print(f"[INFO] 每 {num_steps_per_inference} 步推理一次 ")
step_count = 0
inference_step = 0
x_vel_cmd = 0.0
y_vel_cmd = 0.0
yaw_vel_cmd = 0.0
commands = np.array([x_vel_cmd, y_vel_cmd, yaw_vel_cmd], dtype=np.float32)
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
# 键盘
keyboard_reader = KeyboardReader()
keyboard_reader.init()
view = viewer.launch_passive(model, data)
print("[INFO] 已启动!")
loop_start_time = time.time()
while view.is_running() and not g_exit_requested:
# 键盘命令
if keyboard_reader.is_key_pressed(' '):
x_vel_cmd = 0.0
y_vel_cmd = 0.0
yaw_vel_cmd = 0.0
if keyboard_reader.is_key_held('w'):
x_vel_cmd = MAX_LIN_VEL_X
elif keyboard_reader.is_key_held('s'):
x_vel_cmd = -MAX_LIN_VEL_X
else:
x_vel_cmd = 0.0
if keyboard_reader.is_key_held('q'):
y_vel_cmd = MAX_LIN_VEL_Y
elif keyboard_reader.is_key_held('e'):
y_vel_cmd = -MAX_LIN_VEL_Y
else:
y_vel_cmd = 0.0
if keyboard_reader.is_key_held('a'):
yaw_vel_cmd = MAX_ANG_VEL
elif keyboard_reader.is_key_held('d'):
yaw_vel_cmd = -MAX_ANG_VEL
else:
yaw_vel_cmd = 0.0
if keyboard_reader.is_key_pressed('r'):
# 重置机器人到初始位置
data.qpos[0:3] = np.array([0.0, 0.0, 0.42])
data.qpos[3:7] = np.array([1.0, 0.0, 0.0, 0.0])
data.qpos[7:19] = DEFAULT_JOINT_ANGLES
data.qvel[:] = 0.0
data.ctrl[:] = 0.0
last_actions = np.zeros(NUM_ACTIONS, dtype=np.float32)
mujoco.mj_forward(model, data)
print("[RESET] 机器人已重置")
if keyboard_reader.is_key_pressed('escape'):
break
# 推理 (每 N 步一次)
if inference_step == 0:
commands[0] = x_vel_cmd
commands[1] = y_vel_cmd
commands[2] = yaw_vel_cmd
obs = compute_observations_motrix(model, data, commands, last_actions)
# 推理
action = session.run(None, {'observations': obs.reshape(1, -1).astype(np.float32)})[0][0]
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
last_actions = action.copy()
# PD 控制
# joint_targets = action * action_scale + default_angles
joint_targets = DEFAULT_JOINT_ANGLES + action * ACTION_SCALE
current_pos = data.qpos[7:19]
current_vel = data.qvel[6:18]
torques = KP * (joint_targets - current_pos) - KD * current_vel
torques = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
data.ctrl[:] = torques
mujoco.mj_step(model, data)
view.sync()
expected_time = step_count * ctrl_dt
elapsed = time.time() - loop_start_time
sleep_time = expected_time - elapsed
if sleep_time > 0:
time.sleep(sleep_time)
step_count += 1
inference_step = (inference_step + 1) % num_steps_per_inference
if step_count % 200 == 0:
trunk_z = data.qpos[2]
lin_vel = np.linalg.norm(data.qvel[0:3])
print(f"\n========== Step {step_count} ==========")
print(f"[CMD] x={x_vel_cmd:.2f}, y={y_vel_cmd:.2f}, yaw={yaw_vel_cmd:.2f}")
print(f"[OBS] gyro={obs[0:3]}, grav={obs[3:6]}")
print(f"[ACTION] raw={action[:4]}... scaled={action[:4]*ACTION_SCALE}...")
print(f"[TARGET] {joint_targets[:4]}...")
print(f"[TORQUE] {torques[:4]}...")
print(f"[STATE] z={trunk_z:.3f}m, vel={lin_vel:.3f}m/s")
print(f"==========================================\n")
keyboard_reader.restore()
view.close()
if __name__ == "__main__":
main()

17
scripts/play.py
View File

@@ -36,6 +36,7 @@ _RAND_SEED = flags.DEFINE_bool("rand-seed", False, "Generate random seed")
_RLLIB = flags.DEFINE_string(
"rllib", None, "The RL framework (skrl/rslrl). Auto-discovered from latest training if not specified."
)
_FORCE_PHASE = flags.DEFINE_integer("force-phase", None, "Lock terrain phase (0=flat,1=rough,2=stairs,3=mixed)")
def get_inference_backend(policy_path: Path | str, rllib: str):
@@ -179,13 +180,21 @@ def main(argv):
backend = get_inference_backend(policy_path, rllib)
# Build env config overrides
env_cfg_override = {}
if _FORCE_PHASE.present:
env_cfg_override["force_phase"] = _FORCE_PHASE.value
if not env_cfg_override:
env_cfg_override = None
if rllib == "rslrl":
# RSLRL evaluation flow (always uses torch backend)
assert device_supports.torch, "PyTorch is not available on your device"
from motrix_rl.rslrl.torch.train import ppo
config.torch.backend = "torch"
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override,
enable_render=enable_render, env_cfg_override=env_cfg_override)
trainer.play(policy_path)
elif backend == "jax":
@@ -193,7 +202,8 @@ def main(argv):
from motrix_rl.skrl.jax.train import ppo
config.jax.backend = "jax" # or "numpy"
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override,
enable_render=enable_render, env_cfg_override=env_cfg_override)
trainer.play(policy_path)
elif backend == "torch":
@@ -201,7 +211,8 @@ def main(argv):
from motrix_rl.skrl.torch.train import ppo
config.torch.backend = "torch"
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override,
enable_render=enable_render, env_cfg_override=env_cfg_override)
trainer.play(policy_path)

147
scripts/play_dreamwaq.py Normal file
View File

@@ -0,0 +1,147 @@
#!/usr/bin/env python3
"""DreamWaQ play — renders env with trained policy in MotrixSim.
Usage:
uv run scripts/play_dreamwaq.py
uv run scripts/play_dreamwaq.py --num-envs 16
"""
import argparse, os, time, sys
# CRITICAL: disable JAX GPU memory preallocation BEFORE importing jax.
# Otherwise JAX grabs 75% of GPU memory and starves the MotrixSim (Vulkan)
# renderer → "Couldn't get swap chain texture" crash. Must be set first.
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import motrix_envs.locomotion.go1.dreamwaq # noqa
import motrix_rl.tasks.go1_dreamwaq # noqa
import numpy as np
import jax, jax.numpy as jnp
import pickle, msgpack
from motrix_envs import registry as env_registry
from motrix_envs.np.renderer import NpRenderer
from motrix_rl.skrl.jax.train.dreamwaq_ppo import DreamWaQWrapper, CENet
def _decode_arr(ext):
if not hasattr(ext, "code"): return None
parts = msgpack.unpackb(ext.data, raw=False)
if not isinstance(parts, list) or len(parts) < 3: return None
shape = []
def _flatten(s):
if isinstance(s, list):
for x in s: _flatten(x)
elif isinstance(s, int): shape.append(s)
_flatten(parts[0])
return np.frombuffer(parts[2], dtype=np.dtype(parts[1])).reshape(shape)
def load_params(ckpt_path):
with open(ckpt_path, 'rb') as f:
ckpt = pickle.load(f)
raw = msgpack.unpackb(ckpt['policy'])['params']
params = {}
for name, val in raw.items():
if isinstance(val, dict):
params[name] = {k: _decode_arr(v) for k, v in val.items()}
else:
params[name] = _decode_arr(val)
# State-preprocessor stats for first 64 dims (REQUIRED: policy trained on normalized obs)
mean64 = std64 = None
if 'state_preprocessor' in ckpt:
sp = msgpack.unpackb(ckpt['state_preprocessor'], raw=False)
mean64 = _decode_arr(sp['running_mean'])[:64].astype(np.float32)
std64 = np.sqrt(_decode_arr(sp['running_variance'])[:64]).astype(np.float32)
return params, mean64, std64
CLIP_ACT = 23.7
CLIP_OBS = 100.0
def policy_forward(x, p, mean64=None, std64=None):
x = jnp.array(x[:, :64])
# Apply state-preprocessor normalization (clip((x-mean)/(std+eps), -5, 5))
if mean64 is not None:
x = jnp.clip((x - jnp.array(mean64)) / (jnp.array(std64) + 1e-8), -5.0, 5.0)
else:
x = jnp.clip(x, -CLIP_OBS, CLIP_OBS)
x = jax.nn.elu(x @ jnp.array(p['Dense_0']['kernel']) + jnp.array(p['Dense_0']['bias']))
x = jax.nn.elu(x @ jnp.array(p['Dense_1']['kernel']) + jnp.array(p['Dense_1']['bias']))
x = jax.nn.elu(x @ jnp.array(p['Dense_2']['kernel']) + jnp.array(p['Dense_2']['bias']))
return np.clip(np.array(x @ jnp.array(p['Dense_3']['kernel']) + jnp.array(p['Dense_3']['bias'])),
-CLIP_ACT, CLIP_ACT)
def main():
p = argparse.ArgumentParser()
p.add_argument("--num-envs", type=int, default=9)
p.add_argument("--checkpoint", default=None)
args = p.parse_args()
# Auto-find checkpoint
if args.checkpoint is None:
run_dir = "runs/go1-dreamwaq-walk/skrl"
runs = sorted([d for d in os.listdir(run_dir) if os.path.isdir(os.path.join(run_dir, d)) and d.startswith("26-")])
args.checkpoint = os.path.join(run_dir, runs[-1], "checkpoints", "best_agent.pickle")
# Load policy + state-preprocessor normalization
policy_params, mean64, std64 = load_params(args.checkpoint)
print(f"[Play] Policy: {args.checkpoint}")
print(f"[Play] State normalization: {'ON' if mean64 is not None else 'OFF'}")
# Load VAE (saved in skrl/ base dir, not run subdir)
run_dir = os.path.dirname(os.path.dirname(os.path.dirname(args.checkpoint))) # skrl/ base
vae_path = os.path.join(run_dir, "cenet_params.pkl")
if not os.path.exists(vae_path):
vae_files = sorted([f for f in os.listdir(run_dir) if f.startswith("vae_")],
key=lambda x: int(x.split("_")[1].split(".")[0]))
if vae_files:
vae_path = os.path.join(run_dir, vae_files[-1])
with open(vae_path, 'rb') as f:
vae_params = pickle.load(f)
print(f"[Play] VAE: {vae_path}")
# Create env + renderer (like view.py)
raw_env = env_registry.make("go1-dreamwaq-walk", num_envs=args.num_envs)
renderer = NpRenderer(raw_env)
# CENet for policy inference
cenet = CENet()
rng = jax.random.PRNGKey(42)
wrapper = DreamWaQWrapper(raw_env, cenet, vae_params, rng=rng)
# Init env
raw_env.init_state()
wrapper._vae_buf = []
n = raw_env._num_envs
print(f"[Play] {n} envs, Ctrl+C to stop")
from motrixsim.render import RenderClosedError
try:
while True:
# CENet inference (mean mode)
hist = jnp.array(raw_env._state.info.get("obs_history",
np.zeros((n, 5, 45), dtype=np.float32)))
z, vel = cenet.apply(vae_params, hist, method=cenet.inference)
code = np.concatenate([np.array(vel), np.array(z)], axis=-1)
obs_arr = raw_env._state.obs
priv = raw_env._state.info.get("privileged_obs", np.zeros((n, 235), dtype=np.float32))
heights = priv[:, 48:] if priv.shape[1] > 48 else np.zeros((n, 187), dtype=np.float32)
base_vel = raw_env._state.info.get("base_vel", np.zeros((n, 3), dtype=np.float32))
base_vel_n = base_vel * np.array([2.0, 2.0, 1.0], dtype=np.float32)
aug_obs = np.concatenate([code, obs_arr, base_vel_n, heights], axis=-1)
actions = policy_forward(aug_obs, policy_params, mean64, std64)
wrapper.step(actions)
renderer.render()
time.sleep(0.01)
except (KeyboardInterrupt, RenderClosedError):
pass
try: renderer.close()
except: pass
print("[Play] Done")
if __name__ == "__main__":
main()

141
scripts/play_dreamwaq_rsl.py Normal file
View File

@@ -0,0 +1,141 @@
#!/usr/bin/env python3
"""DreamWaQ rsl_rl play — render the trained ActorCritic_DWAQ policy in NATIVE MotrixSim.
Loads the PyTorch checkpoint DIRECTLY (no ONNX). ONNX is only for cross-sim
deployment (e.g. MuJoCo sim2sim); the native MotrixSim env runs the torch policy.
Deterministic inference: mean CENet code + actor.
Usage:
uv run scripts/play_dreamwaq_rsl.py # auto-find latest, walk forward
uv run scripts/play_dreamwaq_rsl.py --checkpoint runs/.../model_1100.pt --vx 0.5
uv run scripts/play_dreamwaq_rsl.py --vx 0 --num-envs 1 # stand still, single robot
"""
import argparse, glob, os, sys, time
# avoid JAX grabbing GPU memory and starving the MotrixSim (Vulkan) renderer
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
os.environ.setdefault("JAX_PLATFORMS", "cpu")
# --terrain / --level / --flat-stairs / --stairs: pick hfield scene (before import).
terrain_type = "pyramid"
if "--flat-stairs" in sys.argv: terrain_type = "flat_stairs"
elif "--stairs" in sys.argv: terrain_type = "stairs"
if "--terrain" in sys.argv or "--level" in sys.argv or "--flat-stairs" in sys.argv or "--stairs" in sys.argv:
os.environ["DREAMWAQ_TERRAIN"] = terrain_type
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import numpy as np
import torch
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401 register env
from motrix_envs import registry as env_registry
from motrix_envs.np.renderer import NpRenderer
from motrix_rl.dwaq_rsl.actor_critic_dwaq import ActorCritic_DWAQ
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
NUM_OBS, NUM_PRIV, NUM_HIST, NUM_ACT, CENET_OUT = 45, 235, 5, 12, 19
CLIP_ACT = 23.7
def _iter_of(path):
try:
return int(os.path.basename(path).split("_")[1].split(".")[0])
except Exception:
return -1
def find_latest():
models = glob.glob(os.path.join(PROJECT, "runs", "go1-dreamwaq-walk", "rsl_dwaq", "*", "model_*.pt"))
if not models:
print("[ERROR] no rsl_dwaq checkpoints found"); sys.exit(1)
return max(models, key=_iter_of) # highest iteration (flat model_1100 > terrain early models)
def main():
p = argparse.ArgumentParser()
p.add_argument("--checkpoint", default=None)
p.add_argument("--num-envs", type=int, default=4)
p.add_argument("--vx", type=float, default=0.5, help="forward velocity command [m/s]")
p.add_argument("--vy", type=float, default=0.0, help="lateral velocity command [m/s]")
p.add_argument("--wz", type=float, default=0.0, help="yaw rate command [rad/s]")
p.add_argument("--terrain", action="store_true",
help="view the training pyramid hfield (default: flat plane)")
p.add_argument("--flat-stairs", action="store_true",
help="view the 2-level flat+stairs terrain (implies terrain)")
p.add_argument("--stairs", action="store_true",
help="view the stairs terrain scene (implies terrain)")
p.add_argument("--level", type=int, default=None,
help="force ALL spawns at this terrain level (implies --terrain)")
p.add_argument("--spawn-height", type=float, default=None,
help="spawn clearance above terrain in meters (default 0.45; try 1-2 to experiment)")
args = p.parse_args()
ckpt = args.checkpoint or find_latest()
ac = ActorCritic_DWAQ(NUM_OBS + CENET_OUT, NUM_PRIV, NUM_ACT, NUM_HIST * NUM_OBS, CENET_OUT)
ac.load_state_dict(torch.load(ckpt, map_location="cpu")["model_state_dict"])
ac.eval()
print(f"[Play-rsl] policy (native torch): {ckpt}")
env = env_registry.make("go1-dreamwaq-walk", num_envs=args.num_envs)
if args.level is not None:
env._force_level = args.level # pin all spawns to this level (read in reset)
print(f"[Play-rsl] forcing ALL spawns at terrain level {args.level}")
if args.spawn_height is not None:
env._spawn_absolute = args.spawn_height # absolute world z, no offset
print(f"[Play-rsl] spawn absolute z = {args.spawn_height}m")
renderer = NpRenderer(env)
env.init_state()
n = env._num_envs
cmd = np.array([args.vx, args.vy, args.wz], dtype=np.float32)
print(f"[Play-rsl] {n} envs | cmd=(vx={args.vx}, vy={args.vy}, wz={args.wz}) | Ctrl+C to stop")
@torch.no_grad()
def act_fn(obs, hist):
obs_t = torch.from_numpy(obs)
h = ac.encoder(torch.from_numpy(hist).reshape(obs.shape[0], -1)) # (n,225)->(n,64)
code = torch.cat([ac.encode_mean_vel(h), ac.encode_mean_latent(h)], dim=-1) # (n,19)
return ac.actor(torch.cat([code, obs_t], dim=-1)).numpy() # (n,12)
from motrixsim.render import RenderClosedError
show_heights = False
try:
while True:
if renderer._render.input.is_key_just_pressed("r"):
env.init_state()
print("[R] Reset all envs")
if renderer._render.input.is_key_just_pressed("h"):
show_heights = not show_heights
print(f"[H] Height points: {'ON' if show_heights else 'OFF'}")
env._state.info["commands"][:] = cmd
obs = env._state.obs.astype(np.float32)
hist = env._state.info.get("obs_history",
np.zeros((n, NUM_HIST, NUM_OBS), np.float32)).astype(np.float32)
act = act_fn(obs, hist)
env.step(np.clip(act, -CLIP_ACT, CLIP_ACT).astype(np.float32))
if show_heights:
from motrix_envs.math import quaternion
pose = env._body.get_pose(env._state.data)
bp = pose[0, :3]
yaw = quaternion.get_yaw(pose[0:1, 3:7])[0]
cos_y, sin_y = np.cos(yaw), np.sin(yaw)
for gy in env._hy:
for gx in env._hx:
wx = bp[0] + cos_y*gx - sin_y*gy
wy = bp[1] + sin_y*gx + cos_y*gy
wz = float(env._sample_terrain_height(np.array([[wx,wy]]))[0])
g = renderer._render.gizmos
g.draw_sphere(0.02, (np.float32(wx), np.float32(wy), np.float32(wz)))
renderer.render()
time.sleep(0.01)
except (KeyboardInterrupt, RenderClosedError):
pass
try:
renderer.close()
except Exception:
pass
print("[Play-rsl] done")
if __name__ == "__main__":
main()

430
scripts/terrain_editor.py Normal file
View File

@@ -0,0 +1,430 @@
#!/usr/bin/env python3
"""Terrain editor GUI — draw pyramids, mark spawn zones, export PNG + coordinates.
Usage:
uv run scripts/terrain_editor.py
"""
import tkinter as tk
from tkinter import ttk, messagebox
import numpy as np, os, cv2
# ═══ defaults ═══
HS = 0.05; VS = 0.005
CELL_M = 8.0; BORDER_M = 5.0
PLATFORM_M = 1.0 # platform 1m
SPAWN_RADIUS_M = 0.5 # spawn zone ±0.5m around center
DEFAULT_STEP_H = 0.20; DEFAULT_STEP_D = 0.20; DEFAULT_NUM_STEPS = 10
DEFAULT_REF_PLANE_CM = 200 # all cells start from same reference height
CELL_PX = int(CELL_M / HS); BORDER_PX = int(BORDER_M / HS)
PLATFORM_PX = int(PLATFORM_M / HS); SPAWN_RADIUS_PX = int(SPAWN_RADIUS_M / HS)
class TerrainEditor:
def __init__(self, root):
self.root = root
self.root.title("Terrain Editor")
self.rows = 2; self.cols = 4
self.cell_types = {}
self._init_defaults()
self.selected = (0, 0)
self._dragging = False
self._build_ui()
self._sync_params()
self._redraw_all()
def _init_defaults(self):
for r in range(self.rows):
for c in range(self.cols):
if r == 0:
self.cell_types[(r, c)] = {"type": "flat", "spawn": True, "level": 0,
"ref_plane_cm": DEFAULT_REF_PLANE_CM}
else:
self.cell_types[(r, c)] = {
"type": "convex" if c % 2 == 0 else "concave",
"step_h": DEFAULT_STEP_H, "step_d": DEFAULT_STEP_D,
"num_steps": DEFAULT_NUM_STEPS, "spawn": True, "level": 1,
"ref_plane_cm": DEFAULT_REF_PLANE_CM,
}
# ═══ UI ═══
def _build_ui(self):
paned = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
paned.pack(fill=tk.BOTH, expand=True)
left = ttk.Frame(paned); paned.add(left, weight=2)
right = ttk.Frame(paned); paned.add(right, weight=1)
self._build_preview_ui(left)
self._build_params_ui(right)
def _build_preview_ui(self, parent):
ttk.Label(parent, text="Terrain Preview (click to select cell, right-click toggle spawn)", font=("", 10)).pack(pady=2)
self.info_label = ttk.Label(parent, text="")
self.info_label.pack()
self.preview = tk.Canvas(parent, bg="#333", width=600, height=400)
self.preview.pack(fill=tk.BOTH, expand=True, padx=5, pady=5)
self.preview.bind("<Button-1>", self._on_click)
self.preview.bind("<B1-Motion>", self._on_drag)
self.preview.bind("<Button-3>", self._on_right_click)
ttk.Label(parent, text="Left-click: select | Right-click: toggle spawn | Drag: select").pack()
ctrl = ttk.Frame(parent)
ctrl.pack(pady=5)
ttk.Label(ctrl, text="Rows:").pack(side=tk.LEFT)
self.rows_var = tk.IntVar(value=self.rows)
ttk.Spinbox(ctrl, from_=1, to=10, width=4, textvariable=self.rows_var,
command=self._on_grid_size).pack(side=tk.LEFT, padx=2)
ttk.Label(ctrl, text="Cols:").pack(side=tk.LEFT, padx=(10,0))
self.cols_var = tk.IntVar(value=self.cols)
ttk.Spinbox(ctrl, from_=1, to=10, width=4, textvariable=self.cols_var,
command=self._on_grid_size).pack(side=tk.LEFT, padx=2)
ttk.Button(parent, text="Export PNG + Coords", command=self._export).pack(pady=5)
def _build_params_ui(self, parent):
f = ttk.Frame(parent); f.pack(padx=10, pady=5, fill=tk.X)
ttk.Label(f, text="Cell Type:").grid(row=0, column=0, sticky=tk.W)
self.type_var = tk.StringVar(value="flat")
ttk.Combobox(f, textvariable=self.type_var, values=["flat", "convex", "concave"],
state="readonly", width=10).grid(row=0, column=1, padx=5)
self.type_var.trace("w", lambda *a: self._on_param_change())
ttk.Label(f, text="Level:").grid(row=0, column=2, sticky=tk.W, padx=(20,0))
self.level_var = tk.IntVar(value=0)
ttk.Spinbox(f, from_=0, to=9, width=3, textvariable=self.level_var,
command=self._on_param_change).grid(row=0, column=3)
self.spawn_var = tk.BooleanVar(value=True)
ttk.Checkbutton(f, text="Spawn", variable=self.spawn_var,
command=self._on_param_change).grid(row=0, column=4, padx=10)
ttk.Separator(parent, orient=tk.HORIZONTAL).pack(fill=tk.X, pady=5, padx=10)
ttk.Label(parent, text="Pyramid Params").pack()
g = ttk.Frame(parent); g.pack(padx=10, pady=5, fill=tk.X)
ttk.Label(g, text="Step height (m):").grid(row=0, column=0, sticky=tk.W)
self.sh_var = tk.StringVar(value=str(DEFAULT_STEP_H))
ttk.Entry(g, textvariable=self.sh_var, width=7).grid(row=0, column=1, padx=5)
self.sh_var.trace("w", lambda *a: self._on_param_change())
ttk.Label(g, text="Step tread (m):").grid(row=1, column=0, sticky=tk.W)
self.sd_var = tk.StringVar(value=str(DEFAULT_STEP_D))
ttk.Entry(g, textvariable=self.sd_var, width=7).grid(row=1, column=1, padx=5)
self.sd_var.trace("w", lambda *a: self._on_param_change())
ttk.Label(g, text="Num steps:").grid(row=2, column=0, sticky=tk.W)
self.ns_var = tk.StringVar(value=str(DEFAULT_NUM_STEPS))
ttk.Entry(g, textvariable=self.ns_var, width=7).grid(row=2, column=1, padx=5)
ttk.Label(g, text="Ref plane (cm):").grid(row=3, column=0, sticky=tk.W)
self.ref_var = tk.StringVar(value=str(DEFAULT_REF_PLANE_CM))
ttk.Entry(g, textvariable=self.ref_var, width=7).grid(row=3, column=1, padx=5)
self.ns_var.trace("w", lambda *a: self._on_param_change())
ttk.Separator(parent, orient=tk.HORIZONTAL).pack(fill=tk.X, pady=5, padx=10)
ttk.Label(parent, text="Selected Cell").pack()
self.cell_label = ttk.Label(parent, text="")
self.cell_label.pack()
# ═══ events ═══
def _cell_at(self, ex, ey):
m = 5; pw = self.preview.winfo_width(); ph = self.preview.winfo_height()
cw = (pw - 2*m) // max(self.cols,1); ch = (ph - 2*m) // max(self.rows,1)
col = (ex - m) // cw; row = (ey - m) // ch
if 0 <= col < self.cols and 0 <= row < self.rows:
return row, col, m + col*cw, m + row*ch, cw, ch
return None
def _on_click(self, event):
v = self._cell_at(event.x, event.y)
if v:
self.selected = (v[0], v[1])
self._sync_params(); self._redraw_all()
def _on_drag(self, event):
v = self._cell_at(event.x, event.y)
if v:
self.selected = (v[0], v[1])
self._sync_params(); self._redraw_all()
def _on_right_click(self, event):
v = self._cell_at(event.x, event.y)
if v:
r, c = v[0], v[1]
ct = self.cell_types.get((r, c), {"type": "flat", "spawn": True, "level": r})
ct = dict(ct) # copy before modifying
ct["spawn"] = not ct.get("spawn", True)
self.cell_types[(r, c)] = ct
if (r, c) == self.selected:
self._sync_params()
self._redraw_all()
def _on_grid_size(self):
try: nr = self.rows_var.get()
except: nr = self.rows
try: nc = self.cols_var.get()
except: nc = self.cols
if nr == self.rows and nc == self.cols: return
old = self.cell_types
self.rows, self.cols = nr, nc
self.cell_types = {}
for r in range(nr):
for c in range(nc):
self.cell_types[(r,c)] = old.get((r,c), {"type": "flat", "spawn": True, "level": r})
self._redraw_all()
# ═══ sync ═══
def _sync_params(self):
ct = self.cell_types.get(self.selected, {"type": "flat", "spawn": True, "level": self.selected[0]})
self.type_var.set(ct.get("type", "flat"))
self.spawn_var.set(ct.get("spawn", True))
r = self.selected[0]
self.level_var.set(ct.get("level", r))
self.ref_var.set(str(ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)))
self.sh_var.set(str(ct.get("step_h", DEFAULT_STEP_H)))
self.sd_var.set(str(ct.get("step_d", DEFAULT_STEP_D)))
self.ns_var.set(str(ct.get("num_steps", DEFAULT_NUM_STEPS)))
r, c = self.selected
half_x = BORDER_M + self.cols * CELL_M / 2
half_y = BORDER_M + self.rows * CELL_M / 2
cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
self.cell_label.config(text=f"({r},{c}) center: x={cx:+.1f} y={cy:+.1f} type={ct['type']}")
def _on_param_change(self):
r, c = self.selected
try: sh = float(self.sh_var.get()); sd = float(self.sd_var.get())
except: return
try: ns = int(self.ns_var.get())
except: return
try: ref_cm = float(self.ref_var.get())
except: ref_cm = DEFAULT_REF_PLANE_CM
cell = {"type": self.type_var.get(), "spawn": self.spawn_var.get(),
"level": self.level_var.get(), "ref_plane_cm": ref_cm}
if cell["type"] != "flat":
cell.update({"step_h": sh, "step_d": sd, "num_steps": ns})
self.cell_types[(r, c)] = cell
self._redraw_all()
# ═══ draw ═══
def _redraw_all(self):
w = self.preview.winfo_width(); h = self.preview.winfo_height()
if w < 10: w = 600
if h < 10: h = 400
self._draw_preview(w, h)
half_x = BORDER_M + self.cols * CELL_M / 2
half_y = BORDER_M + self.rows * CELL_M / 2
self.info_label.config(
text=f"{self.rows}×{self.cols} "
f"{self.cols*CELL_M+2*BORDER_M:.0f}×{self.rows*CELL_M+2*BORDER_M:.0f}m "
f"spawn_cy = {half_y-BORDER_M-CELL_M/2:.0f} - row*{CELL_M:.0f}")
def _draw_preview(self, pw, ph):
cv = self.preview; cv.delete("all")
m = 5; cw = (pw - 2*m) // max(self.cols, 1); ch = (ph - 2*m) // max(self.rows, 1)
cw = max(cw, 30); ch = max(ch, 30)
colors = {"flat": "#5b8c5a", "convex": "#c0392b", "concave": "#2471a3"}
for r in range(self.rows):
for c in range(self.cols):
x1, y1 = m + c*cw, m + r*ch
x2, y2 = x1 + cw, y1 + ch
ct = self.cell_types.get((r,c), {"type": "flat", "spawn": True, "level": r})
cv.create_rectangle(x1, y1, x2, y2, fill=colors.get(ct["type"], "#555"),
outline="#888", width=1)
# Cell center dot
cx = (x1+x2)//2; cy = (y1+y2)//2
cv.create_oval(cx-3, cy-3, cx+3, cy+3, fill="white", outline="")
# Pyramid stairs rings
if ct["type"] != "flat":
sh = ct.get("step_h", DEFAULT_STEP_H)
sd = ct.get("step_d", DEFAULT_STEP_D)
ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
concave = ct["type"] == "concave"
p2 = max(2, cw // 16)
step_px = max(1, (cw//2 - p2) // max(ns, 1))
h_max = int(sh * ns / VS)
for i in range(ns + 1):
half = p2 + (ns - i) * step_px
if concave:
frac = (ns - i) / max(ns, 1)
else:
frac = i / max(ns, 1)
g = int(180 - frac * 100)
clr = f"#{g:02x}{g:02x}{g:02x}"
cv.create_rectangle(cx - half, cy - half, cx + half, cy + half,
fill=clr, outline="")
# Spawn zone (green rect)
if ct.get("spawn", True):
sz = max(2, int(SPAWN_RADIUS_M / CELL_M * cw))
cv.create_rectangle(cx - sz, cy - sz, cx + sz, cy + sz,
outline="#00ff00", width=2)
# Level + type label
lvl = ct.get("level", r)
lbl = f"L{lvl} {ct['type'][:3]}"
if ct["type"] != "flat":
lbl = f"L{lvl} {ct['type'][:3]}-{sh*100:.0f}cm"
cv.create_text(x1 + 20, y1 + 10, text=lbl, fill="white",
font=("", 8), anchor=tk.NW)
# Highlight selected cell
r, c = self.selected
x1, y1 = m + c*cw, m + r*ch
x2, y2 = x1 + cw, y1 + ch
cv.create_rectangle(x1, y1, x2, y2, outline="yellow", width=3)
# Level labels on right
for r in range(self.rows):
y = m + r*ch + ch//2
cv.create_text(pw - 15, y, text=f"L{r}", fill="white", font=("", 12, "bold"))
# ═══ generate + export ═══
def _generate_png(self):
tot_rows = self.rows * CELL_PX + 2 * BORDER_PX
tot_cols = self.cols * CELL_PX + 2 * BORDER_PX
canvas = np.zeros((tot_rows, tot_cols), dtype=np.uint16)
for r in range(self.rows):
for c in range(self.cols):
x0 = BORDER_PX + c * CELL_PX; y0 = BORDER_PX + r * CELL_PX
ct = self.cell_types.get((r,c), {"type": "flat", "level": r})
if ct["type"] == "flat": continue
sh = ct.get("step_h", DEFAULT_STEP_H)
sd = ct.get("step_d", DEFAULT_STEP_D)
ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
concave = ct["type"] == "concave"
ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
ref_vs = int(ref_cm / 100.0 / VS)
h_vs = int(sh / VS); d_px = int(sd / HS)
p2 = PLATFORM_PX // 2
cx = x0 + CELL_PX // 2; cy = y0 + CELL_PX // 2
cv2.rectangle(canvas, (x0, y0), (x0+CELL_PX, y0+CELL_PX), int(ref_vs), -1)
for i in range(ns + 1):
half = p2 + (ns - i) * d_px
x1, y1 = cx - half, cy - half; x2, y2 = cx + half, cy + half
if concave:
h = ref_vs - h_vs * i
else:
h = ref_vs + h_vs * i
cv2.rectangle(canvas, (x1, y1), (x2, y2), int(h), -1)
hf_m = canvas.astype(np.float32) * VS
z_min, z_max = float(hf_m.min()), float(hf_m.max())
z_range = max(z_max - z_min, 0.001)
png = ((hf_m - z_min) / z_range * 65535).astype(np.uint16)
return png, z_range, z_min
def _export(self):
png, z_range, z_min = self._generate_png()
out_dir = os.path.join(os.path.dirname(__file__), "..",
"motrix_envs", "src", "motrix_envs", "locomotion",
"go1", "xmls", "assets")
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, "flat_stairs.png")
cv2.imwrite(out_path, png)
w_m = (self.cols * CELL_PX + 2 * BORDER_PX) * HS
h_m = (self.rows * CELL_PX + 2 * BORDER_PX) * HS
half_x = BORDER_M + self.cols * CELL_M / 2
half_y = BORDER_M + self.rows * CELL_M / 2
lines = [
f"# Terrain: {self.rows}×{self.cols} {w_m:.0f}×{h_m:.0f}m",
f"XML: size=\"{w_m/2:.1f} {h_m/2:.1f} {z_range:.3f} {max(z_min,0.001):.3f}\"",
f"dreamwaq.py: terrain_rows={self.rows} terrain_cols={self.cols}",
f"",
f"# === Cell centers ===",
]
for r in range(self.rows):
for c in range(self.cols):
cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
ct = self.cell_types.get((r,c), {"type": "flat", "level": r})
lines.append(f" ({r},{c}): x={cx:+.1f} y={cy:+.1f} {ct['type']}")
lines.append("")
lines.append("# === Level boundaries (robot out of bounds → reset) ===")
level_bounds = {}
for r in range(self.rows):
for c in range(self.cols):
lv = self.cell_types.get((r,c), {"level": r})["level"]
if lv not in level_bounds:
level_bounds[lv] = {"rmin": r, "rmax": r, "cmin": c, "cmax": c}
else:
b = level_bounds[lv]
b["rmin"] = min(b["rmin"], r)
b["rmax"] = max(b["rmax"], r)
b["cmin"] = min(b["cmin"], c)
b["cmax"] = max(b["cmax"], c)
for lv in sorted(level_bounds):
b = level_bounds[lv]
x_min = -half_x + BORDER_M + b["cmin"] * CELL_M
x_max = -half_x + BORDER_M + (b["cmax"] + 1) * CELL_M
y_min = half_y - BORDER_M - (b["rmax"] + 1) * CELL_M
y_max = half_y - BORDER_M - b["rmin"] * CELL_M
lines.append(f" level {lv}: x=[{x_min:+.1f}, {x_max:+.1f}] "
f"y=[{y_min:+.1f}, {y_max:+.1f}] "
f"({b['rmax']-b['rmin']+1}×{b['cmax']-b['cmin']+1} cells)")
lines.append("")
lines.append("# === Spawn positions ===")
for lv in sorted(level_bounds):
b = level_bounds[lv]
spawn_cells = [(r,c) for r in range(b["rmin"], b["rmax"]+1)
for c in range(b["cmin"], b["cmax"]+1)
if self.cell_types.get((r,c), {}).get("spawn", True)]
if spawn_cells:
lines.append(f" level {lv}: {len(spawn_cells)} spawn cells")
for (rr, cc) in spawn_cells:
cx = -half_x + BORDER_M + cc * CELL_M + CELL_M / 2
cy = half_y - BORDER_M - rr * CELL_M - CELL_M / 2
ct = self.cell_types.get((rr,cc), {})
if ct.get("type") == "flat":
z_plat = 0
else:
ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
sh = ct.get("step_h", 0)
ns = ct.get("num_steps", 0)
concave = ct["type"] == "concave"
z_plat = (ref_cm - ns * sh * 100) / 100.0 if concave else (ref_cm + ns * sh * 100) / 100.0
lines.append(f" ({rr},{cc}) x={cx:+.1f} y={cy:+.1f} {ct['type']} "
f"z_plat={z_plat*100:.0f}cm")
lines.append("")
lines.append("# === Pyramid tread details ===")
for r in range(self.rows):
for c in range(self.cols):
ct = self.cell_types.get((r,c), {"type": "flat"})
if ct["type"] == "flat":
continue
cx = -half_x + BORDER_M + c * CELL_M + CELL_M / 2
cy = half_y - BORDER_M - r * CELL_M - CELL_M / 2
sh = ct.get("step_h", DEFAULT_STEP_H)
sd = ct.get("step_d", DEFAULT_STEP_D)
ns = ct.get("num_steps", DEFAULT_NUM_STEPS)
concave = ct["type"] == "concave"
ref_cm = ct.get("ref_plane_cm", DEFAULT_REF_PLANE_CM)
ref_z = ref_cm / 100.0
h_vs = int(sh / VS)
p2 = PLATFORM_PX // 2
d_px = int(sd / HS)
plat_z = (ref_cm - ns * sh * 100) / 100.0 if concave else (ref_cm + ns * sh * 100) / 100.0
lines.append(f" ({r},{c}) {ct['type']} center=({cx:+.1f}, {cy:+.1f}) "
f"ref_plane={ref_z*100:.0f}cm platform={plat_z*100:.0f}cm "
f"step_h={sh*100:.0f}cm tread={sd*100:.0f}cm steps={ns}")
for i in range(ns + 1):
half_m = (p2 + (ns - i) * d_px) * HS
if concave:
z = ref_z - (h_vs * i) * VS
else:
z = ref_z + (h_vs * i) * VS
ring_type = "platform" if i == ns else "ring"
lines.append(f" {ring_type} {i}: z={z*100:5.0f}cm "
f"half={half_m:.2f}m "
f"x=[{cx-half_m:+.1f},{cx+half_m:+.1f}] "
f"y=[{cy-half_m:+.1f},{cy+half_m:+.1f}]")
lines.append("")
lines.append(f"# Training: DREAMWAQ_TERRAIN=flat_stairs "
f"uv run scripts/train_dreamwaq_rsl.py --level N")
info = "\n".join(lines)
print(info)
messagebox.showinfo("Exported", f"{out_path}\n\n{info}")
if __name__ == "__main__":
root = tk.Tk()
root.geometry("900x550")
TerrainEditor(root)
root.mainloop()

23
scripts/train.py
View File

@@ -35,6 +35,9 @@ _TRAIN_BACKEND = flags.DEFINE_string("train-backend", None, "The learning backen
_SEED = flags.DEFINE_integer("seed", None, "Random seed for reproducibility")
_RAND_SEED = flags.DEFINE_bool("rand-seed", False, "Generate random seed")
_RLLIB = flags.DEFINE_string("rllib", "skrl", "The RL framework (skrl/rslrl)")
_CHECKPOINT = flags.DEFINE_string("checkpoint", None, "Resume training from a checkpoint (.pickle/.pt)")
_FORCE_PHASE = flags.DEFINE_integer("force-phase", None, "Lock curriculum to a specific phase (0=flat,1=rough,2=stairs,3=mixed)")
_TRACKING_LINVEL_SCALE = flags.DEFINE_float("tracking-linvel-scale", None, "Override tracking_lin_vel reward scale")
def get_train_backend(supports: utils.DeviceSupports, train_backend_arg: str | None, rllib: str):
@@ -104,6 +107,15 @@ def main(argv):
# Determine the training backend
train_backend = get_train_backend(device_supports, _TRAIN_BACKEND.value, rllib)
# Build env config overrides from command-line flags
env_cfg_override = {}
if _FORCE_PHASE.present:
env_cfg_override["force_phase"] = _FORCE_PHASE.value
if _TRACKING_LINVEL_SCALE.present:
env_cfg_override["tracking_lin_vel_scale"] = _TRACKING_LINVEL_SCALE.value
if not env_cfg_override:
env_cfg_override = None
trainer = None
if rllib == "rslrl":
# RSLRL training flow
@@ -111,22 +123,25 @@ def main(argv):
assert train_backend == "torch", "RSLRL only supports PyTorch backend"
from motrix_rl.rslrl.torch.train import ppo
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override,
enable_render=enable_render, env_cfg_override=env_cfg_override)
elif train_backend == "jax":
from motrix_rl.skrl.jax.train import ppo
config.jax.backend = "jax" # or "numpy"
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override,
enable_render=enable_render, env_cfg_override=env_cfg_override)
elif train_backend == "torch":
from motrix_rl.skrl.torch.train import ppo
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override,
enable_render=enable_render, env_cfg_override=env_cfg_override)
else:
raise Exception(f"Unknown train backend: {train_backend}")
trainer.train()
trainer.train(checkpoint=_CHECKPOINT.value)
if __name__ == "__main__":

42
scripts/train_cts.py Normal file
View File

@@ -0,0 +1,42 @@
#!/usr/bin/env python3
"""Train CTS (Concurrent Teacher-Student) Go1 locomotion.
Usage:
uv run scripts/train_cts.py
uv run scripts/train_cts.py --num-envs 512
"""
import logging
import motrix_rl.tasks.go1_go2style # noqa: triggers env + rlcfg registration
from absl import app, flags
from skrl import config as skrl_config
from motrix_rl import utils
from motrix_rl.skrl.jax.train.cts_ppo import CTSTrainer
logger = logging.getLogger(__name__)
_ENV = flags.DEFINE_string("env", "go1-cts-flat-walk-go2style", "CTS env to train")
_NUM_ENVS = flags.DEFINE_integer("num-envs", 1024, "Number of environments")
_SEED = flags.DEFINE_integer("seed", None, "Random seed")
def main(argv):
supports = utils.get_device_supports()
logger.info(supports)
env_name = _ENV.value
override = {}
if _NUM_ENVS.present:
override["num_envs"] = _NUM_ENVS.value
if _SEED.present:
override["runner.seed"] = _SEED.value
skrl_config.jax.backend = "jax"
trainer = CTSTrainer(env_name=env_name, cfg_override=override)
trainer.train()
if __name__ == "__main__":
app.run(main)

42
scripts/train_dreamwaq.py Normal file
View File

@@ -0,0 +1,42 @@
#!/usr/bin/env python3
"""DreamWaQ training — Manaro-Alpha aligned.
Usage:
uv run scripts/train_dreamwaq.py # default (2048 envs, 100M steps)
uv run scripts/train_dreamwaq.py --num-envs 4096 --timesteps 150M
"""
import argparse
# Register env + config
import motrix_envs.locomotion.go1.dreamwaq # noqa
import motrix_rl.tasks.go1_dreamwaq # noqa
from motrix_rl.skrl.jax.train.dreamwaq_ppo import DreamWaQTrainer
def main():
p = argparse.ArgumentParser()
p.add_argument("--num-envs", type=int, default=2048)
p.add_argument("--timesteps", type=str, default="100M")
p.add_argument("--seed", type=int, default=42)
args = p.parse_args()
ts = args.timesteps
if ts.endswith("M"): ts = int(float(ts[:-1]) * 1_000_000)
elif ts.endswith("K"): ts = int(float(ts[:-1]) * 1_000)
else: ts = int(ts)
# SKRL timesteps = env.step() calls, NOT individual env steps
skrl_ts = ts // args.num_envs
override = {
"num_envs": args.num_envs,
"runner.seed": args.seed,
"runner.trainer.timesteps": skrl_ts,
}
trainer = DreamWaQTrainer(env_name="go1-dreamwaq-walk", cfg_override=override)
trainer.train()
if __name__ == "__main__":
main()

122
scripts/train_dreamwaq_rsl.py Normal file
View File

@@ -0,0 +1,122 @@
#!/usr/bin/env python3
"""DreamWaQ training via the faithful rsl_rl-1.0.2 port (upstream-aligned).
Uses upstream's ActorCritic_DWAQ + PPO (joint VAE training) + OnPolicyRunner,
with MotrixLab's DreamWaQ env. Config matches upstream Go1RoughCfgPPO.
Usage:
uv run scripts/train_dreamwaq_rsl.py --num-envs 2048 --iterations 3000
"""
import argparse
import os
import sys
import torch
import motrix_envs.locomotion.go1.dreamwaq # noqa: F401 (registers env)
from motrix_envs import registry as env_registry
from motrix_rl.dwaq_rsl import OnPolicyRunner, DwaqVecEnv
# Upstream Go1RoughCfgPPO (legged_robot_config.py LeggedRobotCfgPPO + Go1 overrides)
TRAIN_CFG = {
"runner": {
"policy_class_name": "ActorCritic_DWAQ",
"algorithm_class_name": "PPO",
"num_steps_per_env": 24,
"save_interval": 50,
},
"algorithm": {
"value_loss_coef": 1.0,
"use_clipped_value_loss": True,
"clip_param": 0.2,
"entropy_coef": 0.01,
"num_learning_epochs": 5,
"num_mini_batches": 4,
"learning_rate": 1.0e-3,
"schedule": "adaptive",
"gamma": 0.99,
"lam": 0.95,
"desired_kl": 0.01,
"max_grad_norm": 1.0,
},
"policy": {
"init_noise_std": 1.0,
},
}
def main():
p = argparse.ArgumentParser()
p.add_argument("--num-envs", type=int, default=2048)
p.add_argument("--iterations", type=int, default=3000)
p.add_argument("--seed", type=int, default=1)
p.add_argument("--init-noise-std", type=float, default=1.0)
p.add_argument("--force-std", action="store_true",
help="force reset action std to --init-noise-std even when resuming")
p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
p.add_argument("--resume", default=None,
help="checkpoint .pt to warm-start from (e.g. the flat policy before "
"the pyramid-terrain curriculum). Loads model + optimizer; the "
"action std comes from the checkpoint, not --init-noise-std.")
p.add_argument("--level", type=int, default=None,
help="force ALL envs to this terrain level (0-9), skip curriculum")
args = p.parse_args()
torch.manual_seed(args.seed)
import numpy as np
np.random.seed(args.seed)
TRAIN_CFG["policy"]["init_noise_std"] = args.init_noise_std
raw_env = env_registry.make("go1-dreamwaq-walk", num_envs=args.num_envs)
if args.level is not None:
raw_env._force_level = args.level
print(f"[DreamWaQ-rsl] forcing ALL envs at terrain level {args.level}")
env = DwaqVecEnv(raw_env, device=args.device)
print(f"[DreamWaQ-rsl] {args.num_envs} envs | obs={env.num_obs} priv={env.num_privileged_obs} "
f"hist={env.num_obs_hist} act={env.num_actions} | device={args.device}")
print(f"[DreamWaQ-rsl] actor_in={env.num_obs + 19} critic_in={env.num_privileged_obs} "
f"cenet_in={env.num_obs_hist * env.num_obs}")
log_dir = os.path.join("runs", "go1-dreamwaq-walk", "rsl_dwaq",
__import__("datetime").datetime.now().strftime("%m-%d_%H-%M-%S"))
os.makedirs(log_dir, exist_ok=True)
# Save config snapshot for later reference
import json
cfg_snapshot = {
"command_line": sys.argv,
"kp": raw_env.cfg.control_config.stiffness,
"kd": raw_env.cfg.control_config.damping,
"action_scale": raw_env.cfg.control_config.action_scale,
"rewards": dict(raw_env.cfg.reward_config.scales),
"sigma": raw_env.cfg.reward_config.tracking_sigma,
"only_positive": raw_env.cfg.reward_config.only_positive_rewards,
"force_level": args.level,
"init_noise_std": TRAIN_CFG["policy"]["init_noise_std"],
"entropy_coef": TRAIN_CFG["algorithm"]["entropy_coef"],
"terrain_rows": raw_env.cfg.terrain_rows,
"terrain_cols": raw_env.cfg.terrain_cols,
"scene": raw_env.cfg.model_file,
}
with open(os.path.join(log_dir, "config.json"), "w") as f:
json.dump(cfg_snapshot, f, indent=2, default=str)
runner = OnPolicyRunner(env, TRAIN_CFG, log_dir=log_dir, device=args.device)
if args.resume:
runner.load(args.resume)
runner.current_learning_iteration = 0
if args.force_std:
runner.alg.actor_critic.std.data.fill_(args.init_noise_std)
print(f"[DreamWaQ-rsl] warm-start from {args.resume} "
f"(std FORCED to {args.init_noise_std})")
else:
print(f"[DreamWaQ-rsl] warm-start from {args.resume} "
f"(model+optimizer; std from checkpoint)")
print(f"[DreamWaQ-rsl] log_dir={log_dir} | training {args.iterations} iterations...")
runner.learn(args.iterations, init_at_random_ep_len=True)
print("[DreamWaQ-rsl] done.")
if __name__ == "__main__":
main()

95
scripts/train_go2style.py Normal file
View File

@@ -0,0 +1,95 @@
#!/usr/bin/env python3
"""Train Go1 go2style flat-terrain locomotion.
Usage:
uv run scripts/train_go2style.py
uv run scripts/train_go2style.py --rllib rslrl
"""
import logging
# IMPORTANT: trigger registration of go2style env + rl config
import motrix_rl.tasks.go1_go2style # noqa: F401
# Now run the standard training pipeline
from absl import app, flags
from skrl import config
from motrix_rl import utils
logger = logging.getLogger(__name__)
_ENV = flags.DEFINE_string("env", "go1-flat-terrain-walk-go2style", "The env to train")
_SIM_BACKEND = flags.DEFINE_string("sim-backend", None, "Simulation backend")
_NUM_ENVS = flags.DEFINE_integer("num-envs", 2048, "Number of envs")
_RENDER = flags.DEFINE_bool("render", False, "Render the env")
_TRAIN_BACKEND = flags.DEFINE_string("train-backend", None, "learning backend (jax/torch)")
_SEED = flags.DEFINE_integer("seed", None, "Random seed")
_RAND_SEED = flags.DEFINE_bool("rand-seed", False, "Generate random seed")
_RLLIB = flags.DEFINE_string("rllib", "skrl", "RL framework (skrl/rslrl)")
def get_train_backend(supports, train_backend_arg, rllib):
if rllib == "rslrl":
if train_backend_arg is not None and train_backend_arg != "torch":
raise Exception("RSLRL only supports PyTorch backend.")
if not supports.torch:
raise Exception("RSLRL requires PyTorch.")
return "torch"
if train_backend_arg is not None:
backend = train_backend_arg
if backend == "jax" and not supports.jax:
raise Exception("JAX not available.")
if backend == "torch" and not supports.torch:
raise Exception("PyTorch not available.")
return backend
if supports.jax and supports.jax_gpu:
return "jax"
elif supports.torch and supports.torch_gpu:
return "torch"
elif supports.jax:
return "jax"
elif supports.torch:
return "torch"
else:
raise Exception("Neither JAX nor PyTorch available.")
def main(argv):
device_supports = utils.get_device_supports()
logger.info(device_supports)
env_name = _ENV.value
enable_render = _RENDER.value
rl_override = {}
if _NUM_ENVS.present:
rl_override["num_envs"] = _NUM_ENVS.value
if _RAND_SEED.value:
rl_override["runner.seed"] = None
elif _SEED.present:
rl_override["runner.seed"] = _SEED.value
sim_backend = _SIM_BACKEND.value
rllib = _RLLIB.value
train_backend = get_train_backend(device_supports, _TRAIN_BACKEND.value, rllib)
if rllib == "rslrl":
assert device_supports.torch
from motrix_rl.rslrl.torch.train import ppo
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
elif train_backend == "jax":
from motrix_rl.skrl.jax.train import ppo
config.jax.backend = "jax"
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
elif train_backend == "torch":
from motrix_rl.skrl.torch.train import ppo
config.torch.backend = "torch"
trainer = ppo.Trainer(env_name, sim_backend, cfg_override=rl_override, enable_render=enable_render)
else:
raise Exception(f"Unknown train backend: {train_backend}")
trainer.train()
if __name__ == "__main__":
app.run(main)

135
scripts/view_go2style.py Normal file
View File

@@ -0,0 +1,135 @@
#!/usr/bin/env python3
"""MuJoCo sim2sim visualization for go2style policy (45-dim obs, no linvel).
Controls:
W/S: forward/back Q/E: left/right A/D: rotate Space: stop R: reset
"""
import numpy as np
import mujoco
from mujoco import viewer
import onnxruntime as ort
import os, sys, time, threading, queue
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ONNX_PATH = os.path.join(PROJECT, "exports_go1_go2style", "policy.onnx")
XML_DIR = os.path.join(PROJECT, "motrix_envs", "src", "motrix_envs", "locomotion", "go1", "xmls")
# go2style params
NUM_OBS = 45
NUM_ACTIONS = 12
OBS_SCALES = {'ang_vel': 0.25, 'dof_pos': 1.0, 'dof_vel': 0.05}
ACTION_SCALE = 0.25
KP = 20.0
KD = 0.0 # MuJoCo joint自带damping=0.5, PD kd=0避免过阻尼
CLIP_ACTIONS = 23.7
CLIP_OBS = 100.0
MAX_VX, MAX_VY, MAX_WZ = 1.0, 1.0, 1.0
DEFAULT_ANGLES = np.array([
-0.0, 0.9, -1.8, 0.0, 0.9, -1.8, -0.0, 0.9, -1.8, 0.0, 0.9, -1.8,
], dtype=np.float32)
from pynput import keyboard
class KB:
def __init__(self):
self._q = queue.Queue(); self.running = True
self.held = set(); self._t = None; self._l = None
def _n(self, k):
try:
if hasattr(k,'char') and k.char: return k.char.lower()
except: pass
return str(k).lower()
def _w(self):
while self.running:
try:
et, k = self._q.get(timeout=0.05)
n = self._n(k)
if et == 'press': self.held.add(n)
elif et == 'release': self.held.discard(n)
except queue.Empty: pass
def init(self):
def op(k): self._q.put(('press',k))
def or_(k): self._q.put(('release',k))
self._l = keyboard.Listener(on_press=op, on_release=or_)
self._l.start()
self._t = threading.Thread(target=self._w, daemon=True); self._t.start()
print("[KB] 键盘就绪")
def held_keys(self): return self.held.copy()
def stop(self): self.running = False; self._l.stop()
def get_sensor(m, d, name):
sid = mujoco.mj_name2id(m, mujoco.mjtObj.mjOBJ_SENSOR, name)
if sid < 0: return None
adr = m.sensor_adr[sid]; dim = m.sensor_dim[sid]
return d.sensordata[adr:adr+dim].copy()
def compute_obs(model, data, commands, last_action):
obs = np.zeros(NUM_OBS, dtype=np.float32)
# gyro [0:3]
g = get_sensor(model, data, "gyro")
obs[0:3] = (g if g is not None else data.qvel[3:6]) * OBS_SCALES['ang_vel']
# gravity [3:6]
R = data.xmat[1].reshape(3,3)
obs[3:6] = (R.T @ np.array([0.,0.,-1.])).astype(np.float32)
# joint pos [6:18]
obs[6:18] = (data.qpos[7:19] - DEFAULT_ANGLES) * OBS_SCALES['dof_pos']
# joint vel [18:30]
obs[18:30] = data.qvel[6:18] * OBS_SCALES['dof_vel']
# last action [30:42]
obs[30:42] = last_action
# commands [42:45]
obs[42:45] = commands * np.array([2.0, 2.0, 0.25], dtype=np.float32)
return np.clip(obs, -CLIP_OBS, CLIP_OBS)
def main():
os.chdir(XML_DIR)
xml = open("scene_motor_actuator.xml").read()
model = mujoco.MjModel.from_xml_string(xml)
data = mujoco.MjData(model)
data.qpos[0:3] = [0,0,0.42]; data.qpos[3:7] = [1,0,0,0]; data.qpos[7:19] = DEFAULT_ANGLES
mujoco.mj_forward(model, data)
session = ort.InferenceSession(ONNX_PATH, providers=['CPUExecutionProvider'])
print(f"[ONNX] {ONNX_PATH}")
print(f"[CTRL] W/S前后 Q/E左右 A/D旋转 Space停 R重置 Esc退出")
kb = KB(); kb.init()
view = viewer.launch_passive(model, data)
step, vx, vy, wz = 0, 0.0, 0.0, 0.0
last_action = np.zeros(NUM_ACTIONS, dtype=np.float32)
action = np.zeros(NUM_ACTIONS, dtype=np.float32)
decimation = 2 # MuJoCo dt=0.005, policy dt=0.01 → 2 steps per inference
while view.is_running():
keys = kb.held_keys()
if 'escape' in keys: break
if 'r' in keys:
data.qpos[0:3]=[0,0,0.42]; data.qpos[3:7]=[1,0,0,0]; data.qpos[7:19]=DEFAULT_ANGLES
data.qvel[:]=0; last_action[:]=0; mujoco.mj_forward(model,data); print("[R] 重置")
if ' ' in keys: vx=vy=wz=0.0
vx = MAX_VX if 'w' in keys else (-MAX_VX if 's' in keys else 0.0)
vy = MAX_VY if 'q' in keys else (-MAX_VY if 'e' in keys else 0.0)
wz = MAX_WZ if 'a' in keys else (-MAX_WZ if 'd' in keys else 0.0)
if step % decimation == 0:
cmd = np.array([vx, vy, wz], dtype=np.float32)
obs = compute_obs(model, data, cmd, last_action)
action = session.run(None, {'observations': obs.reshape(1,-1).astype(np.float32)})[0][0]
action = np.clip(action, -CLIP_ACTIONS, CLIP_ACTIONS)
last_action = action.copy()
target = DEFAULT_ANGLES + action * ACTION_SCALE
torques = KP*(target - data.qpos[7:19]) - KD*data.qvel[6:18]
data.ctrl[:] = np.clip(torques, -CLIP_ACTIONS, CLIP_ACTIONS)
mujoco.mj_step(model, data)
view.sync()
step += 1
time.sleep(0.001)
kb.stop(); view.close()
if __name__ == "__main__":
main()

89
scripts/view_orig.py Normal file
View File

@@ -0,0 +1,89 @@
#!/usr/bin/env python3
"""MuJoCo viewer for original Go1 (45-dim, PD 80/1.0, action_scale=0.05)"""
import numpy as np, mujoco, onnxruntime as ort, os, time, threading, queue
from mujoco import viewer
from pynput import keyboard
PROJECT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
ONNX = os.path.join(PROJECT, "exports_go1_orig", "policy.onnx")
XML_DIR = os.path.join(PROJECT, "motrix_envs/src/motrix_envs/locomotion/go1/xmls")
# Original params
NUM_OBS = 57
KP, KD = 80.0, 0.5 # KD=0.5 + joint_damping(0.5) = 1.0 = training kd
ACTION_SCALE = 0.05
CLIP = 23.7
DEFAULT = np.array([-0.0,0.9,-1.8,0.0,0.9,-1.8,-0.0,0.9,-1.8,0.0,0.9,-1.8], dtype=np.float32)
class KB:
def __init__(self):
self._q=queue.Queue(); self.running=True; self.held=set()
def _n(self,k):
try:
if hasattr(k,'char') and k.char: return k.char.lower()
except: pass
return str(k).lower()
def _w(self):
while self.running:
try:
et,k=self._q.get(timeout=0.05); n=self._n(k)
if et=='press': self.held.add(n)
elif et=='release': self.held.discard(n)
except queue.Empty: pass
def init(self):
self._l=keyboard.Listener(on_press=lambda k:self._q.put(('press',k)), on_release=lambda k:self._q.put(('release',k)))
self._l.start(); self._t=threading.Thread(target=self._w,daemon=True); self._t.start()
def held_keys(self): return self.held.copy()
def stop(self): self.running=False; self._l.stop()
def get_sensor(m,d,name):
sid=mujoco.mj_name2id(m,mujoco.mjtObj.mjOBJ_SENSOR,name)
if sid<0: return None
adr=m.sensor_adr[sid]; return d.sensordata[adr:adr+m.sensor_dim[sid]].copy()
def compute_obs(model,data,cmd,last_a):
obs=np.zeros(NUM_OBS,dtype=np.float32)
g=get_sensor(model,data,"gyro")
obs[0:3]=(g if g is not None else data.qvel[3:6])*0.25
R=data.xmat[1].reshape(3,3)
obs[3:6]=(R.T@np.array([0.,0.,-1.])).astype(np.float32)
obs[6:18]=(data.qpos[7:19]-DEFAULT)*1.0
obs[18:30]=data.qvel[6:18]*0.05
obs[30:42]=last_a
obs[42:45]=cmd; obs[45:57]=0.0 # contact_force*np.array([2.,2.,0.25],dtype=np.float32)
return np.clip(obs,-100.,100.)
def main():
os.chdir(XML_DIR)
model=mujoco.MjModel.from_xml_string(open("scene_motor_actuator.xml").read())
data=mujoco.MjData(model)
data.qpos[0:3]=[0,0,0.42]; data.qpos[3:7]=[1,0,0,0]; data.qpos[7:19]=DEFAULT
mujoco.mj_forward(model,data)
session=ort.InferenceSession(ONNX,providers=['CPUExecutionProvider'])
print(f"[ORIG] kp={KP} kd={KD+0.5} action_scale={ACTION_SCALE} | W/S前后 Q/E左右 A/D旋转")
kb=KB(); kb.init()
view=viewer.launch_passive(model,data)
step,vx,vy,wz=0,0.,0.,0.
last_a=np.zeros(12,dtype=np.float32); action=np.zeros(12,dtype=np.float32)
dec=2
while view.is_running():
keys=kb.held_keys()
if 'escape' in keys: break
if 'r' in keys:
data.qpos[0:3]=[0,0,0.42]; data.qpos[3:7]=[1,0,0,0]; data.qpos[7:19]=DEFAULT
data.qvel[:]=0; last_a[:]=0; mujoco.mj_forward(model,data)
if ' ' in keys: vx=vy=wz=0.
vx=1.0 if 'w' in keys else (-1.0 if 's' in keys else 0.)
vy=1.0 if 'q' in keys else (-1.0 if 'e' in keys else 0.)
wz=1.0 if 'a' in keys else (-1.0 if 'd' in keys else 0.)
if step%dec==0:
obs=compute_obs(model,data,np.array([vx,vy,wz],dtype=np.float32),last_a)
action=session.run(None,{'observations':obs.reshape(1,-1).astype(np.float32)})[0][0]
action=np.clip(action,-CLIP,CLIP); last_a=action.copy()
target=DEFAULT+action*ACTION_SCALE
t=KP*(target-data.qpos[7:19])-KD*data.qvel[6:18]
data.ctrl[:]=np.clip(t,-CLIP,CLIP)
mujoco.mj_step(model,data); view.sync(); step+=1; time.sleep(0.001)
kb.stop(); view.close()
if __name__=="__main__": main()

53
uv.lock generated
View File

@@ -2,9 +2,6 @@ version = 1
revision = 3
requires-python = "==3.10.*"
[options]
prerelease-mode = "allow"
[manifest]
members = [
"motrix-envs",
@@ -94,7 +91,7 @@ name = "certifi"
version = "2022.12.7"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/certifi-2022.12.7-py3-none-any.whl", hash = "sha256:4ad3232f5e926d6718ec31cfc1fcadfde020920e278684144551c91769c7bc18" },
{ url = "https://download.pytorch.org/whl/certifi-2022.12.7-py3-none-any.whl", hash = "sha256:4ad3232f5e926d6718ec31cfc1fcadfde020920e278684144551c91769c7bc18", upload-time = "2023-10-03T17:33:19Z" },
]
[[package]]
@@ -102,7 +99,7 @@ name = "charset-normalizer"
version = "2.1.1"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/charset_normalizer-2.1.1-py3-none-any.whl", hash = "sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f" },
{ url = "https://download.pytorch.org/whl/charset_normalizer-2.1.1-py3-none-any.whl", hash = "sha256:83e9a75d1911279afd89352c68b45348559d1fc0506b054b346651b5e7fee29f", upload-time = "2023-10-03T17:33:19Z" },
]
[[package]]
@@ -136,7 +133,7 @@ name = "colorama"
version = "0.4.6"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6" },
{ url = "https://download.pytorch.org/whl/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", upload-time = "2023-10-05T23:50:34Z" },
]
[[package]]
@@ -346,7 +343,7 @@ name = "idna"
version = "3.4"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/idna-3.4-py3-none-any.whl", hash = "sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2" },
{ url = "https://download.pytorch.org/whl/idna-3.4-py3-none-any.whl", hash = "sha256:90b77e79eaa3eba6de819a0c442c0b4ceefc341a7a2ab77d7562bf49f425c5c2", upload-time = "2023-10-06T21:35:40Z" },
]
[[package]]
@@ -366,7 +363,7 @@ dependencies = [
{ name = "zipp" },
]
wheels = [
{ url = "https://download.pytorch.org/whl/importlib_metadata-7.1.0-py3-none-any.whl" },
{ url = "https://download.pytorch.org/whl/importlib_metadata-7.1.0-py3-none-any.whl", upload-time = "2025-08-05T21:34:24Z" },
]
[[package]]
@@ -466,7 +463,7 @@ dependencies = [
{ name = "markupsafe" },
]
wheels = [
{ url = "https://download.pytorch.org/whl/jinja2-3.1.6-py3-none-any.whl" },
{ url = "https://download.pytorch.org/whl/jinja2-3.1.6-py3-none-any.whl", upload-time = "2025-10-14T18:38:59Z" },
]
[[package]]
@@ -526,10 +523,10 @@ name = "markupsafe"
version = "3.0.3"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-win_amd64.whl" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-macosx_11_0_arm64.whl", upload-time = "2026-03-27T13:54:27Z" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", upload-time = "2026-03-27T13:54:27Z" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", upload-time = "2026-03-27T13:54:27Z" },
{ url = "https://download.pytorch.org/whl/markupsafe-3.0.3-cp310-cp310-win_amd64.whl", upload-time = "2026-03-27T13:54:29Z" },
]
[[package]]
@@ -1058,7 +1055,7 @@ name = "packaging"
version = "24.1"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124" },
{ url = "https://download.pytorch.org/whl/packaging-24.1-py3-none-any.whl", hash = "sha256:5b8f2217dbdbd2f7f384c41c628544e6d52f2d0f53c6d0c3ea61aa5d1d7ff124", upload-time = "2024-10-29T23:48:01Z" },
]
[[package]]
@@ -1191,7 +1188,7 @@ dependencies = [
{ name = "urllib3" },
]
wheels = [
{ url = "https://download.pytorch.org/whl/requests-2.28.1-py3-none-any.whl", hash = "sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349" },
{ url = "https://download.pytorch.org/whl/requests-2.28.1-py3-none-any.whl", hash = "sha256:8fefa2a1a1365bf5520aac41836fbee479da67864514bdb821f31ce07ce65349", upload-time = "2023-10-06T21:36:51Z" },
]
[[package]]
@@ -1245,7 +1242,7 @@ name = "setuptools"
version = "70.2.0"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/setuptools-70.2.0-py3-none-any.whl", hash = "sha256:b8b8060bb426838fbe942479c90296ce976249451118ef566a5a0b7d8b78fb05" },
{ url = "https://download.pytorch.org/whl/setuptools-70.2.0-py3-none-any.whl", hash = "sha256:b8b8060bb426838fbe942479c90296ce976249451118ef566a5a0b7d8b78fb05", upload-time = "2025-01-30T19:44:58Z" },
]
[[package]]
@@ -1624,9 +1621,9 @@ dependencies = [
{ name = "typing-extensions" },
]
wheels = [
{ url = "https://download.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b1f0cdd0720ad60536deb5baa427b782fd920dd4fcf72e244d32974caafa3b9e" },
{ url = "https://download.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:ac1849553ee673dfafb44c610c60cb60a2890f0e117f43599a526cf777eb8b8c" },
{ url = "https://download.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:c52c4b869742f00b12cb34521d1381be6119fa46244791704b00cc4a3cb06850" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:b1f0cdd0720ad60536deb5baa427b782fd920dd4fcf72e244d32974caafa3b9e", upload-time = "2025-05-14T03:33:41Z" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:ac1849553ee673dfafb44c610c60cb60a2890f0e117f43599a526cf777eb8b8c", upload-time = "2025-04-22T18:19:25Z" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torch-2.7.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:c52c4b869742f00b12cb34521d1381be6119fa46244791704b00cc4a3cb06850", upload-time = "2025-04-22T18:19:44Z" },
]
[[package]]
@@ -1637,8 +1634,8 @@ dependencies = [
{ name = "torch" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/cu128/torchaudio-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:d62b6e9b792ad37af6d1289ba283e1029e71b4ff9cd3c6cf7f0e7776f23254b2" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torchaudio-2.7.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:f96c2be8aff6c827e76fd3a85e69a54ba5b9a37090853ed886f056ddfbca09a4" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torchaudio-2.7.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:d62b6e9b792ad37af6d1289ba283e1029e71b4ff9cd3c6cf7f0e7776f23254b2", upload-time = "2025-04-22T18:30:28Z" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torchaudio-2.7.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:f96c2be8aff6c827e76fd3a85e69a54ba5b9a37090853ed886f056ddfbca09a4", upload-time = "2025-04-22T18:30:29Z" },
]
[[package]]
@@ -1651,8 +1648,8 @@ dependencies = [
{ name = "torch" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.22.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:59df5a550113a80ce523047066eaaedb168c69482da88c3ab246716ab45ba092" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.22.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:cdd90b768b01b0d638cb06a6c211b550b275c0c207b5210b7cbb5cea8dde11db" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.22.0%2Bcu128-cp310-cp310-manylinux_2_28_x86_64.whl", hash = "sha256:59df5a550113a80ce523047066eaaedb168c69482da88c3ab246716ab45ba092", upload-time = "2025-04-22T18:30:20Z" },
{ url = "https://download-r2.pytorch.org/whl/cu128/torchvision-0.22.0%2Bcu128-cp310-cp310-win_amd64.whl", hash = "sha256:cdd90b768b01b0d638cb06a6c211b550b275c0c207b5210b7cbb5cea8dde11db", upload-time = "2025-04-22T18:30:20Z" },
]
[[package]]
@@ -1663,7 +1660,7 @@ dependencies = [
{ name = "colorama", marker = "sys_platform == 'win32'" },
]
wheels = [
{ url = "https://download.pytorch.org/whl/tqdm-4.66.5-py3-none-any.whl", hash = "sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd" },
{ url = "https://download.pytorch.org/whl/tqdm-4.66.5-py3-none-any.whl", hash = "sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd", upload-time = "2024-10-30T00:09:55Z" },
]
[[package]]
@@ -1686,8 +1683,8 @@ dependencies = [
{ name = "setuptools" },
]
wheels = [
{ url = "https://download-r2.pytorch.org/whl/triton-3.3.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6cbb1b3c214d17e06e2c12399747f645fb7e8b72cb24271557cdce0893c764c6" },
{ url = "https://download-r2.pytorch.org/whl/triton-3.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4198996c9fa3fd811e3bc007f0fc9853c784be3dae6d30714f579c5106d70616" },
{ url = "https://download-r2.pytorch.org/whl/triton-3.3.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6cbb1b3c214d17e06e2c12399747f645fb7e8b72cb24271557cdce0893c764c6", upload-time = "2026-01-22T22:41:17Z" },
{ url = "https://download-r2.pytorch.org/whl/triton-3.3.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4198996c9fa3fd811e3bc007f0fc9853c784be3dae6d30714f579c5106d70616", upload-time = "2026-01-22T22:41:22Z" },
]
[[package]]
@@ -1703,7 +1700,7 @@ name = "urllib3"
version = "1.26.13"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/urllib3-1.26.13-py2.py3-none-any.whl", hash = "sha256:47cc05d99aaa09c9e72ed5809b60e7ba354e64b59c9c173ac3018642d8bb41fc" },
{ url = "https://download.pytorch.org/whl/urllib3-1.26.13-py2.py3-none-any.whl", hash = "sha256:47cc05d99aaa09c9e72ed5809b60e7ba354e64b59c9c173ac3018642d8bb41fc", upload-time = "2023-10-06T21:48:51Z" },
]
[[package]]
@@ -1747,5 +1744,5 @@ name = "zipp"
version = "3.19.2"
source = { registry = "https://download.pytorch.org/whl/cu128" }
wheels = [
{ url = "https://download.pytorch.org/whl/zipp-3.19.2-py3-none-any.whl" },
{ url = "https://download.pytorch.org/whl/zipp-3.19.2-py3-none-any.whl", upload-time = "2025-08-05T22:17:38Z" },
]