| """Canonical adapters shared by episode fitting and batched inference.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import TYPE_CHECKING |
|
|
| import numpy as np |
|
|
| from ._so3 import ( |
| conjugate_quaternion, |
| multiply_quaternions, |
| normalize_quaternion, |
| quaternion_to_rotvec, |
| ) |
| from .native_backend import native_kernel |
| from .semantic import ( |
| difference_additive, |
| eef_pose_absolute, |
| eef_pose_components, |
| eef_pose_deltas, |
| eef_pose_from_components, |
| integrate_additive, |
| integrate_eef_deltas, |
| normalize_gripper_actions, |
| resample_uniform, |
| wrap_to_pi, |
| ) |
|
|
| if TYPE_CHECKING: |
| from .presets import ComponentSpec, PresetConfig |
|
|
|
|
| def _resample(values, source_hz, target_hz, method, *, kind="linear", target_length=None): |
| source_hz, target_hz = float(source_hz), float(target_hz) |
| if source_hz == target_hz and target_length in (None, values.shape[1]): |
| return values |
| kernel = native_kernel("resample_uniform_batch") |
| if kernel is not None: |
| return kernel( |
| values, |
| source_hz, |
| target_hz, |
| kind, |
| -1 if target_length is None else target_length, |
| method, |
| ) |
| return np.stack( |
| [ |
| resample_uniform( |
| row, |
| source_hz=source_hz, |
| target_hz=target_hz, |
| kind=kind, |
| target_length=target_length, |
| method=method, |
| ) |
| for row in values |
| ] |
| ) |
|
|
|
|
| def _pose_map(values, kernel_name, reference): |
| kernel = native_kernel(kernel_name) |
| return kernel(values) if kernel is not None else np.stack([reference(row) for row in values]) |
|
|
|
|
| def _rotvec_poses(values): |
| kernel = native_kernel("eef_pose_from_rotvec_batch") |
| if kernel is not None: |
| return kernel(values) |
| flat = values.reshape(-1, 6) |
| return eef_pose_from_components(flat[:, :3], flat[:, 3:], "rotation_vector").reshape( |
| *values.shape[:2], 7 |
| ) |
|
|
|
|
| def _poses(values: np.ndarray, component: ComponentSpec) -> np.ndarray: |
| position = np.take(values, component.position_dims, axis=-1) |
| rotation = np.take(values, component.rotation_dims, axis=-1) |
| if component.rotation_format == "rotation_vector": |
| return _rotvec_poses(np.concatenate((position, rotation), axis=-1)) |
| return eef_pose_from_components( |
| position.reshape(-1, 3), rotation.reshape(-1, rotation.shape[-1]), component.rotation_format |
| ).reshape(*values.shape[:2], 7) |
|
|
|
|
| def _local(poses, anchor): |
| quaternion = normalize_quaternion(poses[..., 3:7]) |
| inverse = conjugate_quaternion(normalize_quaternion(anchor[..., 3:7])) |
| relative = multiply_quaternions(inverse, quaternion) |
| return np.concatenate( |
| (poses[..., :3].astype(np.float64) - anchor[..., :3], quaternion_to_rotvec(relative)), |
| axis=-1, |
| ).astype(np.float32) |
|
|
|
|
| def _joint(preset, values, component, source_hz, current_state): |
| source = np.take(values, component.joint_dims, axis=-1) |
| if preset.input == "delta": |
| source = integrate_additive(source, 1) |
| anchor = np.zeros((len(source), 1, source.shape[-1]), np.float32) |
| else: |
| anchor = ( |
| source[:, :1] |
| if current_state is None |
| else np.take(current_state, component.joint_dims, axis=-1)[:, None] |
| ) |
| joined = np.concatenate((anchor, source), axis=1) |
| if component.joint_wrap == "pi": |
| outside = np.abs(joined) > np.pi + 1e-5 |
| if outside.any(): |
| row = int(np.flatnonzero(outside.any(axis=(1, 2)))[0]) |
| column = component.joint_dims[int(np.argmax(outside[row].any(axis=0)))] |
| raise ValueError( |
| f"action space {preset.name!r} declares joint_wrap='pi' for component " |
| f"{component.name!r}, so its absolute joint columns must lie in (-pi, pi]; " |
| f"source column {column} reaches {joined[row][outside[row]][0]:.6g} rad. Wrap the " |
| "inputs, or register a layout with joint_wrap='none' for coordinates that " |
| "are continuous past one turn, such as a multi-turn or prismatic axis." |
| ) |
| batch, steps, width = joined.shape |
| kernel = native_kernel("unwrap_angles") |
| unwrapped = ( |
| ( |
| kernel(joined.transpose(1, 0, 2).reshape(steps, batch * width), True) |
| .reshape(steps, batch, width) |
| .transpose(1, 0, 2) |
| ) |
| if kernel is not None |
| else np.unwrap(joined, axis=1) |
| ) |
| anchor, source = unwrapped[:, :1], unwrapped[:, 1:] |
| return np.asarray( |
| _resample(source, source_hz, preset.target_hz, preset.interpolation) - anchor, |
| dtype=np.float32, |
| ) |
|
|
|
|
| def _eef(preset, values, component, source_hz, current_state): |
| if preset.input == "delta": |
| delta = np.concatenate( |
| ( |
| np.take(values, component.position_dims, axis=-1), |
| np.take(values, component.rotation_dims, axis=-1), |
| ), |
| axis=-1, |
| ) |
| poses = _pose_map(delta, "integrate_eef_deltas_batch", integrate_eef_deltas) |
| sampled = _resample( |
| poses, source_hz, preset.target_hz, preset.interpolation, kind="pose_xyzw" |
| ) |
| return _pose_map(sampled, "eef_pose_absolute_batch", eef_pose_absolute) |
| poses = _poses(values, component) |
| anchor = poses[:, :1] if current_state is None else _poses(current_state[:, None], component) |
| sampled = _resample(poses, source_hz, preset.target_hz, preset.interpolation, kind="pose_xyzw") |
| return _local(sampled, anchor) |
|
|
|
|
| def canonicalize_batch(preset: PresetConfig, values, *, source_hz, current_state): |
| """Return component ``(B,T_codec,D_active)`` arrays from validated input.""" |
| return { |
| component.name: canonicalize_component_batch( |
| preset, values, component, source_hz=source_hz, current_state=current_state |
| ) |
| for component in preset.components |
| } |
|
|
|
|
| def canonicalize_component_batch(preset, values, component, *, source_hz, current_state): |
| """Convert one component of validated ``(B,T,D)`` source actions.""" |
| motion = (_joint if component.space == "joint" else _eef)( |
| preset, values, component, source_hz, current_state |
| ) |
| if component.gripper_dim is None: |
| return motion |
| gripper = values[..., component.gripper_dim : component.gripper_dim + 1] |
| if component.gripper_threshold is not None: |
| gripper = ((gripper >= component.gripper_threshold) if component.gripper_open_is_high |
| else (gripper <= component.gripper_threshold)).astype(np.float32) |
| else: |
| gripper = normalize_gripper_actions( |
| gripper.reshape(-1, 1), |
| closed_value=component.gripper_closed_value, |
| open_value=component.gripper_open_value, |
| ).reshape(*values.shape[:2], 1) |
| gripper = _resample(gripper, source_hz, preset.target_hz, preset.interpolation, kind="hold") |
| return np.concatenate((motion, gripper), axis=-1) |
|
|
|
|
| def _external_eef(motion, component, current_state): |
| poses = _rotvec_poses(motion) |
| if current_state is None: |
| return _pose_map(poses, "eef_pose_deltas_batch", eef_pose_deltas) |
| anchor = _poses(current_state[:, None], component) |
| quaternion = multiply_quaternions( |
| normalize_quaternion(anchor[..., 3:7]), normalize_quaternion(poses[..., 3:7]) |
| ) |
| attached = np.concatenate( |
| (poses[..., :3].astype(np.float64) + anchor[..., :3], quaternion), axis=-1 |
| ).astype(np.float32) |
| return eef_pose_components(attached.reshape(-1, 7), component.rotation_format).reshape( |
| *attached.shape[:2], 3 + len(component.rotation_dims) |
| ) |
|
|
|
|
| def externalize_batch( |
| values, component: ComponentSpec, *, source_hz, target_hz, target_length, method, current_state |
| ): |
| """Resample decoded paths and return ``(B,T_output,D_component)`` arrays.""" |
| sampled = resample_canonical_batch( |
| values, |
| component, |
| source_hz=source_hz, |
| target_hz=target_hz, |
| target_length=target_length, |
| method=method, |
| ) |
| return externalize_canonical_batch(sampled, component, current_state=current_state) |
|
|
|
|
| def resample_canonical_batch( |
| values, component: ComponentSpec, *, source_hz, target_hz, target_length, method |
| ): |
| """Rate-convert local ``(B,T,D_active)`` paths, holding gripper commands.""" |
| motion_width = len(component.joint_dims or ()) if component.space == "joint" else 6 |
| motion = values[..., :motion_width] |
| if component.space == "eef": |
| poses = _rotvec_poses(motion) |
| poses = _resample( |
| poses, source_hz, target_hz, method, kind="pose_xyzw", target_length=target_length |
| ) |
| motion = _pose_map(poses, "eef_pose_absolute_batch", eef_pose_absolute) |
| else: |
| motion = _resample(motion, source_hz, target_hz, method, target_length=target_length) |
| if component.gripper_dim is None: |
| return np.asarray(motion, dtype=np.float32) |
| gripper = _resample( |
| values[..., motion_width : motion_width + 1], |
| source_hz, |
| target_hz, |
| method, |
| kind="hold", |
| target_length=target_length, |
| ) |
| return np.asarray(np.concatenate((motion, gripper), axis=-1), dtype=np.float32) |
|
|
|
|
| def externalize_canonical_batch(values, component: ComponentSpec, *, current_state): |
| """Restore source coordinates or deltas from ``(B,T,D_active)`` local paths.""" |
| motion_width = len(component.joint_dims or ()) if component.space == "joint" else 6 |
| motion = values[..., :motion_width] |
| if component.space == "eef": |
| motion = _external_eef(motion, component, current_state) |
| elif current_state is None: |
| motion = difference_additive(motion, 1) |
| else: |
| motion = motion + np.take(current_state, component.joint_dims, axis=-1)[:, None] |
| if component.joint_wrap == "pi": |
| motion = wrap_to_pi(motion) |
| if component.gripper_dim is None: |
| return np.asarray(motion, dtype=np.float32) |
| gripper = values[..., motion_width : motion_width + 1] |
| if component.gripper_threshold is not None: |
| gripper = gripper if component.gripper_open_is_high else 1.0 - gripper |
| return np.asarray(np.concatenate((motion, gripper), axis=-1), dtype=np.float32) |
| gripper = component.gripper_closed_value + gripper * ( |
| component.gripper_open_value - component.gripper_closed_value |
| ) |
| return np.asarray(np.concatenate((motion, gripper), axis=-1), dtype=np.float32) |
|
|
