refactor deploy scripts.
This commit is contained in:
wertyuilife
@@ -1,5 +1,10 @@
|
||||
policy_path: "{ROOT_DIR}/deploy/pretrain/go2/go2_moe_cts_1500.pt"
|
||||
xml_path: "{ROOT_DIR}/resources/go2/stairs.xml"
|
||||
policy_path: "{ROOT_DIR}/deploy/pre_train/wyh/go2_moe_cts_v1_49500.pt"
|
||||
|
||||
xml_path: "{ROOT_DIR}/resource/stairs.xml"
|
||||
|
||||
render_fps: 240
|
||||
video_fps: 60
|
||||
save_video: false
|
||||
|
||||
# Total simulation time
|
||||
simulation_duration: 60000000.0
|
||||
@@ -10,7 +15,7 @@ control_decimation: 10
|
||||
base_init_pos: [0.0, 0.0, 0.57]
|
||||
base_init_quat: [1.0, 0.0, 0.0, 0.0]
|
||||
|
||||
# For IsaacLab DelayedPDActuatorCfg, if not used, set all delay to 0.
|
||||
# Based on the implementation of IsaacLab's DelayedPDActuatorCfg.
|
||||
actuator_delay_min: 0
|
||||
actuator_delay_max: 0
|
||||
actuator_delay_seed: 0
|
||||
|
||||
@@ -1,387 +1,134 @@
|
||||
"""CTS Policy deployment for Unitree Go2 in MuJoCo."""
|
||||
"""Simplified 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
|
||||
|
||||
from utils import (
|
||||
build_delay_buffers,
|
||||
display_current_command,
|
||||
gravity_from_quat,
|
||||
infer_action,
|
||||
init_joystick,
|
||||
latest_obs_frame,
|
||||
load_config,
|
||||
open_video_writer,
|
||||
pd_control,
|
||||
push_obs_history,
|
||||
read_joystick_command,
|
||||
sample_delayed_targets,
|
||||
set_initial_state,
|
||||
setup_tracking_camera,
|
||||
MujocoRenderUtils,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 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
|
||||
CONFIG_NAME = "go2.yaml"
|
||||
VIDEO_DIR = Path(__file__).with_name("videos")
|
||||
ACTUATOR_GROUPS = (
|
||||
np.array([0, 3, 6, 9], dtype=np.int64),
|
||||
np.array([1, 4, 7, 10], dtype=np.int64),
|
||||
np.array([2, 5, 8, 11], dtype=np.int64),
|
||||
)
|
||||
|
||||
|
||||
@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
|
||||
|
||||
def display_current_command(cmd: np.ndarray) -> None:
|
||||
"""Refresh the current command on the same terminal line."""
|
||||
cmd_text = f"\rCurrent command | vx: {cmd[0]: .3f} vy: {cmd[1]: .3f} wz: {cmd[2]: .3f}"
|
||||
print(cmd_text, end="", flush=True)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# 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"],
|
||||
def build_features(data, action: np.ndarray, cmd: np.ndarray, cfg):
|
||||
joint_pos = (data.qpos[7:] - cfg.default_angles) * cfg.dof_pos_scale
|
||||
joint_vel = data.qvel[6:] * cfg.dof_vel_scale
|
||||
return {
|
||||
"ang_vel": data.qvel[3:6] * cfg.ang_vel_scale,
|
||||
"gravity": gravity_from_quat(data.qpos[3:7]),
|
||||
"cmd": cmd * cfg.cmd_scale,
|
||||
"joint_pos": joint_pos[cfg.idx_mj2model],
|
||||
"joint_vel": joint_vel[cfg.idx_mj2model],
|
||||
"last_action": action[cfg.idx_mj2model],
|
||||
}
|
||||
|
||||
# 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),
|
||||
}
|
||||
def action_to_target(action: np.ndarray, cfg):
|
||||
return cfg.default_angles + action * cfg.action_pos_scale
|
||||
|
||||
# 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()
|
||||
def main() -> None:
|
||||
cfg = load_config(CONFIG_NAME)
|
||||
layout = [
|
||||
("ang_vel", 3),
|
||||
("gravity", 3),
|
||||
("cmd", 3),
|
||||
("joint_pos", cfg.num_actions),
|
||||
("joint_vel", cfg.num_actions),
|
||||
("last_action", cfg.num_actions),
|
||||
]
|
||||
joystick = init_joystick()
|
||||
cmd = cfg.cmd_init.copy()
|
||||
display_current_command(cmd)
|
||||
|
||||
model = mujoco.MjModel.from_xml_path(str(cfg.xml_path))
|
||||
data = mujoco.MjData(model)
|
||||
model.opt.timestep = cfg.dt
|
||||
set_initial_state(data, cfg.base_init_pos, cfg.base_init_quat, cfg.default_angles)
|
||||
mujoco.mj_forward(model, data)
|
||||
|
||||
# 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]}"
|
||||
renderer = mujoco.Renderer(model, height=360, width=640)
|
||||
policy = torch.jit.load(str(cfg.policy_path))
|
||||
writer, frame_skip, video_path = open_video_writer(
|
||||
cfg.save_video, policy_path=cfg.policy_path, cmd=cmd, dt=cfg.dt, video_dir=VIDEO_DIR, video_fps=cfg.video_fps
|
||||
)
|
||||
|
||||
# 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)
|
||||
action = np.zeros(cfg.num_actions, dtype=np.float32)
|
||||
target_pos = cfg.default_angles.copy()
|
||||
target_vel = np.zeros(cfg.num_actions, dtype=np.float32)
|
||||
obs = np.zeros(cfg.num_obs * cfg.history_len, dtype=np.float32)
|
||||
pos_history = build_delay_buffers(target_pos, delay_max=cfg.delay_max)
|
||||
delay_rng = np.random.default_rng(cfg.delay_seed)
|
||||
render_substeps = max(1, int((1.0 / cfg.render_fps) / cfg.dt))
|
||||
mujoco_render_utils = MujocoRenderUtils()
|
||||
|
||||
with mujoco.viewer.launch_passive(model, data) as viewer:
|
||||
setup_tracking_camera(viewer)
|
||||
start_time = time.time()
|
||||
counter = 0
|
||||
|
||||
while viewer.is_running() and time.time() - start_time < sim_cfg["duration"]:
|
||||
while viewer.is_running() and time.time() - start_time < cfg.duration:
|
||||
step_start = time.time()
|
||||
if joystick and counter % cfg.decimation == 0:
|
||||
cmd = read_joystick_command(joystick, cfg.max_cmd)
|
||||
|
||||
# Update command from joystick
|
||||
if use_joystick and counter % sim_cfg["decimation"] == 0:
|
||||
cmd = get_joystick_command(joystick, max_cmd)
|
||||
data.ctrl[:] = pd_control(target_pos, data.qpos[7:], cfg.kps, target_vel, data.qvel[6:], cfg.kds)
|
||||
mujoco.mj_step(model, data)
|
||||
mujoco_render_utils.update(cmd, data)
|
||||
|
||||
# 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)
|
||||
renderer.update_scene(data, camera=viewer.cam)
|
||||
mujoco_render_utils.update_external_rendering(renderer, ctype='renderer')
|
||||
writer.append_data(renderer.render())
|
||||
except Exception as e:
|
||||
print(f"Render error: {e}")
|
||||
except Exception as exc:
|
||||
print(f"Render error: {exc}")
|
||||
|
||||
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)
|
||||
if counter % cfg.decimation == 0:
|
||||
push_obs_history(obs, build_features(data, action, cmd, cfg), layout, cfg.history_len)
|
||||
action = infer_action(policy, obs, latest_obs_frame(obs, layout, cfg.history_len), cfg.idx_model2mj)
|
||||
pos_history.append(action_to_target(action, cfg).copy())
|
||||
target_pos = sample_delayed_targets(pos_history, ACTUATOR_GROUPS, cfg.delay_min, cfg.delay_max, delay_rng)
|
||||
display_current_command(cmd)
|
||||
|
||||
# Sync viewer
|
||||
if counter % render_substeps == 0:
|
||||
viewer.sync()
|
||||
|
||||
# Time management
|
||||
sleep_time = sim_cfg["dt"] - (time.time() - step_start) - 0.1
|
||||
if counter % render_substeps == 0:
|
||||
mujoco_render_utils.update_external_rendering(viewer, ctype='viewer')
|
||||
viewer.sync()
|
||||
sleep_time = cfg.dt - (time.time() - step_start) - 0.1
|
||||
if sleep_time > 0:
|
||||
time.sleep(sleep_time)
|
||||
print()
|
||||
|
||||
# Cleanup
|
||||
print()
|
||||
if writer:
|
||||
print(f"Video saved: {video_path}")
|
||||
writer.close()
|
||||
print(f"Video saved: {video_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
328
deploy/deploy_mujoco/utils.py
Normal file
328
deploy/deploy_mujoco/utils.py
Normal file
@@ -0,0 +1,328 @@
|
||||
"""Shared helpers for MuJoCo deployment scripts."""
|
||||
|
||||
from collections import deque
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
from typing import NamedTuple
|
||||
|
||||
import imageio
|
||||
import numpy as np
|
||||
import pygame
|
||||
import torch
|
||||
import yaml
|
||||
|
||||
from typing import Union, Literal
|
||||
import mujoco
|
||||
import mujoco.viewer
|
||||
|
||||
ROOT_DIR = Path(__file__).resolve().parents[2]
|
||||
CONFIG_DIR = Path(__file__).with_name("configs")
|
||||
|
||||
|
||||
class CTSPolicyInputs(NamedTuple):
|
||||
policy: torch.Tensor
|
||||
single_obs: torch.Tensor
|
||||
|
||||
|
||||
def load_config(config_name: str):
|
||||
with (CONFIG_DIR / config_name).open("r", encoding="utf-8") as file:
|
||||
raw = yaml.safe_load(file)
|
||||
|
||||
def path_value(name: str) -> Path:
|
||||
return Path(raw[name].replace("{ROOT_DIR}", str(ROOT_DIR)))
|
||||
|
||||
data = {
|
||||
"policy_path": path_value("policy_path"),
|
||||
"xml_path": path_value("xml_path"),
|
||||
"duration": float(raw["simulation_duration"]),
|
||||
"dt": float(raw["simulation_dt"]),
|
||||
"decimation": int(raw["control_decimation"]),
|
||||
"history_len": int(raw.get("history_len", 1)),
|
||||
"num_actions": int(raw["num_actions"]),
|
||||
"num_obs": int(raw["num_obs"]),
|
||||
"delay_min": int(raw.get("actuator_delay_min", 0)),
|
||||
"delay_max": int(raw.get("actuator_delay_max", 0)),
|
||||
"delay_seed": int(raw.get("actuator_delay_seed", 0)),
|
||||
"save_video": bool(raw.get("save_video", False)),
|
||||
"render_fps": int(raw.get("render_fps", 60)),
|
||||
"video_fps": int(raw.get("video_fps", 50)),
|
||||
}
|
||||
if data["delay_min"] < 0 or data["delay_max"] < data["delay_min"]:
|
||||
raise ValueError(f"Invalid actuator delay range: min={data['delay_min']}, max={data['delay_max']}.")
|
||||
|
||||
for name in (
|
||||
"kps",
|
||||
"kds",
|
||||
"default_angles",
|
||||
"torque_limit",
|
||||
"base_init_pos",
|
||||
"base_init_quat",
|
||||
"max_cmd",
|
||||
"cmd_init",
|
||||
"cmd_scale",
|
||||
):
|
||||
if name in raw:
|
||||
data[name] = np.asarray(raw[name], dtype=np.float32)
|
||||
|
||||
for name in (
|
||||
"lin_vel_scale",
|
||||
"ang_vel_scale",
|
||||
"dof_pos_scale",
|
||||
"dof_vel_scale",
|
||||
"action_pos_scale",
|
||||
"action_vel_scale",
|
||||
):
|
||||
if name in raw:
|
||||
data[name] = float(raw[name])
|
||||
|
||||
idx_model2mj = idx_mj2model = np.arange(data["num_actions"], dtype=np.int64)
|
||||
if "mujoco_joint_names" in raw and "model_joint_names" in raw:
|
||||
mj_names = raw["mujoco_joint_names"]
|
||||
model_names = raw["model_joint_names"]
|
||||
idx_model2mj = np.asarray([model_names.index(name) for name in mj_names], dtype=np.int64)
|
||||
idx_mj2model = np.asarray([mj_names.index(name) for name in model_names], dtype=np.int64)
|
||||
data["idx_model2mj"] = idx_model2mj
|
||||
data["idx_mj2model"] = idx_mj2model
|
||||
return SimpleNamespace(**data)
|
||||
|
||||
|
||||
def gravity_from_quat(quaternion: np.ndarray) -> np.ndarray:
|
||||
qw, qx, qy, qz = quaternion
|
||||
return np.array(
|
||||
[
|
||||
2 * (-qz * qx + qw * qy),
|
||||
-2 * (qz * qy + qw * qx),
|
||||
1 - 2 * (qw * qw + qz * qz),
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
|
||||
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:
|
||||
return (target_q - q) * kp + (target_dq - dq) * kd
|
||||
|
||||
|
||||
def init_joystick():
|
||||
pygame.init()
|
||||
if pygame.joystick.get_count() == 0:
|
||||
print("No Joystick detected. Using default commands from config.")
|
||||
return None
|
||||
joystick = pygame.joystick.Joystick(0)
|
||||
joystick.init()
|
||||
print(f"Detected Joystick: {joystick.get_name()}")
|
||||
return joystick
|
||||
|
||||
|
||||
def read_joystick_command(joystick, max_cmd: np.ndarray) -> np.ndarray:
|
||||
pygame.event.pump()
|
||||
axes = np.array([joystick.get_axis(i) for i in (0, 1, 3)], dtype=np.float32)
|
||||
axes[np.abs(axes) < 0.1] = 0.0
|
||||
return np.array(
|
||||
[-axes[1] * max_cmd[0], -axes[0] * max_cmd[1], -axes[2] * max_cmd[2]],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
|
||||
def display_current_command(cmd: np.ndarray) -> None:
|
||||
print(f"\rCurrent command | vx: {cmd[0]: .3f} vy: {cmd[1]: .3f} wz: {cmd[2]: .3f}", end="", flush=True)
|
||||
|
||||
|
||||
def set_initial_state(data, base_pos: np.ndarray, base_quat: np.ndarray, joint_pos: np.ndarray) -> None:
|
||||
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 setup_tracking_camera(viewer, *, trackbodyid: int = 1, distance: float = 3.0,
|
||||
elevation: float = -30.0, azimuth: float = 0.0) -> None:
|
||||
|
||||
viewer.cam.type = mujoco.mjtCamera.mjCAMERA_TRACKING
|
||||
viewer.cam.trackbodyid = trackbodyid
|
||||
viewer.cam.distance = distance
|
||||
viewer.cam.elevation = elevation
|
||||
viewer.cam.azimuth = azimuth
|
||||
|
||||
|
||||
def open_video_writer(enabled: bool, *, policy_path: Path, cmd: np.ndarray,
|
||||
dt: float, video_dir: Path, video_fps: int):
|
||||
if not enabled:
|
||||
return None, 1, None
|
||||
video_dir.mkdir(parents=True, exist_ok=True)
|
||||
video_path = video_dir / f"{policy_path.stem}_cmd_{cmd[0]}_{cmd[1]}_{cmd[2]}.mp4"
|
||||
frame_skip = max(1, int((1.0 / dt) / video_fps))
|
||||
writer = imageio.get_writer(video_path, fps=video_fps)
|
||||
print(f"Recording: {video_path} (Sim FPS: {1.0 / dt:.1f}, Video FPS: {video_fps})")
|
||||
return writer, frame_skip, video_path
|
||||
|
||||
|
||||
def push_obs_history(obs: np.ndarray, features: dict, layout: list[tuple[str, int]], history_len: int) -> None:
|
||||
offset = 0
|
||||
for name, dim in layout:
|
||||
block = obs[offset:offset + dim * history_len]
|
||||
block[:-dim] = block[dim:]
|
||||
block[-dim:] = features[name]
|
||||
offset += dim * history_len
|
||||
|
||||
|
||||
def latest_obs_frame(obs: np.ndarray, layout: list[tuple[str, int]], history_len: int) -> np.ndarray:
|
||||
frames = []
|
||||
offset = 0
|
||||
for _, dim in layout:
|
||||
end = offset + dim * history_len
|
||||
frames.append(obs[end - dim:end])
|
||||
offset = end
|
||||
return np.concatenate(frames, axis=0)
|
||||
|
||||
|
||||
def infer_action(policy, obs: np.ndarray, single_obs: np.ndarray, idx_model2mj: np.ndarray) -> np.ndarray:
|
||||
result = policy(
|
||||
CTSPolicyInputs(
|
||||
policy=torch.from_numpy(obs).unsqueeze(0),
|
||||
single_obs=torch.from_numpy(single_obs).unsqueeze(0),
|
||||
)
|
||||
)
|
||||
return result.detach().cpu().numpy().squeeze()[idx_model2mj]
|
||||
|
||||
|
||||
def build_delay_buffers(*defaults: np.ndarray, delay_max: int):
|
||||
size = delay_max + 1
|
||||
buffers = tuple(deque((value.copy() for _ in range(size)), maxlen=size) for value in defaults)
|
||||
return buffers[0] if len(buffers) == 1 else buffers
|
||||
|
||||
|
||||
def sample_delayed_targets(histories, actuator_groups, delay_min: int, delay_max: int, rng: np.random.Generator):
|
||||
if isinstance(histories, deque):
|
||||
histories = (histories,)
|
||||
delayed = [history[-1].copy() for history in histories]
|
||||
for joint_ids in actuator_groups:
|
||||
delay_steps = int(rng.integers(delay_min, delay_max + 1)) if delay_max > 0 else 0
|
||||
for i, history in enumerate(histories):
|
||||
delayed[i][joint_ids] = history[-1 - delay_steps][joint_ids]
|
||||
return delayed[0] if len(delayed) == 1 else tuple(delayed)
|
||||
|
||||
|
||||
class MujocoRenderUtils:
|
||||
def __init__(self):
|
||||
self.target_velocity = None
|
||||
|
||||
self.vis_smooth_factor = 1.0
|
||||
self.ren_smooth_factor = 1.0
|
||||
|
||||
self.vis_cur_vel = np.zeros(3)
|
||||
self.ren_cur_vel = np.zeros(3)
|
||||
|
||||
self.mj_data = None
|
||||
self._renderer_capacity_warned = False
|
||||
self._viewer_capacity_warned = False
|
||||
|
||||
def update(self, target_velocity, mj_data):
|
||||
self.target_velocity = target_velocity
|
||||
self.mj_data = mj_data
|
||||
|
||||
def update_external_rendering(self,
|
||||
handle: Union[mujoco.viewer.Handle, mujoco.Renderer],
|
||||
ctype: Literal['viewer', 'renderer'],
|
||||
):
|
||||
""" Update external rendering handle (viewer or renderer). """
|
||||
|
||||
def has_geom_slot(container, index: int, warned_attr: str) -> bool:
|
||||
capacity = len(container.geoms)
|
||||
if index < capacity:
|
||||
return True
|
||||
if not getattr(self, warned_attr):
|
||||
print(f"Warning: {ctype} geom capacity exceeded ({capacity}); skipping velocity arrows.")
|
||||
setattr(self, warned_attr, True)
|
||||
return False
|
||||
|
||||
def add_thick_arrow(geom_elem, pos, vec, rgba, scale=0.7):
|
||||
vel_norm = np.linalg.norm(vec)
|
||||
display_norm = min(vel_norm * scale, 1.0)
|
||||
|
||||
if display_norm < 0.10:
|
||||
mujoco.mjv_initGeom(
|
||||
geom_elem,
|
||||
type=mujoco.mjtGeom.mjGEOM_NONE,
|
||||
size=[0,0,0], pos=pos, mat=np.eye(3).flatten(), rgba=[0,0,0,0]
|
||||
)
|
||||
return
|
||||
|
||||
mat = np.zeros(9)
|
||||
target_quat = np.zeros(4)
|
||||
vec_normalized = vec / vel_norm
|
||||
mujoco.mju_quatZ2Vec(target_quat, vec_normalized)
|
||||
mujoco.mju_quat2Mat(mat, target_quat)
|
||||
|
||||
mat = mat.reshape(3, 3)
|
||||
mat[:, 2] *= display_norm
|
||||
|
||||
mujoco.mjv_initGeom(
|
||||
geom_elem,
|
||||
type=mujoco.mjtGeom.mjGEOM_ARROW,
|
||||
size=[0.02, 0.02, display_norm], # [height, width, length]
|
||||
pos=pos,
|
||||
mat=mat.flatten(),
|
||||
rgba=rgba
|
||||
)
|
||||
|
||||
viewer_geom_idx = 0
|
||||
if ctype == 'viewer':
|
||||
handle.user_scn.ngeom = 0 # reset user scene geometry
|
||||
|
||||
if self.target_velocity is not None:
|
||||
base_pos_world = self.mj_data.qpos[:3]
|
||||
base_quat = self.mj_data.qpos[3:7]
|
||||
|
||||
# rendering arrows start position
|
||||
offset_body = np.array([0.0, 0.0, 0.2])
|
||||
offset_world = np.zeros(3)
|
||||
mujoco.mju_rotVecQuat(offset_world, offset_body, base_quat)
|
||||
start_pos = base_pos_world + offset_world
|
||||
|
||||
tgt_vel_body = np.array([self.target_velocity[0], self.target_velocity[1], 0.0])
|
||||
|
||||
raw_cur_vel_world = self.mj_data.qvel[:3]
|
||||
raw_cur_vel = np.zeros(3)
|
||||
neg_quat = np.zeros(4)
|
||||
mujoco.mju_negQuat(neg_quat, base_quat)
|
||||
mujoco.mju_rotVecQuat(raw_cur_vel, raw_cur_vel_world, neg_quat)
|
||||
cur_vel_body = np.array([raw_cur_vel[0], raw_cur_vel[1], 0.0])
|
||||
|
||||
# EMA: v_smooth = alpha * v_new + (1 - alpha) * v_old
|
||||
# alpha = self.vis_smooth_factor if ctype == 'viewer' else self.ren_smooth_factor
|
||||
self.vis_cur_vel = cur_vel_body
|
||||
self.ren_cur_vel = cur_vel_body
|
||||
|
||||
tgt_vel_world = np.zeros(3)
|
||||
cur_vel_world = np.zeros(3)
|
||||
mujoco.mju_rotVecQuat(tgt_vel_world, tgt_vel_body, base_quat)
|
||||
if ctype == 'viewer':
|
||||
mujoco.mju_rotVecQuat(cur_vel_world, self.vis_cur_vel, base_quat)
|
||||
else:
|
||||
mujoco.mju_rotVecQuat(cur_vel_world, self.ren_cur_vel, base_quat)
|
||||
|
||||
COLOR_CMD = [0, 1, 0, 1] # Green 0x00ff00
|
||||
COLOR_REAL = [0, 0, 1, 1] # Blue 0x0000ff
|
||||
|
||||
if ctype == 'viewer':
|
||||
# Cmd Arrow
|
||||
if has_geom_slot(handle.user_scn, viewer_geom_idx, "_viewer_capacity_warned"):
|
||||
add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, tgt_vel_world, COLOR_CMD)
|
||||
viewer_geom_idx += 1
|
||||
# Real Arrow
|
||||
if has_geom_slot(handle.user_scn, viewer_geom_idx, "_viewer_capacity_warned"):
|
||||
add_thick_arrow(handle.user_scn.geoms[viewer_geom_idx], start_pos, cur_vel_world, COLOR_REAL)
|
||||
viewer_geom_idx += 1
|
||||
else:
|
||||
scene = handle.scene
|
||||
if has_geom_slot(scene, scene.ngeom, "_renderer_capacity_warned"):
|
||||
add_thick_arrow(scene.geoms[scene.ngeom], start_pos, tgt_vel_world, COLOR_CMD)
|
||||
scene.ngeom += 1
|
||||
if has_geom_slot(scene, scene.ngeom, "_renderer_capacity_warned"):
|
||||
add_thick_arrow(scene.geoms[scene.ngeom], start_pos, cur_vel_world, COLOR_REAL)
|
||||
scene.ngeom += 1
|
||||
|
||||
if ctype == 'viewer':
|
||||
handle.user_scn.ngeom = viewer_geom_idx
|
||||
Reference in New Issue
Block a user