refactor deploy scripts.
This commit is contained in:
wertyuilife
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