47372443ff
Upstream: https://github.com/pollen-robotics/microduck_rl Upstream-Commit: d424a0c899f6b33cbd3daeb279913134349c0b63 Upstream-Branch: develop
1384 lines
69 KiB
Python
1384 lines
69 KiB
Python
#!/usr/bin/env python3
|
|
"""Simple script to run ONNX policy inference in MuJoCo with rendering."""
|
|
|
|
import argparse
|
|
import csv
|
|
import math
|
|
import os
|
|
import pickle
|
|
import queue
|
|
import select
|
|
import sys
|
|
import termios
|
|
import threading
|
|
import time
|
|
import tty
|
|
import numpy as np
|
|
import mujoco
|
|
import mujoco.viewer
|
|
import onnxruntime as ort
|
|
|
|
MICRODUCK_XML = "src/mjlab_microduck/robot/microduck/scene.xml"
|
|
# MICRODUCK_XML = "src/mjlab_microduck/robot/microduck/scene_ramps.xml"
|
|
# MICRODUCK_XML = "src/mjlab_microduck/robot/microduck/scene_floor_objects.xml"
|
|
# MICRODUCK_XML = "src/mjlab_microduck/robot/microduck/scene_robot_walk.xml"
|
|
MICRODUCK_ROLLERS_XML = "src/mjlab_microduck/robot/microduck/scene_rollers.xml"
|
|
MICRODUCK_BALL_XML = "src/mjlab_microduck/robot/microduck/scene_ball.xml"
|
|
|
|
# Body pose command constants (must match training constants)
|
|
BODY_CMD_MAX_Z = 0.03 # ±30 mm
|
|
BODY_CMD_MAX_XY = 0.02 # ±20 mm
|
|
BODY_CMD_MAX_ANGLE = math.radians(30) # ±30°
|
|
|
|
# Ball placement for kick behaviors (must match microduck_ball_kick_env_cfg's
|
|
# reset_ball_in_front_of_foot params: ball center in the robot's yaw frame).
|
|
BALL_OFFSET_X = 0.09
|
|
BALL_OFFSET_ABS_Y = 0.042
|
|
BALL_RADIUS = 0.035
|
|
|
|
# Default pose used by the policy (legs flexed, standing position)
|
|
# This is the reference pose that:
|
|
# - Actions are offsets from (motor_target = DEFAULT_POSE + action * scale)
|
|
# - Joint observations are relative to (obs_joint_pos = current_pos - DEFAULT_POSE)
|
|
# STAND2 pose (matches HOME_FRAME in microduck_constants.py): trunk shifted
|
|
# ~5mm forward so the CoM sits over the ankle axis. Leg pitch chain leaned
|
|
# forward vs the old pose: hip_pitch 30°→26.24°, ankle 30°→25.95°, knee 0°→0.28°.
|
|
DEFAULT_POSE = np.array([
|
|
0.0, # left_hip_yaw
|
|
-0.0873, # left_hip_roll
|
|
-0.4579, # left_hip_pitch
|
|
-0.0049, # left_knee
|
|
0.4530, # left_ankle
|
|
0.3491, # neck_pitch
|
|
0.3491, # head_pitch
|
|
0.0, # head_yaw
|
|
0.0, # head_roll
|
|
0.0, # right_hip_yaw
|
|
0.0873, # right_hip_roll
|
|
0.4579, # right_hip_pitch
|
|
0.0049, # right_knee
|
|
-0.4530, # right_ankle
|
|
], dtype=np.float32)
|
|
|
|
|
|
class TerminalInput:
|
|
"""Single-keypress reader on stdin (cbreak mode, background thread).
|
|
|
|
Replaces the MuJoCo viewer key_callback: keypresses in the viewer window
|
|
also fire the viewer's built-in visualization shortcuts (frames, labels,
|
|
rendering toggles…), so commands are read from the terminal instead.
|
|
Arrow keys arrive as ESC [ A/B/C/D escape sequences and are translated to
|
|
symbolic names ("up"/"down"/"left"/"right"); letters are lowercased.
|
|
cbreak (not raw) mode keeps ISIG enabled, so Ctrl+C still works.
|
|
"""
|
|
|
|
_ARROWS = {"A": "up", "B": "down", "C": "right", "D": "left"}
|
|
|
|
def __init__(self):
|
|
self._queue = queue.Queue()
|
|
self.enabled = sys.stdin.isatty()
|
|
self._fd = sys.stdin.fileno() if self.enabled else -1
|
|
self._old_attrs = None
|
|
self._stop = threading.Event()
|
|
|
|
def __enter__(self):
|
|
if not self.enabled:
|
|
print("WARNING: stdin is not a TTY — keyboard control disabled")
|
|
return self
|
|
self._old_attrs = termios.tcgetattr(self._fd)
|
|
tty.setcbreak(self._fd)
|
|
threading.Thread(target=self._reader, daemon=True).start()
|
|
return self
|
|
|
|
def __exit__(self, *exc):
|
|
self._stop.set()
|
|
if self._old_attrs is not None:
|
|
termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs)
|
|
|
|
def _read1(self, timeout):
|
|
"""Read one byte from stdin, or None on timeout. os.read (unbuffered):
|
|
buffered sys.stdin.read would swallow escape-sequence bytes past what
|
|
select reported ready."""
|
|
r, _, _ = select.select([self._fd], [], [], timeout)
|
|
if not r:
|
|
return None
|
|
data = os.read(self._fd, 1)
|
|
return data.decode(errors="ignore") if data else None
|
|
|
|
def _reader(self):
|
|
while not self._stop.is_set():
|
|
ch = self._read1(0.1)
|
|
if not ch:
|
|
continue
|
|
if ch == "\x1b": # possible arrow-key escape sequence
|
|
if self._read1(0.05) == "[":
|
|
final = self._read1(0.05)
|
|
name = self._ARROWS.get(final) if final else None
|
|
if name:
|
|
self._queue.put(name)
|
|
continue # bare ESC / unknown sequence: ignore
|
|
self._queue.put(ch.lower() if ch.isalpha() else ch)
|
|
|
|
def get_keys(self):
|
|
"""Drain and return all pending keys (symbolic names / characters)."""
|
|
keys = []
|
|
while True:
|
|
try:
|
|
keys.append(self._queue.get_nowait())
|
|
except queue.Empty:
|
|
return keys
|
|
|
|
|
|
class PolicyInference:
|
|
def __init__(self, model, data, walking_onnx_path=None, action_scale=1.0,
|
|
delay_min_lag=0, delay_max_lag=0,
|
|
standing_onnx_path=None, switch_threshold=0.05,
|
|
use_projected_gravity=False, ground_pick_onnx_path=None, ground_pick_period=4.0,
|
|
sit_onnx_path=None, new_cmd_obs=False, slope_onnx_path=None,
|
|
sitstand_onnx_path=None,
|
|
kick_left_onnx_path=None, kick_right_onnx_path=None,
|
|
roulade_onnx_path=None,
|
|
kick_duration=3.0, roulade_duration=2.0):
|
|
self.model = model
|
|
self.data = data
|
|
self.action_scale = action_scale
|
|
self.use_projected_gravity = use_projected_gravity
|
|
self.delay_min_lag = delay_min_lag
|
|
self.delay_max_lag = delay_max_lag
|
|
self.switch_threshold = switch_threshold
|
|
# When True: emit the unified 13D command vector and treat head_offset /
|
|
# body_cmd as policy COMMANDS (no add to ctrl, no joint_pos correction).
|
|
# When False: legacy behaviour (3D command, head_offset added to ctrl[5:9]).
|
|
self.new_cmd_obs = new_cmd_obs
|
|
|
|
# Load walking policy
|
|
self.walking_session = None
|
|
self.default_gait_period_from_onnx = None
|
|
if walking_onnx_path:
|
|
print(f"Loading walking policy from: {walking_onnx_path}")
|
|
self.walking_session = ort.InferenceSession(walking_onnx_path)
|
|
w_input_shape = self.walking_session.get_inputs()[0].shape
|
|
w_output_shape = self.walking_session.get_outputs()[0].shape
|
|
print(f"Walking policy input: {self.walking_session.get_inputs()[0].name}, shape: {w_input_shape}")
|
|
print(f"Walking policy output: {self.walking_session.get_outputs()[0].name}, shape: {w_output_shape}")
|
|
|
|
# Try to read gait period from ONNX metadata
|
|
try:
|
|
model_metadata = self.walking_session.get_modelmeta()
|
|
if hasattr(model_metadata, 'custom_metadata_map') and 'gait_period' in model_metadata.custom_metadata_map:
|
|
self.default_gait_period_from_onnx = float(model_metadata.custom_metadata_map['gait_period'])
|
|
print(f"Found gait period in ONNX metadata: {self.default_gait_period_from_onnx:.4f}s")
|
|
except Exception as e:
|
|
print(f"Could not read gait period from ONNX metadata: {e}")
|
|
|
|
# Load standing policy
|
|
self.standing_session = None
|
|
if standing_onnx_path:
|
|
print(f"\nLoading standing policy from: {standing_onnx_path}")
|
|
self.standing_session = ort.InferenceSession(standing_onnx_path)
|
|
s_input_shape = self.standing_session.get_inputs()[0].shape
|
|
s_output_shape = self.standing_session.get_outputs()[0].shape
|
|
print(f"Standing policy input: {self.standing_session.get_inputs()[0].name}, shape: {s_input_shape}")
|
|
print(f"Standing policy output: {self.standing_session.get_outputs()[0].name}, shape: {s_output_shape}")
|
|
if self.walking_session:
|
|
print(f"Policy switching threshold: {switch_threshold} (vel command magnitude)")
|
|
|
|
# Load ground pick policy
|
|
self.ground_pick_session = None
|
|
self.ground_pick_mode = False
|
|
self.ground_pick_phase = 0.0
|
|
self.ground_pick_period = ground_pick_period
|
|
if ground_pick_onnx_path:
|
|
print(f"\nLoading ground pick policy from: {ground_pick_onnx_path}")
|
|
self.ground_pick_session = ort.InferenceSession(ground_pick_onnx_path)
|
|
gp_input_shape = self.ground_pick_session.get_inputs()[0].shape
|
|
print(f"Ground pick policy input shape: {gp_input_shape}")
|
|
|
|
# Load sit policy. Two flavours share the Y key and self.sit_session:
|
|
# - --sit (is_sitstand=False): the OLD one-way sit policy. Sits
|
|
# unconditionally on a zero twist command; standing back up is done
|
|
# by switching back to the standing/walking session.
|
|
# - --sitstand (is_sitstand=True): the commanded sit↔stand policy.
|
|
# twist[0] is a posture flag (0=stand, 1=sit); the SAME policy sits,
|
|
# holds, and stands back up — Y just flips the flag.
|
|
self.sit_session = None
|
|
self.sit_mode = False
|
|
self.is_sitstand = False
|
|
if sit_onnx_path and sitstand_onnx_path:
|
|
raise ValueError("Provide only one of --sit / --sitstand")
|
|
if sit_onnx_path:
|
|
print(f"\nLoading sit policy from: {sit_onnx_path}")
|
|
self.sit_session = ort.InferenceSession(sit_onnx_path)
|
|
sit_input_shape = self.sit_session.get_inputs()[0].shape
|
|
print(f"Sit policy input shape: {sit_input_shape}")
|
|
elif sitstand_onnx_path:
|
|
if not self.new_cmd_obs:
|
|
raise ValueError(
|
|
"--sitstand policies use the unified 13D command obs (61D); run with --new-cmd-obs"
|
|
)
|
|
print(f"\nLoading sitstand policy from: {sitstand_onnx_path}")
|
|
self.sit_session = ort.InferenceSession(sitstand_onnx_path)
|
|
self.is_sitstand = True
|
|
ss_input_shape = self.sit_session.get_inputs()[0].shape
|
|
print(f"Sitstand policy input shape: {ss_input_shape}")
|
|
|
|
# Load slope policy (passive descent, runs with zero twist command)
|
|
self.slope_session = None
|
|
self.slope_mode = False
|
|
if slope_onnx_path:
|
|
print(f"\nLoading slope policy from: {slope_onnx_path}")
|
|
self.slope_session = ort.InferenceSession(slope_onnx_path)
|
|
sl_input_shape = self.slope_session.get_inputs()[0].shape
|
|
print(f"Slope policy input shape: {sl_input_shape}")
|
|
|
|
# Episodic behavior policies (kick left/right, roulade). All three use
|
|
# the unified 61D obs layout with an ALL-ZERO 13D command (twist forced
|
|
# ~0 in training, head/body slots zero-padded), so triggering one is a
|
|
# plain session swap; after `duration` seconds control hands back to
|
|
# walking/standing (the behavior policies end standing on their own).
|
|
self.behavior_sessions = {}
|
|
self.behavior_durations = {}
|
|
self.behavior_mode = None # name of the running behavior, or None
|
|
self.behavior_time_left = 0.0
|
|
for name, path, duration in (
|
|
("kick_left", kick_left_onnx_path, kick_duration),
|
|
("kick_right", kick_right_onnx_path, kick_duration),
|
|
("roulade", roulade_onnx_path, roulade_duration),
|
|
):
|
|
if not path:
|
|
continue
|
|
if not self.new_cmd_obs:
|
|
raise ValueError(
|
|
f"--{name.replace('_', '-')} policies use the unified 13D "
|
|
"command obs (61D); run with --new-cmd-obs"
|
|
)
|
|
print(f"\nLoading {name} policy from: {path}")
|
|
self.behavior_sessions[name] = ort.InferenceSession(path)
|
|
self.behavior_durations[name] = duration
|
|
print(f"{name} policy input shape: {self.behavior_sessions[name].get_inputs()[0].shape}"
|
|
f" (auto-return after {duration:.1f}s)")
|
|
|
|
# Validate at least one policy loaded. A sitstand policy can run alone
|
|
# (it holds the stand at flag=0), unlike the old one-way sit policy.
|
|
if not self.walking_session and not self.standing_session and not self.is_sitstand:
|
|
raise ValueError("At least one of --walking, --standing or --sitstand must be provided")
|
|
|
|
# Determine initial active session and policy
|
|
if self.walking_session:
|
|
self.current_policy = "walking"
|
|
self.ort_session = self.walking_session
|
|
elif self.standing_session:
|
|
self.current_policy = "standing"
|
|
self.ort_session = self.standing_session
|
|
else:
|
|
# sitstand-only: start standing (posture flag 0).
|
|
self.current_policy = "sit"
|
|
self.ort_session = self.sit_session
|
|
|
|
# Get input/output names from active session
|
|
self.input_name = self.ort_session.get_inputs()[0].name
|
|
self.output_name = self.ort_session.get_outputs()[0].name
|
|
|
|
# Get sensor IDs and body IDs
|
|
self.imu_ang_vel_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_SENSOR, "imu_ang_vel")
|
|
self.trunk_base_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "trunk_base")
|
|
|
|
# Trunk freejoint qpos address (needed to place the ball in the robot's
|
|
# yaw frame) and optional ball freejoint (present in scene_ball.xml).
|
|
_trunk_jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint")
|
|
self._trunk_qpos_adr = int(model.jnt_qposadr[_trunk_jid])
|
|
_ball_jid = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "ball_free")
|
|
if _ball_jid >= 0:
|
|
self.ball_qpos_adr = int(model.jnt_qposadr[_ball_jid])
|
|
self.ball_qvel_adr = int(model.jnt_dofadr[_ball_jid])
|
|
else:
|
|
self.ball_qpos_adr = None
|
|
self.ball_qvel_adr = None
|
|
|
|
print(f"Sensors found:")
|
|
print(f" imu_ang_vel: id={self.imu_ang_vel_id}")
|
|
print(f"Body IDs:")
|
|
print(f" trunk_base: id={self.trunk_base_id}")
|
|
|
|
# Joint information
|
|
self.n_joints = model.nu
|
|
|
|
# For robots with passive/interspersed joints (e.g. roller skates), the actuated
|
|
# joints are not contiguous in qpos/qvel. Compute the correct indices from the
|
|
# actuator transmission joint IDs so extraction works for any joint ordering.
|
|
self.joint_qpos_indices = [
|
|
int(model.jnt_qposadr[model.actuator_trnid[i, 0]]) for i in range(model.nu)
|
|
]
|
|
self.joint_qvel_indices = [
|
|
int(model.jnt_dofadr[model.actuator_trnid[i, 0]]) for i in range(model.nu)
|
|
]
|
|
|
|
# Default pose for the policy (flexed legs)
|
|
self.default_pose = DEFAULT_POSE[:self.n_joints]
|
|
print(f"Number of actuators: {self.n_joints}")
|
|
print(f"Default pose: {self.default_pose}")
|
|
print(f"Action scale: {self.action_scale}")
|
|
|
|
# Last action (for observation history)
|
|
self.last_action = np.zeros(self.n_joints, dtype=np.float32)
|
|
|
|
# Velocity command [lin_vel_x, lin_vel_y, ang_vel_z] — controls walking / policy switching
|
|
self.vel_cmd = np.zeros(3, dtype=np.float32)
|
|
# Key-press step sizes and limits (overridden per mode in main())
|
|
self.vel_step_x = 0.05
|
|
self.vel_step_y = 0.05
|
|
self.vel_step_ang = 0.3
|
|
self.vel_max_x = 0.3
|
|
self.vel_min_x = -0.3
|
|
self.vel_max_y = 0.3
|
|
self.vel_min_y = -0.3
|
|
self.vel_max_ang = 1.5
|
|
# Body pose command. In new_cmd_obs mode this is 6D
|
|
# [x, y, z, roll, pitch, yaw] (m, m, m, rad, rad, rad)
|
|
# In legacy mode only [z, pitch, roll] (first 3 indices reused as
|
|
# [z, pitch, roll] to keep the legacy normalization path working).
|
|
self.body_cmd = np.zeros(6 if self.new_cmd_obs else 3, dtype=np.float32)
|
|
# Obs command vector (3D in legacy mode, 13D when new_cmd_obs=True).
|
|
self.command = np.zeros(13 if self.new_cmd_obs else 3, dtype=np.float32)
|
|
|
|
# Body pose mode (like head mode but for standing body pose control)
|
|
self.body_pose_mode = False
|
|
self.body_cmd_step_xy = 0.005 # 5 mm per keypress (4 to max)
|
|
self.body_cmd_step_z = 0.01 # 10 mm per keypress (3 to max)
|
|
self.body_cmd_step_angle = math.radians(10) # 10° per keypress (3 to max)
|
|
|
|
# Head control mode. In legacy mode head_offset is added on top of
|
|
# ctrl[5:9]; in new_cmd_obs mode it's a *command* fed to the policy.
|
|
# Final per-joint training caps: neck/head_pitch ±1.1, head_yaw ±1.4,
|
|
# head_roll ±0.31. Slider max = widest joint cap; head_roll naturally
|
|
# gets clipped by the policy since it was never trained beyond 0.31.
|
|
self.head_mode = False
|
|
self.head_offset = np.zeros(4, dtype=np.float32)
|
|
if self.new_cmd_obs:
|
|
self.head_max = 1.4
|
|
self.head_step = 0.1
|
|
else:
|
|
self.head_max = 2.5
|
|
self.head_step = 0.83
|
|
|
|
# Action delay buffer
|
|
self.use_delay = self.delay_max_lag > 0
|
|
if self.use_delay:
|
|
buffer_size = self.delay_max_lag + 1
|
|
self.action_buffer = [np.zeros(self.n_joints, dtype=np.float32) for _ in range(buffer_size)]
|
|
self.buffer_index = 0
|
|
self.current_lag = np.random.randint(self.delay_min_lag, self.delay_max_lag + 1)
|
|
print(f"\nActuator delay enabled:")
|
|
print(f" Min lag: {self.delay_min_lag} timesteps")
|
|
print(f" Max lag: {self.delay_max_lag} timesteps")
|
|
print(f" Sampled lag: {self.current_lag} timesteps")
|
|
print(f" Buffer size: {buffer_size}")
|
|
else:
|
|
self.action_buffer = None
|
|
self.current_lag = 0
|
|
|
|
def _update_command(self):
|
|
"""Update self.command (fed into obs) based on current policy and commands.
|
|
|
|
Legacy mode (new_cmd_obs=False): self.command is 3D.
|
|
New mode (new_cmd_obs=True): self.command is 13D:
|
|
[vx, vy, vtheta, ← twist
|
|
neck_pitch, head_pitch, head_yaw, head_roll, ← head_pose deltas
|
|
body_x, body_y, body_z, body_roll, body_pitch, body_yaw] ← body_pose
|
|
We keep the existing keyboard mappings: head_offset (4D) drives the head
|
|
slots; body_cmd[0..2] currently mean (Δz, Δpitch, Δroll) and are routed
|
|
into body_pose slots [z, pitch, roll]; x/y/yaw stay 0 (not exposed on
|
|
keyboard yet). ground_pick still owns slots [0..2] for phase encoding.
|
|
"""
|
|
if self.new_cmd_obs:
|
|
if self.behavior_mode is not None:
|
|
# Kick/roulade were trained with an all-zero 13D command
|
|
# (twist ~0, head/body slots zero-padded) — feeding stale
|
|
# head/body commands would be out-of-distribution.
|
|
self.command = np.zeros(13, dtype=np.float32)
|
|
return
|
|
cmd = np.zeros(13, dtype=np.float32)
|
|
# twist slot (or phase encoding for ground_pick — overwritten there)
|
|
if self.current_policy == "walking":
|
|
cmd[0:3] = self.vel_cmd
|
|
elif self.current_policy == "sit" and self.is_sitstand:
|
|
# Sitstand posture flag: 1 = sit, 0 = stand. NOT zeros — the
|
|
# all-zero twist is the STAND command for this policy, which is
|
|
# why feeding it the old sit-policy zero command did nothing.
|
|
cmd[0] = 1.0 if self.sit_mode else 0.0
|
|
# else standing/old-sit/ground_pick: leave twist 0 (ground_pick
|
|
# writes its phase encoding later)
|
|
cmd[3:7] = self.head_offset
|
|
cmd[7:13] = self.body_cmd # [x, y, z, roll, pitch, yaw]
|
|
self.command = cmd
|
|
return
|
|
|
|
# Legacy 3D command
|
|
if self.current_policy == "walking":
|
|
self.command = self.vel_cmd.copy()
|
|
elif self.current_policy == "sit":
|
|
# Sit was trained with a near-zero twist command.
|
|
self.command = np.zeros(3, dtype=np.float32)
|
|
elif self.current_policy == "standing":
|
|
# Normalize body pose cmd to match training's body_pose_cmd_obs
|
|
self.command = np.array([
|
|
self.body_cmd[0] / BODY_CMD_MAX_Z,
|
|
self.body_cmd[1] / BODY_CMD_MAX_ANGLE,
|
|
self.body_cmd[2] / BODY_CMD_MAX_ANGLE,
|
|
], dtype=np.float32)
|
|
elif self.current_policy == "slope":
|
|
# Passive descent: zero command (like standing coast)
|
|
self.command = np.zeros(3, dtype=np.float32)
|
|
# ground_pick: command is set directly by update_ground_pick_phase
|
|
|
|
def _update_policy_session(self):
|
|
"""Switch between walking and standing sessions based on vel_cmd magnitude."""
|
|
if not (self.walking_session and self.standing_session):
|
|
return # Only one policy loaded, no switching
|
|
if self.ground_pick_mode:
|
|
return # Don't switch during ground pick
|
|
if self.sit_mode:
|
|
return # Don't switch while sitting
|
|
if self.slope_mode:
|
|
return # Don't switch during slope mode
|
|
if self.behavior_mode is not None:
|
|
return # Don't switch during a kick/roulade
|
|
|
|
magnitude = float(np.linalg.norm(self.vel_cmd))
|
|
new_policy = "standing" if magnitude <= self.switch_threshold else "walking"
|
|
if new_policy != self.current_policy:
|
|
self.current_policy = new_policy
|
|
self.ort_session = self.standing_session if new_policy == "standing" else self.walking_session
|
|
print(f"Switched to {self.current_policy} policy (vel magnitude: {magnitude:.3f})")
|
|
self._update_command()
|
|
|
|
def set_vel_cmd(self, lin_vel_x=0.0, lin_vel_y=0.0, ang_vel_z=0.0):
|
|
"""Set velocity command (used for walking / policy switching)."""
|
|
self.vel_cmd = np.array([lin_vel_x, lin_vel_y, ang_vel_z], dtype=np.float32)
|
|
self._update_policy_session()
|
|
self._update_command()
|
|
print(f"Vel cmd: [{lin_vel_x:.2f}, {lin_vel_y:.2f}, {ang_vel_z:.2f}] [{self.current_policy}]")
|
|
|
|
def toggle_body_pose_mode(self):
|
|
"""Toggle body pose control mode on/off."""
|
|
self.body_pose_mode = not self.body_pose_mode
|
|
if self.body_pose_mode:
|
|
print("Body pose mode: ON")
|
|
print(f" UP/DOWN: Δz ±{self.body_cmd_step_z*1000:.0f}mm (max ±{BODY_CMD_MAX_Z*1000:.0f}mm)")
|
|
print(f" LEFT/RIGHT: Δpitch ±{math.degrees(self.body_cmd_step_angle):.0f}° (max ±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
print(f" A/E: Δroll ±{math.degrees(self.body_cmd_step_angle):.0f}° (max ±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
if self.new_cmd_obs:
|
|
print(f" Z/S: Δyaw ±{math.degrees(self.body_cmd_step_angle):.0f}° (max ±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
print(f" SPACE: reset body pose to zero")
|
|
self._print_body_cmd()
|
|
else:
|
|
print("Body pose mode: OFF")
|
|
|
|
def toggle_slope_mode(self):
|
|
"""Toggle slope policy mode on/off (passive descent, zero twist command)."""
|
|
if self.slope_session is None:
|
|
print("Slope unavailable: no --slope policy loaded")
|
|
return
|
|
if self.behavior_mode is not None:
|
|
print(f"Cannot toggle slope mode during {self.behavior_mode}")
|
|
return
|
|
self.slope_mode = not self.slope_mode
|
|
if self.slope_mode:
|
|
self.ort_session = self.slope_session
|
|
self.current_policy = "slope"
|
|
self.set_vel_cmd(0.0, 0.0, 0.0) # passive descent: zero command
|
|
print("Slope mode: ON (passive descent)")
|
|
else:
|
|
self.vel_cmd = np.zeros(3, dtype=np.float32)
|
|
if self.walking_session:
|
|
self.current_policy = "walking"
|
|
self.ort_session = self.walking_session
|
|
else:
|
|
self.current_policy = "standing"
|
|
self.ort_session = self.standing_session
|
|
self._update_command()
|
|
print("Slope mode: OFF")
|
|
|
|
def _print_body_cmd(self):
|
|
if self.new_cmd_obs:
|
|
x, y, z, roll, pitch, yaw = self.body_cmd
|
|
print(
|
|
f"Body cmd: x={x*1000:5.1f}mm y={y*1000:5.1f}mm z={z*1000:5.1f}mm "
|
|
f"roll={math.degrees(roll):5.1f}° pitch={math.degrees(pitch):5.1f}° "
|
|
f"yaw={math.degrees(yaw):5.1f}°"
|
|
)
|
|
else:
|
|
print(
|
|
f"Body cmd: z={self.body_cmd[0]*1000:.1f}mm "
|
|
f"pitch={math.degrees(self.body_cmd[1]):.1f}° "
|
|
f"roll={math.degrees(self.body_cmd[2]):.1f}°"
|
|
)
|
|
|
|
# --- body command bumpers (index differs between legacy 3D and new 6D) ---
|
|
def _body_idx(self, axis: str) -> int:
|
|
"""Map an axis name to the body_cmd index, depending on the active mode."""
|
|
if self.new_cmd_obs:
|
|
return {"x": 0, "y": 1, "z": 2, "roll": 3, "pitch": 4, "yaw": 5}[axis]
|
|
return {"z": 0, "pitch": 1, "roll": 2}[axis]
|
|
|
|
def bump_body(self, axis: str, delta: float):
|
|
idx = self._body_idx(axis)
|
|
cap = BODY_CMD_MAX_Z if axis == "z" else BODY_CMD_MAX_XY if axis in ("x", "y") else BODY_CMD_MAX_ANGLE
|
|
self.body_cmd[idx] = float(np.clip(self.body_cmd[idx] + delta, -cap, cap))
|
|
self._update_command()
|
|
self._print_body_cmd()
|
|
|
|
def quat_rotate_inverse(self, quat, vec):
|
|
"""Rotate a vector by the inverse of a quaternion [w, x, y, z]."""
|
|
w = quat[0]
|
|
xyz = quat[1:4]
|
|
t = np.cross(xyz, vec) * 2
|
|
return vec - w * t + np.cross(xyz, t)
|
|
|
|
def get_raw_accelerometer(self):
|
|
"""Get raw accelerometer reading from MuJoCo sensor."""
|
|
sensor_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SENSOR, "imu_accel")
|
|
if sensor_id < 0:
|
|
raise ValueError("Sensor 'imu_accel' not found in model")
|
|
|
|
sensor_adr = self.model.sensor_adr[sensor_id]
|
|
accel_raw = self.data.sensordata[sensor_adr:sensor_adr+3].copy().astype(np.float32)
|
|
accel_negated = -accel_raw
|
|
mag = np.linalg.norm(accel_negated)
|
|
if mag > 0.1:
|
|
return accel_negated / mag
|
|
else:
|
|
quat = self.data.xquat[self.trunk_base_id].copy().astype(np.float32)
|
|
world_gravity = np.array([0.0, 0.0, -1.0], dtype=np.float32)
|
|
return self.quat_rotate_inverse(quat, world_gravity)
|
|
|
|
def get_projected_gravity(self):
|
|
"""Get projected gravity in body frame."""
|
|
quat = self.data.xquat[self.trunk_base_id].copy().astype(np.float32)
|
|
world_gravity = np.array([0.0, 0.0, -1.0], dtype=np.float32)
|
|
return self.quat_rotate_inverse(quat, world_gravity)
|
|
|
|
def get_base_ang_vel(self):
|
|
"""Get base angular velocity from IMU gyro sensor."""
|
|
sensor_adr = self.model.sensor_adr[self.imu_ang_vel_id]
|
|
return self.data.sensordata[sensor_adr:sensor_adr + 3].copy().astype(np.float32)
|
|
|
|
def get_joint_pos_relative(self):
|
|
"""Get joint positions relative to default pose."""
|
|
current_pos = self.data.qpos[self.joint_qpos_indices].copy().astype(np.float32)
|
|
return current_pos - self.default_pose
|
|
|
|
def get_joint_vel(self):
|
|
"""Get joint velocities."""
|
|
return self.data.qvel[self.joint_qvel_indices].copy().astype(np.float32)
|
|
|
|
def get_observations(self):
|
|
"""Collect observations matching policy input.
|
|
|
|
Order for velocity/standing task:
|
|
1. base_ang_vel (3D)
|
|
2. raw_accelerometer OR projected_gravity (3D)
|
|
3. joint_pos (14D) - relative to default
|
|
4. joint_vel (14D)
|
|
5. actions (14D) - last action
|
|
6. command (3D) - vel cmd (walking) or normalized body pose cmd (standing)
|
|
Total: 51D
|
|
"""
|
|
obs = []
|
|
|
|
obs.append(self.get_base_ang_vel())
|
|
|
|
if self.use_projected_gravity:
|
|
obs.append(self.get_projected_gravity())
|
|
else:
|
|
obs.append(self.get_raw_accelerometer())
|
|
|
|
obs.append(self.get_joint_pos_relative())
|
|
obs.append(self.get_joint_vel())
|
|
obs.append(self.last_action)
|
|
obs.append(self.command)
|
|
|
|
return np.concatenate(obs).astype(np.float32)
|
|
|
|
def trigger_ground_pick(self):
|
|
"""Start one ground pick cycle. Automatically returns to walking when done."""
|
|
if self.ground_pick_session is None:
|
|
print("Ground pick unavailable: no --ground-pick policy loaded")
|
|
return
|
|
if self.ground_pick_mode:
|
|
print("Ground pick already in progress")
|
|
return
|
|
if self.sit_mode:
|
|
print("Cannot ground pick while sitting (press Y to stand up first)")
|
|
return
|
|
if self.behavior_mode is not None:
|
|
print(f"Cannot ground pick during {self.behavior_mode}")
|
|
return
|
|
self.ground_pick_mode = True
|
|
self.ground_pick_phase = 0.0
|
|
self.ort_session = self.ground_pick_session
|
|
self.current_policy = "ground_pick"
|
|
print(f"Ground pick: started (period={self.ground_pick_period:.1f}s)")
|
|
|
|
def _end_ground_pick(self):
|
|
"""Switch back after a ground pick cycle completes."""
|
|
self.ground_pick_mode = False
|
|
self.vel_cmd = np.zeros(3, dtype=np.float32)
|
|
if self.walking_session:
|
|
self.current_policy = "walking"
|
|
self.ort_session = self.walking_session
|
|
else:
|
|
self.current_policy = "standing"
|
|
self.ort_session = self.standing_session
|
|
self._update_command()
|
|
print(f"Ground pick: done → back to {self.current_policy}")
|
|
|
|
def update_ground_pick_phase(self, dt: float):
|
|
"""Advance the ground pick phase; auto-exit when one full cycle completes."""
|
|
if not self.ground_pick_mode:
|
|
return
|
|
new_phase = self.ground_pick_phase + dt / self.ground_pick_period
|
|
if new_phase >= 0.7:
|
|
self._end_ground_pick()
|
|
return
|
|
self.ground_pick_phase = new_phase
|
|
# ground_pick policies use the first 3 slots (twist) as phase encoding.
|
|
# Higher slots (head/body) stay at whatever _update_command set them to.
|
|
self.command[0] = np.cos(2 * np.pi * self.ground_pick_phase)
|
|
self.command[1] = np.sin(2 * np.pi * self.ground_pick_phase)
|
|
self.command[2] = 0.0
|
|
|
|
def trigger_behavior(self, name):
|
|
"""Start an episodic behavior (kick_left / kick_right / roulade).
|
|
|
|
The behavior policies were trained to run from a standing start with an
|
|
all-zero command and end standing, so triggering is a session swap; a
|
|
timer hands control back to walking/standing afterwards.
|
|
"""
|
|
session = self.behavior_sessions.get(name)
|
|
if session is None:
|
|
print(f"{name} unavailable: no --{name.replace('_', '-')} policy loaded")
|
|
return
|
|
if self.behavior_mode is not None:
|
|
print(f"Cannot start {name}: {self.behavior_mode} already in progress")
|
|
return
|
|
if self.ground_pick_mode:
|
|
print(f"Cannot start {name} during ground pick")
|
|
return
|
|
if self.sit_mode:
|
|
print(f"Cannot start {name} while sitting (press Y to stand up first)")
|
|
return
|
|
if self.slope_mode:
|
|
print(f"Cannot start {name} during slope mode")
|
|
return
|
|
if name in ("kick_left", "kick_right"):
|
|
self._place_ball(name)
|
|
self.behavior_mode = name
|
|
self.behavior_time_left = self.behavior_durations[name]
|
|
self.vel_cmd = np.zeros(3, dtype=np.float32)
|
|
self.current_policy = name
|
|
self.ort_session = session
|
|
self._update_command()
|
|
print(f"{name}: started (auto-return in {self.behavior_time_left:.1f}s)")
|
|
|
|
def _place_ball(self, behavior):
|
|
"""Teleport the ball in front of the kicking foot, matching training's
|
|
reset_ball_in_front_of_foot (offset in the robot's yaw frame)."""
|
|
if self.ball_qpos_adr is None or self.ball_qvel_adr is None:
|
|
print("No ball in scene (kick will swing at air)")
|
|
return
|
|
adr = self._trunk_qpos_adr
|
|
x, y = float(self.data.qpos[adr]), float(self.data.qpos[adr + 1])
|
|
qw, qx, qy, qz = self.data.qpos[adr + 3:adr + 7]
|
|
yaw = math.atan2(2.0 * (qw * qz + qx * qy), 1.0 - 2.0 * (qy * qy + qz * qz))
|
|
off_y = -BALL_OFFSET_ABS_Y if behavior == "kick_right" else BALL_OFFSET_ABS_Y
|
|
bx = x + math.cos(yaw) * BALL_OFFSET_X - math.sin(yaw) * off_y
|
|
by = y + math.sin(yaw) * BALL_OFFSET_X + math.cos(yaw) * off_y
|
|
self.data.qpos[self.ball_qpos_adr:self.ball_qpos_adr + 7] = [bx, by, BALL_RADIUS, 1, 0, 0, 0]
|
|
self.data.qvel[self.ball_qvel_adr:self.ball_qvel_adr + 6] = 0.0
|
|
foot = behavior.split("_")[1]
|
|
print(f"Ball placed at ({bx:.3f}, {by:.3f}) in front of the {foot} foot")
|
|
|
|
def update_behavior(self, dt: float):
|
|
"""Advance the behavior timer; hand back to walking/standing when done."""
|
|
if self.behavior_mode is None:
|
|
return
|
|
self.behavior_time_left -= dt
|
|
if self.behavior_time_left <= 0.0:
|
|
self._end_behavior()
|
|
|
|
def _end_behavior(self):
|
|
name = self.behavior_mode
|
|
self.behavior_mode = None
|
|
self.vel_cmd = np.zeros(3, dtype=np.float32)
|
|
if self.walking_session:
|
|
self.current_policy = "walking"
|
|
self.ort_session = self.walking_session
|
|
elif self.standing_session:
|
|
self.current_policy = "standing"
|
|
self.ort_session = self.standing_session
|
|
else:
|
|
# sitstand-only setup: the sitstand policy holds the stand (flag 0).
|
|
self.current_policy = "sit"
|
|
self.ort_session = self.sit_session
|
|
self._update_command()
|
|
print(f"{name}: done → back to {self.current_policy}")
|
|
|
|
def toggle_sit(self):
|
|
"""Toggle sitting on/off (Y key).
|
|
|
|
Old one-way sit policy (--sit): Y off switches back to the standing/
|
|
walking session, which does the standing back up.
|
|
Sitstand policy (--sitstand): Y just flips the posture flag — the SAME
|
|
policy sits, holds the sit, and stands back up gently (trained response
|
|
to a flag flip is a ~2 s glide). The session stays active after
|
|
standing (it holds the stand); a velocity command switches back to
|
|
walking/standing as usual.
|
|
"""
|
|
if self.sit_session is None:
|
|
print("Sit unavailable: no --sit/--sitstand policy loaded")
|
|
return
|
|
if self.ground_pick_mode:
|
|
print("Cannot sit during ground pick")
|
|
return
|
|
if self.behavior_mode is not None:
|
|
print(f"Cannot sit during {self.behavior_mode}")
|
|
return
|
|
self.sit_mode = not self.sit_mode
|
|
if self.sit_mode:
|
|
self.vel_cmd = np.zeros(3, dtype=np.float32)
|
|
self.current_policy = "sit"
|
|
self.ort_session = self.sit_session
|
|
print("Sit: ON" + (" (sitstand flag=1; Y again to stand up)" if self.is_sitstand else ""))
|
|
elif self.is_sitstand:
|
|
# Stay on the sitstand session — it stands up itself (flag → 0).
|
|
# Do NOT swap to the standing policy here: it would take over
|
|
# mid-rise from a seated state it wasn't trained on.
|
|
print("Sit: OFF → sitstand policy standing up (flag=0)")
|
|
else:
|
|
if self.standing_session:
|
|
self.current_policy = "standing"
|
|
else:
|
|
self.current_policy = "walking"
|
|
self.ort_session = self.standing_session if self.current_policy == "standing" else self.walking_session
|
|
print(f"Sit: OFF → back to {self.current_policy}")
|
|
self._update_command()
|
|
|
|
def toggle_head_mode(self):
|
|
"""Toggle head control mode on/off."""
|
|
self.head_mode = not self.head_mode
|
|
if self.head_mode:
|
|
print("Head mode: ON")
|
|
print(f" Z/S: neck_pitch | UP/DOWN: head_pitch | LEFT/RIGHT: head_yaw | A/E: head_roll | SPACE: reset (max ±{self.head_max:.2f} rad)")
|
|
else:
|
|
print("Head mode: OFF")
|
|
|
|
def infer(self):
|
|
"""Run policy inference and return action."""
|
|
obs = self.get_observations()
|
|
obs_batch = obs.reshape(1, -1)
|
|
action = self.ort_session.run([self.output_name], {self.input_name: obs_batch})[0]
|
|
action = action.squeeze(0).astype(np.float32)
|
|
self.last_action = action.copy()
|
|
return action
|
|
|
|
def apply_action(self, action):
|
|
"""Apply action to MuJoCo controls with optional delay."""
|
|
if self.use_delay:
|
|
self.action_buffer[self.buffer_index] = action.copy()
|
|
delayed_index = (self.buffer_index - self.current_lag) % len(self.action_buffer)
|
|
delayed_action = self.action_buffer[delayed_index]
|
|
self.buffer_index = (self.buffer_index + 1) % len(self.action_buffer)
|
|
target_positions = self.default_pose + delayed_action * self.action_scale
|
|
else:
|
|
target_positions = self.default_pose + action * self.action_scale
|
|
|
|
self.data.ctrl[:] = target_positions
|
|
# Legacy mode: head_offset is an external perturbation added on top of
|
|
# the policy output. New mode: head_offset is a COMMAND fed into the
|
|
# policy's obs, so the policy itself produces the offset head pose.
|
|
if not self.new_cmd_obs:
|
|
self.data.ctrl[5:9] += self.head_offset
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Run ONNX policy in MuJoCo")
|
|
parser.add_argument("--roller", action="store_true", help="Use roller skate robot XML (robot_walk_rollers.xml)")
|
|
parser.add_argument("--walking", type=str, default=None, help="Path to walking policy ONNX file")
|
|
parser.add_argument("--standing", "-s", type=str, default=None, help="Path to standing policy ONNX file")
|
|
parser.add_argument("--ground-pick", type=str, default=None, help="Path to ground pick policy ONNX file (press G to activate)")
|
|
parser.add_argument("--sit", type=str, default=None, help="Path to OLD one-way sitting policy ONNX file (press Y to sit, Y again switches back to standing/walking policy)")
|
|
parser.add_argument("--sitstand", type=str, default=None, help="Path to sitstand policy ONNX (commanded sit<->stand; press Y to sit, Y again the SAME policy stands back up). Requires --new-cmd-obs. Can run standalone.")
|
|
parser.add_argument("--slope", type=str, default=None, help="Path to slope policy ONNX file (press Y to toggle)")
|
|
parser.add_argument("--kick-left", type=str, default=None, help="Path to LEFT-foot ball kick policy ONNX (press K to trigger). Requires --new-cmd-obs. Loads a scene with a ball.")
|
|
parser.add_argument("--kick-right", type=str, default=None, help="Path to RIGHT-foot ball kick policy ONNX (press L to trigger). Requires --new-cmd-obs. Loads a scene with a ball.")
|
|
parser.add_argument("--roulade", type=str, default=None, help="Path to roulade (forward roll) policy ONNX (press R to trigger). Requires --new-cmd-obs.")
|
|
parser.add_argument("--kick-duration", type=float, default=3.0, help="Seconds a kick policy stays active before handing back to standing/walking (default: 3.0)")
|
|
parser.add_argument("--roulade-duration", type=float, default=2.0, help="Seconds the roulade policy stays active before handing back to standing/walking (default: 2.0, ~the roll itself; the standing/walking policy takes over for the settle)")
|
|
parser.add_argument("--lin-vel-x", type=float, default=0.0, help="Initial linear velocity X command (m/s)")
|
|
parser.add_argument("--lin-vel-y", type=float, default=0.0, help="Initial linear velocity Y command (m/s)")
|
|
parser.add_argument("--ang-vel-z", type=float, default=0.0, help="Initial angular velocity Z command (rad/s)")
|
|
parser.add_argument("--action-scale", type=float, default=1.0, help="Action scale (default: 1.0)")
|
|
parser.add_argument("--raw-accelerometer", action="store_true", help="Use raw accelerometer instead of projected gravity")
|
|
parser.add_argument("--delay", type=int, nargs='*', default=None, help="Enable actuator delay: --delay MIN MAX or --delay LAG")
|
|
parser.add_argument("--debug", action="store_true", help="Print observations and actions")
|
|
parser.add_argument("--save-csv", type=str, default=None, help="Save observations and actions to CSV file")
|
|
parser.add_argument("--record", type=str, default=None, help="Enable recording mode: save observations to pickle file on Ctrl+C")
|
|
parser.add_argument("--switch-threshold", type=float, default=0.05, help="Vel command magnitude threshold for walking/standing switch (default: 0.05)")
|
|
parser.add_argument("--ground-pick-period", type=float, default=4.0, help="Ground pick phase period in seconds (default: 4.0)")
|
|
parser.add_argument("--new-cmd-obs", action="store_true",
|
|
help="Use the unified 13D command obs layout (twist+head_pose+body_pose). "
|
|
"Required for policies trained with the new pose-command-tracking setup. "
|
|
"Old policies (51D obs, head_offset added to ctrl) need this flag OFF.")
|
|
parser.add_argument("--current-limit", type=float, default=1.75,
|
|
help="XL330 firmware current limit [A]. Actuator torque is clipped to "
|
|
"+/- current_limit * kt (kt from the bam package), matching the "
|
|
"current saturation modeled in training. <=0 disables.")
|
|
parser.add_argument("--foot-friction", type=float, default=None,
|
|
help="Override the foot sliding friction (mu) to emulate the real grippy "
|
|
"PU sole. Training used mu~1.0 (range 0.7-1.3); real PU is likely "
|
|
"~1.5-2.5. e.g. --foot-friction 2.0")
|
|
parser.add_argument("--foot-solref", type=float, default=None,
|
|
help="Soften foot contact: solref time constant (s) for the foot geoms "
|
|
"(default sim ~0.02 = stiff/rigid). Larger = softer, to emulate the "
|
|
"compliant PU sole. e.g. --foot-solref 0.04")
|
|
args = parser.parse_args()
|
|
|
|
if not args.walking and not args.standing and not args.sitstand:
|
|
parser.error("At least one of --walking, --standing or --sitstand must be provided")
|
|
if args.sitstand and not args.new_cmd_obs:
|
|
parser.error("--sitstand policies use the unified 13D command obs (61D); add --new-cmd-obs")
|
|
if (args.kick_left or args.kick_right or args.roulade) and not args.new_cmd_obs:
|
|
parser.error("--kick-left/--kick-right/--roulade policies use the unified 13D command obs (61D); add --new-cmd-obs")
|
|
if (args.kick_left or args.kick_right or args.roulade) and args.roller:
|
|
parser.error("kick/roulade policies are trained on the walking robot, not the roller model")
|
|
|
|
# Parse delay arguments
|
|
delay_min_lag = 0
|
|
delay_max_lag = 0
|
|
if args.delay is not None:
|
|
if len(args.delay) == 0:
|
|
delay_min_lag = 1
|
|
delay_max_lag = 2
|
|
elif len(args.delay) == 1:
|
|
delay_min_lag = args.delay[0]
|
|
delay_max_lag = args.delay[0]
|
|
elif len(args.delay) == 2:
|
|
delay_min_lag = args.delay[0]
|
|
delay_max_lag = args.delay[1]
|
|
else:
|
|
print("Error: --delay accepts 0, 1, or 2 arguments")
|
|
return
|
|
|
|
# Load MuJoCo model. Kick policies get a scene with a ball to kick.
|
|
if args.roller:
|
|
xml_path = MICRODUCK_ROLLERS_XML
|
|
elif args.kick_left or args.kick_right:
|
|
xml_path = MICRODUCK_BALL_XML
|
|
else:
|
|
xml_path = MICRODUCK_XML
|
|
print(f"Loading MuJoCo model from: {xml_path}")
|
|
model = mujoco.MjModel.from_xml_path(xml_path)
|
|
model.opt.timestep = 0.005
|
|
data = mujoco.MjData(model)
|
|
|
|
# XL330 firmware current limit. The motors saturate current at ~1.75 A; since
|
|
# torque = kt * current, this caps the actuator force at +/- kt * I_max. The
|
|
# MuJoCo position actuators here are not the BAM voltage model, but clipping
|
|
# their output force reproduces the same current saturation the policy was
|
|
# trained against (see BamActuator.max_current). kt comes from the bam package.
|
|
if args.current_limit and args.current_limit > 0:
|
|
from bam.model import load_model
|
|
kt = load_model(motor_name="xl330", model="m6").kt.value
|
|
torque_limit = kt * args.current_limit
|
|
model.actuator_forcerange[:, 0] = -torque_limit
|
|
model.actuator_forcerange[:, 1] = torque_limit
|
|
model.actuator_forcelimited[:] = 1
|
|
print(f"Current limit: {args.current_limit:.2f} A -> torque limit "
|
|
f"+/-{torque_limit:.4f} Nm (kt={kt:.4f})")
|
|
|
|
# Foot contact override — emulate the real grippy + soft PU sole to check
|
|
# whether it reproduces the on-robot forward-fall-at-speed. Training used
|
|
# rigid feet at mu~1.0; the real sole is grippier (higher mu) and compliant
|
|
# (softer solref). Applied to the foot collision geoms only.
|
|
if args.foot_friction is not None or args.foot_solref is not None:
|
|
import re as _re
|
|
n_feet = 0
|
|
for g in range(model.ngeom):
|
|
gname = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_GEOM, g)
|
|
if gname and _re.match(r"^(left|right)_foot_collision$", gname):
|
|
if args.foot_friction is not None:
|
|
model.geom_friction[g, 0] = args.foot_friction # tangential mu
|
|
if args.foot_solref is not None:
|
|
model.geom_solref[g, 0] = args.foot_solref # softer contact
|
|
model.geom_solref[g, 1] = 1.0
|
|
n_feet += 1
|
|
print(f"Foot override on {n_feet} geoms: "
|
|
f"mu={args.foot_friction if args.foot_friction is not None else 'default'}, "
|
|
f"solref={args.foot_solref if args.foot_solref is not None else 'default'}")
|
|
|
|
# Initialize policy
|
|
policy = PolicyInference(
|
|
model, data,
|
|
walking_onnx_path=args.walking,
|
|
action_scale=args.action_scale,
|
|
delay_min_lag=delay_min_lag,
|
|
delay_max_lag=delay_max_lag,
|
|
standing_onnx_path=args.standing,
|
|
switch_threshold=args.switch_threshold,
|
|
use_projected_gravity=not args.raw_accelerometer,
|
|
ground_pick_onnx_path=args.ground_pick,
|
|
ground_pick_period=args.ground_pick_period,
|
|
sit_onnx_path=args.sit,
|
|
new_cmd_obs=args.new_cmd_obs,
|
|
slope_onnx_path=args.slope,
|
|
sitstand_onnx_path=args.sitstand,
|
|
kick_left_onnx_path=args.kick_left,
|
|
kick_right_onnx_path=args.kick_right,
|
|
roulade_onnx_path=args.roulade,
|
|
kick_duration=args.kick_duration,
|
|
roulade_duration=args.roulade_duration,
|
|
)
|
|
policy.set_vel_cmd(args.lin_vel_x, args.lin_vel_y, args.ang_vel_z)
|
|
|
|
# Set realistic wheel bearing friction for roller inference (must be done
|
|
# programmatically — non-zero frictionloss in the XML breaks training)
|
|
if args.roller:
|
|
import re
|
|
for j in range(model.njnt):
|
|
name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_JOINT, j)
|
|
if name and re.match(r"^passive_.*", name):
|
|
dof_adr = model.jnt_dofadr[j]
|
|
model.dof_frictionloss[dof_adr] = 0.003
|
|
|
|
# Per-mode velocity command limits matching training ranges
|
|
if args.roller:
|
|
policy.vel_step_x = 0.05 # lin_vel_x step (range -0.5..0.6)
|
|
policy.vel_step_y = 0.0 # no lateral command for rollers
|
|
policy.vel_step_ang = 0.1 # heading error step (range ±1.0 rad)
|
|
policy.vel_max_x = 0.6
|
|
policy.vel_min_x = -0.5 # negative = brake
|
|
policy.vel_max_y = 0.0
|
|
policy.vel_min_y = 0.0
|
|
policy.vel_max_ang = 1.0 # ±1.0 rad heading error
|
|
else:
|
|
policy.vel_max_x = 0.3
|
|
policy.vel_min_x = -0.3
|
|
policy.vel_max_y = 0.2
|
|
policy.vel_min_y = -0.2
|
|
policy.vel_max_ang = 1.5
|
|
|
|
# Set initial position to default pose
|
|
freejoint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint")
|
|
qpos_adr = model.jnt_qposadr[freejoint_id]
|
|
data.qpos[qpos_adr + 0] = 0.0
|
|
data.qpos[qpos_adr + 1] = 0.0
|
|
data.qpos[qpos_adr + 2] = 0.1385 if args.roller else 0.125 # rollers add 13.5mm height
|
|
data.qpos[qpos_adr + 3:qpos_adr + 7] = [1, 0, 0, 0]
|
|
for i, qpos_idx in enumerate(policy.joint_qpos_indices):
|
|
data.qpos[qpos_idx] = policy.default_pose[i]
|
|
data.ctrl[:] = policy.default_pose
|
|
mujoco.mj_forward(model, data)
|
|
|
|
# Verify observation size
|
|
test_obs = policy.get_observations()
|
|
cmd_dim = 13 if policy.new_cmd_obs else 3
|
|
expected_obs_size = 3 + 3 + policy.n_joints + policy.n_joints + policy.n_joints + cmd_dim
|
|
breakdown = (
|
|
f"3(ang_vel) + 3(proj_grav) + {policy.n_joints}(joint_pos) + "
|
|
f"{policy.n_joints}(joint_vel) + {policy.n_joints}(last_action) + {cmd_dim}(command)"
|
|
)
|
|
|
|
if test_obs.size != expected_obs_size:
|
|
print(f"\nWARNING: Observation size mismatch!")
|
|
print(f" Expected: {expected_obs_size}")
|
|
print(f" Got: {test_obs.size}")
|
|
print(f" Breakdown: {breakdown}")
|
|
print()
|
|
|
|
print("\n" + "="*80)
|
|
print("MicroDuck Policy Inference")
|
|
print("="*80)
|
|
print(f"Control frequency: 50 Hz (decimation: 4)")
|
|
print(f"Simulation timestep: {model.opt.timestep}s")
|
|
print(f"Observation size: {test_obs.size} (expected: {expected_obs_size})")
|
|
if policy.walking_session:
|
|
print(f"Walking policy: loaded")
|
|
if policy.standing_session:
|
|
print(f"Standing policy: loaded (body pose: z=±{BODY_CMD_MAX_Z*1000:.0f}mm, pitch/roll=±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
if policy.walking_session and policy.standing_session:
|
|
print(f" Switch threshold: {policy.switch_threshold} (vel cmd magnitude)")
|
|
if policy.ground_pick_session:
|
|
print(f"Ground pick policy: loaded (press G)")
|
|
if policy.sit_session:
|
|
kind = "Sitstand" if policy.is_sitstand else "Sit"
|
|
print(f"{kind} policy: loaded (press Y to toggle)")
|
|
if policy.slope_session:
|
|
print(f"Slope policy: loaded (press Y to toggle, passive descent)")
|
|
_behavior_keys = {"kick_left": "K", "kick_right": "L", "roulade": "R"}
|
|
for _name in policy.behavior_sessions:
|
|
print(f"{_name} policy: loaded (press {_behavior_keys[_name]}, "
|
|
f"auto-return after {policy.behavior_durations[_name]:.1f}s)")
|
|
print(f"Active policy: {policy.current_policy}")
|
|
print("Close viewer window to exit")
|
|
print()
|
|
|
|
decimation = 4
|
|
control_step_count = 0
|
|
control_dt = decimation * model.opt.timestep
|
|
|
|
# Rolling buffer of trunk world-frame xy velocity over the last 1 s, used
|
|
# to print a running average so we can compare commanded vs achieved speed.
|
|
from collections import deque
|
|
_vel_window_steps = max(1, int(round(1.0 / control_dt))) # ≈ 50 @ 50 Hz
|
|
vel_history = deque(maxlen=_vel_window_steps)
|
|
|
|
csv_data = [] if args.save_csv else None
|
|
recorded_observations = [] if args.record else None
|
|
policy_enabled = not args.record
|
|
policy_enable_time = None
|
|
original_kp = None
|
|
if args.record:
|
|
original_kp = model.actuator_gainprm[:, 0].copy()
|
|
|
|
# Cache the trunk freejoint qvel address so the push handler can write to
|
|
# the trunk's world-frame linear velocity directly (qvel[0..3]).
|
|
_freejoint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, "trunk_base_freejoint")
|
|
_trunk_qvel_adr = int(model.jnt_dofadr[_freejoint_id])
|
|
PUSH_MAX = 1.0 # matches the final velstand push_magnitude curriculum cap
|
|
|
|
def random_push():
|
|
"""Set the trunk's world-frame xy velocity to a random vector of
|
|
magnitude PUSH_MAX, simulating the push_by_setting_velocity training
|
|
event. Doesn't accumulate — overwrites current linear velocity."""
|
|
import random
|
|
angle = random.uniform(0, 2 * np.pi)
|
|
vx = PUSH_MAX * np.cos(angle)
|
|
vy = PUSH_MAX * np.sin(angle)
|
|
data.qvel[_trunk_qvel_adr + 0] = vx
|
|
data.qvel[_trunk_qvel_adr + 1] = vy
|
|
print(f"PUSH applied: v=[{vx:.2f}, {vy:.2f}, 0] m/s (angle={np.degrees(angle):.0f}°)")
|
|
|
|
# Keys come from the TERMINAL (raw stdin, see TerminalInput) — not from the
|
|
# MuJoCo viewer window, whose keypresses also fire built-in visualization
|
|
# shortcuts. `key` is a symbolic name: "up"/"down"/"left"/"right", " ", or
|
|
# a lowercase letter.
|
|
quit_requested = False
|
|
|
|
def handle_key(key):
|
|
nonlocal policy_enabled, quit_requested
|
|
try:
|
|
if key == "up":
|
|
if policy.head_mode:
|
|
policy.head_offset[1] = np.clip(policy.head_offset[1] + policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode:
|
|
policy.bump_body("z", policy.body_cmd_step_z)
|
|
else:
|
|
policy.set_vel_cmd(policy.vel_max_x, policy.vel_cmd[1], policy.vel_cmd[2])
|
|
elif key == "down":
|
|
if policy.head_mode:
|
|
policy.head_offset[1] = np.clip(policy.head_offset[1] - policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode:
|
|
policy.bump_body("z", -policy.body_cmd_step_z)
|
|
else:
|
|
policy.set_vel_cmd(policy.vel_min_x, policy.vel_cmd[1], policy.vel_cmd[2])
|
|
elif key == "right":
|
|
if policy.head_mode:
|
|
policy.head_offset[2] = np.clip(policy.head_offset[2] - policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode:
|
|
policy.bump_body("pitch", -policy.body_cmd_step_angle)
|
|
elif args.roller:
|
|
new_ang = np.clip(policy.vel_cmd[2] - policy.vel_step_ang, -policy.vel_max_ang, policy.vel_max_ang)
|
|
policy.set_vel_cmd(policy.vel_cmd[0], policy.vel_cmd[1], new_ang)
|
|
else:
|
|
policy.set_vel_cmd(policy.vel_cmd[0], policy.vel_min_y, policy.vel_cmd[2])
|
|
elif key == "left":
|
|
if policy.head_mode:
|
|
policy.head_offset[2] = np.clip(policy.head_offset[2] + policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode:
|
|
policy.bump_body("pitch", policy.body_cmd_step_angle)
|
|
elif args.roller:
|
|
new_ang = np.clip(policy.vel_cmd[2] + policy.vel_step_ang, -policy.vel_max_ang, policy.vel_max_ang)
|
|
policy.set_vel_cmd(policy.vel_cmd[0], policy.vel_cmd[1], new_ang)
|
|
else:
|
|
policy.set_vel_cmd(policy.vel_cmd[0], policy.vel_max_y, policy.vel_cmd[2])
|
|
elif key == " ":
|
|
if policy.head_mode:
|
|
policy.head_offset[:] = 0.0
|
|
policy._update_command()
|
|
print("Head offset reset to zero")
|
|
elif policy.body_pose_mode:
|
|
policy.body_cmd[:] = 0.0
|
|
policy._update_command()
|
|
print("Body pose cmd reset to zero")
|
|
else:
|
|
policy.set_vel_cmd(0.0, 0.0, 0.0)
|
|
elif key == "t":
|
|
# Toggle policy inference on/off. When OFF the controller stops
|
|
# querying the ONNX policy and the motors hold the last applied
|
|
# target (no fresh ctrl writes).
|
|
policy_enabled = not policy_enabled
|
|
print(f"Policy inference: {'ON' if policy_enabled else 'OFF (paused)'}")
|
|
elif key == "g":
|
|
policy.trigger_ground_pick()
|
|
elif key == "k":
|
|
policy.trigger_behavior("kick_left")
|
|
elif key == "l":
|
|
policy.trigger_behavior("kick_right")
|
|
elif key == "r":
|
|
policy.trigger_behavior("roulade")
|
|
elif key == "q":
|
|
quit_requested = True
|
|
print("Quit requested")
|
|
elif key == "y":
|
|
# Y toggles whichever aux policy is loaded (--sit or --slope).
|
|
if policy.sit_session is not None:
|
|
policy.toggle_sit()
|
|
else:
|
|
policy.toggle_slope_mode()
|
|
elif key == "h":
|
|
policy.toggle_head_mode()
|
|
elif key == "b":
|
|
policy.toggle_body_pose_mode()
|
|
elif key == "p":
|
|
random_push()
|
|
elif key == "a":
|
|
if policy.head_mode:
|
|
policy.head_offset[3] = np.clip(policy.head_offset[3] + policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode:
|
|
policy.bump_body("roll", policy.body_cmd_step_angle)
|
|
else:
|
|
policy.set_vel_cmd(policy.vel_cmd[0], policy.vel_cmd[1], policy.vel_max_ang)
|
|
elif key == "e":
|
|
if policy.head_mode:
|
|
policy.head_offset[3] = np.clip(policy.head_offset[3] - policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode:
|
|
policy.bump_body("roll", -policy.body_cmd_step_angle)
|
|
else:
|
|
policy.set_vel_cmd(policy.vel_cmd[0], policy.vel_cmd[1], -policy.vel_max_ang)
|
|
elif key == "z":
|
|
if policy.head_mode:
|
|
policy.head_offset[0] = np.clip(policy.head_offset[0] + policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode and policy.new_cmd_obs:
|
|
policy.bump_body("yaw", policy.body_cmd_step_angle)
|
|
elif key == "s":
|
|
if policy.head_mode:
|
|
policy.head_offset[0] = np.clip(policy.head_offset[0] - policy.head_step, -policy.head_max, policy.head_max)
|
|
policy._update_command()
|
|
print(f"Head offset: neck={policy.head_offset[0]:.2f} pitch={policy.head_offset[1]:.2f} yaw={policy.head_offset[2]:.2f} roll={policy.head_offset[3]:.2f}")
|
|
elif policy.body_pose_mode and policy.new_cmd_obs:
|
|
policy.bump_body("yaw", -policy.body_cmd_step_angle)
|
|
except Exception as e:
|
|
print(f"Key press error: {e}")
|
|
|
|
print("\nKeyboard controls (type in THIS terminal — the viewer window no longer captures keys):")
|
|
print(" [ Velocity mode (default) ]")
|
|
print(" UP arrow: increase lin_vel_x (push/accelerate)")
|
|
print(" DOWN arrow: decrease lin_vel_x (0=coast, negative=brake)")
|
|
if args.roller:
|
|
print(" LEFT/RIGHT arrow: turn left/right (ang_vel_z heading error)")
|
|
print(" A / E: turn left/right (ang_vel_z, incremental)")
|
|
else:
|
|
print(" LEFT/RIGHT arrow: strafe left/right (lin_vel_y)")
|
|
print(" A / E: turn left/right (ang_vel_z)")
|
|
print(" SPACE: coast (zero all commands)")
|
|
print(" T: toggle policy inference on/off (paused = motors hold last target)")
|
|
print(" G: trigger ground pick (requires --ground-pick)")
|
|
print(" Y: toggle sit (with --sit/--sitstand) or slope mode (with --slope)")
|
|
print(" K: kick with LEFT foot (requires --kick-left)")
|
|
print(" L: kick with RIGHT foot (requires --kick-right)")
|
|
print(" R: roulade / forward roll (requires --roulade)")
|
|
print(f" P: random push (trunk vel = {PUSH_MAX:.1f} m/s in random direction)")
|
|
print(" Q: quit")
|
|
print(" [ Body pose mode — press B to toggle ]")
|
|
print(f" UP/DOWN arrow: Δz ±10mm (max ±{BODY_CMD_MAX_Z*1000:.0f}mm)")
|
|
print(f" LEFT/RIGHT arrow: Δpitch ±10° (max ±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
print(f" A / E: Δroll ±10° (max ±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
if args.new_cmd_obs:
|
|
print(f" Z / S: Δyaw ±10° (new_cmd_obs only, max ±{math.degrees(BODY_CMD_MAX_ANGLE):.0f}°)")
|
|
print(" SPACE: reset body pose to zero")
|
|
print(" [ Head mode — press H to toggle ]")
|
|
print(" Z / S: neck_pitch ±step")
|
|
print(" UP/DOWN arrow: head_pitch ±step")
|
|
print(" LEFT/RIGHT arrow: head_yaw ±step")
|
|
print(" A / E: head_roll ±step")
|
|
print(" SPACE: reset head offset to zero")
|
|
|
|
with TerminalInput() as term, \
|
|
mujoco.viewer.launch_passive(model, data, show_left_ui=False, show_right_ui=False) as viewer:
|
|
viewer.sync()
|
|
start_time = time.time()
|
|
|
|
if args.record:
|
|
policy_enable_time = start_time + 1.0
|
|
print("Recording mode: policy will be enabled after 1 second standby")
|
|
for i in range(model.nu):
|
|
model.actuator_gainprm[i, 0] = 2.0
|
|
model.actuator_biasprm[i, 1] = -2.0
|
|
print(" Standby mode: kp set to 2.0")
|
|
|
|
try:
|
|
prev_step_time = time.time()
|
|
|
|
while viewer.is_running() and not quit_requested:
|
|
step_start = time.time()
|
|
|
|
for key in term.get_keys():
|
|
handle_key(key)
|
|
|
|
if not policy_enabled and policy_enable_time is not None:
|
|
if step_start >= policy_enable_time:
|
|
policy_enabled = True
|
|
if original_kp is not None:
|
|
for i in range(model.nu):
|
|
kp = original_kp[i]
|
|
model.actuator_gainprm[i, 0] = kp
|
|
model.actuator_biasprm[i, 1] = -kp
|
|
print("Policy inference enabled (after 1s standby)")
|
|
print(f" Restored original kp gains (range: [{original_kp.min():.2f}, {original_kp.max():.2f}])")
|
|
|
|
actual_dt = step_start - prev_step_time
|
|
prev_step_time = step_start
|
|
|
|
policy.update_ground_pick_phase(actual_dt)
|
|
policy.update_behavior(actual_dt)
|
|
|
|
if policy_enabled:
|
|
action = policy.infer()
|
|
policy.apply_action(action)
|
|
else:
|
|
# Paused: keep last ctrl, don't query the policy. Motors
|
|
# hold position. Use a zero action just so downstream
|
|
# logging (csv/debug) sees something consistent.
|
|
action = np.zeros(policy.n_joints, dtype=np.float32)
|
|
|
|
control_step_count += 1
|
|
|
|
# Track BODY-frame forward/lateral velocity + yaw rate, print the
|
|
# 1-second moving average once per second vs the commanded values.
|
|
# Body frame so "forward" / "turn" are directly comparable to the
|
|
# command (which is in the robot frame): lets us see if the policy
|
|
# actually achieves commanded forward speed and turn rate.
|
|
quat = data.qpos[qpos_adr + 3:qpos_adr + 7].astype(np.float32)
|
|
v_world = np.array([
|
|
data.qvel[_trunk_qvel_adr + 0],
|
|
data.qvel[_trunk_qvel_adr + 1],
|
|
data.qvel[_trunk_qvel_adr + 2],
|
|
], dtype=np.float32)
|
|
v_body = policy.quat_rotate_inverse(quat, v_world)
|
|
yaw_rate = float(data.qvel[_trunk_qvel_adr + 5]) # body-frame wz
|
|
vel_history.append((float(v_body[0]), float(v_body[1]), yaw_rate))
|
|
if control_step_count % _vel_window_steps == 0 and len(vel_history) > 0:
|
|
n = len(vel_history)
|
|
avg_fwd = sum(v[0] for v in vel_history) / n
|
|
avg_lat = sum(v[1] for v in vel_history) / n
|
|
avg_yaw = sum(v[2] for v in vel_history) / n
|
|
cmd_x, cmd_y, cmd_yaw = policy.vel_cmd[0], policy.vel_cmd[1], policy.vel_cmd[2]
|
|
trunk_z = float(data.qpos[qpos_adr + 2])
|
|
print(
|
|
f"[vel 1s avg] achieved/cmd fwd={avg_fwd:+.2f}/{cmd_x:+.2f} "
|
|
f"lat={avg_lat:+.2f}/{cmd_y:+.2f} m/s "
|
|
f"yaw={avg_yaw:+.2f}/{cmd_yaw:+.2f} rad/s "
|
|
f"trunk_z={trunk_z*1000:.1f} mm"
|
|
)
|
|
|
|
if csv_data is not None:
|
|
obs = policy.get_observations()
|
|
row = {'step': control_step_count, 'time': control_step_count * control_dt}
|
|
for i in range(obs.size):
|
|
row[f'obs_{i}'] = obs[i]
|
|
for i in range(action.size):
|
|
row[f'action_{i}'] = action[i]
|
|
csv_data.append(row)
|
|
|
|
if recorded_observations is not None:
|
|
obs = policy.get_observations()
|
|
timestamp = time.time() - start_time
|
|
recorded_observations.append({'timestamp': timestamp, 'observation': obs.tolist()})
|
|
|
|
if args.debug:
|
|
should_print = control_step_count <= 10 or control_step_count % 50 == 0
|
|
if should_print:
|
|
obs = policy.get_observations()
|
|
pos = data.qpos[qpos_adr:qpos_adr + 3]
|
|
quat = data.qpos[qpos_adr + 3:qpos_adr + 7]
|
|
com_height = pos[2]
|
|
|
|
print(f"\n{'='*70}")
|
|
print(f"Step {control_step_count} DEBUG:")
|
|
print(f"{'='*70}")
|
|
print(f"Active policy: {policy.current_policy}")
|
|
print(f"Base state:")
|
|
print(f" Position: [{pos[0]:7.4f}, {pos[1]:7.4f}, {pos[2]:7.4f}]")
|
|
print(f" CoM height: {com_height:7.4f}")
|
|
print(f" Quaternion: [{quat[0]:7.4f}, {quat[1]:7.4f}, {quat[2]:7.4f}, {quat[3]:7.4f}]")
|
|
print(f"\nObservation (shape {obs.shape}, total {obs.size}):")
|
|
print(f" Ang vel [0:3]: {obs[0:3]}")
|
|
print(f" Proj grav [3:6]: {obs[3:6]}")
|
|
print(f" Joint pos [6:{6+policy.n_joints}]: {obs[6:6+policy.n_joints]}")
|
|
print(f" Joint vel [{6+policy.n_joints}:{6+2*policy.n_joints}]: {obs[6+policy.n_joints:6+2*policy.n_joints]}")
|
|
print(f" Last action [{6+2*policy.n_joints}:{6+3*policy.n_joints}]: {obs[6+2*policy.n_joints:6+3*policy.n_joints]}")
|
|
cmd_end = 6+3*policy.n_joints+3
|
|
print(f" Command [{6+3*policy.n_joints}:{cmd_end}]: {obs[6+3*policy.n_joints:cmd_end]}")
|
|
if policy.current_policy == "standing":
|
|
print(f" Body cmd (raw): z={policy.body_cmd[0]*1000:.1f}mm pitch={math.degrees(policy.body_cmd[1]):.1f}° roll={math.degrees(policy.body_cmd[2]):.1f}°")
|
|
print(f"\nAction output:")
|
|
print(f" Raw action: {action}")
|
|
print(f" Action min/max: [{action.min():.4f}, {action.max():.4f}]")
|
|
if policy.use_delay:
|
|
print(f" Delay: {policy.current_lag} timesteps (buffered)")
|
|
print(f" Applied ctrl (first 5): {data.ctrl[:5]}")
|
|
print(f" Applied ctrl (last 5): {data.ctrl[-5:]}")
|
|
|
|
for _ in range(decimation):
|
|
mujoco.mj_step(model, data)
|
|
|
|
viewer.sync()
|
|
|
|
elapsed = time.time() - step_start
|
|
sleep_time = control_dt - elapsed
|
|
if sleep_time > 0:
|
|
time.sleep(sleep_time)
|
|
|
|
except KeyboardInterrupt:
|
|
print("\n\nKeyboardInterrupt received (Ctrl+C). Saving data...")
|
|
|
|
print("\nInference stopped.")
|
|
|
|
if csv_data is not None and len(csv_data) > 0:
|
|
print(f"\nSaving {len(csv_data)} steps to: {args.save_csv}")
|
|
with open(args.save_csv, 'w', newline='') as csvfile:
|
|
fieldnames = csv_data[0].keys()
|
|
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
writer.writerows(csv_data)
|
|
print(f"CSV file saved successfully!")
|
|
print(f" Columns: {len(fieldnames)}")
|
|
print(f" Rows: {len(csv_data)}")
|
|
|
|
if recorded_observations is not None and len(recorded_observations) > 0:
|
|
print(f"\nSaving {len(recorded_observations)} recorded observations to: {args.record}")
|
|
with open(args.record, 'wb') as f:
|
|
pickle.dump(recorded_observations, f)
|
|
print(f"Recorded observations saved to {args.record}")
|
|
print(f" Observations: {len(recorded_observations)}")
|
|
print(f" Duration: {recorded_observations[-1]['timestamp']:.2f}s")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|