diff --git a/src/unilab/managers/_noise/noise_cfg.py b/src/unilab/managers/_noise/noise_cfg.py index d5cf2d8ec..13383b018 100644 --- a/src/unilab/managers/_noise/noise_cfg.py +++ b/src/unilab/managers/_noise/noise_cfg.py @@ -68,14 +68,21 @@ def apply(self, data: np.ndarray, *, rng: np.random.Generator | None = None) -> n_min = self._as_array(self.n_min, data.dtype) n_max = self._as_array(self.n_max, data.dtype) - # Generate uniform noise in [0, 1) and scale to [n_min, n_max). + # Generate uniform noise in [0, 1) and transform the generated array + # in place. The pre-existing expression allocated one array for each + # multiply, add, and final data operation; keeping the same float32 + # cast and ufunc order preserves bit-level results while returning the + # transformed noise buffer itself. noise = rng.random(data.shape).astype(data.dtype, copy=False) - noise = noise * (n_max - n_min) + n_min + np.multiply(noise, n_max - n_min, out=noise) + np.add(noise, n_min, out=noise) if self.operation == "add": - return data + noise + np.add(data, noise, out=noise) + return noise elif self.operation == "scale": - return data * noise + np.multiply(data, noise, out=noise) + return noise elif self.operation == "abs": return noise else: diff --git a/src/unilab/managers/observation_manager.py b/src/unilab/managers/observation_manager.py index 7bf1b4dae..8cfa6af60 100644 --- a/src/unilab/managers/observation_manager.py +++ b/src/unilab/managers/observation_manager.py @@ -335,7 +335,7 @@ def _row_env_ids(mask: np.ndarray) -> list[int]: ) # Sanitize (applies to both "warn" and "sanitize" policies). - return np.nan_to_num(tensor, nan=0.0, posinf=0.0, neginf=0.0) + return np.nan_to_num(tensor, copy=False, nan=0.0, posinf=0.0, neginf=0.0) def compute( self, @@ -381,6 +381,17 @@ def compute_group( group_term_names = self._group_obs_term_names[group_name] group_obs: dict[str, np.ndarray] = {} obs_terms = zip(group_term_names, self._group_obs_term_cfgs[group_name], strict=False) + # In the strict default policy a finite result is by far the common + # case. For concatenated groups, scan the assembled output once and + # only inspect individual slices when an error is actually found; this + # retains the per-term diagnostic while removing repeated full scans + # from the hot path. + defer_error_nan_check = ( + self._group_obs_concatenate[group_name] + and not self._group_obs_temporal[group_name] + and group_cfg.nan_check_per_term + and group_cfg.nan_policy == "error" + ) # Reset path (issue #1259 R2): when no term in this group uses delay or # history buffers, everything downstream of the term call is row # independent, so only the reset rows are processed. Term calls and @@ -409,10 +420,28 @@ def compute_group( # NoiseModel.__call__ likewise returns a new array. obs = self._group_obs_class_instances[group_name][term_name](obs) fresh = True - if not row_scoped and not fresh: - # Terms may return backend/command-owned buffers; copy before the - # in-place clip/scale below. Skipped when noise already produced - # a fresh array (issue #1296). + sanitizes_per_term = group_cfg.nan_check_per_term and group_cfg.nan_policy in ( + "warn", + "sanitize", + ) + exposes_term_output = ( + not self._group_obs_concatenate[group_name] + and term_cfg.delay_max_lag == 0 + and term_cfg.history_length == 0 + ) + if ( + not row_scoped + and not fresh + and ( + term_cfg.clip is not None + or term_cfg.scale is not None + or sanitizes_per_term + or exposes_term_output + ) + ): + # Concatenation and temporal buffers already copy their inputs. + # Only take a defensive copy when this pipeline may mutate the + # term or expose it directly to callers. obs = obs.copy() if row_scoped: # Fresh row copy; safe for the in-place clip/scale below. @@ -425,7 +454,11 @@ def compute_group( np.multiply(obs, scale, out=obs) # Check for NaN/Inf before delay/history buffers (per-term checking). - if group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": + if ( + group_cfg.nan_check_per_term + and group_cfg.nan_policy != "disabled" + and not defer_error_nan_check + ): obs = self._check_and_handle_nans( obs, context=f"{group_name}/{term_name}", @@ -474,6 +507,33 @@ def compute_group( result = np.concatenate( list(group_obs.values()), axis=self._group_obs_concatenate_dim[group_name] ) + if defer_error_nan_check: + finite = np.isfinite(result) + if not finite.all(): + axis = self._group_obs_concatenate_dim[group_name] + axis = axis if axis >= 0 else result.ndim + axis + offset = 0 + for term_name, term_dims in zip( + group_term_names, + self._group_obs_term_dim[group_name], + strict=True, + ): + width = int(term_dims[axis - 1]) + selectors = [slice(None)] * result.ndim + selectors[axis] = slice(offset, offset + width) + selector = tuple(selectors) + if not finite[selector].all(): + # Reuse the established diagnostic path so the + # error still names the first offending term and + # reports reset-row IDs when applicable. + self._check_and_handle_nans( + result[selector], + context=f"{group_name}/{term_name}", + policy=group_cfg.nan_policy, + env_ids=env_ids if row_scoped else None, + ) + break + offset += width # Final check for concatenated result (non-per-term checking). if not group_cfg.nan_check_per_term and group_cfg.nan_policy != "disabled": result = self._check_and_handle_nans( diff --git a/src/unilab/tasks/motion_tracking/common/kernels.py b/src/unilab/tasks/motion_tracking/common/kernels.py index c046bb052..d722b76c8 100644 --- a/src/unilab/tasks/motion_tracking/common/kernels.py +++ b/src/unilab/tasks/motion_tracking/common/kernels.py @@ -3,6 +3,7 @@ from __future__ import annotations import os +from math import atan2, cos, sin, sqrt import numpy as np from numba import config, njit, prange, set_num_threads @@ -28,6 +29,461 @@ 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(inline="always") +def _quat_mul_components( + w1: float, + x1: float, + y1: float, + z1: float, + w2: float, + x2: float, + y2: float, + z2: float, +) -> tuple[float, float, float, float]: + """Hamilton product for scalar quaternion components.""" + return ( + 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, + ) + + +@njit(inline="always") +def _quat_apply_inverse_components( + w: float, + x: float, + y: float, + z: float, + vx: float, + vy: float, + vz: float, +) -> tuple[float, float, float]: + """Rotate one vector by the inverse of a unit quaternion.""" + qx = -x + qy = -y + qz = -z + tx = 2.0 * (qy * vz - qz * vy) + ty = 2.0 * (qz * vx - qx * vz) + tz = 2.0 * (qx * vy - qy * vx) + return ( + vx + w * tx + qy * tz - qz * ty, + vy + w * ty + qz * tx - qx * tz, + vz + w * tz + qx * ty - qy * tx, + ) + + +@njit(inline="always") +def _quat_to_rot6d_components( + w: float, + x: float, + y: float, + z: float, +) -> tuple[float, float, float, float, float, float]: + """Flatten the first two rotation-matrix columns.""" + 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 + return ( + 1.0 - 2.0 * (yy + zz), + 2.0 * (xy - wz), + 2.0 * (xy + wz), + 1.0 - 2.0 * (xx + zz), + 2.0 * (xz - wy), + 2.0 * (yz + wx), + ) + + +@njit(cache=True, nogil=True, parallel=True) +def update_motion_relative_state_kernel( + env_ids: np.ndarray, + anchor_body_idx: int, + motion_body_pos_local_w: np.ndarray, + motion_body_pos_w: np.ndarray, + motion_body_quat_w: np.ndarray, + robot_body_pos_w: np.ndarray, + robot_body_quat_w: np.ndarray, + body_pos_relative_w: np.ndarray, + body_quat_relative_w: np.ndarray, + motion_anchor_pos_b: np.ndarray, + motion_anchor_ori_b: np.ndarray, + robot_body_pos_b: np.ndarray, + robot_body_ori_b: np.ndarray, +) -> None: + """Write all MotionCommand relative transforms for selected rows.""" + num_bodies = motion_body_pos_local_w.shape[1] + for row in prange(env_ids.shape[0]): + env_idx = env_ids[row] + motion_anchor_x = motion_body_pos_local_w[env_idx, anchor_body_idx, 0] + motion_anchor_y = motion_body_pos_local_w[env_idx, anchor_body_idx, 1] + motion_anchor_z = motion_body_pos_local_w[env_idx, anchor_body_idx, 2] + motion_anchor_w = motion_body_quat_w[env_idx, anchor_body_idx, 0] + motion_anchor_qx = motion_body_quat_w[env_idx, anchor_body_idx, 1] + motion_anchor_qy = motion_body_quat_w[env_idx, anchor_body_idx, 2] + motion_anchor_qz = motion_body_quat_w[env_idx, anchor_body_idx, 3] + robot_anchor_x = robot_body_pos_w[env_idx, anchor_body_idx, 0] + robot_anchor_y = robot_body_pos_w[env_idx, anchor_body_idx, 1] + robot_anchor_z = robot_body_pos_w[env_idx, anchor_body_idx, 2] + robot_anchor_w = robot_body_quat_w[env_idx, anchor_body_idx, 0] + robot_anchor_qx = robot_body_quat_w[env_idx, anchor_body_idx, 1] + robot_anchor_qy = robot_body_quat_w[env_idx, anchor_body_idx, 2] + robot_anchor_qz = robot_body_quat_w[env_idx, anchor_body_idx, 3] + + yaw_w, yaw_x, yaw_y, yaw_z = _quat_mul_components( + robot_anchor_w, + robot_anchor_qx, + robot_anchor_qy, + robot_anchor_qz, + motion_anchor_w, + -motion_anchor_qx, + -motion_anchor_qy, + -motion_anchor_qz, + ) + half_yaw = 0.5 * atan2( + 2.0 * (yaw_w * yaw_z + yaw_x * yaw_y), + 1.0 - 2.0 * (yaw_y * yaw_y + yaw_z * yaw_z), + ) + delta_w = cos(half_yaw) + delta_z = sin(half_yaw) + yaw_cross = 2.0 * delta_w * delta_z + yaw_z2 = 2.0 * delta_z * delta_z + + anchor_vx = motion_body_pos_w[env_idx, anchor_body_idx, 0] - robot_anchor_x + anchor_vy = motion_body_pos_w[env_idx, anchor_body_idx, 1] - robot_anchor_y + anchor_vz = motion_body_pos_w[env_idx, anchor_body_idx, 2] - robot_anchor_z + anchor_pos_x, anchor_pos_y, anchor_pos_z = _quat_apply_inverse_components( + robot_anchor_w, + robot_anchor_qx, + robot_anchor_qy, + robot_anchor_qz, + anchor_vx, + anchor_vy, + anchor_vz, + ) + motion_anchor_pos_b[env_idx, 0] = anchor_pos_x + motion_anchor_pos_b[env_idx, 1] = anchor_pos_y + motion_anchor_pos_b[env_idx, 2] = anchor_pos_z + rel_w, rel_x, rel_y, rel_z = _quat_mul_components( + robot_anchor_w, + -robot_anchor_qx, + -robot_anchor_qy, + -robot_anchor_qz, + motion_anchor_w, + motion_anchor_qx, + motion_anchor_qy, + motion_anchor_qz, + ) + anchor_rot6d = _quat_to_rot6d_components(rel_w, rel_x, rel_y, rel_z) + for component in range(6): + motion_anchor_ori_b[env_idx, component] = anchor_rot6d[component] + + for body_idx in range(num_bodies): + body_motion_w = motion_body_quat_w[env_idx, body_idx, 0] + body_motion_x = motion_body_quat_w[env_idx, body_idx, 1] + body_motion_y = motion_body_quat_w[env_idx, body_idx, 2] + body_motion_z = motion_body_quat_w[env_idx, body_idx, 3] + out_w, out_x, out_y, out_z = _quat_mul_components( + delta_w, + 0.0, + 0.0, + delta_z, + body_motion_w, + body_motion_x, + body_motion_y, + body_motion_z, + ) + body_quat_relative_w[env_idx, body_idx, 0] = out_w + body_quat_relative_w[env_idx, body_idx, 1] = out_x + body_quat_relative_w[env_idx, body_idx, 2] = out_y + body_quat_relative_w[env_idx, body_idx, 3] = out_z + + vx = motion_body_pos_local_w[env_idx, body_idx, 0] - motion_anchor_x + vy = motion_body_pos_local_w[env_idx, body_idx, 1] - motion_anchor_y + vz = motion_body_pos_local_w[env_idx, body_idx, 2] - motion_anchor_z + relative_x = vx + relative_x -= yaw_cross * vy + relative_x -= yaw_z2 * vx + relative_x += robot_anchor_x + relative_y = vy + relative_y += yaw_cross * vx + relative_y -= yaw_z2 * vy + relative_y += robot_anchor_y + body_pos_relative_w[env_idx, body_idx, 0] = relative_x + body_pos_relative_w[env_idx, body_idx, 1] = relative_y + body_pos_relative_w[env_idx, body_idx, 2] = vz + motion_anchor_z + + robot_vx = robot_body_pos_w[env_idx, body_idx, 0] - robot_anchor_x + robot_vy = robot_body_pos_w[env_idx, body_idx, 1] - robot_anchor_y + robot_vz = robot_body_pos_w[env_idx, body_idx, 2] - robot_anchor_z + robot_pos_x, robot_pos_y, robot_pos_z = _quat_apply_inverse_components( + robot_anchor_w, + robot_anchor_qx, + robot_anchor_qy, + robot_anchor_qz, + robot_vx, + robot_vy, + robot_vz, + ) + robot_body_pos_b[env_idx, body_idx, 0] = robot_pos_x + robot_body_pos_b[env_idx, body_idx, 1] = robot_pos_y + robot_body_pos_b[env_idx, body_idx, 2] = robot_pos_z + body_robot_w = robot_body_quat_w[env_idx, body_idx, 0] + body_robot_x = robot_body_quat_w[env_idx, body_idx, 1] + body_robot_y = robot_body_quat_w[env_idx, body_idx, 2] + body_robot_z = robot_body_quat_w[env_idx, body_idx, 3] + robot_rel_w, robot_rel_x, robot_rel_y, robot_rel_z = _quat_mul_components( + robot_anchor_w, + -robot_anchor_qx, + -robot_anchor_qy, + -robot_anchor_qz, + body_robot_w, + body_robot_x, + body_robot_y, + body_robot_z, + ) + robot_rot6d = _quat_to_rot6d_components( + robot_rel_w, + robot_rel_x, + robot_rel_y, + robot_rel_z, + ) + for component in range(6): + robot_body_ori_b[env_idx, body_idx, component] = robot_rot6d[component] + + +@njit(cache=True, nogil=True, parallel=True) +def update_object_relative_state_kernel( + env_ids: np.ndarray, + robot_anchor_pos_w: np.ndarray, + robot_anchor_quat_w: np.ndarray, + object_pos_w: np.ndarray, + object_quat_w: np.ndarray, + object_lin_vel_w: np.ndarray, + object_state_b: np.ndarray, +) -> None: + """Write BoxMotionCommand object pose and velocity for selected rows.""" + for row in prange(env_ids.shape[0]): + env_idx = env_ids[row] + anchor_x = robot_anchor_pos_w[env_idx, 0] + anchor_y = robot_anchor_pos_w[env_idx, 1] + anchor_z = robot_anchor_pos_w[env_idx, 2] + anchor_w = robot_anchor_quat_w[env_idx, 0] + anchor_qx = robot_anchor_quat_w[env_idx, 1] + anchor_qy = robot_anchor_quat_w[env_idx, 2] + anchor_qz = robot_anchor_quat_w[env_idx, 3] + pos_x, pos_y, pos_z = _quat_apply_inverse_components( + anchor_w, + anchor_qx, + anchor_qy, + anchor_qz, + object_pos_w[env_idx, 0] - anchor_x, + object_pos_w[env_idx, 1] - anchor_y, + object_pos_w[env_idx, 2] - anchor_z, + ) + object_state_b[env_idx, 0] = pos_x + object_state_b[env_idx, 1] = pos_y + object_state_b[env_idx, 2] = pos_z + rel_w, rel_x, rel_y, rel_z = _quat_mul_components( + anchor_w, + -anchor_qx, + -anchor_qy, + -anchor_qz, + object_quat_w[env_idx, 0], + object_quat_w[env_idx, 1], + object_quat_w[env_idx, 2], + object_quat_w[env_idx, 3], + ) + rot6d = _quat_to_rot6d_components(rel_w, rel_x, rel_y, rel_z) + for component in range(6): + object_state_b[env_idx, 3 + component] = rot6d[component] + vel_x, vel_y, vel_z = _quat_apply_inverse_components( + anchor_w, + anchor_qx, + anchor_qy, + anchor_qz, + object_lin_vel_w[env_idx, 0], + object_lin_vel_w[env_idx, 1], + object_lin_vel_w[env_idx, 2], + ) + object_state_b[env_idx, 9] = vel_x + object_state_b[env_idx, 10] = vel_y + object_state_b[env_idx, 11] = vel_z + + +@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, @@ -97,16 +553,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) @@ -167,4 +614,7 @@ 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", + "update_motion_relative_state_kernel", + "update_object_relative_state_kernel", ] diff --git a/src/unilab/tasks/motion_tracking/common/manager_terms.py b/src/unilab/tasks/motion_tracking/common/manager_terms.py index 782072e54..26e85651d 100644 --- a/src/unilab/tasks/motion_tracking/common/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/common/manager_terms.py @@ -5,7 +5,6 @@ import dataclasses import math from dataclasses import dataclass, field -from types import SimpleNamespace from typing import TYPE_CHECKING, Any, Literal, cast import numpy as np @@ -13,9 +12,6 @@ from unilab.envs.mdp.actions import JointPositionAction, JointPositionActionCfg from unilab.managers import CommandTerm, CommandTermCfg, ManagerTermBase, ManagerTermBaseCfg from unilab.managers.scene_entity_config import SceneEntityCfg -from unilab.utils.geometry import ( - np_write_relative_anchor_transform_pos_rot6d, -) from unilab.utils.rotation import ( np_quat_apply_inverse, np_quat_error_magnitude_squared_batched, @@ -30,10 +26,10 @@ reward_motion_body_ori_kernel, reward_motion_body_pos_kernel, termination_anchor_pos_kernel, + update_motion_metrics_kernel, + update_motion_relative_state_kernel, ) from .motion_loader import MotionData, MotionLoader, MotionSampler -from .observations import write_body_ori6_in_anchor_frame, write_body_pos_in_anchor_frame -from .transforms import update_relative_transforms if TYPE_CHECKING: from unilab.base.entity import Entity @@ -190,12 +186,9 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self.robot_body_pos_b = np.empty_like(self._body_pos_w) self.robot_body_ori_b = np.empty((self.num_envs, num_bodies, 6), dtype=dtype) self.joint_default_bias = np.zeros((self.num_envs, num_joints), dtype=dtype) - self._delta_pos_w = np.empty((self.num_envs, 3), dtype=dtype) - self._delta_ori_w = np.empty((self.num_envs, 4), dtype=dtype) - self._body_vec_error = np.empty_like(self._body_pos_w) - 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. @@ -223,7 +216,11 @@ def __init__(self, cfg: MotionCommandCfg, env: ManagerBasedRlEnv): self.metrics[name] = np.zeros(self.num_envs, dtype=dtype) self._refresh_motion() self._refresh_robot_state(force=True) + # Configure and compile both fused kernels on the cold path so the first + # measured manager step contains no Numba worker/JIT initialization. + configure_motion_kernel_runtime() self._refresh_relative_state() + self._update_metrics() def _make_motion_loader( self, @@ -421,139 +418,55 @@ def _refresh_robot_state( self._robot_cache_step = step def _refresh_relative_state(self, env_ids: np.ndarray | None = None) -> None: - if env_ids is not None: - self._refresh_relative_state_rows(env_ids) - return - update_relative_transforms( - self, - self._motion_data, + rows = self._all_env_ids if env_ids is None else env_ids + update_motion_relative_state_kernel( + rows, + self.anchor_body_idx, + self._motion_data.body_pos_w, + self._body_pos_w, + self._motion_data.body_quat_w, self._robot_body_pos_w, self._robot_body_quat_w, - ) - np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self.anchor_pos_w, - self.anchor_quat_w, + self.body_pos_relative_w, + self.body_quat_relative_w, self.motion_anchor_pos_b, self.motion_anchor_ori_b, - ) - write_body_pos_in_anchor_frame( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self._robot_body_pos_w, self.robot_body_pos_b, - body_vec_error=self._body_vec_error, - ) - write_body_ori6_in_anchor_frame( - self.robot_anchor_quat_w, - self._robot_body_quat_w, self.robot_body_ori_b, ) - def _refresh_relative_state_rows(self, env_ids: np.ndarray) -> None: - """Row-scoped variant of `_refresh_relative_state` for partial resets. - - Computes the same transforms on the gathered reset rows and scatters the - results back; untouched rows keep their per-step values. - """ - num_rows = len(env_ids) - num_bodies = self._robot_body_pos_w.shape[1] - dtype = self._body_pos_w.dtype - robot_pos_rows = self._robot_body_pos_w[env_ids] - robot_quat_rows = self._robot_body_quat_w[env_ids] - motion_rows = SimpleNamespace( - body_pos_w=self._motion_data.body_pos_w[env_ids], - body_quat_w=self._motion_data.body_quat_w[env_ids], - ) - # `update_relative_transforms` reads/writes these attributes on the env; - # a namespace with row-sized scratch keeps the shared buffers untouched. - scratch = SimpleNamespace( - anchor_body_idx=self.anchor_body_idx, - _delta_pos_w=np.empty((num_rows, 3), dtype=dtype), - _delta_ori_w=np.empty((num_rows, 4), dtype=dtype), - body_quat_relative_w=np.empty((num_rows, num_bodies, 4), dtype=dtype), - body_pos_relative_w=np.empty((num_rows, num_bodies, 3), dtype=dtype), - _body_vec_error=np.empty((num_rows, num_bodies, 3), dtype=dtype), - _env_error=np.empty(num_rows, dtype=dtype), - _reward_term=np.empty(num_rows, dtype=dtype), - ) - update_relative_transforms(scratch, motion_rows, robot_pos_rows, robot_quat_rows) - self.body_pos_relative_w[env_ids] = scratch.body_pos_relative_w - self.body_quat_relative_w[env_ids] = scratch.body_quat_relative_w - - anchor_idx = self.anchor_body_idx - robot_anchor_pos_rows = robot_pos_rows[:, anchor_idx] - robot_anchor_quat_rows = robot_quat_rows[:, anchor_idx] - anchor_pos_rows = self._body_pos_w[env_ids][:, anchor_idx] - anchor_quat_rows = motion_rows.body_quat_w[:, anchor_idx] - motion_anchor_pos_b = np.empty((num_rows, 3), dtype=dtype) - motion_anchor_ori_b = np.empty((num_rows, 6), dtype=dtype) - np_write_relative_anchor_transform_pos_rot6d( - robot_anchor_pos_rows, - robot_anchor_quat_rows, - anchor_pos_rows, - anchor_quat_rows, - motion_anchor_pos_b, - motion_anchor_ori_b, - ) - self.motion_anchor_pos_b[env_ids] = motion_anchor_pos_b - self.motion_anchor_ori_b[env_ids] = motion_anchor_ori_b - robot_body_pos_b = np.empty((num_rows, num_bodies, 3), dtype=dtype) - write_body_pos_in_anchor_frame( - robot_anchor_pos_rows, - robot_anchor_quat_rows, - robot_pos_rows, - robot_body_pos_b, - body_vec_error=scratch._body_vec_error, - ) - self.robot_body_pos_b[env_ids] = robot_body_pos_b - robot_body_ori_b = np.empty((num_rows, num_bodies, 6), dtype=dtype) - write_body_ori6_in_anchor_frame( - robot_anchor_quat_rows, - robot_quat_rows, - robot_body_ori_b, - ) - 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) diff --git a/src/unilab/tasks/motion_tracking/common/observations.py b/src/unilab/tasks/motion_tracking/common/observations.py deleted file mode 100644 index 6332c4ac5..000000000 --- a/src/unilab/tasks/motion_tracking/common/observations.py +++ /dev/null @@ -1,66 +0,0 @@ -"""In-place body-frame writers shared by motion manager terms.""" - -from __future__ import annotations - -import numpy as np - - -def write_body_pos_in_anchor_frame( - anchor_pos: np.ndarray, - anchor_quat: np.ndarray, - body_pos: np.ndarray, - out: np.ndarray, - *, - body_vec_error: np.ndarray, -) -> None: - aw = anchor_quat[:, None, 0] - ax = anchor_quat[:, None, 1] - ay = anchor_quat[:, None, 2] - az = anchor_quat[:, None, 3] - - num_envs, num_bodies = body_pos.shape[:2] - rel_pos = body_vec_error[:num_envs, :num_bodies] - vx = rel_pos[..., 0] - vy = rel_pos[..., 1] - vz = rel_pos[..., 2] - np.subtract(body_pos[..., 0], anchor_pos[:, None, 0], out=vx) - np.subtract(body_pos[..., 1], anchor_pos[:, None, 1], out=vy) - np.subtract(body_pos[..., 2], anchor_pos[:, None, 2], out=vz) - - tx = 2 * (az * vy - ay * vz) - ty = 2 * (ax * vz - az * vx) - tz = 2 * (ay * vx - ax * vy) - - out[..., 0] = vx + aw * tx + az * ty - ay * tz - out[..., 1] = vy + aw * ty + ax * tz - az * tx - out[..., 2] = vz + aw * tz + ay * tx - ax * ty - - -def write_body_ori6_in_anchor_frame( - anchor_quat: np.ndarray, - body_quat: np.ndarray, - out: np.ndarray, -) -> None: - aw = anchor_quat[:, None, 0] - ax = anchor_quat[:, None, 1] - ay = anchor_quat[:, None, 2] - az = anchor_quat[:, None, 3] - bw = body_quat[..., 0] - bx = body_quat[..., 1] - by = body_quat[..., 2] - bz = body_quat[..., 3] - - rw = aw * bw + ax * bx + ay * by + az * bz - rx = aw * bx - ax * bw - ay * bz + az * by - ry = aw * by + ax * bz - ay * bw - az * bx - rz = aw * bz - ax * by + ay * bx - az * bw - - out[..., 0] = 1 - 2 * (ry * ry + rz * rz) - out[..., 1] = 2 * (rx * ry - rw * rz) - out[..., 2] = 2 * (rx * ry + rw * rz) - out[..., 3] = 1 - 2 * (rx * rx + rz * rz) - out[..., 4] = 2 * (rx * rz - rw * ry) - out[..., 5] = 2 * (ry * rz + rw * rx) - - -__all__ = ["write_body_ori6_in_anchor_frame", "write_body_pos_in_anchor_frame"] diff --git a/src/unilab/tasks/motion_tracking/common/transforms.py b/src/unilab/tasks/motion_tracking/common/transforms.py deleted file mode 100644 index 31b4f31a6..000000000 --- a/src/unilab/tasks/motion_tracking/common/transforms.py +++ /dev/null @@ -1,103 +0,0 @@ -"""Shared relative body-transform computation for motion tracking. - -Fills the environment's ``body_pos_relative_w`` / ``body_quat_relative_w`` -reference buffers each step. The op order and in-place ``out=`` usage are -load-bearing for stable observation and reward semantics. -""" - -from __future__ import annotations - -from typing import Any - -import numpy as np - - -def update_relative_transforms( - env: Any, - motion_data: Any, - robot_body_pos_w: np.ndarray, - robot_body_quat_w: np.ndarray, -) -> None: - """Update relative body transforms for tracking.""" - # Get anchor states - anchor_pos_w = motion_data.body_pos_w[:, env.anchor_body_idx] - anchor_quat_w = motion_data.body_quat_w[:, env.anchor_body_idx] - robot_anchor_pos_w = robot_body_pos_w[:, env.anchor_body_idx] - robot_anchor_quat_w = robot_body_quat_w[:, env.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 = env._delta_pos_w - delta_pos_w[:] = robot_anchor_pos_w - delta_pos_w[:, 2] = anchor_pos_w[:, 2] - - # Compute yaw-only rotation difference, equivalent to - # np_yaw_quat(np_quat_mul(robot_anchor_quat_w, np_quat_inv(anchor_quat_w))). - delta_ori_w = env._delta_ori_w - rw, rx, ry, rz = ( - robot_anchor_quat_w[:, 0], - robot_anchor_quat_w[:, 1], - robot_anchor_quat_w[:, 2], - robot_anchor_quat_w[:, 3], - ) - aw, ax, ay, az = ( - anchor_quat_w[:, 0], - anchor_quat_w[:, 1], - anchor_quat_w[:, 2], - anchor_quat_w[:, 3], - ) - qw = rw * aw + rx * ax + ry * ay + rz * az - qx = -rw * ax + rx * aw - ry * az + rz * ay - qy = -rw * ay + rx * az + ry * aw - rz * ax - qz = -rw * az - rx * ay + ry * ax + rz * aw - half_yaw = 0.5 * np.arctan2(2 * (qw * qz + qx * qy), 1 - 2 * (qy * qy + qz * qz)) - np.cos(half_yaw, out=delta_ori_w[:, 0]) - delta_ori_w[:, 1:3] = 0.0 - np.sin(half_yaw, out=delta_ori_w[:, 3]) - - dw1 = delta_ori_w[:, 0] - dz1 = delta_ori_w[:, 3] - dw = dw1[:, None] - dz = dz1[:, None] - mw = motion_data.body_quat_w[..., 0] - mx = motion_data.body_quat_w[..., 1] - my = motion_data.body_quat_w[..., 2] - mz = motion_data.body_quat_w[..., 3] - out_quat = env.body_quat_relative_w - out_quat[..., 0] = dw * mw - out_quat[..., 0] -= dz * mz - out_quat[..., 1] = dw * mx - out_quat[..., 1] -= dz * my - out_quat[..., 2] = dw * my - out_quat[..., 2] += dz * mx - out_quat[..., 3] = dw * mz - out_quat[..., 3] += dz * mw - - rel_pos = env._body_vec_error - vx = rel_pos[..., 0] - vy = rel_pos[..., 1] - vz = rel_pos[..., 2] - np.subtract(motion_data.body_pos_w[..., 0], anchor_pos_w[:, None, 0], out=vx) - np.subtract(motion_data.body_pos_w[..., 1], anchor_pos_w[:, None, 1], out=vy) - np.subtract(motion_data.body_pos_w[..., 2], anchor_pos_w[:, None, 2], out=vz) - - yaw_cross = env._env_error - yaw_z2 = env._reward_term - np.multiply(dw1, dz1, out=yaw_cross) - yaw_cross *= 2.0 - np.square(dz1, out=yaw_z2) - yaw_z2 *= 2.0 - yaw_cross_2d = yaw_cross[:, None] - yaw_z2_2d = yaw_z2[:, None] - - out_pos = env.body_pos_relative_w - out_pos[..., 0] = vx - out_pos[..., 0] -= yaw_cross_2d * vy - out_pos[..., 0] -= yaw_z2_2d * vx - out_pos[..., 0] += delta_pos_w[:, None, 0] - out_pos[..., 1] = vy - out_pos[..., 1] += yaw_cross_2d * vx - out_pos[..., 1] -= yaw_z2_2d * vy - out_pos[..., 1] += delta_pos_w[:, None, 1] - out_pos[..., 2] = vz - out_pos[..., 2] += delta_pos_w[:, None, 2] diff --git a/src/unilab/tasks/motion_tracking/g1/manager_terms.py b/src/unilab/tasks/motion_tracking/g1/manager_terms.py index 32e34cbb8..f9764680b 100644 --- a/src/unilab/tasks/motion_tracking/g1/manager_terms.py +++ b/src/unilab/tasks/motion_tracking/g1/manager_terms.py @@ -9,16 +9,13 @@ from unilab.managers import ManagerTermBase, ManagerTermBaseCfg from unilab.managers.scene_entity_config import SceneEntityCfg +from unilab.tasks.motion_tracking.common.kernels import update_object_relative_state_kernel from unilab.tasks.motion_tracking.common.manager_terms import ( MotionCommand, MotionCommandCfg, MotionJointPositionAction, ) -from unilab.utils.geometry import np_write_relative_anchor_transform_pos_rot6d -from unilab.utils.rotation import ( - np_quat_apply_inverse, - np_quat_error_magnitude_squared_batched, -) +from unilab.utils.rotation import np_quat_error_magnitude_squared_batched from .motion_box_loader import BoxMotionData, BoxMotionLoader @@ -115,37 +112,15 @@ def _resample_command(self, env_ids: np.ndarray) -> None: self.object.write_root_state_to_sim(object_state, env_ids=env_ids) def _refresh_object_state(self, env_ids: np.ndarray | None = None) -> None: - if env_ids is None: - np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w, - self.robot_anchor_quat_w, - self.object.data.root_link_pos_w, - self.object.data.root_link_quat_w, - self._object_obs_b[:, :3], - self._object_obs_b[:, 3:9], - ) - self._object_obs_b[:, 9:12] = np_quat_apply_inverse( - self.robot_anchor_quat_w, - self.object.data.root_link_lin_vel_w, - ) - return - num_rows = len(env_ids) - dtype = self._object_obs_b.dtype - pos_b = np.empty((num_rows, 3), dtype=dtype) - rot_b = np.empty((num_rows, 6), dtype=dtype) - np_write_relative_anchor_transform_pos_rot6d( - self.robot_anchor_pos_w[env_ids], - self.robot_anchor_quat_w[env_ids], - self.object.data.root_link_pos_w[env_ids], - self.object.data.root_link_quat_w[env_ids], - pos_b, - rot_b, - ) - self._object_obs_b[env_ids, :3] = pos_b - self._object_obs_b[env_ids, 3:9] = rot_b - self._object_obs_b[env_ids, 9:12] = np_quat_apply_inverse( - self.robot_anchor_quat_w[env_ids], - self.object.data.root_link_lin_vel_w[env_ids], + rows = self._all_env_ids if env_ids is None else env_ids + update_object_relative_state_kernel( + rows, + self.robot_anchor_pos_w, + self.robot_anchor_quat_w, + self.object.data.root_link_pos_w, + self.object.data.root_link_quat_w, + self.object.data.root_link_lin_vel_w, + self._object_obs_b, ) def post_compute(self) -> None: diff --git a/src/unilab/utils/geometry.py b/src/unilab/utils/geometry.py index 5dc933d3f..5358f4f6d 100644 --- a/src/unilab/utils/geometry.py +++ b/src/unilab/utils/geometry.py @@ -148,66 +148,3 @@ def np_cartesian_to_spherical(cart: np.ndarray) -> np.ndarray: phi = np.arctan2(cart[..., 2:3], cart[..., 0:1]) theta = np.arcsin(np.clip(cart[..., 1:2] / length, -1.0, 1.0)) return np.concatenate([length, phi, theta], axis=-1) - - -def np_write_relative_anchor_transform_pos_rot6d( - source_anchor_pos_w: np.ndarray, - source_anchor_quat_w: np.ndarray, - target_anchor_pos_w: np.ndarray, - target_anchor_quat_w: np.ndarray, - out_pos: np.ndarray, - out_rot6d: np.ndarray, -) -> None: - """Fused frame-transform + 6D rotation flatten, writing in place. - - Computes the position of ``target_anchor`` in ``source_anchor``'s frame - (written to ``out_pos`` of shape ``(N, 3)``) and the relative rotation - ``conj(source_anchor) * target_anchor`` expressed as the flattened - first two columns of its rotation matrix (written to ``out_rot6d`` of - shape ``(N, 6)``). No intermediate quaternion arrays are allocated. - - This helper computes the relative anchor transform used by motion tracking. - """ - aw = source_anchor_quat_w[:, 0] - ax = source_anchor_quat_w[:, 1] - ay = source_anchor_quat_w[:, 2] - az = source_anchor_quat_w[:, 3] - - vx = target_anchor_pos_w[:, 0] - source_anchor_pos_w[:, 0] - vy = target_anchor_pos_w[:, 1] - source_anchor_pos_w[:, 1] - vz = target_anchor_pos_w[:, 2] - source_anchor_pos_w[:, 2] - - qx = -ax - qy = -ay - qz = -az - tx = 2 * (qy * vz - qz * vy) - ty = 2 * (qz * vx - qx * vz) - tz = 2 * (qx * vy - qy * vx) - out_pos[:, 0] = vx + aw * tx + qy * tz - qz * ty - out_pos[:, 1] = vy + aw * ty + qz * tx - qx * tz - out_pos[:, 2] = vz + aw * tz + qx * ty - qy * tx - - bw = target_anchor_quat_w[:, 0] - bx = target_anchor_quat_w[:, 1] - by = target_anchor_quat_w[:, 2] - bz = target_anchor_quat_w[:, 3] - rw = aw * bw + ax * bx + ay * by + az * bz - rx = aw * bx - ax * bw - ay * bz + az * by - ry = aw * by + ax * bz - ay * bw - az * bx - rz = aw * bz - ax * by + ay * bx - az * bw - - xx = rx * rx - yy = ry * ry - zz = rz * rz - xy = rx * ry - xz = rx * rz - yz = ry * rz - wx = rw * rx - wy = rw * ry - wz = rw * rz - out_rot6d[:, 0] = 1 - 2 * (yy + zz) - out_rot6d[:, 1] = 2 * (xy - wz) - out_rot6d[:, 2] = 2 * (xy + wz) - out_rot6d[:, 3] = 1 - 2 * (xx + zz) - out_rot6d[:, 4] = 2 * (xz - wy) - out_rot6d[:, 5] = 2 * (yz + wx) diff --git a/tests/managers/test_observation_buffers_noise.py b/tests/managers/test_observation_buffers_noise.py index 0688dbfe7..818fea0a0 100644 --- a/tests/managers/test_observation_buffers_noise.py +++ b/tests/managers/test_observation_buffers_noise.py @@ -95,6 +95,28 @@ def test_noise_configs_use_supplied_generator() -> None: np.testing.assert_array_equal(ConstantNoiseCfg(bias=2.0, operation="abs").apply(data), 2.0) +@pytest.mark.parametrize("operation", ["add", "scale", "abs"]) +def test_uniform_noise_inplace_matches_reference_expression(operation: str) -> None: + data = np.arange(24, dtype=np.float32).reshape(8, 3) + n_min = np.asarray([-0.2, -0.1, -0.05], dtype=np.float32) + n_max = np.asarray([0.3, 0.4, 0.5], dtype=np.float32) + cfg = UniformNoiseCfg(n_min=tuple(n_min), n_max=tuple(n_max), operation=operation) + + reference_rng = np.random.default_rng(1702) + unit = reference_rng.random(data.shape).astype(data.dtype, copy=False) + noise = unit * (n_max - n_min) + n_min + if operation == "add": + expected = data + noise + elif operation == "scale": + expected = data * noise + else: + expected = noise + + actual = cfg.apply(data, rng=np.random.default_rng(1702)) + np.testing.assert_array_equal(actual, expected) + np.testing.assert_array_equal(data, np.arange(24, dtype=np.float32).reshape(8, 3)) + + def test_additive_bias_noise_supports_scalar_terms() -> None: from unilab.managers._noise import NoiseModelWithAdditiveBias @@ -143,6 +165,86 @@ def test_observation_groups_pipeline_order_and_history(fake_env: FakeEnv) -> Non assert manager.get_active_iterable_terms(0)[0][0] == "policy-state" +def test_concatenated_result_owns_each_result_and_protects_term_buffers( + fake_env: FakeEnv, +) -> None: + source = fake_env.obs + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "source": ObservationTermCfg(func=lambda env: env.obs), + "constant": ObservationTermCfg( + func=lambda env: np.ones((env.num_envs, 1), dtype=np.float32) + ), + } + ) + }, + fake_env, + ) + + first = manager.compute(update_history=True)["policy"] + assert isinstance(first, np.ndarray) + first_address = first.ctypes.data + expected_first = np.concatenate((source.copy(), np.ones((fake_env.num_envs, 1))), axis=1) + np.testing.assert_array_equal(first, expected_first) + + fake_env.obs += 100.0 + second = manager.compute(update_history=True)["policy"] + assert isinstance(second, np.ndarray) + assert second.ctypes.data != first_address + np.testing.assert_array_equal(first, expected_first) + np.testing.assert_array_equal(second[:, :2], fake_env.obs) + assert not np.shares_memory(second, fake_env.obs) + + +def test_concatenated_nan_sanitize_does_not_mutate_term_owned_input( + fake_env: FakeEnv, +) -> None: + source = fake_env.obs.copy() + source[0, 0] = np.nan + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={"state": ObservationTermCfg(func=lambda env: source)}, + nan_policy="sanitize", + nan_check_per_term=False, + ) + }, + fake_env, + ) + + result = manager.compute(update_history=True)["policy"] + assert isinstance(result, np.ndarray) + assert np.isfinite(result).all() + assert np.isnan(source[0, 0]) + + +def test_concatenated_nan_error_still_identifies_offending_term(fake_env: FakeEnv) -> None: + def invalid(env: FakeEnv) -> np.ndarray: + result = np.ones((env.num_envs, 1), dtype=np.float32) + result[2, 0] = np.nan + return result + + manager = ObservationManager( + { + "policy": ObservationGroupCfg( + terms={ + "finite": ObservationTermCfg(func=lambda env: env.obs), + "invalid": ObservationTermCfg(func=invalid), + } + ) + }, + fake_env, + ) + + with pytest.raises( + ValueError, + match=r"NaN detected.*'policy/invalid'.*environments: \[2\]", + ): + manager.compute(update_history=True) + + def test_observation_noise_model_delay_and_seed_reproducibility() -> None: cfg = { "policy": ObservationGroupCfg( diff --git a/tests/tasks/test_motion_term_parity.py b/tests/tasks/test_motion_term_parity.py index 6c56381f0..b0d1e2b18 100644 --- a/tests/tasks/test_motion_term_parity.py +++ b/tests/tasks/test_motion_term_parity.py @@ -13,7 +13,15 @@ from unilab.managers import RewardTermCfg, TerminationTermCfg from unilab.tasks.motion_tracking.common import kernels from unilab.tasks.motion_tracking.common import manager_terms as mt -from unilab.utils.rotation import np_quat_error_magnitude_squared_batched +from unilab.utils.rotation import ( + np_matrix_first_two_cols_from_quat, + np_quat_apply_batched, + np_quat_apply_inverse_batched, + np_quat_error_magnitude_squared_batched, + np_quat_inv, + np_quat_mul_batched, + np_yaw_quat, +) def _make_env(command: Any) -> SimpleNamespace: @@ -320,6 +328,241 @@ def test_motion_hot_kernels_compile_parallel_on_term_construction(body_setup) -> assert get_num_threads() == min(8, config.NUMBA_DEFAULT_NUM_THREADS) +def test_motion_metrics_kernel_matches_numpy_and_scopes_rows() -> None: + rng = np.random.default_rng(1701) + num_envs, num_bodies, num_joints = 257, 12, 29 + anchor_body_idx = 4 + motion_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + relative_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + motion_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + robot_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + relative_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + motion_lin = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_lin = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + motion_ang = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_ang = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + motion_joint_pos = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + robot_joint_pos = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + motion_joint_vel = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + robot_joint_vel = rng.standard_normal((num_envs, num_joints), dtype=np.float32) + inputs = ( + motion_pos, + robot_pos, + motion_quat, + robot_quat, + motion_lin, + robot_lin, + motion_ang, + robot_ang, + relative_pos, + relative_quat, + motion_joint_pos, + robot_joint_pos, + motion_joint_vel, + robot_joint_vel, + ) + snapshots = tuple(value.copy() for value in inputs) + expected = ( + np.linalg.norm(motion_pos[:, anchor_body_idx] - robot_pos[:, anchor_body_idx], axis=-1), + np.sqrt( + np_quat_error_magnitude_squared_batched( + motion_quat[:, anchor_body_idx], robot_quat[:, anchor_body_idx] + ) + ), + np.linalg.norm(motion_lin[:, anchor_body_idx] - robot_lin[:, anchor_body_idx], axis=-1), + np.linalg.norm(motion_ang[:, anchor_body_idx] - robot_ang[:, anchor_body_idx], axis=-1), + np.linalg.norm(relative_pos - robot_pos, axis=-1).mean(axis=-1), + np.sqrt(np_quat_error_magnitude_squared_batched(relative_quat, robot_quat)).mean(axis=-1), + np.linalg.norm(motion_lin - robot_lin, axis=-1).mean(axis=-1), + np.linalg.norm(motion_ang - robot_ang, axis=-1).mean(axis=-1), + np.linalg.norm(motion_joint_pos - robot_joint_pos, axis=-1), + np.linalg.norm(motion_joint_vel - robot_joint_vel, axis=-1), + ) + outputs = tuple(np.full(num_envs, -123.0, dtype=np.float32) for _ in expected) + + def run(rows: np.ndarray) -> None: + kernels.update_motion_metrics_kernel( + rows, + anchor_body_idx, + motion_pos, + robot_pos, + motion_quat, + robot_quat, + motion_lin, + robot_lin, + motion_ang, + robot_ang, + relative_pos, + relative_quat, + motion_joint_pos, + robot_joint_pos, + motion_joint_vel, + robot_joint_vel, + *outputs, + ) + + selected = np.asarray([0, 3, 128, 256], dtype=np.int32) + run(selected) + untouched = np.ones(num_envs, dtype=bool) + untouched[selected] = False + for actual, reference in zip(outputs, expected, strict=True): + np.testing.assert_allclose(actual[selected], reference[selected], rtol=2e-6, atol=1e-5) + np.testing.assert_array_equal(actual[untouched], -123.0) + + run(np.arange(num_envs, dtype=np.int32)) + for actual, reference in zip(outputs, expected, strict=True): + np.testing.assert_allclose(actual, reference, rtol=2e-6, atol=1e-5) + for actual, snapshot in zip(inputs, snapshots, strict=True): + np.testing.assert_array_equal(actual, snapshot) + assert kernels.update_motion_metrics_kernel.targetoptions["nopython"] is True + assert kernels.update_motion_metrics_kernel.targetoptions["nogil"] is True + assert kernels.update_motion_metrics_kernel.targetoptions["parallel"] is True + assert kernels.update_motion_metrics_kernel.signatures + + +def test_motion_relative_state_kernel_matches_numpy_and_scopes_rows() -> None: + rng = np.random.default_rng(1818) + num_envs, num_bodies = 257, 12 + anchor_body_idx = 4 + motion_pos_local = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + env_origins = rng.standard_normal((num_envs, 1, 3), dtype=np.float32) + motion_pos_world = motion_pos_local + env_origins + motion_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + robot_pos = rng.standard_normal((num_envs, num_bodies, 3), dtype=np.float32) + robot_quat = _unit_quat(rng.standard_normal((num_envs, num_bodies, 4), dtype=np.float32)) + inputs = (motion_pos_local, motion_pos_world, motion_quat, robot_pos, robot_quat) + snapshots = tuple(value.copy() for value in inputs) + + motion_anchor_pos_local = motion_pos_local[:, anchor_body_idx] + motion_anchor_quat = motion_quat[:, anchor_body_idx] + robot_anchor_pos = robot_pos[:, anchor_body_idx] + robot_anchor_quat = robot_quat[:, anchor_body_idx] + delta_pos = robot_anchor_pos.copy() + delta_pos[:, 2] = motion_anchor_pos_local[:, 2] + delta_quat = np_yaw_quat( + np_quat_mul_batched(robot_anchor_quat, np_quat_inv(motion_anchor_quat)) + ) + expected_body_pos_relative = np_quat_apply_batched( + delta_quat[:, None], + motion_pos_local - motion_anchor_pos_local[:, None], + ) + expected_body_pos_relative += delta_pos[:, None] + expected_body_quat_relative = np_quat_mul_batched(delta_quat[:, None], motion_quat) + expected_motion_anchor_pos = np_quat_apply_inverse_batched( + robot_anchor_quat, + motion_pos_world[:, anchor_body_idx] - robot_anchor_pos, + ) + expected_motion_anchor_ori = np_matrix_first_two_cols_from_quat( + np_quat_mul_batched(np_quat_inv(robot_anchor_quat), motion_anchor_quat) + ) + expected_robot_body_pos = np_quat_apply_inverse_batched( + robot_anchor_quat[:, None], + robot_pos - robot_anchor_pos[:, None], + ) + expected_robot_body_ori = np_matrix_first_two_cols_from_quat( + np_quat_mul_batched(np_quat_inv(robot_anchor_quat)[:, None], robot_quat) + ) + expected = ( + expected_body_pos_relative, + expected_body_quat_relative, + expected_motion_anchor_pos, + expected_motion_anchor_ori, + expected_robot_body_pos, + expected_robot_body_ori, + ) + outputs = tuple(np.full(value.shape, -123.0, dtype=np.float32) for value in expected) + output_addresses = tuple(value.ctypes.data for value in outputs) + + def run(rows: np.ndarray) -> None: + kernels.update_motion_relative_state_kernel( + rows, + anchor_body_idx, + motion_pos_local, + motion_pos_world, + motion_quat, + robot_pos, + robot_quat, + *outputs, + ) + + selected = np.asarray([0, 3, 128, 256], dtype=np.int32) + run(selected) + untouched = np.ones(num_envs, dtype=bool) + untouched[selected] = False + for actual, reference in zip(outputs, expected, strict=True): + np.testing.assert_allclose(actual[selected], reference[selected], rtol=3e-6, atol=2e-6) + np.testing.assert_array_equal(actual[untouched], -123.0) + + run(np.arange(num_envs, dtype=np.int32)) + for actual, reference in zip(outputs, expected, strict=True): + np.testing.assert_allclose(actual, reference, rtol=3e-6, atol=2e-6) + for actual, snapshot in zip(inputs, snapshots, strict=True): + np.testing.assert_array_equal(actual, snapshot) + assert tuple(value.ctypes.data for value in outputs) == output_addresses + assert kernels.update_motion_relative_state_kernel.targetoptions["nopython"] is True + assert kernels.update_motion_relative_state_kernel.targetoptions["nogil"] is True + assert kernels.update_motion_relative_state_kernel.targetoptions["parallel"] is True + assert kernels.update_motion_relative_state_kernel.signatures + + +def test_object_relative_state_kernel_matches_numpy_and_scopes_rows() -> None: + rng = np.random.default_rng(1819) + num_envs = 257 + anchor_pos = rng.standard_normal((num_envs, 3), dtype=np.float32) + anchor_quat = _unit_quat(rng.standard_normal((num_envs, 4), dtype=np.float32)) + object_pos = rng.standard_normal((num_envs, 3), dtype=np.float32) + object_quat = _unit_quat(rng.standard_normal((num_envs, 4), dtype=np.float32)) + object_lin_vel = rng.standard_normal((num_envs, 3), dtype=np.float32) + inputs = (anchor_pos, anchor_quat, object_pos, object_quat, object_lin_vel) + snapshots = tuple(value.copy() for value in inputs) + expected = np.concatenate( + ( + np_quat_apply_inverse_batched(anchor_quat, object_pos - anchor_pos), + np_matrix_first_two_cols_from_quat( + np_quat_mul_batched(np_quat_inv(anchor_quat), object_quat) + ), + np_quat_apply_inverse_batched(anchor_quat, object_lin_vel), + ), + axis=-1, + ) + output = np.full(expected.shape, -123.0, dtype=np.float32) + output_address = output.ctypes.data + + selected = np.asarray([0, 3, 128, 256], dtype=np.int32) + kernels.update_object_relative_state_kernel( + selected, + anchor_pos, + anchor_quat, + object_pos, + object_quat, + object_lin_vel, + output, + ) + untouched = np.ones(num_envs, dtype=bool) + untouched[selected] = False + np.testing.assert_allclose(output[selected], expected[selected], rtol=3e-6, atol=2e-6) + np.testing.assert_array_equal(output[untouched], -123.0) + + kernels.update_object_relative_state_kernel( + np.arange(num_envs, dtype=np.int32), + anchor_pos, + anchor_quat, + object_pos, + object_quat, + object_lin_vel, + output, + ) + np.testing.assert_allclose(output, expected, rtol=3e-6, atol=2e-6) + for actual, snapshot in zip(inputs, snapshots, strict=True): + np.testing.assert_array_equal(actual, snapshot) + assert output.ctypes.data == output_address + assert kernels.update_object_relative_state_kernel.targetoptions["nopython"] is True + assert kernels.update_object_relative_state_kernel.targetoptions["nogil"] is True + assert kernels.update_object_relative_state_kernel.targetoptions["parallel"] is True + assert kernels.update_object_relative_state_kernel.signatures + + def test_joint_pos_limits_bit_parity() -> None: rng = np.random.default_rng(123) joint_pos = rng.standard_normal((8, 5), dtype=np.float32)