Blame for runtime/backends/actioncodec2/semantic.py

"""Canonical physical preprocessing for routed ActionCodec2 profiles."""

from __future__ import annotations

import math

import numpy as np
from scipy.interpolate import CubicSpline, PchipInterpolator
from scipy.spatial.transform import Rotation, RotationSpline

from ._so3 import (
    euler_xyz_to_quaternion as _euler_xyz_to_quaternion,
)
from ._so3 import (
    multiply_quaternions as _quaternion_multiply,
)
from ._so3 import (
    normalize_quaternion as _normalize_quaternion,
)
from ._so3 import (
    quaternion_to_rotvec as _quaternion_to_rotvec,
)
from ._so3 import (
    rotvec_to_quaternion as _rotvec_to_quaternion,
)
from .native_backend import native_kernel


def integrate_additive(values: np.ndarray, order: int) -> np.ndarray:
    """Integrate along the penultimate (time) axis from a zero boundary."""
    result = np.asarray(values, dtype=np.float32)
    kernel = native_kernel("additive_transform")
    if kernel is not None:
        shape = result.shape
        return kernel(result.reshape(-1, *shape[-2:]), order, True).reshape(shape)
    for _ in range(order):
        result = np.cumsum(result, axis=-2, dtype=np.float32)
    return result


def difference_additive(values: np.ndarray, order: int) -> np.ndarray:
    """Take zero-anchored finite differences along the time axis."""
    result = np.asarray(values, dtype=np.float32)
    kernel = native_kernel("additive_transform")
    if kernel is not None:
        shape = result.shape
        return kernel(result.reshape(-1, *shape[-2:]), order, False).reshape(shape)
    for _ in range(order):
        previous = np.concatenate((np.zeros_like(result[..., :1, :]), result[..., :-1, :]), axis=-2)
        result = result - previous
    return result


def wrap_to_pi(values: np.ndarray) -> np.ndarray:
    """Return angles in radians folded into ``(-pi, pi]``."""
    angles = np.asarray(values, dtype=np.float64)
    kernel = native_kernel("wrap_angles")
    if kernel is not None:
        return kernel(angles)
    return np.asarray(np.pi - np.mod(np.pi - angles, 2.0 * np.pi), dtype=np.float32)


_INTERPOLATION_METHODS = {"linear", "pchip", "cubic"}


def _validate_rate(value: float, name: str) -> float:
    rate = float(value)
    if not math.isfinite(rate) or rate <= 0.0:
        raise ValueError(f"{name} must be positive and finite")
    return rate


def _quaternion_slerp(left: np.ndarray, right: np.ndarray, weight: np.ndarray) -> np.ndarray:
    first = _normalize_quaternion(np.asarray(left, dtype=np.float64))
    second = _normalize_quaternion(np.asarray(right, dtype=np.float64))
    dot = np.sum(first * second, axis=1)
    flip = dot < 0.0
    second[flip] *= -1.0
    dot = np.clip(np.abs(dot), 0.0, 1.0)
    output = np.empty_like(first)
    near = dot > 0.9995
    output[near] = first[near] + weight[near, None] * (second[near] - first[near])
    regular = ~near
    if np.any(regular):
        angle = np.arccos(dot[regular])
        denominator = np.sin(angle)
        w = weight[regular]
        output[regular] = (
            np.sin((1.0 - w) * angle)[:, None] / denominator[:, None] * first[regular]
            + np.sin(w * angle)[:, None] / denominator[:, None] * second[regular]
        )
    return _normalize_quaternion(output)


def resampled_horizon(length: int, *, source_hz: float, target_hz: float) -> int:
    """Return the target control-step count for a regular source horizon."""
    if isinstance(length, bool) or not isinstance(length, (int, np.integer)) or length < 1:
        raise ValueError("a trajectory must contain at least one sample")
    source_rate = _validate_rate(source_hz, "source_hz")
    target_rate = _validate_rate(target_hz, "target_hz")
    scaled = int(math.floor(int(length) * target_rate / source_rate + 0.5))
    return max(1, scaled)


def _target_times(
    length: int,
    source_hz: float,
    target_hz: float,
    *,
    target_length: int | None,
) -> np.ndarray:
    count = (
        resampled_horizon(length, source_hz=source_hz, target_hz=target_hz)
        if target_length is None
        else target_length
    )
    if isinstance(count, bool) or not isinstance(count, (int, np.integer)) or count < 1:
        raise ValueError("target_length must be a positive integer")
    return np.arange(int(count), dtype=np.float64) / target_hz


def _guarded_natural_cubic(
    source_time: np.ndarray,
    values: np.ndarray,
    target_time: np.ndarray,
) -> np.ndarray:
    sampled = np.asarray(
        CubicSpline(source_time, values, axis=0, bc_type="natural")(target_time),
        dtype=np.float64,
    )
    lower = np.min(values, axis=0)
    upper = np.max(values, axis=0)

    tolerance = 1e-9 + 0.05 * (upper - lower)
    fallback_axes = np.any(
        (sampled < lower[None] - tolerance[None]) | (sampled > upper[None] + tolerance[None]),
        axis=0,
    )
    if np.any(fallback_axes):
        sampled[:, fallback_axes] = PchipInterpolator(
            source_time, values[:, fallback_axes], axis=0
        )(target_time)
    return sampled


def resample_uniform(
    values: np.ndarray,
    *,
    source_hz: float,
    target_hz: float,
    kind: str = "linear",
    target_length: int | None = None,
    method: str = "linear",
) -> np.ndarray:
    """Resample one regular control-horizon trajectory onto a target clock."""
    source_rate = _validate_rate(source_hz, "source_hz")
    target_rate = _validate_rate(target_hz, "target_hz")
    array = np.asarray(values, dtype=np.float64)
    if array.ndim != 2 or len(array) < 1 or not np.all(np.isfinite(array)):
        raise ValueError("values must be a non-empty finite [T, D] array")
    if kind not in {"linear", "hold", "pose_xyzw", "pose_euler_xyz"}:
        raise ValueError("unsupported resampling kind")
    if method not in _INTERPOLATION_METHODS:
        raise ValueError(f"method must be one of {sorted(_INTERPOLATION_METHODS)}")
    if kind == "pose_xyzw" and array.shape[1] != 7:
        raise ValueError("pose_xyzw values must have shape [T, 7]")
    if kind == "pose_euler_xyz" and array.shape[1] != 6:
        raise ValueError("pose_euler_xyz values must have shape [T, 6]")
    kernel = native_kernel("resample_uniform_batch")
    if kernel is not None:
        if target_length is not None and (
            isinstance(target_length, bool)
            or not isinstance(target_length, (int, np.integer))
            or target_length < 1
        ):
            raise ValueError("target_length must be a positive integer")
        if kind == "pose_euler_xyz":
            array = np.concatenate((array[:, :3], _euler_xyz_to_quaternion(array[:, 3:6])), axis=1)
            kind = "pose_xyzw"
        return kernel(
            array[None],
            source_rate,
            target_rate,
            kind,
            -1 if target_length is None else int(target_length),
            method,
        )[0]
    standard_length = resampled_horizon(len(array), source_hz=source_rate, target_hz=target_rate)
    if source_rate == target_rate and target_length in {None, standard_length}:
        if kind == "pose_euler_xyz":
            array = np.concatenate((array[:, :3], _euler_xyz_to_quaternion(array[:, 3:6])), axis=1)
        return np.asarray(array, dtype=np.float32)

    target_time = _target_times(len(array), source_rate, target_rate, target_length=target_length)
    source_time = np.arange(len(array), dtype=np.float64) / source_rate
    interpolation_time = np.minimum(target_time, source_time[-1])
    source_position = target_time * source_rate
    left = np.minimum(np.floor(source_position).astype(np.int64), len(array) - 1)
    right = np.minimum(left + 1, len(array) - 1)
    weight = source_position - left
    if kind == "hold":
        result = array[left]
    else:
        pose = kind in {"pose_xyzw", "pose_euler_xyz"}
        continuous = array[:, :3] if pose else array
        if len(array) == 1:
            sampled = np.repeat(continuous, len(target_time), axis=0)
        elif method == "pchip":
            sampled = PchipInterpolator(source_time, continuous, axis=0)(interpolation_time)
        elif method == "cubic":
            sampled = _guarded_natural_cubic(source_time, continuous, interpolation_time)
        else:
            sampled = array[left, : continuous.shape[1]] + weight[:, None] * (
                array[right, : continuous.shape[1]] - array[left, : continuous.shape[1]]
            )
        if not pose:
            result = sampled
        else:
            quaternion = (
                array[:, 3:7] if kind == "pose_xyzw" else _euler_xyz_to_quaternion(array[:, 3:6])
            )
            if len(array) == 1:
                rotation = np.repeat(quaternion, len(target_time), axis=0)
            elif method == "linear":
                rotation = _quaternion_slerp(quaternion[left], quaternion[right], weight)
            else:
                rotation = RotationSpline(source_time, Rotation.from_quat(quaternion))(
                    interpolation_time
                ).as_quat()
            result = np.concatenate((sampled, rotation), axis=1)
    return np.asarray(result, dtype=np.float32)


def eef_pose_absolute(poses: np.ndarray) -> np.ndarray:
    """Convert canonical ``xyz + xyzw`` state to ``xyz + rotation_vector``."""
    values = np.asarray(poses, dtype=np.float64)
    if values.ndim != 2 or values.shape[1] != 7 or not np.all(np.isfinite(values)):
        raise ValueError("poses must be a finite [T, 7] array")
    kernel = native_kernel("eef_pose_absolute_batch")
    if kernel is not None and len(values):
        return kernel(values[None])[0]
    return np.asarray(
        np.concatenate((values[:, :3], _quaternion_to_rotvec(values[:, 3:7])), axis=1),
        dtype=np.float32,
    )


def eef_pose_components(poses: np.ndarray, representation: str) -> np.ndarray:
    """Convert canonical ``xyz + xyzw`` poses to one registered source layout."""
    values = np.asarray(poses, dtype=np.float64)
    if values.ndim != 2 or values.shape[1] != 7 or not np.all(np.isfinite(values)):
        raise ValueError("poses must be a finite [T, 7] array")
    quaternion = _normalize_quaternion(values[:, 3:7])
    if representation == "rotation_vector":
        rotation = _quaternion_to_rotvec(quaternion)
    elif representation == "euler_xyz":
        rotation = Rotation.from_quat(quaternion).as_euler("xyz")
    elif representation == "quaternion_xyzw":
        rotation = quaternion
    else:
        raise ValueError("unsupported EEF rotation representation")
    return np.asarray(np.concatenate((values[:, :3], rotation), axis=1), dtype=np.float32)


def eef_pose_from_components(
    position: np.ndarray,
    rotation: np.ndarray,
    representation: str,
) -> np.ndarray:
    """Build canonical ``xyz + xyzw`` poses from an explicit source layout."""
    translation = np.asarray(position, dtype=np.float64)
    orientation = np.asarray(rotation, dtype=np.float64)
    if translation.ndim != 2 or translation.shape[1] != 3 or len(translation) < 1:
        raise ValueError("position must have non-empty shape [T, 3]")
    if len(orientation) != len(translation) or not np.all(np.isfinite(translation)):
        raise ValueError("position and rotation must have matching finite trajectories")
    if representation == "rotation_vector":
        if orientation.ndim != 2 or orientation.shape[1] != 3:
            raise ValueError("rotation_vector must have shape [T, 3]")
        quaternion = _rotvec_to_quaternion(orientation)
    elif representation == "euler_xyz":
        if orientation.ndim != 2 or orientation.shape[1] != 3:
            raise ValueError("euler_xyz must have shape [T, 3]")
        quaternion = _euler_xyz_to_quaternion(orientation)
    elif representation == "quaternion_xyzw":
        if orientation.ndim != 2 or orientation.shape[1] != 4:
            raise ValueError("quaternion_xyzw must have shape [T, 4]")
        quaternion = _normalize_quaternion(orientation)
    else:
        raise ValueError("unsupported EEF rotation representation")
    if not np.all(np.isfinite(orientation)):
        raise ValueError("position and rotation must be finite")
    return np.asarray(np.concatenate((translation, quaternion), axis=1), dtype=np.float32)


def eef_pose_local_absolute(
    poses: np.ndarray, *, initial_pose: np.ndarray | None = None
) -> np.ndarray:
    """Express measured poses relative to an explicit or inferred initial pose."""
    values = np.asarray(poses, dtype=np.float64)
    if (
        values.ndim != 2
        or values.shape[1] != 7
        or len(values) < 1
        or not np.all(np.isfinite(values))
    ):
        raise ValueError("poses must be a finite non-empty [T, 7] array")
    if initial_pose is None:
        initial = values[:1]
    else:
        initial = np.asarray(initial_pose, dtype=np.float64)
        if initial.shape == (7,):
            initial = initial[None]
        if initial.shape != (1, 7) or not np.all(np.isfinite(initial)):
            raise ValueError("initial_pose must be a finite [7] pose")
    quaternion = _normalize_quaternion(values[:, 3:7])
    initial_quaternion = _normalize_quaternion(initial[:, 3:7])
    inverse_start = initial_quaternion.copy()
    inverse_start[:, :3] *= -1.0
    relative = _quaternion_multiply(np.broadcast_to(inverse_start, quaternion.shape), quaternion)
    return np.asarray(
        np.concatenate((values[:, :3] - initial[:, :3], _quaternion_to_rotvec(relative)), axis=1),
        dtype=np.float32,
    )


def eef_pose_deltas(poses: np.ndarray) -> np.ndarray:
    """Differentiate a local ``xyz + xyzw`` path from zero/identity."""
    values = np.asarray(poses, dtype=np.float64)
    if (
        values.ndim != 2
        or values.shape[1] != 7
        or len(values) < 1
        or not np.all(np.isfinite(values))
    ):
        raise ValueError("poses must be a finite non-empty [T, 7] array")
    kernel = native_kernel("eef_pose_deltas_batch")
    if kernel is not None:
        return kernel(values[None])[0]
    quaternion = _normalize_quaternion(values[:, 3:7])
    previous_position = np.concatenate((np.zeros((1, 3)), values[:-1, :3]), axis=0)
    identity = np.asarray([[0.0, 0.0, 0.0, 1.0]], dtype=np.float64)
    previous_quaternion = np.concatenate((identity, quaternion[:-1]), axis=0)
    inverse_previous = previous_quaternion.copy()
    inverse_previous[:, :3] *= -1.0
    increment = _quaternion_multiply(inverse_previous, quaternion)
    return np.asarray(
        np.concatenate(
            (values[:, :3] - previous_position, _quaternion_to_rotvec(increment)), axis=1
        ),
        dtype=np.float32,
    )


def eef_pose_from_local_absolute(poses: np.ndarray, *, initial_pose: np.ndarray) -> np.ndarray:
    """Attach an external initial pose to a local ``xyz + xyzw`` path."""
    values = np.asarray(poses, dtype=np.float64)
    initial = np.asarray(initial_pose, dtype=np.float64)
    if (
        values.ndim != 2
        or values.shape[1] != 7
        or len(values) < 1
        or not np.all(np.isfinite(values))
    ):
        raise ValueError("poses must be a finite non-empty [T, 7] array")
    if initial.shape == (7,):
        initial = initial[None]
    if initial.shape != (1, 7) or not np.all(np.isfinite(initial)):
        raise ValueError("initial_pose must be a finite [7] pose")
    local_quaternion = _normalize_quaternion(values[:, 3:7])
    initial_quaternion = _normalize_quaternion(initial[:, 3:7])
    quaternion = _quaternion_multiply(
        np.broadcast_to(initial_quaternion, local_quaternion.shape), local_quaternion
    )
    return np.asarray(
        np.concatenate((values[:, :3] + initial[:, :3], quaternion), axis=1),
        dtype=np.float32,
    )


def integrate_eef_deltas(deltas: np.ndarray) -> np.ndarray:
    """Integrate ``[dx,dy,dz,body-rotvec]`` from a zero/identity pose."""
    values = np.asarray(deltas, dtype=np.float64)
    if values.ndim != 2 or values.shape[1] != 6 or not np.all(np.isfinite(values)):
        raise ValueError("deltas must be a finite [T, 6] array")
    if len(values) < 1:
        raise ValueError("deltas must contain at least one step")
    kernel = native_kernel("integrate_eef_deltas_batch")
    if kernel is not None:
        return kernel(values[None])[0]
    translation = np.cumsum(values[:, :3], axis=0)
    increments = _rotvec_to_quaternion(values[:, 3:6])
    orientation = np.empty((len(values), 4), dtype=np.float64)
    current = np.asarray([[0.0, 0.0, 0.0, 1.0]], dtype=np.float64)
    for index, increment in enumerate(increments):
        current = _quaternion_multiply(current, increment[None])
        orientation[index] = current[0]
    return np.asarray(np.concatenate((translation, orientation), axis=1), dtype=np.float32)


def normalize_gripper_actions(
    values: np.ndarray,
    *,
    closed_value: float,
    open_value: float,
) -> np.ndarray:
    """Map one raw controller gripper channel to the canonical open fraction."""
    array = np.asarray(values, dtype=np.float32)
    if array.ndim != 2 or array.shape[1] != 1 or not np.all(np.isfinite(array)):
        raise ValueError("gripper action must be a finite [T, 1] array")
    if not (math.isfinite(closed_value) and math.isfinite(open_value)):
        raise ValueError("gripper anchors must be finite")
    if closed_value == open_value:
        raise ValueError("gripper closed/open anchors must differ")
    kernel = native_kernel("gripper_transform")
    if kernel is not None:
        return kernel(array, "normalize", closed_value, open_value)
    fraction = (array - closed_value) / (open_value - closed_value)
    return np.clip(fraction, 0.0, 1.0).astype(np.float32)


def discretize_gripper_actions(
    values: np.ndarray,
    *,
    mode: str,
    threshold: float = 0.8,
) -> np.ndarray:
    """Quantize the canonical open fraction for the codec."""
    array = np.asarray(values, dtype=np.float32)
    if array.ndim != 2 or array.shape[1] != 1 or not np.all(np.isfinite(array)):
        raise ValueError("gripper fraction must be a finite [T, 1] array")
    if mode == "binary":
        if not math.isfinite(threshold):
            raise ValueError("gripper threshold must be finite")
        kernel = native_kernel("gripper_transform")
        if kernel is not None:
            return kernel(array, mode, threshold, 0)
        return (array >= threshold).astype(np.float32)
    if mode == "eleven":
        kernel = native_kernel("gripper_transform")
        if kernel is not None:
            return kernel(array, mode, threshold, 0)
        return (np.rint(array * 10.0) / 10.0).astype(np.float32)
    if mode == "twenty":
        return (np.rint(array * 20.0) / 20.0).astype(np.float32)
    raise ValueError("gripper mode must be 'binary', 'eleven' or 'twenty'")