onf-c2/src/onf/graph/train/align.py

"""A dynamic-time-warping teacher for the blend weight.

`mse(a_exec, a_corr)` tells alpha how much its blend hurt on THIS chunk. It does not tell it the
thing that generalises: whether the segment being followed is the same path as the one the arm
should be on, up to how fast it is walked. Two segments that trace the same motion at different
speeds have a large pointwise error and a small warped one -- and it is the warped one that decides
whether steering towards the segment is a good idea.

So a second, cheap supervision channel: warp the reference segment the blend is about to follow
onto the true future segment, and label the pair "aligned" when the mean warped distance clears a
threshold. Offline at training time, worth nothing at deploy time, and it gives alpha a target on
rows where the regression term happens to be small for the wrong reason.

Positions only. The rotation block is measured in radians against metres and has no shared scale
with it; mixing them would make the threshold meaningless.
"""

from __future__ import annotations

from typing import Final

import torch
from torch import Tensor

__all__ = ["ALIGN_H", "dtw_cost", "alignment_label"]

# Metres. Mean warped per-step distance below which a segment counts as "the same path". Chosen at
# roughly a gripper width: closer than this and following the reference cannot put the hand
# somewhere the true segment does not also pass through.
ALIGN_H: Final[float] = 0.05

_BIG: Final[float] = 1e9


def dtw_cost(a: Tensor, b: Tensor) -> Tensor:
    """Mean-per-step dynamic-time-warping distance between two batches of position segments.

    Classic |a| x |b| DP with the three standard moves, divided by the number of aligned pairs on
    the realised path. The loop runs over the K x K cells, not over the batch: K is the action-chunk
    length (8), so 64 vectorised [B] updates cost less than one Python pass over a batch.

    Args:
        a: [B, Ka, 3] positions, metres.
        b: [B, Kb, 3] positions, metres.

    Returns:
        [B] mean warped distance, metres. Path length is used as the normaliser, so the value is
        comparable across different K.

    Raises:
        ValueError: Either input is not [B, K, 3] with K >= 1, or the batch sizes disagree.
    """
    for name, tensor in (("a", a), ("b", b)):
        if tensor.dim() != 3 or tensor.shape[-1] != 3 or tensor.shape[1] < 1:
            raise ValueError(f"{name} must be [B, K, 3] with K >= 1, got {tuple(tensor.shape)}")
    if a.shape[0] != b.shape[0]:
        raise ValueError(f"batch mismatch: a has {a.shape[0]} rows, b has {b.shape[0]}")

    local = torch.cdist(a, b)                                  # [B, Ka, Kb]
    n_b, k_a, k_b = local.shape
    big = torch.full((n_b,), _BIG, device=local.device, dtype=local.dtype)
    # (accumulated cost, path length) per cell; only the previous row is ever read.
    prev_cost = [big] * k_b
    prev_len = [torch.zeros_like(big)] * k_b

    for i in range(k_a):
        row_cost: list[Tensor] = []
        row_len: list[Tensor] = []
        for j in range(k_b):
            if i == 0 and j == 0:
                best_cost, best_len = torch.zeros_like(big), torch.zeros_like(big)
            else:
                cands = [(prev_cost[j], prev_len[j])]                       # from above
                if j > 0:
                    cands.append((row_cost[j - 1], row_len[j - 1]))         # from the left
                    cands.append((prev_cost[j - 1], prev_len[j - 1]))       # diagonal
                stacked = torch.stack([c for c, _ in cands])
                lengths = torch.stack([n for _, n in cands])
                pick = stacked.argmin(dim=0)
                best_cost = stacked.gather(0, pick[None])[0]
                best_len = lengths.gather(0, pick[None])[0]
            row_cost.append(best_cost + local[:, i, j])
            row_len.append(best_len + 1.0)
        prev_cost, prev_len = row_cost, row_len

    return prev_cost[-1] / prev_len[-1].clamp_min(1.0)


def alignment_label(ee_used: Tensor, ee_true: Tensor, h: float = ALIGN_H) -> Tensor:
    """Whether the segment the blend will follow traces the same path as the true future one.

    Args:
        ee_used: [B, K, 6] reference segment the blend is built from, concat(position, axis-angle).
        ee_true: [B, K, 6] the label node's own segment, same layout.
        h: Threshold on the mean warped distance, metres.

    Returns:
        [B] float in {0, 1}, ready as a BCE target.
    """
    return (dtw_cost(ee_used[..., :3], ee_true[..., :3]) < h).to(ee_used.dtype)