onf-c2/src/onf/graph/core/segments.py

"""Node reference segments as torch tensors, and the differentiable posterior expectation over them.

The training path for the action-chunk blend: an action-space loss reaches the retriever only
through the expectation's weights. Deploy uses the numpy arms in onf.graph.run.readout instead.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

import numpy as np
import torch

from onf.graph.core import schema

if TYPE_CHECKING:
    from onf.graph.core.nodes import NodeTable

__all__ = ["EXPECT_K", "ee_segment_table", "expected_ee_segment"]

EXPECT_K = 32              # expected_ee_segment support width; matches ReadoutContext.topm
_EPS = 1e-12


def ee_segment_table(nodes: "NodeTable", seg_k: int = schema.SEG_K) -> torch.Tensor:
    """Materialize every node's reference end-effector segment as one tensor.

    Row v equals onf.graph.core.nodes.NodeTable.ee_segment(v, seg_k): the same clip at
    raw_ptr[owner+1] and the same edge padding, for all V nodes at once. Build it ONCE per graph and
    hold it -- on the long suite it is [28476, 8, 6] float32, about 5.5 MB, and rebuilding it per
    call would cost more than the expectation it feeds.

    Args:
        nodes: The demonstration-graph node table.
        seg_k: Segment length in frames.

    Returns:
        [V, seg_k, 6] float32 base-frame concat(position, unnormalized axis-angle).

    Raises:
        ValueError: seg_k < 1, or some node's t_raw falls outside its own demo's raw range.
    """
    if seg_k < 1:
        raise ValueError(f"seg_k must be >= 1, got {seg_k}")
    owner = nodes.owner.astype(np.int64)
    raw_ptr = nodes.raw_ptr.astype(np.int64)
    start = nodes.t_raw.astype(np.int64)
    demo_start, demo_end = raw_ptr[owner], raw_ptr[owner + 1]

    outside = np.flatnonzero((start < demo_start) | (start >= demo_end))
    if outside.size:
        v = int(outside[0])
        raise ValueError(
            f"t_raw[{v}]={start[v]} outside demo {owner[v]}'s raw range "
            f"[{demo_start[v]}, {demo_end[v]})"
        )

    # Clamping the read index at the demo's last frame IS the edge padding: past that frame every
    # further row repeats it, which is what NodeTable._edge_pad produces.
    last = np.minimum(start + seg_k, demo_end) - 1
    offsets = np.arange(seg_k, dtype=np.int64)
    idx = np.minimum(start[:, None] + offsets[None, :], last[:, None])
    return torch.as_tensor(np.asarray(nodes.ee_base[idx], dtype=np.float32))


def expected_ee_segment(p: torch.Tensor, table: torch.Tensor, k: int = EXPECT_K) -> torch.Tensor:
    """Posterior expectation of the end-effector segment, differentiable w.r.t. p.

    Takes the top-k of p, renormalises over that support and contracts the table against it. topk's
    INDICES are a hard selection carrying no gradient; the WEIGHTS carry one, and that is the only
    path by which an action-space loss can reach the retriever.

    One body, both ranks: p is [V] or [B, V], addressed from the end of the shape.

    Note:
        Deliberately a FIXED k, where the numpy Readout._aggregate takes weighted_topk's
        0.99-cumulative-mass support. A mass-based K is data-dependent and would make the batched
        support ragged. The two agree exactly when the realised supports coincide -- pinned by
        tests/test_readout.py::test_expected_ee_segment_matches_the_numpy_aggregate.

    Args:
        p: [V] or [B, V] distribution over nodes; need not already sum to 1.
        table: [V, seg_k, 6] segments from ee_segment_table.
        k: Support width, clamped to V.

    Returns:
        [seg_k, 6] or [B, seg_k, 6], in the promoted dtype of p and table.

    Raises:
        ValueError: k < 1.
    """
    if k < 1:
        raise ValueError(f"k must be >= 1, got {k}")
    top_p, idx = torch.topk(p, min(k, p.shape[-1]), dim=-1)
    w = top_p / top_p.sum(dim=-1, keepdim=True).clamp_min(_EPS)
    return (w[..., None, None] * table[idx]).sum(dim=-3)