fix: clamp std before distribution, lower init_noise to 0.5, NaN guard
This commit is contained in:
8x54zj-m
264
scripts/dreamwaq_sim2sim_mujoco.py
Normal file
264
scripts/dreamwaq_sim2sim_mujoco.py
Normal 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
272
scripts/eval_go1_commands.py
Normal 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
217
scripts/export_dreamwaq_onnx.py
Normal 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
88
scripts/export_dreamwaq_rsl_onnx.py
Normal 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
289
scripts/export_go1_no_linevel_onnx.py
Normal 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
314
scripts/export_go1_onnx.py
Normal 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
119
scripts/gen_flat_stairs.py
Normal 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
95
scripts/gen_stairs_box.py
Normal 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
108
scripts/gen_stairs_test.py
Normal 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
428
scripts/go1_no_linevel_sim2sim_mujoco.py
Normal 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
371
scripts/go1_sim2sim_mujoco.py
Normal 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()
|
||||
@@ -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
147
scripts/play_dreamwaq.py
Normal 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
141
scripts/play_dreamwaq_rsl.py
Normal 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
430
scripts/terrain_editor.py
Normal 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()
|
||||
@@ -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
42
scripts/train_cts.py
Normal 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
42
scripts/train_dreamwaq.py
Normal 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
122
scripts/train_dreamwaq_rsl.py
Normal 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
95
scripts/train_go2style.py
Normal 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
135
scripts/view_go2style.py
Normal 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
89
scripts/view_orig.py
Normal 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()
|
||||
Reference in New Issue
Block a user