380 lines
13 KiB
Python
380 lines
13 KiB
Python
"""CTS Policy deployment for Unitree Go2 in MuJoCo."""
|
|
|
|
import time
|
|
from collections import deque
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import NamedTuple
|
|
|
|
import imageio
|
|
import mujoco
|
|
import mujoco.viewer
|
|
import numpy as np
|
|
import pygame
|
|
import torch
|
|
import yaml
|
|
from argparse import ArgumentParser
|
|
|
|
|
|
# ============================================================================
|
|
# Types & Constants
|
|
# ============================================================================
|
|
|
|
ROOT_DIR = str(Path(__file__).parent.parent.parent)
|
|
CONFIG_DIR = f"{ROOT_DIR}/deploy/deploy_mujoco/configs"
|
|
VIDEO_DIR = Path(__file__).parent / "videos"
|
|
|
|
class CTSPolicyInputs(NamedTuple):
|
|
"""Input format for CTS policy."""
|
|
policy: torch.Tensor
|
|
single_obs: torch.Tensor
|
|
|
|
|
|
@dataclass
|
|
class ObsBlockCfg:
|
|
"""Configuration for observation block."""
|
|
name: str
|
|
dim: int # Single frame dimension
|
|
|
|
|
|
ACTUATOR_GROUPS = {
|
|
"hip": [0, 3, 6, 9],
|
|
"thigh": [1, 4, 7, 10],
|
|
"calf": [2, 5, 8, 11],
|
|
}
|
|
|
|
|
|
# ============================================================================
|
|
# Helper Functions
|
|
# ============================================================================
|
|
|
|
def get_gravity_orientation(quaternion: np.ndarray) -> np.ndarray:
|
|
"""Compute gravity vector in body frame from quaternion."""
|
|
qw, qx, qy, qz = quaternion
|
|
|
|
gravity = np.zeros(3)
|
|
gravity[0] = 2 * (-qz * qx + qw * qy)
|
|
gravity[1] = -2 * (qz * qy + qw * qx)
|
|
gravity[2] = 1 - 2 * (qw * qw + qz * qz)
|
|
|
|
return gravity
|
|
|
|
|
|
def pd_control(target_q: np.ndarray, q: np.ndarray, kp: np.ndarray,
|
|
target_dq: np.ndarray, dq: np.ndarray, kd: np.ndarray) -> np.ndarray:
|
|
"""Compute PD control torques."""
|
|
return (target_q - q) * kp + (target_dq - dq) * kd
|
|
|
|
|
|
def get_joystick_command(joystick, max_cmd: np.ndarray) -> np.ndarray:
|
|
"""Read command from Xbox controller."""
|
|
pygame.event.pump()
|
|
|
|
dead_zone = 0.1
|
|
axes = [joystick.get_axis(i) for i in [0, 1, 3]] # LX, LY, RX
|
|
axes = [0 if abs(a) < dead_zone else a for a in axes]
|
|
|
|
cmd = np.array([-axes[1] * max_cmd[0], -axes[0] * max_cmd[1], -axes[2] * max_cmd[2]], dtype=np.float32)
|
|
return cmd
|
|
|
|
|
|
def load_config(config_file: str) -> dict:
|
|
"""Load and parse YAML configuration."""
|
|
with open(f"{CONFIG_DIR}/{config_file}", "r") as f:
|
|
config = yaml.load(f, Loader=yaml.FullLoader)
|
|
|
|
# Replace path placeholders
|
|
config["policy_path"] = config["policy_path"].replace("{ROOT_DIR}", ROOT_DIR)
|
|
config["xml_path"] = config["xml_path"].replace("{ROOT_DIR}", ROOT_DIR)
|
|
|
|
return config
|
|
|
|
|
|
def build_observation(obs: np.ndarray, features: dict, obs_cfg: list, history_len: int) -> None:
|
|
"""Update stacked observation buffer with new frame features (in-place)."""
|
|
ptr = 0
|
|
for cfg in obs_cfg:
|
|
dim = cfg.dim
|
|
start, end = ptr, ptr + dim * history_len
|
|
|
|
# Roll history and insert new frame
|
|
obs[start:end] = np.roll(obs[start:end], shift=-dim, axis=0)
|
|
obs[end - dim:end] = features[cfg.name]
|
|
|
|
ptr = end
|
|
|
|
|
|
def extract_single_obs(obs: np.ndarray, obs_cfg: list, history_len: int) -> np.ndarray:
|
|
"""Extract most recent single-frame observation from stacked buffer."""
|
|
single = []
|
|
ptr = 0
|
|
for cfg in obs_cfg:
|
|
dim = cfg.dim
|
|
block = obs[ptr:ptr + dim * history_len]
|
|
single.append(block[-dim:]) # Get last dim elements (most recent frame)
|
|
ptr += dim * history_len
|
|
return np.concatenate(single, axis=0)
|
|
|
|
|
|
def apply_action(action: np.ndarray, default_angles: np.ndarray,
|
|
action_pos_scale: float) -> np.ndarray:
|
|
"""Transform policy action to position targets."""
|
|
return action * action_pos_scale + default_angles
|
|
|
|
|
|
def set_initial_state(data: mujoco.MjData, base_pos: np.ndarray, base_quat: np.ndarray,
|
|
joint_pos: np.ndarray) -> None:
|
|
"""Apply the Isaac-style initial pose before the first mj_forward."""
|
|
data.qpos[:] = 0.0
|
|
data.qvel[:] = 0.0
|
|
data.qpos[:3] = base_pos
|
|
data.qpos[3:7] = base_quat
|
|
data.qpos[7:] = joint_pos
|
|
|
|
|
|
def build_delay_buffers(default_pos: np.ndarray, delay_max: int) -> deque:
|
|
"""Create history buffers for delayed control targets."""
|
|
history_len = delay_max + 1
|
|
pos_history = deque((default_pos.copy() for _ in range(history_len)), maxlen=history_len)
|
|
return pos_history
|
|
|
|
|
|
def sample_delayed_targets(pos_history: deque, delay_min: int,
|
|
delay_max: int, rng: np.random.Generator) -> np.ndarray:
|
|
"""Apply per-actuator-group discrete control delays."""
|
|
delayed_pos = pos_history[-1].copy()
|
|
|
|
for joint_ids in ACTUATOR_GROUPS.values():
|
|
delay_steps = int(rng.integers(delay_min, delay_max + 1)) if delay_max > 0 else 0
|
|
pos_source = pos_history[-1 - delay_steps]
|
|
delayed_pos[joint_ids] = pos_source[joint_ids]
|
|
|
|
return delayed_pos
|
|
|
|
|
|
def init_joystick() -> tuple:
|
|
"""Initialize pygame and joystick if available."""
|
|
pygame.init()
|
|
|
|
if pygame.joystick.get_count() > 0:
|
|
joystick = pygame.joystick.Joystick(0)
|
|
joystick.init()
|
|
print(f"Detected Joystick: {joystick.get_name()}")
|
|
return joystick, True
|
|
|
|
print("No Joystick detected. Using default commands from config.")
|
|
return None, False
|
|
|
|
|
|
def setup_viewer_camera(viewer) -> None:
|
|
"""Configure tracking camera."""
|
|
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
|
viewer.cam.trackbodyid = 1
|
|
viewer.cam.distance = 3.0
|
|
viewer.cam.elevation = -30.0
|
|
viewer.cam.azimuth = 0.0
|
|
|
|
|
|
# ============================================================================
|
|
# Main
|
|
# ============================================================================
|
|
|
|
def main():
|
|
parser = ArgumentParser()
|
|
parser.add_argument("--save-video", action="store_true", help="Save video of simulation.")
|
|
args = parser.parse_args()
|
|
|
|
# Configuration
|
|
config_file = "go2.yaml"
|
|
render_fps = 240
|
|
|
|
config = load_config(config_file)
|
|
|
|
# Extract config parameters
|
|
sim_cfg = {
|
|
"duration": config["simulation_duration"],
|
|
"dt": config["simulation_dt"],
|
|
"decimation": config["control_decimation"],
|
|
}
|
|
|
|
# PD controller gains
|
|
kps = np.array(config["kps"], dtype=np.float32)
|
|
kds = np.array(config["kds"], dtype=np.float32)
|
|
default_angles = np.array(config["default_angles"], dtype=np.float32)
|
|
base_init_pos = np.array(config["base_init_pos"], dtype=np.float32)
|
|
base_init_quat = np.array(config["base_init_quat"], dtype=np.float32)
|
|
delay_min = int(config.get("actuator_delay_min", 0))
|
|
delay_max = int(config.get("actuator_delay_max", 0))
|
|
delay_seed = int(config.get("actuator_delay_seed", 0))
|
|
if delay_min < 0 or delay_max < delay_min:
|
|
raise ValueError(
|
|
f"Invalid actuator delay range: min={delay_min}, max={delay_max}."
|
|
)
|
|
delay_rng = np.random.default_rng(delay_seed)
|
|
|
|
# Scaling factors
|
|
scales = {
|
|
"lin_vel": config["lin_vel_scale"],
|
|
"ang_vel": config["ang_vel_scale"],
|
|
"dof_pos": config["dof_pos_scale"],
|
|
"dof_vel": config["dof_vel_scale"],
|
|
"action_pos": config["action_pos_scale"],
|
|
"cmd": np.array(config["cmd_scale"], dtype=np.float32),
|
|
}
|
|
|
|
# Policy dimensions
|
|
num_actions = config["num_actions"]
|
|
num_obs = config["num_obs"]
|
|
history_len = config.get("history_len", 1)
|
|
max_cmd = np.array(config["max_cmd"], dtype=np.float32)
|
|
cmd = np.array(config["cmd_init"], dtype=np.float32)
|
|
|
|
# Joint name mapping
|
|
idx_model2mj = idx_mj2model = list(range(num_actions))
|
|
if "mujoco_joint_names" in config and "model_joint_names" in config:
|
|
mj_names = config["mujoco_joint_names"]
|
|
model_names = config["model_joint_names"]
|
|
idx_model2mj = [model_names.index(j) for j in mj_names]
|
|
idx_mj2model = [mj_names.index(j) for j in model_names]
|
|
|
|
# Initialize joystick
|
|
joystick, use_joystick = init_joystick()
|
|
|
|
# Prepare video output
|
|
VIDEO_DIR.mkdir(parents=True, exist_ok=True)
|
|
model_name = Path(config["policy_path"]).stem
|
|
cmd_str = f"cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}"
|
|
|
|
# Initialize state
|
|
action = np.zeros(num_actions, dtype=np.float32)
|
|
target_dof_pos = default_angles.copy()
|
|
target_dof_vel = np.zeros(num_actions, dtype=np.float32)
|
|
obs = np.zeros(num_obs * history_len, dtype=np.float32)
|
|
pos_history = build_delay_buffers(target_dof_pos, delay_max)
|
|
|
|
# Build observation config dynamically
|
|
num_joints = num_actions
|
|
obs_cfg = [
|
|
ObsBlockCfg("ang_vel", 3),
|
|
ObsBlockCfg("gravity", 3),
|
|
ObsBlockCfg("cmd", 3),
|
|
ObsBlockCfg("joint_pos", num_joints),
|
|
ObsBlockCfg("joint_vel", num_actions),
|
|
ObsBlockCfg("last_action", num_actions),
|
|
]
|
|
|
|
# Load MuJoCo model
|
|
m = mujoco.MjModel.from_xml_path(config["xml_path"])
|
|
d = mujoco.MjData(m)
|
|
m.opt.timestep = sim_cfg["dt"]
|
|
set_initial_state(d, base_init_pos, base_init_quat, default_angles)
|
|
mujoco.mj_forward(m, d)
|
|
|
|
renderer = mujoco.Renderer(m, height=360, width=640)
|
|
|
|
# Load policy
|
|
policy = torch.jit.load(config["policy_path"])
|
|
|
|
# Setup video recording
|
|
writer = None
|
|
if args.save_video:
|
|
video_path = VIDEO_DIR / f"{model_name}_{cmd_str}.mp4"
|
|
video_fps = 50
|
|
sim_fps = 1.0 / sim_cfg["dt"]
|
|
frame_skip = max(1, int(sim_fps / video_fps))
|
|
writer = imageio.get_writer(video_path, fps=video_fps)
|
|
print(f"Recording: {video_path} (Sim FPS: {sim_fps:.1f}, Video FPS: {video_fps})")
|
|
|
|
render_substeps = int((1.0 / render_fps) / sim_cfg["dt"])
|
|
|
|
# Run simulation
|
|
with mujoco.viewer.launch_passive(m, d) as viewer:
|
|
setup_viewer_camera(viewer)
|
|
|
|
start_time = time.time()
|
|
counter = 0
|
|
|
|
while viewer.is_running() and time.time() - start_time < sim_cfg["duration"]:
|
|
step_start = time.time()
|
|
|
|
# Update command from joystick
|
|
if use_joystick and counter % sim_cfg["decimation"] == 0:
|
|
cmd = get_joystick_command(joystick, max_cmd)
|
|
|
|
# Compute and apply control
|
|
tau = pd_control(target_dof_pos, d.qpos[7:], kps, target_dof_vel, d.qvel[6:], kds)
|
|
d.ctrl[:] = tau
|
|
|
|
mujoco.mj_step(m, d)
|
|
|
|
# Record frame
|
|
if writer and counter % frame_skip == 0:
|
|
try:
|
|
renderer.update_scene(d, camera=viewer.cam)
|
|
writer.append_data(renderer.render())
|
|
except Exception as e:
|
|
print(f"Render error: {e}")
|
|
|
|
counter += 1
|
|
|
|
# Policy update at control frequency
|
|
if counter % sim_cfg["decimation"] == 0:
|
|
# Extract sensor data
|
|
qj = d.qpos[7:]
|
|
dqj = d.qvel[6:]
|
|
quat = d.qpos[3:7]
|
|
ang_vel = d.qvel[3:6]
|
|
|
|
# Scale observations
|
|
qj = (qj - default_angles) * scales["dof_pos"]
|
|
dqj = dqj * scales["dof_vel"]
|
|
ang_vel = ang_vel * scales["ang_vel"]
|
|
|
|
gravity = get_gravity_orientation(quat)
|
|
|
|
# Build observation features
|
|
features = {
|
|
"ang_vel": ang_vel,
|
|
"gravity": gravity,
|
|
"cmd": cmd * scales["cmd"],
|
|
"joint_pos": qj[idx_mj2model],
|
|
"joint_vel": dqj[idx_mj2model],
|
|
"last_action": action[idx_mj2model],
|
|
}
|
|
|
|
# Update stacked observation
|
|
build_observation(obs, features, obs_cfg, history_len)
|
|
|
|
# Extract single-frame observation
|
|
single_obs = extract_single_obs(obs, obs_cfg, history_len)
|
|
|
|
# Run policy
|
|
obs_tensor = torch.from_numpy(obs).unsqueeze(0)
|
|
single_tensor = torch.from_numpy(single_obs).unsqueeze(0)
|
|
|
|
result = policy(CTSPolicyInputs(policy=obs_tensor, single_obs=single_tensor))
|
|
action = result.detach().cpu().numpy().squeeze()[idx_model2mj]
|
|
|
|
# Apply action
|
|
latest_target_pos = apply_action(action, default_angles, scales["action_pos"])
|
|
pos_history.append(latest_target_pos.copy())
|
|
target_dof_pos = sample_delayed_targets(pos_history, delay_min, delay_max, delay_rng)
|
|
|
|
# Sync viewer
|
|
if counter % render_substeps == 0:
|
|
viewer.sync()
|
|
|
|
# Time management
|
|
sleep_time = sim_cfg["dt"] - (time.time() - step_start) - 0.1
|
|
if sleep_time > 0:
|
|
time.sleep(sleep_time)
|
|
|
|
# Cleanup
|
|
if writer:
|
|
print(f"Video saved: {video_path}")
|
|
writer.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|