no_linevel_no_foot
This commit is contained in:
cyy_mac
BIN
deploy_45dim/better3.onnx
Normal file
BIN
deploy_45dim/better3.onnx
Normal file
Binary file not shown.
403
deploy_45dim/go1_no_linevel_sim2sim_mujoco.py
Normal file
403
deploy_45dim/go1_no_linevel_sim2sim_mujoco.py
Normal file
@@ -0,0 +1,403 @@
|
||||
#!/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
|
||||
# ============================================================
|
||||
HERE = os.path.dirname(os.path.abspath(__file__))
|
||||
DEFAULT_ONNX_PATH = os.path.join(HERE, "better3.onnx")
|
||||
GO1_XML = os.path.join(HERE, "..", "sim2sim_mujoco_example", "data", "go1", "xml", "go1.xml")
|
||||
|
||||
# ============================================================
|
||||
# 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()
|
||||
|
||||
# Use local Go1 XML (terrain switching disabled on macOS)
|
||||
os.chdir(os.path.dirname(GO1_XML)) # mesh paths are relative
|
||||
with open(GO1_XML, "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}")
|
||||
|
||||
# Use simple spawn
|
||||
spawn_xyz = np.array([0.0, 0.0, 0.42], dtype=np.float64)
|
||||
|
||||
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()
|
||||
|
||||
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
|
||||
|
||||
# Reset
|
||||
if keyboard_reader.is_key_pressed("r"):
|
||||
spawn_xyz[:] = [0.0, 0.0, 0.42]
|
||||
|
||||
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)
|
||||
print(f"[RESET] 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")
|
||||
|
||||
keyboard_reader.restore()
|
||||
viewer_handle.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user