From 8cdc8fd2b463c696fb578822fceec89b036e9e66 Mon Sep 17 00:00:00 2001 From: caozx1110 Date: Mon, 16 Mar 2026 18:30:02 +0800 Subject: [PATCH] feat: add motion imitation tracking for G1 humanoid Migrate motion tracking implementation from mjlab to UniLab architecture. Adapts BeyondMimic-style motion imitation to UniLab's non-manager based environment system with CPU simulation and GPU training. Key components: - Math utilities for quaternion operations and transforms - Motion loader with adaptive/uniform/start sampling modes - G1 motion tracking environment with exponential reward functions - MuJoCo backend extensions for body state access - CSV to NPZ motion preprocessing script Architecture differences from mjlab: - Direct NpEnv inheritance vs manager-based system - CPU MuJoCo simulation vs GPU MuJoCo-Warp - Decorator-based registration vs automatic config registration - Inline reward/observation computation vs manager terms Training config: 4096 envs, 30k iterations, 50 Hz control Co-Authored-By: Claude Opus 4.6 --- scripts/motion/README.md | 68 ++ scripts/motion/csv_to_npz.py | 356 +++++++++++ unilab/base/backend/mujoco_backend.py | 142 ++++ unilab/config/locomotion_params.py | 12 +- unilab/envs/locomotion/g1/__init__.py | 1 + .../locomotion/g1/motion_tracking/__init__.py | 5 + .../g1/motion_tracking/math_utils.py | 304 +++++++++ .../g1/motion_tracking/motion_loader.py | 261 ++++++++ .../locomotion/g1/motion_tracking/tracking.py | 604 ++++++++++++++++++ 9 files changed, 1750 insertions(+), 3 deletions(-) create mode 100644 scripts/motion/README.md create mode 100644 scripts/motion/csv_to_npz.py create mode 100644 unilab/envs/locomotion/g1/motion_tracking/__init__.py create mode 100644 unilab/envs/locomotion/g1/motion_tracking/math_utils.py create mode 100644 unilab/envs/locomotion/g1/motion_tracking/motion_loader.py create mode 100644 unilab/envs/locomotion/g1/motion_tracking/tracking.py diff --git a/scripts/motion/README.md b/scripts/motion/README.md new file mode 100644 index 000000000..965d62c6e --- /dev/null +++ b/scripts/motion/README.md @@ -0,0 +1,68 @@ +# Motion Preprocessing + +This directory contains scripts for preprocessing motion data for motion tracking tasks. + +## CSV to NPZ Conversion + +The `csv_to_npz.py` script converts motion data from CSV format to NPZ format with precomputed forward kinematics. + +### Input Format + +CSV files should contain motion data in Unitree's generalized coordinate convention: +- Columns 0-2: Base position (x, y, z) +- Columns 3-6: Base quaternion (x, y, z, w) - will be converted to wxyz internally +- Columns 7+: Joint angles (29 joints for G1) + +### Output Format + +NPZ files contain: +- `fps`: Frame rate (integer) +- `joint_pos`: Joint positions (N_frames × N_joints) +- `joint_vel`: Joint velocities (N_frames × N_joints) +- `body_pos_w`: Body positions in world frame (N_frames × N_bodies × 3) +- `body_quat_w`: Body quaternions in world frame (N_frames × N_bodies × 4, wxyz) +- `body_lin_vel_w`: Body linear velocities (N_frames × N_bodies × 3) +- `body_ang_vel_w`: Body angular velocities (N_frames × N_bodies × 3) + +### Usage + +```bash +# Basic usage +uv run python scripts/motion/csv_to_npz.py \ + --input_file path/to/motion.csv \ + --output_file path/to/motion.npz \ + --input_fps 30 \ + --output_fps 50 + +# With custom model file +uv run python scripts/motion/csv_to_npz.py \ + --input_file path/to/motion.csv \ + --output_file path/to/motion.npz \ + --input_fps 30 \ + --output_fps 50 \ + --model_file path/to/model.xml + +# Process specific line range +uv run python scripts/motion/csv_to_npz.py \ + --input_file path/to/motion.csv \ + --output_file path/to/motion.npz \ + --input_fps 30 \ + --output_fps 50 \ + --line_range 100 500 +``` + +### Parameters + +- `--input_file`: Path to input CSV file (required) +- `--output_file`: Path to output NPZ file (required) +- `--input_fps`: Frame rate of input CSV (default: 30) +- `--output_fps`: Desired output frame rate (default: 50) +- `--model_file`: MuJoCo model file (default: G1 flat scene) +- `--line_range`: Line range to process [start, end] (optional) + +### Notes + +- The script uses LERP for position interpolation and SLERP for quaternion interpolation +- Velocities are computed using numerical differentiation +- Forward kinematics is computed using MuJoCo for all bodies +- The output FPS should match the control frequency of your training environment (typically 50 Hz) diff --git a/scripts/motion/csv_to_npz.py b/scripts/motion/csv_to_npz.py new file mode 100644 index 000000000..97e94fd7c --- /dev/null +++ b/scripts/motion/csv_to_npz.py @@ -0,0 +1,356 @@ +"""Convert CSV motion files to NPZ format with forward kinematics. + +This script converts motion data from CSV format (Unitree convention) to NPZ format +with precomputed forward kinematics for all bodies. + +Input CSV format: +- Base position (3): x, y, z +- Base quaternion (4): x, y, z, w (will be converted to w, x, y, z internally) +- Joint angles (29): all joint positions + +Output NPZ format: +- fps: Frame rate (integer) +- joint_pos: Joint positions (N_frames × N_joints) +- joint_vel: Joint velocities (N_frames × N_joints) +- body_pos_w: Body positions in world frame (N_frames × N_bodies × 3) +- body_quat_w: Body quaternions in world frame (N_frames × N_bodies × 4, wxyz) +- body_lin_vel_w: Body linear velocities (N_frames × N_bodies × 3) +- body_ang_vel_w: Body angular velocities (N_frames × N_bodies × 3) +""" + +import argparse +from pathlib import Path + +import mujoco +import numpy as np +from tqdm import tqdm + + +def quat_slerp(q1: np.ndarray, q2: np.ndarray, t: float) -> np.ndarray: + """Spherical linear interpolation between two quaternions (wxyz format).""" + # Ensure shortest path + dot = np.dot(q1, q2) + if dot < 0: + q2 = -q2 + dot = -dot + + # If quaternions are very close, use linear interpolation + if dot > 0.9995: + result = q1 + t * (q2 - q1) + return result / np.linalg.norm(result) + + # Compute angle + theta = np.arccos(np.clip(dot, -1, 1)) + sin_theta = np.sin(theta) + + # Compute interpolation weights + w1 = np.sin((1 - t) * theta) / sin_theta + w2 = np.sin(t * theta) / sin_theta + + return w1 * q1 + w2 * q2 + + +def quat_conjugate(q: np.ndarray) -> np.ndarray: + """Compute quaternion conjugate (wxyz format).""" + return np.array([q[0], -q[1], -q[2], -q[3]]) + + +def quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: + """Multiply two quaternions (wxyz format).""" + w1, x1, y1, z1 = q1 + w2, x2, y2, z2 = q2 + return np.array( + [ + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ] + ) + + +def axis_angle_from_quat(q: np.ndarray) -> np.ndarray: + """Convert quaternion to axis-angle representation (wxyz format).""" + # Normalize + q = q / np.linalg.norm(q) + w, x, y, z = q + + # Compute angle + angle = 2 * np.arccos(np.clip(w, -1, 1)) + + # Compute axis + sin_half = np.sin(angle / 2) + if abs(sin_half) < 1e-6: + # No rotation + return np.zeros(3) + + axis = np.array([x, y, z]) / sin_half + return axis * angle + + +class MotionLoader: + """Load and interpolate motion from CSV file.""" + + def __init__( + self, + motion_file: str, + input_fps: int, + output_fps: int, + line_range: tuple[int, int] | None = None, + ): + self.motion_file = motion_file + self.input_fps = input_fps + self.output_fps = output_fps + self.input_dt = 1.0 / self.input_fps + self.output_dt = 1.0 / self.output_fps + self.line_range = line_range + self._load_motion() + self._interpolate_motion() + self._compute_velocities() + + def _load_motion(self): + """Load motion from CSV file.""" + if self.line_range is None: + motion = np.loadtxt(self.motion_file, delimiter=",", dtype=np.float32) + else: + motion = np.loadtxt( + self.motion_file, + delimiter=",", + skiprows=self.line_range[0] - 1, + max_rows=self.line_range[1] - self.line_range[0] + 1, + dtype=np.float32, + ) + + self.motion_base_poss_input = motion[:, :3] + # Convert quaternion from xyzw to wxyz + self.motion_base_rots_input = motion[:, 3:7][:, [3, 0, 1, 2]] + self.motion_dof_poss_input = motion[:, 7:] + + self.input_frames = motion.shape[0] + self.duration = (self.input_frames - 1) * self.input_dt + + def _interpolate_motion(self): + """Interpolate motion to output FPS.""" + times = np.arange(0, self.duration, self.output_dt, dtype=np.float32) + self.output_frames = times.shape[0] + index_0, index_1, blend = self._compute_frame_blend(times) + + # Linear interpolation for positions + self.motion_base_poss = ( + self.motion_base_poss_input[index_0] * (1 - blend[:, None]) + + self.motion_base_poss_input[index_1] * blend[:, None] + ) + + # Spherical linear interpolation for quaternions + self.motion_base_rots = np.zeros((self.output_frames, 4), dtype=np.float32) + for i in range(self.output_frames): + self.motion_base_rots[i] = quat_slerp( + self.motion_base_rots_input[index_0[i]], + self.motion_base_rots_input[index_1[i]], + blend[i], + ) + + # Linear interpolation for joint positions + self.motion_dof_poss = ( + self.motion_dof_poss_input[index_0] * (1 - blend[:, None]) + + self.motion_dof_poss_input[index_1] * blend[:, None] + ) + + print( + f"Motion interpolated: {self.input_frames} frames @ {self.input_fps} Hz " + f"→ {self.output_frames} frames @ {self.output_fps} Hz" + ) + + def _compute_frame_blend(self, times: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Compute frame indices and blend weights for interpolation.""" + phase = times / self.duration + index_0 = np.floor(phase * (self.input_frames - 1)).astype(np.int32) + index_1 = np.minimum(index_0 + 1, self.input_frames - 1) + blend = phase * (self.input_frames - 1) - index_0 + return index_0, index_1, blend + + def _compute_velocities(self): + """Compute velocities using numerical differentiation.""" + # Linear velocities + self.motion_base_lin_vels = np.gradient(self.motion_base_poss, self.output_dt, axis=0) + self.motion_dof_vels = np.gradient(self.motion_dof_poss, self.output_dt, axis=0) + + # Angular velocities from quaternion derivatives + self.motion_base_ang_vels = self._so3_derivative(self.motion_base_rots, self.output_dt) + + def _so3_derivative(self, rotations: np.ndarray, dt: float) -> np.ndarray: + """Compute angular velocity from quaternion sequence.""" + # Use central differences + q_prev = rotations[:-2] + q_next = rotations[2:] + + # Compute relative quaternion + omega = np.zeros((rotations.shape[0], 3), dtype=np.float32) + for i in range(len(q_prev)): + q_rel = quat_mul(q_next[i], quat_conjugate(q_prev[i])) + omega[i + 1] = axis_angle_from_quat(q_rel) / (2.0 * dt) + + # Repeat first and last + omega[0] = omega[1] + omega[-1] = omega[-2] + + return omega + + +def run_simulation( + motion_loader: MotionLoader, + model_file: str, + joint_names: list[str], + output_file: str, +): + """Run MuJoCo simulation to compute forward kinematics.""" + # Load model + model = mujoco.MjModel.from_xml_path(model_file) + data = mujoco.MjData(model) + + # Get joint indices + joint_indices = [] + for name in joint_names: + jnt_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, name) + if jnt_id < 0: + raise ValueError(f"Joint '{name}' not found in model") + joint_indices.append(jnt_id) + + # Prepare output arrays + num_frames = motion_loader.output_frames + num_joints = len(joint_indices) + num_bodies = model.nbody + + joint_pos = np.zeros((num_frames, num_joints), dtype=np.float32) + joint_vel = np.zeros((num_frames, num_joints), dtype=np.float32) + body_pos_w = np.zeros((num_frames, num_bodies, 3), dtype=np.float32) + body_quat_w = np.zeros((num_frames, num_bodies, 4), dtype=np.float32) + body_lin_vel_w = np.zeros((num_frames, num_bodies, 3), dtype=np.float32) + body_ang_vel_w = np.zeros((num_frames, num_bodies, 3), dtype=np.float32) + + print(f"\nProcessing {num_frames} frames...") + for i in tqdm(range(num_frames)): + # Set root state + data.qpos[0:3] = motion_loader.motion_base_poss[i] + data.qpos[3:7] = motion_loader.motion_base_rots[i] + data.qvel[0:3] = motion_loader.motion_base_lin_vels[i] + data.qvel[3:6] = motion_loader.motion_base_ang_vels[i] + + # Set joint states + for j, jnt_id in enumerate(joint_indices): + qpos_adr = model.jnt_qposadr[jnt_id] + qvel_adr = model.jnt_dofadr[jnt_id] + data.qpos[qpos_adr] = motion_loader.motion_dof_poss[i, j] + data.qvel[qvel_adr] = motion_loader.motion_dof_vels[i, j] + + # Run forward kinematics + mujoco.mj_kinematics(model, data) + mujoco.mj_comVel(model, data) + + # Extract joint states + for j, jnt_id in enumerate(joint_indices): + qpos_adr = model.jnt_qposadr[jnt_id] + qvel_adr = model.jnt_dofadr[jnt_id] + joint_pos[i, j] = data.qpos[qpos_adr] + joint_vel[i, j] = data.qvel[qvel_adr] + + # Extract body states + for body_id in range(num_bodies): + body_pos_w[i, body_id] = data.xpos[body_id] + body_quat_w[i, body_id] = data.xquat[body_id] + body_lin_vel_w[i, body_id] = data.cvel[body_id, 3:6] + body_ang_vel_w[i, body_id] = data.cvel[body_id, 0:3] + + # Save to NPZ + print(f"\nSaving to {output_file}...") + np.savez( + output_file, + fps=np.array([motion_loader.output_fps], dtype=np.int32), + joint_pos=joint_pos, + joint_vel=joint_vel, + body_pos_w=body_pos_w, + body_quat_w=body_quat_w, + body_lin_vel_w=body_lin_vel_w, + body_ang_vel_w=body_ang_vel_w, + ) + print("Done!") + + +def main(): + parser = argparse.ArgumentParser( + description="Convert CSV motion to NPZ with forward kinematics" + ) + parser.add_argument("--input_file", type=str, required=True, help="Input CSV file") + parser.add_argument("--output_file", type=str, required=True, help="Output NPZ file") + parser.add_argument("--input_fps", type=float, default=30.0, help="Input frame rate") + parser.add_argument("--output_fps", type=float, default=50.0, help="Output frame rate") + parser.add_argument( + "--model_file", + type=str, + default=None, + help="MuJoCo model file (default: G1 flat scene)", + ) + parser.add_argument( + "--line_range", + type=int, + nargs=2, + default=None, + help="Line range to process (start, end)", + ) + + args = parser.parse_args() + + # Default model file + if args.model_file is None: + script_dir = Path(__file__).parent.parent + args.model_file = str( + script_dir / "unilab" / "envs" / "locomotion" / "g1" / "xml" / "scene_flat.xml" + ) + + # G1 joint names (in order) + joint_names = [ + "left_hip_pitch_joint", + "left_hip_roll_joint", + "left_hip_yaw_joint", + "left_knee_joint", + "left_ankle_pitch_joint", + "left_ankle_roll_joint", + "right_hip_pitch_joint", + "right_hip_roll_joint", + "right_hip_yaw_joint", + "right_knee_joint", + "right_ankle_pitch_joint", + "right_ankle_roll_joint", + "waist_yaw_joint", + "waist_roll_joint", + "waist_pitch_joint", + "left_shoulder_pitch_joint", + "left_shoulder_roll_joint", + "left_shoulder_yaw_joint", + "left_elbow_joint", + "left_wrist_roll_joint", + "left_wrist_pitch_joint", + "left_wrist_yaw_joint", + "right_shoulder_pitch_joint", + "right_shoulder_roll_joint", + "right_shoulder_yaw_joint", + "right_elbow_joint", + "right_wrist_roll_joint", + "right_wrist_pitch_joint", + "right_wrist_yaw_joint", + ] + + # Load and interpolate motion + motion_loader = MotionLoader( + args.input_file, + int(args.input_fps), + int(args.output_fps), + tuple(args.line_range) if args.line_range else None, + ) + + # Run simulation + run_simulation(motion_loader, args.model_file, joint_names, args.output_file) + + +if __name__ == "__main__": + main() diff --git a/unilab/base/backend/mujoco_backend.py b/unilab/base/backend/mujoco_backend.py index 6ba0747b1..6087f2443 100644 --- a/unilab/base/backend/mujoco_backend.py +++ b/unilab/base/backend/mujoco_backend.py @@ -129,3 +129,145 @@ def model(self): def get_physics_state(self) -> np.ndarray: return self._physics_state + + def get_body_pos(self, body_ids: np.ndarray) -> np.ndarray: + """Get body positions in world frame. + + Args: + body_ids: Body indices (num_bodies,) + + Returns: + Body positions (num_envs, num_bodies, 3) + """ + # Need to run forward kinematics to get body positions + # Use batch_forward to compute xpos for all environments + state_out, _ = self._forward_runner.forward( + model=self._model, + data=self._worker_data, + initial_state=self._physics_state.astype(np.float64), + chunk_size=max(1, self._num_envs // self._n_threads), + skipsensor=True, + out_dtype=np.float64, + return_state=False, + ) + + # Extract body positions from MjData + # xpos is stored in data, need to extract from worker_data after forward + body_pos = np.zeros((self._num_envs, len(body_ids), 3), dtype=self._np_dtype) + + # Run forward kinematics for each environment + for i in range(self._num_envs): + worker_idx = i % self._n_threads + data = self._worker_data[worker_idx] + + # Set state + qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] + qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] + data.qpos[:] = qpos + data.qvel[:] = qvel + + # Run forward kinematics + mujoco.mj_kinematics(self._model, data) + + # Extract body positions + for j, body_id in enumerate(body_ids): + body_pos[i, j] = data.xpos[body_id] + + return body_pos + + def get_body_quat(self, body_ids: np.ndarray) -> np.ndarray: + """Get body quaternions in world frame (wxyz format). + + Args: + body_ids: Body indices (num_bodies,) + + Returns: + Body quaternions (num_envs, num_bodies, 4) + """ + body_quat = np.zeros((self._num_envs, len(body_ids), 4), dtype=self._np_dtype) + + # Run forward kinematics for each environment + for i in range(self._num_envs): + worker_idx = i % self._n_threads + data = self._worker_data[worker_idx] + + # Set state + qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] + qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] + data.qpos[:] = qpos + data.qvel[:] = qvel + + # Run forward kinematics + mujoco.mj_kinematics(self._model, data) + + # Extract body quaternions (MuJoCo uses wxyz format) + for j, body_id in enumerate(body_ids): + body_quat[i, j] = data.xquat[body_id] + + return body_quat + + def get_body_lin_vel(self, body_ids: np.ndarray) -> np.ndarray: + """Get body linear velocities in world frame. + + Args: + body_ids: Body indices (num_bodies,) + + Returns: + Body linear velocities (num_envs, num_bodies, 3) + """ + body_lin_vel = np.zeros((self._num_envs, len(body_ids), 3), dtype=self._np_dtype) + + # Run forward kinematics and velocity computation + for i in range(self._num_envs): + worker_idx = i % self._n_threads + data = self._worker_data[worker_idx] + + # Set state + qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] + qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] + data.qpos[:] = qpos + data.qvel[:] = qvel + + # Run forward kinematics and velocity computation + mujoco.mj_kinematics(self._model, data) + mujoco.mj_comVel(self._model, data) + + # Extract body linear velocities + for j, body_id in enumerate(body_ids): + # cvel contains [angular_vel (3), linear_vel (3)] for each body + body_lin_vel[i, j] = data.cvel[body_id, 3:6] + + return body_lin_vel + + def get_body_ang_vel(self, body_ids: np.ndarray) -> np.ndarray: + """Get body angular velocities in world frame. + + Args: + body_ids: Body indices (num_bodies,) + + Returns: + Body angular velocities (num_envs, num_bodies, 3) + """ + body_ang_vel = np.zeros((self._num_envs, len(body_ids), 3), dtype=self._np_dtype) + + # Run forward kinematics and velocity computation + for i in range(self._num_envs): + worker_idx = i % self._n_threads + data = self._worker_data[worker_idx] + + # Set state + qpos = self._physics_state[i, self._idx_qpos : self._idx_qpos + self.nq] + qvel = self._physics_state[i, self._idx_qvel : self._idx_qvel + self.nv] + data.qpos[:] = qpos + data.qvel[:] = qvel + + # Run forward kinematics and velocity computation + mujoco.mj_kinematics(self._model, data) + mujoco.mj_comVel(self._model, data) + + # Extract body angular velocities + for j, body_id in enumerate(body_ids): + # cvel contains [angular_vel (3), linear_vel (3)] for each body + body_ang_vel[i, j] = data.cvel[body_id, 0:3] + + return body_ang_vel diff --git a/unilab/config/locomotion_params.py b/unilab/config/locomotion_params.py index d10acec48..1e5544076 100644 --- a/unilab/config/locomotion_params.py +++ b/unilab/config/locomotion_params.py @@ -32,14 +32,20 @@ def ppo_config(env_name: str) -> config_dict.ConfigDict: obs_group_set = {"default": ["policy"]} if env_name == "Go1JoystickFlatTerrain": cfg.max_iterations = 151 - obs_group_set = {"default": ["policy"], "actor": ['actor']} + obs_group_set = {"default": ["policy"], "actor": ["actor"]} elif env_name == "G1JoystickFlatTerrain": - obs_group_set = {"default": ["policy"], "actor": ['actor']} + obs_group_set = {"default": ["policy"], "actor": ["actor"]} cfg.num_envs = 2048 cfg.max_iterations = 220 elif env_name == "Go2JoystickFlatTerrain": - obs_group_set = {"default": ["policy"], "actor": ['actor']} + obs_group_set = {"default": ["policy"], "actor": ["actor"]} pass + elif env_name == "G1MotionTracking": + obs_group_set = {"default": ["policy"], "actor": ["actor"]} + cfg.num_envs = 4096 + cfg.max_iterations = 30000 + cfg.save_interval = 500 + cfg.entropy_coef = 0.005 return config_dict.create( algo="ppo", diff --git a/unilab/envs/locomotion/g1/__init__.py b/unilab/envs/locomotion/g1/__init__.py index 200f0fe75..0fbf1576d 100644 --- a/unilab/envs/locomotion/g1/__init__.py +++ b/unilab/envs/locomotion/g1/__init__.py @@ -1,2 +1,3 @@ from .joystick import G1JoystickPPO, G1JoystickPPOCfg from .joystick_sac import G1JoystickSACCfg, G1WalkTaskMjSAC +from .motion_tracking import G1MotionTrackingCfg, G1MotionTrackingEnv, G1MotionTrackingEnvCfg diff --git a/unilab/envs/locomotion/g1/motion_tracking/__init__.py b/unilab/envs/locomotion/g1/motion_tracking/__init__.py new file mode 100644 index 000000000..de066f696 --- /dev/null +++ b/unilab/envs/locomotion/g1/motion_tracking/__init__.py @@ -0,0 +1,5 @@ +"""Motion tracking module for G1 humanoid.""" + +from .tracking import G1MotionTrackingCfg, G1MotionTrackingEnv, G1MotionTrackingEnvCfg + +__all__ = ["G1MotionTrackingCfg", "G1MotionTrackingEnv", "G1MotionTrackingEnvCfg"] diff --git a/unilab/envs/locomotion/g1/motion_tracking/math_utils.py b/unilab/envs/locomotion/g1/motion_tracking/math_utils.py new file mode 100644 index 000000000..88e552939 --- /dev/null +++ b/unilab/envs/locomotion/g1/motion_tracking/math_utils.py @@ -0,0 +1,304 @@ +"""Math utilities for motion tracking - quaternion and rotation operations.""" + +from __future__ import annotations + +import numpy as np + + +def quat_mul(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: + """Multiply two quaternions (wxyz format). + + Args: + q1: First quaternion(s) (N, 4) or (4,) + q2: Second quaternion(s) (N, 4) or (4,) + + Returns: + Product quaternion(s) (N, 4) or (4,) + """ + if q1.ndim == 1: + q1 = q1[None, :] + q2 = q2[None, :] + squeeze = True + else: + squeeze = False + + w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3] + w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3] + + result = np.stack( + [ + w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2, + w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2, + w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2, + w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2, + ], + axis=1, + ) + + return result[0] if squeeze else result + + +def quat_conjugate(q: np.ndarray) -> np.ndarray: + """Compute quaternion conjugate (wxyz format). + + Args: + q: Quaternion(s) (N, 4) or (4,) + + Returns: + Conjugate quaternion(s) (N, 4) or (4,) + """ + if q.ndim == 1: + return np.array([q[0], -q[1], -q[2], -q[3]]) + else: + result = q.copy() + result[:, 1:] *= -1 + return result + + +def quat_inv(q: np.ndarray) -> np.ndarray: + """Compute quaternion inverse (wxyz format). + + For unit quaternions, inverse equals conjugate. + + Args: + q: Quaternion(s) (N, 4) or (4,) + + Returns: + Inverse quaternion(s) (N, 4) or (4,) + """ + return quat_conjugate(q) + + +def quat_apply(q: np.ndarray, v: np.ndarray) -> np.ndarray: + """Rotate vector(s) by quaternion(s) (wxyz format). + + Args: + q: Quaternion(s) (N, 4) or (4,) + v: Vector(s) (N, 3) or (3,) + + Returns: + Rotated vector(s) (N, 3) or (3,) + """ + if q.ndim == 1: + q = q[None, :] + v = v[None, :] + squeeze = True + else: + squeeze = False + + # Convert to (w, x, y, z) + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + vx, vy, vz = v[:, 0], v[:, 1], v[:, 2] + + # Compute rotation using quaternion formula + # v' = v + 2 * cross(q.xyz, cross(q.xyz, v) + q.w * v) + t = 2 * np.stack( + [ + y * vz - z * vy, + z * vx - x * vz, + x * vy - y * vx, + ], + axis=1, + ) + + t += 2 * w[:, None] * v + + result = v + np.stack( + [ + y * t[:, 2] - z * t[:, 1], + z * t[:, 0] - x * t[:, 2], + x * t[:, 1] - y * t[:, 0], + ], + axis=1, + ) + + return result[0] if squeeze else result + + +def quat_apply_inverse(q: np.ndarray, v: np.ndarray) -> np.ndarray: + """Rotate vector(s) by inverse quaternion(s) (wxyz format). + + Args: + q: Quaternion(s) (N, 4) or (4,) + v: Vector(s) (N, 3) or (3,) + + Returns: + Rotated vector(s) (N, 3) or (3,) + """ + return quat_apply(quat_inv(q), v) + + +def quat_error_magnitude(q1: np.ndarray, q2: np.ndarray) -> np.ndarray: + """Compute angular error magnitude between quaternions. + + Args: + q1: First quaternion(s) (N, 4) or (4,) + q2: Second quaternion(s) (N, 4) or (4,) + + Returns: + Error magnitude(s) (N,) or scalar + """ + if q1.ndim == 1: + q1 = q1[None, :] + q2 = q2[None, :] + squeeze = True + else: + squeeze = False + + # Compute relative quaternion + q_rel = quat_mul(q2, quat_inv(q1)) + + # Error magnitude is 2 * arcsin(||xyz||) + xyz_norm = np.linalg.norm(q_rel[:, 1:], axis=1) + error = 2 * np.arcsin(np.clip(xyz_norm, -1, 1)) + + return error[0] if squeeze else error + + +def quat_from_euler_xyz(roll: np.ndarray, pitch: np.ndarray, yaw: np.ndarray) -> np.ndarray: + """Convert Euler angles (XYZ order) to quaternion (wxyz format). + + Args: + roll: Roll angle(s) (N,) or scalar + pitch: Pitch angle(s) (N,) or scalar + yaw: Yaw angle(s) (N,) or scalar + + Returns: + Quaternion(s) (N, 4) or (4,) + """ + roll = np.atleast_1d(roll) + pitch = np.atleast_1d(pitch) + yaw = np.atleast_1d(yaw) + squeeze = roll.shape[0] == 1 + + cr = np.cos(roll * 0.5) + sr = np.sin(roll * 0.5) + cp = np.cos(pitch * 0.5) + sp = np.sin(pitch * 0.5) + cy = np.cos(yaw * 0.5) + sy = np.sin(yaw * 0.5) + + w = cr * cp * cy + sr * sp * sy + x = sr * cp * cy - cr * sp * sy + y = cr * sp * cy + sr * cp * sy + z = cr * cp * sy - sr * sp * cy + + result = np.stack([w, x, y, z], axis=1) + return result[0] if squeeze else result + + +def yaw_quat(q: np.ndarray) -> np.ndarray: + """Extract yaw-only quaternion from full quaternion (wxyz format). + + Args: + q: Quaternion(s) (N, 4) or (4,) + + Returns: + Yaw-only quaternion(s) (N, 4) or (4,) + """ + if q.ndim == 1: + q = q[None, :] + squeeze = True + else: + squeeze = False + + # Extract yaw angle + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + yaw = np.arctan2(2 * (w * z + x * y), 1 - 2 * (y * y + z * z)) + + # Create yaw-only quaternion + half_yaw = yaw * 0.5 + result = np.stack( + [ + np.cos(half_yaw), + np.zeros_like(half_yaw), + np.zeros_like(half_yaw), + np.sin(half_yaw), + ], + axis=1, + ) + + return result[0] if squeeze else result + + +def matrix_from_quat(q: np.ndarray) -> np.ndarray: + """Convert quaternion(s) to rotation matrix (wxyz format). + + Args: + q: Quaternion(s) (N, 4) or (4,) + + Returns: + Rotation matrix (N, 3, 3) or (3, 3) + """ + if q.ndim == 1: + q = q[None, :] + squeeze = True + else: + squeeze = False + + w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3] + + # Compute rotation matrix elements + xx = x * x + yy = y * y + zz = z * z + xy = x * y + xz = x * z + yz = y * z + wx = w * x + wy = w * y + wz = w * z + + result = np.stack( + [ + np.stack([1 - 2 * (yy + zz), 2 * (xy - wz), 2 * (xz + wy)], axis=1), + np.stack([2 * (xy + wz), 1 - 2 * (xx + zz), 2 * (yz - wx)], axis=1), + np.stack([2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy)], axis=1), + ], + axis=1, + ) + + return result[0] if squeeze else result + + +def subtract_frame_transforms( + pos1: np.ndarray, quat1: np.ndarray, pos2: np.ndarray, quat2: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Compute relative transform from frame 1 to frame 2. + + Args: + pos1: Position of frame 1 (N, 3) or (3,) + quat1: Orientation of frame 1 (N, 4) or (4,) in wxyz format + pos2: Position of frame 2 (N, 3) or (3,) + quat2: Orientation of frame 2 (N, 4) or (4,) in wxyz format + + Returns: + Tuple of (relative_pos, relative_quat) in frame 1 coordinates + """ + # Relative position in frame 1 coordinates + rel_pos = quat_apply_inverse(quat1, pos2 - pos1) + + # Relative orientation + rel_quat = quat_mul(quat_inv(quat1), quat2) + + return rel_pos, rel_quat + + +def sample_uniform( + lower: float | np.ndarray, + upper: float | np.ndarray, + size: tuple[int, ...], + dtype=np.float32, +) -> np.ndarray: + """Sample uniformly from [lower, upper]. + + Args: + lower: Lower bound(s) + upper: Upper bound(s) + size: Output shape + dtype: Output dtype + + Returns: + Sampled values + """ + return np.random.uniform(lower, upper, size).astype(dtype) diff --git a/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py b/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py new file mode 100644 index 000000000..1090885d4 --- /dev/null +++ b/unilab/envs/locomotion/g1/motion_tracking/motion_loader.py @@ -0,0 +1,261 @@ +"""Motion loading and sampling for motion tracking tasks.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Literal + +import numpy as np + +from .math_utils import sample_uniform + + +@dataclass +class MotionData: + """Container for motion data at specific frame(s).""" + + joint_pos: np.ndarray # (N, num_joints) + joint_vel: np.ndarray # (N, num_joints) + body_pos_w: np.ndarray # (N, num_bodies, 3) + body_quat_w: np.ndarray # (N, num_bodies, 4) + body_lin_vel_w: np.ndarray # (N, num_bodies, 3) + body_ang_vel_w: np.ndarray # (N, num_bodies, 3) + + +class MotionLoader: + """Loads and provides access to motion data from NPZ files.""" + + def __init__(self, motion_file: str, body_indices: np.ndarray | None = None): + """Initialize motion loader. + + Args: + motion_file: Path to NPZ file containing motion data + body_indices: Optional indices to select subset of bodies + """ + data = np.load(motion_file) + + self.fps = int(data["fps"][0]) + self.joint_pos = data["joint_pos"].astype(np.float32) + self.joint_vel = data["joint_vel"].astype(np.float32) + + # Load body data + body_pos_w = data["body_pos_w"].astype(np.float32) + body_quat_w = data["body_quat_w"].astype(np.float32) + body_lin_vel_w = data["body_lin_vel_w"].astype(np.float32) + body_ang_vel_w = data["body_ang_vel_w"].astype(np.float32) + + # Select subset of bodies if specified + if body_indices is not None: + self.body_pos_w = body_pos_w[:, body_indices] + self.body_quat_w = body_quat_w[:, body_indices] + self.body_lin_vel_w = body_lin_vel_w[:, body_indices] + self.body_ang_vel_w = body_ang_vel_w[:, body_indices] + else: + self.body_pos_w = body_pos_w + self.body_quat_w = body_quat_w + self.body_lin_vel_w = body_lin_vel_w + self.body_ang_vel_w = body_ang_vel_w + + self.num_frames = self.joint_pos.shape[0] + self.num_joints = self.joint_pos.shape[1] + self.num_bodies = self.body_pos_w.shape[1] + + def get_motion_at_frame(self, frame_idx: np.ndarray) -> MotionData: + """Get motion data at specified frame indices. + + Args: + frame_idx: Frame indices (N,) + + Returns: + MotionData at specified frames + """ + return MotionData( + joint_pos=self.joint_pos[frame_idx], + joint_vel=self.joint_vel[frame_idx], + body_pos_w=self.body_pos_w[frame_idx], + body_quat_w=self.body_quat_w[frame_idx], + body_lin_vel_w=self.body_lin_vel_w[frame_idx], + body_ang_vel_w=self.body_ang_vel_w[frame_idx], + ) + + +class MotionSampler: + """Handles motion frame sampling with different strategies.""" + + def __init__( + self, + motion_loader: MotionLoader, + mode: Literal["start", "uniform", "adaptive"], + num_envs: int, + bin_count: int | None = None, + adaptive_lambda: float = 0.8, + adaptive_kernel_size: int = 1, + adaptive_uniform_ratio: float = 0.1, + adaptive_alpha: float = 0.001, + ): + """Initialize motion sampler. + + Args: + motion_loader: Motion loader instance + mode: Sampling mode ("start", "uniform", "adaptive") + num_envs: Number of parallel environments + bin_count: Number of bins for adaptive sampling (auto if None) + adaptive_lambda: Decay factor for adaptive kernel + adaptive_kernel_size: Kernel size for adaptive sampling + adaptive_uniform_ratio: Uniform sampling ratio for adaptive mode + adaptive_alpha: EMA alpha for failure count updates + """ + self.motion_loader = motion_loader + self.mode = mode + self.num_envs = num_envs + + # Current frame indices for each environment + self.current_frames = np.zeros(num_envs, dtype=np.int32) + + # Adaptive sampling parameters + if bin_count is None: + # Auto-compute bin count based on motion length and FPS + self.bin_count = int(motion_loader.num_frames // motion_loader.fps) + 1 + else: + self.bin_count = bin_count + + self.adaptive_lambda = adaptive_lambda + self.adaptive_kernel_size = adaptive_kernel_size + self.adaptive_uniform_ratio = adaptive_uniform_ratio + self.adaptive_alpha = adaptive_alpha + + # Failure tracking for adaptive sampling + self.bin_failed_count = np.zeros(self.bin_count, dtype=np.float32) + self._current_bin_failed = np.zeros(self.bin_count, dtype=np.float32) + + # Precompute adaptive kernel + self.kernel = np.array( + [adaptive_lambda**i for i in range(adaptive_kernel_size)], dtype=np.float32 + ) + self.kernel = self.kernel / self.kernel.sum() + + # Metrics + self.sampling_entropy = 0.0 + self.sampling_top1_prob = 0.0 + self.sampling_top1_bin = 0.0 + + def sample_frames(self, env_ids: np.ndarray) -> np.ndarray: + """Sample motion frames for specified environments. + + Args: + env_ids: Environment indices to sample for + + Returns: + Sampled frame indices + """ + if self.mode == "start": + return self._sample_start(env_ids) + elif self.mode == "uniform": + return self._sample_uniform(env_ids) + elif self.mode == "adaptive": + return self._sample_adaptive(env_ids) + else: + raise ValueError(f"Unknown sampling mode: {self.mode}") + + def _sample_start(self, env_ids: np.ndarray) -> np.ndarray: + """Always start from beginning.""" + frames = np.zeros(len(env_ids), dtype=np.int32) + self.current_frames[env_ids] = frames + return frames + + def _sample_uniform(self, env_ids: np.ndarray) -> np.ndarray: + """Sample uniformly across motion.""" + frames = np.random.randint(0, self.motion_loader.num_frames, len(env_ids), dtype=np.int32) + self.current_frames[env_ids] = frames + + # Update metrics + self.sampling_entropy = 1.0 # Maximum entropy for uniform + self.sampling_top1_prob = 1.0 / self.bin_count + self.sampling_top1_bin = 0.5 # No specific bin preference + + return frames + + def _sample_adaptive(self, env_ids: np.ndarray) -> np.ndarray: + """Sample adaptively based on failure statistics.""" + # Compute sampling probabilities + sampling_probs = self.bin_failed_count + self.adaptive_uniform_ratio / float(self.bin_count) + + # Apply smoothing kernel (non-causal convolution) + if self.adaptive_kernel_size > 1: + # Pad and convolve + padded = np.pad(sampling_probs, (0, self.adaptive_kernel_size - 1), mode="edge") + sampling_probs = np.convolve(padded, self.kernel, mode="valid") + + # Normalize to probabilities + sampling_probs = sampling_probs / sampling_probs.sum() + + # Sample bins + sampled_bins = np.random.choice(self.bin_count, size=len(env_ids), p=sampling_probs) + + # Add random offset within bin + bin_offsets = np.random.uniform(0.0, 1.0, len(env_ids)) + frames = ( + (sampled_bins + bin_offsets) / self.bin_count * (self.motion_loader.num_frames - 1) + ).astype(np.int32) + + self.current_frames[env_ids] = frames + + # Update metrics + H = -(sampling_probs * np.log(sampling_probs + 1e-12)).sum() + H_norm = H / math.log(self.bin_count) if self.bin_count > 1 else 1.0 + pmax_idx = np.argmax(sampling_probs) + pmax = sampling_probs[pmax_idx] + + self.sampling_entropy = H_norm + self.sampling_top1_prob = pmax + self.sampling_top1_bin = pmax_idx / self.bin_count + + return frames + + def update_failure_stats( + self, terminated: np.ndarray, current_frames: np.ndarray | None = None + ): + """Update failure statistics for adaptive sampling. + + Args: + terminated: Boolean array indicating which environments terminated + current_frames: Optional current frame indices (uses internal if None) + """ + if self.mode != "adaptive": + return + + if current_frames is None: + current_frames = self.current_frames + + # Find which bins failed + if np.any(terminated): + bin_indices = np.clip( + (current_frames * self.bin_count) // max(self.motion_loader.num_frames, 1), + 0, + self.bin_count - 1, + ) + failed_bins = bin_indices[terminated] + + # Count failures per bin + self._current_bin_failed[:] = 0 + for bin_idx in failed_bins: + self._current_bin_failed[bin_idx] += 1 + + # Update EMA of failure counts + self.bin_failed_count = ( + self.adaptive_alpha * self._current_bin_failed + + (1 - self.adaptive_alpha) * self.bin_failed_count + ) + + def step(self): + """Advance all frames by one step.""" + self.current_frames += 1 + + # Find environments that reached end of motion + done_mask = self.current_frames >= self.motion_loader.num_frames + return np.where(done_mask)[0] + + def get_current_motion(self) -> MotionData: + """Get motion data at current frames for all environments.""" + return self.motion_loader.get_motion_at_frame(self.current_frames) diff --git a/unilab/envs/locomotion/g1/motion_tracking/tracking.py b/unilab/envs/locomotion/g1/motion_tracking/tracking.py new file mode 100644 index 000000000..defef339e --- /dev/null +++ b/unilab/envs/locomotion/g1/motion_tracking/tracking.py @@ -0,0 +1,604 @@ +"""G1 Motion Tracking Environment - Motion imitation task.""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import gymnasium as gym +import mujoco +import numpy as np +from etils import epath + +from unilab.base import registry +from unilab.base.backend import create_backend +from unilab.base.dtype_config import get_global_dtype +from unilab.base.np_env import NpEnvState +from unilab.envs.locomotion.g1.base import G1BaseCfg, G1BaseEnv + +from .math_utils import ( + matrix_from_quat, + quat_apply, + quat_error_magnitude, + quat_from_euler_xyz, + quat_inv, + quat_mul, + sample_uniform, + subtract_frame_transforms, + yaw_quat, +) +from .motion_loader import MotionLoader, MotionSampler + + +@dataclass +class RewardConfig: + """Reward configuration for motion tracking.""" + + scales: dict[str, float] = field( + default_factory=lambda: { + "motion_global_root_pos": 0.5, + "motion_global_root_ori": 0.5, + "motion_body_pos": 1.0, + "motion_body_ori": 1.0, + "motion_body_lin_vel": 1.0, + "motion_body_ang_vel": 1.0, + "action_rate_l2": -0.1, + "joint_limit": -10.0, + } + ) + # Standard deviations for exponential rewards + std_root_pos: float = 0.3 + std_root_ori: float = 0.4 + std_body_pos: float = 0.3 + std_body_ori: float = 0.4 + std_body_lin_vel: float = 1.0 + std_body_ang_vel: float = 3.14 + + +@dataclass +class PoseRandomization: + """Pose randomization ranges for reset.""" + + x: tuple[float, float] = (-0.05, 0.05) + y: tuple[float, float] = (-0.05, 0.05) + z: tuple[float, float] = (-0.01, 0.01) + roll: tuple[float, float] = (-0.1, 0.1) + pitch: tuple[float, float] = (-0.1, 0.1) + yaw: tuple[float, float] = (-0.2, 0.2) + + +@dataclass +class VelocityRandomization: + """Velocity randomization ranges for reset.""" + + x: tuple[float, float] = (-0.5, 0.5) + y: tuple[float, float] = (-0.5, 0.5) + z: tuple[float, float] = (-0.2, 0.2) + roll: tuple[float, float] = (-0.52, 0.52) + pitch: tuple[float, float] = (-0.52, 0.52) + yaw: tuple[float, float] = (-0.78, 0.78) + + +@dataclass +class G1MotionTrackingCfg(G1BaseCfg): + """Configuration for G1 motion tracking environment.""" + + model_file: str = str(epath.Path(__file__).parent.parent / "xml" / "scene_flat.xml") + motion_file: str = "" # Path to NPZ motion file (required) + anchor_body_name: str = "torso_link" + body_names: tuple[str, ...] = ( + "pelvis", + "left_hip_roll_link", + "left_knee_link", + "left_ankle_roll_link", + "right_hip_roll_link", + "right_knee_link", + "right_ankle_roll_link", + "torso_link", + "left_shoulder_roll_link", + "left_elbow_link", + "left_wrist_yaw_link", + "right_shoulder_roll_link", + "right_elbow_link", + "right_wrist_yaw_link", + ) + sampling_mode: str = "adaptive" # "start", "uniform", "adaptive" + max_episode_seconds: float = 10.0 + reward_config: RewardConfig = field(default_factory=RewardConfig) + pose_randomization: PoseRandomization = field(default_factory=PoseRandomization) + velocity_randomization: VelocityRandomization = field(default_factory=VelocityRandomization) + joint_position_range: tuple[float, float] = (-0.1, 0.1) + # Termination thresholds + anchor_pos_z_threshold: float = 0.25 + anchor_ori_threshold: float = 0.8 + ee_body_pos_z_threshold: float = 0.25 + ee_body_names: tuple[str, ...] = ( + "left_ankle_roll_link", + "right_ankle_roll_link", + "left_wrist_yaw_link", + "right_wrist_yaw_link", + ) + + +@registry.envcfg("G1MotionTracking") +@dataclass +class G1MotionTrackingEnvCfg(G1MotionTrackingCfg): + """Registered configuration for G1 motion tracking.""" + + pass + + +@registry.env("G1MotionTracking", sim_backend="mujoco") +class G1MotionTrackingEnv(G1BaseEnv): + """G1 Motion Tracking Environment.""" + + _cfg: G1MotionTrackingCfg + + def __init__(self, cfg: G1MotionTrackingCfg, num_envs=1, backend_type="mujoco"): + if not cfg.motion_file: + raise ValueError("motion_file must be specified in config") + + backend = create_backend( + backend_type, cfg.model_file, num_envs, cfg.sim_dt, body_name=cfg.asset.body_name + ) + super().__init__(cfg, backend, num_envs) + + # Get body IDs from MuJoCo model + self.body_ids = np.array( + [ + mujoco.mj_name2id(self._backend.model, mujoco.mjtObj.mjOBJ_BODY, name) + for name in cfg.body_names + ], + dtype=np.int32, + ) + self.anchor_body_idx = cfg.body_names.index(cfg.anchor_body_name) + + # Get end-effector body indices for termination + self.ee_body_indices = np.array( + [cfg.body_names.index(name) for name in cfg.ee_body_names], dtype=np.int32 + ) + + # Load motion data + self.motion_loader = MotionLoader(cfg.motion_file, body_indices=self.body_ids) + self.motion_sampler = MotionSampler( + self.motion_loader, mode=cfg.sampling_mode, num_envs=num_envs + ) + + # Buffers for relative body transforms + self.body_pos_relative_w = np.zeros( + (num_envs, len(cfg.body_names), 3), dtype=get_global_dtype() + ) + self.body_quat_relative_w = np.zeros( + (num_envs, len(cfg.body_names), 4), dtype=get_global_dtype() + ) + self.body_quat_relative_w[:, :, 0] = 1.0 # Initialize to identity quaternion + + self._enable_reward_log = True + self._init_obs_space() + self._init_reward_functions() + + def _init_obs_space(self): + # Actor observations: motion_anchor_pos_b (3), motion_anchor_ori_b (6), + # base_lin_vel (3), base_ang_vel (3), joint_pos (num_joints), joint_vel (num_joints), actions (num_joints) + num_actor_obs = 3 + 6 + 3 + 3 + self._num_action * 3 + + # Critic observations: actor + body_pos_b (num_bodies * 3) + body_ori_b (num_bodies * 6) + num_critic_obs = num_actor_obs + len(self._cfg.body_names) * 9 + + self._observation_space = gym.spaces.Box( + low=-float("inf"), high=float("inf"), shape=(num_critic_obs,), dtype=float + ) + self.num_actor_obs = num_actor_obs + + def _init_reward_functions(self): + self._reward_fns = { + "motion_global_root_pos": self._reward_motion_global_root_pos, + "motion_global_root_ori": self._reward_motion_global_root_ori, + "motion_body_pos": self._reward_motion_body_pos, + "motion_body_ori": self._reward_motion_body_ori, + "motion_body_lin_vel": self._reward_motion_body_lin_vel, + "motion_body_ang_vel": self._reward_motion_body_ang_vel, + "action_rate_l2": self._reward_action_rate_l2, + "joint_limit": self._reward_joint_limit, + } + + @property + def observation_space(self) -> gym.spaces.Box: + return self._observation_space # type: ignore[no-any-return] + + def update_state(self, state: NpEnvState) -> NpEnvState: + # Get current motion data + motion_data = self.motion_sampler.get_current_motion() + + # Get robot state + linvel = self.get_local_linvel() + gyro = self.get_gyro() + dof_pos = self.get_dof_pos() + dof_vel = self.get_dof_vel() + qpos = self._backend.get_qpos() + + # Get body states + robot_body_pos_w = self._backend.get_body_pos(self.body_ids) + robot_body_quat_w = self._backend.get_body_quat(self.body_ids) + robot_body_lin_vel_w = self._backend.get_body_lin_vel(self.body_ids) + robot_body_ang_vel_w = self._backend.get_body_ang_vel(self.body_ids) + + # Compute relative body transforms (for observations and rewards) + self._update_relative_transforms(motion_data, robot_body_pos_w, robot_body_quat_w) + + # Compute terminations + terminated = self._compute_terminations( + motion_data, robot_body_pos_w, robot_body_quat_w, qpos + ) + + # Update failure statistics for adaptive sampling + self.motion_sampler.update_failure_stats(terminated) + + # Compute reward + reward = self._compute_reward( + state.info, + motion_data, + robot_body_pos_w, + robot_body_quat_w, + robot_body_lin_vel_w, + robot_body_ang_vel_w, + dof_pos, + ) + + # Compute observations + obs = self._compute_obs( + state.info, + motion_data, + linvel, + gyro, + dof_pos, + dof_vel, + robot_body_pos_w, + robot_body_quat_w, + ) + + # Advance motion frames + done_env_ids = self.motion_sampler.step() + if len(done_env_ids) > 0: + # Resample motion for environments that reached end + self.motion_sampler.sample_frames(done_env_ids) + + return state.replace(obs=obs, reward=reward, terminated=terminated) + + def _update_relative_transforms( + self, motion_data, robot_body_pos_w: np.ndarray, robot_body_quat_w: np.ndarray + ): + """Update relative body transforms for tracking.""" + # Get anchor states + anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] + anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] + + # Compute delta transform: keep robot's XY position, use motion's Z height + # and apply yaw-only rotation difference + delta_pos_w = robot_anchor_pos_w.copy() + delta_pos_w[:, 2] = anchor_pos_w[:, 2] + + # Compute yaw-only rotation difference + quat_diff = quat_mul(robot_anchor_quat_w, quat_inv(anchor_quat_w)) + delta_ori_w = yaw_quat(quat_diff) + + # Transform motion body states to robot-relative coordinates + for i in range(len(self._cfg.body_names)): + self.body_quat_relative_w[:, i] = quat_mul(delta_ori_w, motion_data.body_quat_w[:, i]) + rel_pos = motion_data.body_pos_w[:, i] - anchor_pos_w + self.body_pos_relative_w[:, i] = delta_pos_w + quat_apply(delta_ori_w, rel_pos) + + def _compute_terminations( + self, + motion_data, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + qpos: np.ndarray, + ) -> np.ndarray: + """Compute termination conditions.""" + terminated = np.zeros(self._num_envs, dtype=bool) + + # Anchor position error (Z-axis only) + anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] + anchor_pos_error_z = np.abs(anchor_pos_w[:, 2] - robot_anchor_pos_w[:, 2]) + terminated |= anchor_pos_error_z > self._cfg.anchor_pos_z_threshold + + # Anchor orientation error (gravity direction) + anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] + gravity_vec = np.array([0, 0, -1], dtype=get_global_dtype()) + motion_gravity_b = quat_apply(quat_inv(anchor_quat_w), gravity_vec) + robot_gravity_b = quat_apply(quat_inv(robot_anchor_quat_w), gravity_vec) + gravity_error = np.abs(motion_gravity_b[:, 2] - robot_gravity_b[:, 2]) + terminated |= gravity_error > self._cfg.anchor_ori_threshold + + # End-effector position error (Z-axis only) + for ee_idx in self.ee_body_indices: + ee_pos_error_z = np.abs( + self.body_pos_relative_w[:, ee_idx, 2] - robot_body_pos_w[:, ee_idx, 2] + ) + terminated |= ee_pos_error_z > self._cfg.ee_body_pos_z_threshold + + return terminated + + def _compute_obs( + self, + info: dict, + motion_data, + linvel: np.ndarray, + gyro: np.ndarray, + dof_pos: np.ndarray, + dof_vel: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + ) -> np.ndarray: + """Compute observations (critic observations, actor is subset).""" + # Get anchor states + anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] + anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] + + # Motion anchor position in robot frame + motion_anchor_pos_b, _ = subtract_frame_transforms( + robot_anchor_pos_w, robot_anchor_quat_w, anchor_pos_w, anchor_quat_w + ) + + # Motion anchor orientation in robot frame (as rotation matrix first 2 columns) + _, motion_anchor_ori_rel = subtract_frame_transforms( + robot_anchor_pos_w, robot_anchor_quat_w, anchor_pos_w, anchor_quat_w + ) + motion_anchor_ori_mat = matrix_from_quat(motion_anchor_ori_rel) + motion_anchor_ori_b = motion_anchor_ori_mat[:, :, :2].reshape(self._num_envs, 6) + + # Joint positions and velocities + joint_pos_rel = dof_pos - self.default_angles + last_actions = info.get("current_actions", np.zeros_like(joint_pos_rel)) + + # Actor observations + actor_obs = np.concatenate( + [ + motion_anchor_pos_b, + motion_anchor_ori_b, + linvel, + gyro, + joint_pos_rel, + dof_vel, + last_actions, + ], + axis=1, + dtype=get_global_dtype(), + ) + + # Robot body positions in robot anchor frame + robot_body_pos_b = np.zeros( + (self._num_envs, len(self._cfg.body_names), 3), dtype=get_global_dtype() + ) + robot_body_ori_b = np.zeros( + (self._num_envs, len(self._cfg.body_names), 6), dtype=get_global_dtype() + ) + for i in range(len(self._cfg.body_names)): + pos_b, ori_b = subtract_frame_transforms( + robot_anchor_pos_w, + robot_anchor_quat_w, + robot_body_pos_w[:, i], + robot_body_quat_w[:, i], + ) + robot_body_pos_b[:, i] = pos_b + ori_mat = matrix_from_quat(ori_b) + robot_body_ori_b[:, i] = ori_mat[:, :, :2].reshape(self._num_envs, 6) + + # Critic observations (privileged) + critic_obs = np.concatenate( + [ + actor_obs, + robot_body_pos_b.reshape(self._num_envs, -1), + robot_body_ori_b.reshape(self._num_envs, -1), + ], + axis=1, + dtype=get_global_dtype(), + ) + + return critic_obs + + def _compute_reward( + self, + info: dict, + motion_data, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + robot_body_lin_vel_w: np.ndarray, + robot_body_ang_vel_w: np.ndarray, + dof_pos: np.ndarray, + ) -> np.ndarray: + """Compute reward.""" + dtype = get_global_dtype() + reward = np.zeros((self._num_envs,), dtype=dtype) + cfg = self._cfg.reward_config + + step_count = info.get("steps", np.zeros((self._num_envs,), dtype=np.uint32)) + should_log = self._enable_reward_log and (int(step_count[0]) % 4 == 0) + log = {} if should_log else info.get("log", {}) + + # Store motion and robot states in info for reward functions + info["motion_data"] = motion_data + info["robot_body_pos_w"] = robot_body_pos_w + info["robot_body_quat_w"] = robot_body_quat_w + info["robot_body_lin_vel_w"] = robot_body_lin_vel_w + info["robot_body_ang_vel_w"] = robot_body_ang_vel_w + info["dof_pos"] = dof_pos + + for name, scale in cfg.scales.items(): + if scale == 0 or name not in self._reward_fns: + continue + rew = self._reward_fns[name](info) + weighted_rew = rew * scale + reward += weighted_rew + if should_log: + log[f"reward/{name}"] = float(np.mean(weighted_rew)) + + info["log"] = log + return reward * self._cfg.ctrl_dt + + # Reward functions + def _reward_motion_global_root_pos(self, info: dict) -> np.ndarray: + motion_data = info["motion_data"] + robot_body_pos_w = info["robot_body_pos_w"] + anchor_pos_w = motion_data.body_pos_w[:, self.anchor_body_idx] + robot_anchor_pos_w = robot_body_pos_w[:, self.anchor_body_idx] + error = np.sum(np.square(anchor_pos_w - robot_anchor_pos_w), axis=-1) + return np.exp(-error / self._cfg.reward_config.std_root_pos**2) + + def _reward_motion_global_root_ori(self, info: dict) -> np.ndarray: + motion_data = info["motion_data"] + robot_body_quat_w = info["robot_body_quat_w"] + anchor_quat_w = motion_data.body_quat_w[:, self.anchor_body_idx] + robot_anchor_quat_w = robot_body_quat_w[:, self.anchor_body_idx] + error = quat_error_magnitude(anchor_quat_w, robot_anchor_quat_w) ** 2 + return np.exp(-error / self._cfg.reward_config.std_root_ori**2) + + def _reward_motion_body_pos(self, info: dict) -> np.ndarray: + robot_body_pos_w = info["robot_body_pos_w"] + error = np.sum(np.square(self.body_pos_relative_w - robot_body_pos_w), axis=-1) + return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_pos**2) + + def _reward_motion_body_ori(self, info: dict) -> np.ndarray: + robot_body_quat_w = info["robot_body_quat_w"] + error = quat_error_magnitude(self.body_quat_relative_w, robot_body_quat_w) ** 2 + return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_ori**2) + + def _reward_motion_body_lin_vel(self, info: dict) -> np.ndarray: + motion_data = info["motion_data"] + robot_body_lin_vel_w = info["robot_body_lin_vel_w"] + error = np.sum(np.square(motion_data.body_lin_vel_w - robot_body_lin_vel_w), axis=-1) + return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_lin_vel**2) + + def _reward_motion_body_ang_vel(self, info: dict) -> np.ndarray: + motion_data = info["motion_data"] + robot_body_ang_vel_w = info["robot_body_ang_vel_w"] + error = np.sum(np.square(motion_data.body_ang_vel_w - robot_body_ang_vel_w), axis=-1) + return np.exp(-error.mean(-1) / self._cfg.reward_config.std_body_ang_vel**2) + + def _reward_action_rate_l2(self, info: dict) -> np.ndarray: + action_diff = info["current_actions"] - info["last_actions"] + return np.sum(np.square(action_diff), axis=1) + + def _reward_joint_limit(self, info: dict) -> np.ndarray: + dof_pos = info["dof_pos"] + # Get joint limits from model + model = self._backend.model + joint_range = model.jnt_range[7:] # Skip free joint + lower = joint_range[:, 0] + upper = joint_range[:, 1] + + # Compute violation + lower_violation = np.maximum(0, lower - dof_pos) + upper_violation = np.maximum(0, dof_pos - upper) + violation = lower_violation + upper_violation + return np.sum(np.square(violation), axis=1) + + def reset(self, env_indices: np.ndarray): + """Reset specified environments.""" + dtype = get_global_dtype() + num_reset = len(env_indices) + + # Sample motion frames + motion_frames = self.motion_sampler.sample_frames(env_indices) + motion_data = self.motion_loader.get_motion_at_frame(motion_frames) + + # Get anchor body state from motion + root_pos = motion_data.body_pos_w[:, 0].copy() + root_ori = motion_data.body_quat_w[:, 0].copy() + root_lin_vel = motion_data.body_lin_vel_w[:, 0].copy() + root_ang_vel = motion_data.body_ang_vel_w[:, 0].copy() + joint_pos = motion_data.joint_pos.copy() + joint_vel = motion_data.joint_vel.copy() + + # Apply pose randomization + pose_rand = self._cfg.pose_randomization + range_list = [ + (pose_rand.x[0], pose_rand.x[1]), + (pose_rand.y[0], pose_rand.y[1]), + (pose_rand.z[0], pose_rand.z[1]), + (pose_rand.roll[0], pose_rand.roll[1]), + (pose_rand.pitch[0], pose_rand.pitch[1]), + (pose_rand.yaw[0], pose_rand.yaw[1]), + ] + rand_samples = np.array( + [[np.random.uniform(r[0], r[1]) for r in range_list] for _ in range(num_reset)], + dtype=dtype, + ) + root_pos += rand_samples[:, 0:3] + ori_delta = quat_from_euler_xyz(rand_samples[:, 3], rand_samples[:, 4], rand_samples[:, 5]) + root_ori = quat_mul(ori_delta, root_ori) + + # Apply velocity randomization + vel_rand = self._cfg.velocity_randomization + range_list = [ + (vel_rand.x[0], vel_rand.x[1]), + (vel_rand.y[0], vel_rand.y[1]), + (vel_rand.z[0], vel_rand.z[1]), + (vel_rand.roll[0], vel_rand.roll[1]), + (vel_rand.pitch[0], vel_rand.pitch[1]), + (vel_rand.yaw[0], vel_rand.yaw[1]), + ] + rand_samples = np.array( + [[np.random.uniform(r[0], r[1]) for r in range_list] for _ in range(num_reset)], + dtype=dtype, + ) + root_lin_vel += rand_samples[:, :3] + root_ang_vel += rand_samples[:, 3:] + + # Apply joint position randomization + joint_pos += sample_uniform( + self._cfg.joint_position_range[0], + self._cfg.joint_position_range[1], + joint_pos.shape, + dtype=dtype, + ) + + # Clip joint positions to limits + model = self._backend.model + joint_range = model.jnt_range[7:] # Skip free joint + joint_pos = np.clip(joint_pos, joint_range[:, 0], joint_range[:, 1]) + + # Construct qpos and qvel + qpos = np.tile(self._init_qpos, (num_reset, 1)) + qvel = np.tile(self._init_qvel, (num_reset, 1)) + + qpos[:, 0:3] = root_pos + qpos[:, 3:7] = root_ori + qpos[:, 7:] = joint_pos + + qvel[:, 0:3] = root_lin_vel + qvel[:, 3:6] = root_ang_vel + qvel[:, 6:] = joint_vel + + # Set state + self._backend.set_state(env_indices, qpos, qvel) + + # Initialize info + info = { + "current_actions": np.zeros((num_reset, self._num_action), dtype=dtype), + "last_actions": np.zeros((num_reset, self._num_action), dtype=dtype), + } + + # Compute initial observations + motion_data = self.motion_sampler.get_current_motion() + linvel = self.get_local_linvel()[env_indices] + gyro = self.get_gyro()[env_indices] + dof_pos = self.get_dof_pos()[env_indices] + dof_vel = self.get_dof_vel()[env_indices] + robot_body_pos_w = self._backend.get_body_pos(self.body_ids)[env_indices] + robot_body_quat_w = self._backend.get_body_quat(self.body_ids)[env_indices] + + obs = self._compute_obs( + info, motion_data, linvel, gyro, dof_pos, dof_vel, robot_body_pos_w, robot_body_quat_w + ) + + # Actor observations are first num_actor_obs elements + actor_obs = obs[:, : self.num_actor_obs] + + return actor_obs, obs, info