Skip to content

perf(motion): fuse command metrics in numba #1320

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TATP-233 merged 1 commit into from
Aug 26, 2026
Merged

perf(motion): fuse command metrics in numba #1320

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
181 changes: 171 additions & 10 deletions src/unilab/tasks/motion_tracking/common/kernels.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
from __future__ import annotations

import os
from math import atan2, sqrt

import numpy as np
from numba import config, njit, prange, set_num_threads
Expand All @@ -28,6 +29,174 @@ def configure_motion_kernel_runtime() -> None:
_runtime_configured = True


@njit(inline="always")
def _quat_error_squared(
q1_w: float,
q1_x: float,
q1_y: float,
q1_z: float,
q2_w: float,
q2_x: float,
q2_y: float,
q2_z: float,
) -> float:
"""Return the squared shortest-path angular error for two quaternions."""
rel_w = abs(q2_w * q1_w + q2_x * q1_x + q2_y * q1_y + q2_z * q1_z)
rel_x = -q2_w * q1_x + q2_x * q1_w - q2_y * q1_z + q2_z * q1_y
rel_y = -q2_w * q1_y + q2_x * q1_z + q2_y * q1_w - q2_z * q1_x
rel_z = -q2_w * q1_z - q2_x * q1_y + q2_y * q1_x + q2_z * q1_w
xyz_norm = sqrt(rel_x * rel_x + rel_y * rel_y + rel_z * rel_z)
clipped_w = min(max(rel_w, -1.0), 1.0)
angle = 2.0 * atan2(xyz_norm, clipped_w)
return angle * angle


@njit(cache=True, nogil=True, parallel=True)
def update_motion_metrics_kernel(
env_ids: np.ndarray,
anchor_body_idx: int,
motion_body_pos_w: np.ndarray,
robot_body_pos_w: np.ndarray,
motion_body_quat_w: np.ndarray,
robot_body_quat_w: np.ndarray,
motion_body_lin_vel_w: np.ndarray,
robot_body_lin_vel_w: np.ndarray,
motion_body_ang_vel_w: np.ndarray,
robot_body_ang_vel_w: np.ndarray,
body_pos_relative_w: np.ndarray,
body_quat_relative_w: np.ndarray,
motion_joint_pos: np.ndarray,
robot_joint_pos: np.ndarray,
motion_joint_vel: np.ndarray,
robot_joint_vel: np.ndarray,
error_anchor_pos: np.ndarray,
error_anchor_rot: np.ndarray,
error_anchor_lin_vel: np.ndarray,
error_anchor_ang_vel: np.ndarray,
error_body_pos: np.ndarray,
error_body_rot: np.ndarray,
error_body_lin_vel: np.ndarray,
error_body_ang_vel: np.ndarray,
error_joint_pos: np.ndarray,
error_joint_vel: np.ndarray,
) -> None:
"""Write all MotionCommand error metrics for the selected environment rows.

``env_ids`` is a cold-path-owned all-row index buffer for normal steps and
the reset row list for partial resets. Keeping one kernel for both paths
avoids a second production mathematical implementation while preserving the
manager's row-scoped reset contract.
"""
num_bodies = motion_body_pos_w.shape[1]
num_joints = motion_joint_pos.shape[1]
for row in prange(env_ids.shape[0]):
env_idx = env_ids[row]
anchor_dx = (
motion_body_pos_w[env_idx, anchor_body_idx, 0]
- robot_body_pos_w[env_idx, anchor_body_idx, 0]
)
anchor_dy = (
motion_body_pos_w[env_idx, anchor_body_idx, 1]
- robot_body_pos_w[env_idx, anchor_body_idx, 1]
)
anchor_dz = (
motion_body_pos_w[env_idx, anchor_body_idx, 2]
- robot_body_pos_w[env_idx, anchor_body_idx, 2]
)
error_anchor_pos[env_idx] = sqrt(
anchor_dx * anchor_dx + anchor_dy * anchor_dy + anchor_dz * anchor_dz
)
error_anchor_rot[env_idx] = sqrt(
_quat_error_squared(
motion_body_quat_w[env_idx, anchor_body_idx, 0],
motion_body_quat_w[env_idx, anchor_body_idx, 1],
motion_body_quat_w[env_idx, anchor_body_idx, 2],
motion_body_quat_w[env_idx, anchor_body_idx, 3],
robot_body_quat_w[env_idx, anchor_body_idx, 0],
robot_body_quat_w[env_idx, anchor_body_idx, 1],
robot_body_quat_w[env_idx, anchor_body_idx, 2],
robot_body_quat_w[env_idx, anchor_body_idx, 3],
)
)

anchor_lin_sq = 0.0
anchor_ang_sq = 0.0
for component in range(3):
lin_delta = (
motion_body_lin_vel_w[env_idx, anchor_body_idx, component]
- robot_body_lin_vel_w[env_idx, anchor_body_idx, component]
)
ang_delta = (
motion_body_ang_vel_w[env_idx, anchor_body_idx, component]
- robot_body_ang_vel_w[env_idx, anchor_body_idx, component]
)
anchor_lin_sq += lin_delta * lin_delta
anchor_ang_sq += ang_delta * ang_delta
error_anchor_lin_vel[env_idx] = sqrt(anchor_lin_sq)
error_anchor_ang_vel[env_idx] = sqrt(anchor_ang_sq)

body_pos_sum = 0.0
body_rot_sum = 0.0
body_lin_sum = 0.0
body_ang_sum = 0.0
for body_idx in range(num_bodies):
pos_sq = 0.0
lin_sq = 0.0
ang_sq = 0.0
for component in range(3):
pos_delta = (
body_pos_relative_w[env_idx, body_idx, component]
- robot_body_pos_w[env_idx, body_idx, component]
)
lin_delta = (
motion_body_lin_vel_w[env_idx, body_idx, component]
- robot_body_lin_vel_w[env_idx, body_idx, component]
)
ang_delta = (
motion_body_ang_vel_w[env_idx, body_idx, component]
- robot_body_ang_vel_w[env_idx, body_idx, component]
)
pos_sq += pos_delta * pos_delta
lin_sq += lin_delta * lin_delta
ang_sq += ang_delta * ang_delta
body_pos_sum += sqrt(pos_sq)
body_lin_sum += sqrt(lin_sq)
body_ang_sum += sqrt(ang_sq)
body_rot_sum += sqrt(
_quat_error_squared(
body_quat_relative_w[env_idx, body_idx, 0],
body_quat_relative_w[env_idx, body_idx, 1],
body_quat_relative_w[env_idx, body_idx, 2],
body_quat_relative_w[env_idx, body_idx, 3],
robot_body_quat_w[env_idx, body_idx, 0],
robot_body_quat_w[env_idx, body_idx, 1],
robot_body_quat_w[env_idx, body_idx, 2],
robot_body_quat_w[env_idx, body_idx, 3],
)
)
if num_bodies == 0:
error_body_pos[env_idx] = np.nan
error_body_rot[env_idx] = np.nan
error_body_lin_vel[env_idx] = np.nan
error_body_ang_vel[env_idx] = np.nan
else:
inv_bodies = 1.0 / num_bodies
error_body_pos[env_idx] = body_pos_sum * inv_bodies
error_body_rot[env_idx] = body_rot_sum * inv_bodies
error_body_lin_vel[env_idx] = body_lin_sum * inv_bodies
error_body_ang_vel[env_idx] = body_ang_sum * inv_bodies

joint_pos_sq = 0.0
joint_vel_sq = 0.0
for joint_idx in range(num_joints):
pos_delta = motion_joint_pos[env_idx, joint_idx] - robot_joint_pos[env_idx, joint_idx]
vel_delta = motion_joint_vel[env_idx, joint_idx] - robot_joint_vel[env_idx, joint_idx]
joint_pos_sq += pos_delta * pos_delta
joint_vel_sq += vel_delta * vel_delta
error_joint_pos[env_idx] = sqrt(joint_pos_sq)
error_joint_vel[env_idx] = sqrt(joint_vel_sq)


@njit(cache=True, nogil=True, parallel=True)
def termination_anchor_pos_kernel(
motion_body_pos_w: np.ndarray,
Expand Down Expand Up @@ -97,16 +266,7 @@ def reward_motion_body_ori_kernel(
y2 = actual[env_idx, body_idx, 2]
z2 = actual[env_idx, body_idx, 3]

# Relative rotation actual * conjugate(reference), matching
# np_quat_error_magnitude_squared_batched without materializing arrays.
rel_w = abs(w2 * w1 + x2 * x1 + y2 * y1 + z2 * z1)
rel_x = -w2 * x1 + x2 * w1 - y2 * z1 + z2 * y1
rel_y = -w2 * y1 + x2 * z1 + y2 * w1 - z2 * x1
rel_z = -w2 * z1 - x2 * y1 + y2 * x1 + z2 * w1
xyz_norm = np.sqrt(rel_x * rel_x + rel_y * rel_y + rel_z * rel_z)
clipped_w = min(max(rel_w, -1.0), 1.0)
angle = 2.0 * np.arctan2(xyz_norm, clipped_w)
error += angle * angle
error += _quat_error_squared(w1, x1, y1, z1, w2, x2, y2, z2)
out[env_idx] = np.exp(error / denominator)


Expand Down Expand Up @@ -167,4 +327,5 @@ def reward_motion_body_ang_vel_kernel(
"reward_motion_body_ori_kernel",
"reward_motion_body_pos_kernel",
"termination_anchor_pos_kernel",
"update_motion_metrics_kernel",
]
75 changes: 38 additions & 37 deletions src/unilab/tasks/motion_tracking/common/manager_terms.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
reward_motion_body_ori_kernel,
reward_motion_body_pos_kernel,
termination_anchor_pos_kernel,
update_motion_metrics_kernel,
)
from .motion_loader import MotionData, MotionLoader, MotionSampler
from .observations import write_body_ori6_in_anchor_frame, write_body_pos_in_anchor_frame
Expand Down Expand Up @@ -196,6 +197,8 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv):
self._env_error = np.empty(self.num_envs, dtype=dtype)
self._reward_term = np.empty(self.num_envs, dtype=dtype)
self._robot_cache_step = -1
self._all_env_ids = np.arange(self.num_envs, dtype=np.int32)
self._all_env_ids.setflags(write=False)
# Env ids of the most recent scoped (reset-path) compute; None after a
# per-step compute. Written by `_update_command`, consumed by
# `post_compute` to restrict refresh work to the reset rows.
Expand Down Expand Up @@ -224,6 +227,10 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv):
self._refresh_motion()
self._refresh_robot_state(force=True)
self._refresh_relative_state()
# Compile and execute the metrics kernel on the cold path so the first
# measured command-manager step contains no Numba dispatch/JIT latency.
configure_motion_kernel_runtime()
self._update_metrics()

def _make_motion_loader(
self,
Expand Down Expand Up @@ -517,43 +524,37 @@ def _refresh_relative_state_rows(self, env_ids: np.ndarray) -> None:
self.robot_body_ori_b[env_ids] = robot_body_ori_b

def _update_metrics(self, env_ids: np.ndarray | None = None) -> None:
# All error metrics are row-wise functions of the motion/robot buffers;
# on the reset path only the reset rows are recomputed since other rows
# are unchanged since the per-step update.
sel: np.ndarray | slice = slice(None) if env_ids is None else env_ids
self.metrics["error_anchor_pos"][sel] = np.linalg.norm(
self.anchor_pos_w[sel] - self.robot_anchor_pos_w[sel], axis=-1
)
self.metrics["error_anchor_rot"][sel] = np.sqrt(
np_quat_error_magnitude_squared_batched(
self.anchor_quat_w[sel], self.robot_anchor_quat_w[sel]
)
)
self.metrics["error_anchor_lin_vel"][sel] = np.linalg.norm(
self.anchor_lin_vel_w[sel] - self.robot_anchor_lin_vel_w[sel], axis=-1
)
self.metrics["error_anchor_ang_vel"][sel] = np.linalg.norm(
self.anchor_ang_vel_w[sel] - self.robot_anchor_ang_vel_w[sel], axis=-1
)
self.metrics["error_body_pos"][sel] = np.linalg.norm(
self.body_pos_relative_w[sel] - self.robot_body_pos_w[sel], axis=-1
).mean(axis=-1)
self.metrics["error_body_rot"][sel] = np.sqrt(
np_quat_error_magnitude_squared_batched(
self.body_quat_relative_w[sel], self.robot_body_quat_w[sel]
)
).mean(axis=-1)
self.metrics["error_body_lin_vel"][sel] = np.linalg.norm(
self.body_lin_vel_w[sel] - self.robot_body_lin_vel_w[sel], axis=-1
).mean(axis=-1)
self.metrics["error_body_ang_vel"][sel] = np.linalg.norm(
self.body_ang_vel_w[sel] - self.robot_body_ang_vel_w[sel], axis=-1
).mean(axis=-1)
self.metrics["error_joint_pos"][sel] = np.linalg.norm(
self.joint_pos[sel] - self.robot_joint_pos[sel], axis=-1
)
self.metrics["error_joint_vel"][sel] = np.linalg.norm(
self.joint_vel[sel] - self.robot_joint_vel[sel], axis=-1
# All row-wise metrics are written by one Numba kernel. Passing an
# explicit all-row index buffer for normal steps lets the same kernel
# serve partial-reset rows without retaining a NumPy runtime formula.
rows = self._all_env_ids if env_ids is None else env_ids
update_motion_metrics_kernel(
rows,
self.anchor_body_idx,
self._body_pos_w,
self._robot_body_pos_w,
self._motion_data.body_quat_w,
self._robot_body_quat_w,
self._motion_data.body_lin_vel_w,
self._robot_body_lin_vel_w,
self._motion_data.body_ang_vel_w,
self._robot_body_ang_vel_w,
self.body_pos_relative_w,
self.body_quat_relative_w,
self._motion_data.joint_pos,
self.robot_joint_pos,
self._motion_data.joint_vel,
self.robot_joint_vel,
self.metrics["error_anchor_pos"],
self.metrics["error_anchor_rot"],
self.metrics["error_anchor_lin_vel"],
self.metrics["error_anchor_ang_vel"],
self.metrics["error_body_pos"],
self.metrics["error_body_rot"],
self.metrics["error_body_lin_vel"],
self.metrics["error_body_ang_vel"],
self.metrics["error_joint_pos"],
self.metrics["error_joint_vel"],
)
# Sampler statistics are global scalars, so every row tracks them.
self.metrics["sampling_entropy"].fill(self.sampler.sampling_entropy)
Expand Down
Loading