[ARC 1/7] share rotation travel and first-crossing interpolation primitives

egomimic/rldb/zarr/arc_length_tokenizer.py

@@ -143,6 +143,30 @@ def cumulative_arc_length(pos: np.ndarray) -> np.ndarray:
     return np.concatenate([np.array([0.0]), np.cumsum(step)])
 
 
+def cumulative_rotation_length(ypr: np.ndarray) -> np.ndarray:
+    """Return cumulative SO(3) geodesic angle in radians.
+
+    The returned coordinate is an independent rotation clock.  It deliberately
+    does not depend on translational motion, so an in-place wrist rotation can
+    advance even when the EEF position is stationary.
+    """
+    ypr = np.asarray(ypr, dtype=np.float64).reshape(-1, 3)
+    if len(ypr) < 2:
+        return np.zeros(len(ypr), dtype=np.float64)
+    rotations = _ypr_to_rotation(ypr)
+    relative = rotations[:-1].inv() * rotations[1:]
+    step = np.linalg.norm(relative.as_rotvec(), axis=-1)
+    return np.concatenate([np.array([0.0]), np.cumsum(step)])
+
+
+def cumulative_bimanual_translation_length(actions: np.ndarray) -> np.ndarray:
+    """Joint translation clock: left travel plus right travel."""
+    actions = np.asarray(actions, dtype=np.float64)
+    left = np.linalg.norm(np.diff(actions[:, 0:3], axis=0), axis=-1)
+    right = np.linalg.norm(np.diff(actions[:, 7:10], axis=0), axis=-1)
+    return np.concatenate(([0.0], np.cumsum(left + right)))
+
+
 def _bracket_segment(cumdist: np.ndarray, target_s: float) -> tuple[int, float]:
     """Locate the segment containing arc length target_s.
 
@@ -225,6 +249,51 @@ def _bracket_segments(
     return i, alpha
 
 
+def _first_crossing_brackets(
+    cumdist: np.ndarray, targets: np.ndarray
+) -> tuple[np.ndarray, np.ndarray]:
+    """Vectorized first-crossing bracket, mirroring the per-arm scalar rule.
+
+    Unlike :func:`_bracket_segments` this never clamps a target at or past the
+    end to the last segment's ``alpha = 1``: a trailing hold consumes no
+    translation, so the first frame reaching the distance is the answer.
+    ``TokenizeBimanualArcLengthCartesian._translation_bracket`` is the scalar
+    statement of the same rule, and the two must agree exactly.
+    """
+    cumdist = np.asarray(cumdist, dtype=np.float64)
+    targets = np.asarray(targets, dtype=np.float64)
+    n = len(cumdist)
+    capped = np.minimum(targets, cumdist[-1])
+    upper = np.clip(np.searchsorted(cumdist, capped, side="left"), 1, n - 1)
+    interval = cumdist[upper] - cumdist[upper - 1]
+    ok = interval > 1e-12
+    # The numerator keeps the uncapped target, exactly as the scalar path does;
+    # the clip below is what bounds an overshooting target.
+    alpha = np.where(
+        ok, (targets - cumdist[upper - 1]) / np.where(ok, interval, 1.0), 0.0
+    )
+    alpha = np.clip(alpha, 0.0, 1.0)
+    lo = targets <= cumdist[0]
+    return np.where(lo, 0, upper - 1), np.where(lo, 0.0, alpha)
+
+
+def _source_times_at_targets(
+    cumdist: np.ndarray, targets: np.ndarray, dt: float
+) -> np.ndarray:
+    """Map arc targets to fractional source times with scalar-identical math."""
+    indices, alpha = _bracket_segments(cumdist, targets)
+    return (indices + alpha) * dt
+
+
+def _linear_at_targets(
+    values: np.ndarray, cumdist: np.ndarray, targets: np.ndarray
+) -> np.ndarray:
+    """Interpolate all targets while preserving the scalar operation order."""
+    indices, alpha = _bracket_segments(cumdist, targets)
+    blend = alpha[:, None]
+    return (1.0 - blend) * values[indices] + blend * values[indices + 1]
+
+
 def _slerp_segments_ypr(
     ypr: np.ndarray, i: np.ndarray, alpha: np.ndarray
 ) -> np.ndarray:
@@ -1061,6 +1130,22 @@ def transform(self, batch: dict) -> dict:
 # consecutive waypoints are joined by straight segments, so a single interval's
 # chord IS its arc length. That discrepancy only arises when one rate spans the
 # whole token.
+ARC_CHUNKING_MODES = ("race", "multistream", "joint_distance")
+
+
+def resolve_arc_chunking_mode(mode=None, rotation_distance_unit=None):
+    # Source-window and geometric utilities retain their existing defaults.
+    # The Cartesian codec does not use this selector.
+    if mode is None:
+        return "joint_distance" if rotation_distance_unit is not None else "multistream"
+    value = str(mode).strip().lower()
+    if value not in ARC_CHUNKING_MODES:
+        raise ValueError(
+            f"Unknown arc_chunking_mode geometric progress measure: {value!r}"
+        )
+    return value
+
+
 BIMANUAL_VELOCITY_MODES = ("mean", "per_waypoint", "duration")