onf-c2/src/onf/graph/net/modules.py

"""Standalone building blocks of onf.graph.net.gnn.GraphRetrieverNet: sinusoidal phase
encoding, query-window pooling, and query-gated relation embeddings.

Both modules are rank-agnostic — one forward serves the single-query deploy path and the
batched training path, because every reduction is addressed from the end of the shape.
"""

from __future__ import annotations

from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor

from onf.graph.core import schema

__all__ = ["QueryEncoder", "QueryEncoding", "RelationGate", "sinusoidal_psi"]

REL_INIT_STD = 0.02
LOG_W_EPS = 1e-12


def sinusoidal_psi(t_frac: Tensor, n_psi: int) -> Tensor:
    """Fixed (non-learned) sin/cos encoding of a phase fraction.

    A free function rather than a module with a frequency buffer: it holds no state, and a buffer
    would pin the octaves to the module's device while this reads t_frac's own.

    Args:
        t_frac: Phase fractions in [0, 1], any shape [...].
        n_psi: Number of sin/cos octaves.

    Returns:
        Encoding of shape [..., 2*n_psi], dtype and device of t_frac.
    """
    freqs = 2.0 ** torch.arange(n_psi, dtype=t_frac.dtype, device=t_frac.device)   # [n_psi]
    ang = torch.pi * t_frac[..., None] * freqs                                    # [..., n_psi]
    return torch.cat([torch.sin(ang), torch.cos(ang)], dim=-1)                    # [..., 2*n_psi]


@dataclass(frozen=True, eq=False)
class QueryEncoding:
    """One window's encoding, in the three pieces the pipeline reads separately.

    Attributes:
        e_q: Pooled query embedding, [hidden] or [B, hidden].
        z: Per-step states BEFORE pooling, [H, hidden] or [B, H, hidden] -- the learned seed's
            tier-1 keys, which must stay per row because its mixture is per row.
        pool_w: Cleanliness pooling weights, [H] or [B, H], summing to 1.
    """

    e_q: Tensor
    z: Tensor
    pool_w: Tensor


class QueryEncoder(nn.Module):
    """Pool a window of recent joint states into one query embedding e_Q.

    Pooling weights are softmax(log w) — proportional to per-step CLEANLINESS, not to recency:
    when recovery is needed the most recent states are the worst ones, and the older clean states
    are what identify which demonstration strand we were on.

    Args:
        dim: Joint-space dimensionality.
        hidden: Width of the hidden and output embeddings.
    """

    def __init__(self, dim: int, hidden: int = schema.HIDDEN):
        super().__init__()
        self.dim = int(dim)
        self.hidden = int(hidden)
        in_dim = 2 * self.dim + 1                                # per-step [q, qdot, grip]
        self.step = nn.Sequential(
            nn.Linear(in_dim, self.hidden), nn.ReLU(), nn.Linear(self.hidden, self.hidden),
        )
        self.out = nn.Linear(self.hidden, self.hidden)

    def encode(self, qh: Tensor, vh: Tensor, gh: Tensor, w: Tensor) -> QueryEncoding:
        """Encode and cleanliness-pool one window, batched or unbatched, keeping every piece.

        Args:
            qh: Joint positions, [H, dim] or [B, H, dim].
            vh: Joint velocities, same shape as qh.
            gh: Gripper states, [H] or [B, H].
            w: Per-step cleanliness weights, same shape as gh. Floored at LOG_W_EPS before
                the log, so a zero weight yields a near-zero (not exactly zero) share.

        Returns:
            The encoding: e_Q, the per-step states and the pooling weights.
        """
        x = torch.cat([qh, vh, gh[..., None]], dim=-1)                     # [..., H, 2*dim+1]
        z = self.step(x)                                                   # [..., H, hidden]
        pool_w = F.softmax(w.clamp_min(LOG_W_EPS).log(), dim=-1)           # cleanliness, NOT recency
        return QueryEncoding(
            e_q=self.out((pool_w[..., None] * z).sum(-2)), z=z, pool_w=pool_w,
        )

    def forward(self, qh: Tensor, vh: Tensor, gh: Tensor, w: Tensor) -> Tensor:
        """Encode one window down to e_Q alone, batched or unbatched.

        Args:
            qh: Joint positions, [H, dim] or [B, H, dim].
            vh: Joint velocities, same shape as qh.
            gh: Gripper states, [H] or [B, H].
            w: Per-step cleanliness weights, same shape as gh.

        Returns:
            Query embedding e_Q, [hidden] or [B, hidden].
        """
        return self.encode(qh, vh, gh, w).e_q

    def forward_batch(self, qh: Tensor, vh: Tensor, gh: Tensor, w: Tensor) -> Tensor:
        """Batched entry point; delegates to the rank-agnostic forward."""
        return self.forward(qh, vh, gh, w)


class RelationGate(nn.Module):
    """Query-gated relation embeddings: g_rel = MLP([h_rel ; e_Q]).

    Recomputed on every propagation layer rather than cached once, so the relation table stays
    query-dependent.

    Args:
        hidden: Embedding width.
        n_rel: Number of relation types.
    """

    def __init__(self, hidden: int, n_rel: int):
        super().__init__()
        self.emb = nn.Parameter(torch.randn(n_rel, hidden) * REL_INIT_STD)
        self.mlp = nn.Sequential(nn.Linear(2 * hidden, hidden), nn.ReLU(), nn.Linear(hidden, hidden))

    @property
    def n_rel(self) -> int:
        """Number of relation types in the embedding table."""
        return int(self.emb.shape[0])

    def forward(self, e_q: Tensor) -> Tensor:
        """Gate the relation table by a query embedding, batched or unbatched.

        Batching must not collapse the gate to one shared table: every batch element gets its own
        gated copy from its own e_Q. Both expand calls are views, not copies.

        Args:
            e_q: Query embedding, [hidden] or [B, hidden].

        Returns:
            Gated relation embeddings, [n_rel, hidden] or [B, n_rel, hidden].
        """
        target = (*e_q.shape[:-1], *self.emb.shape)                        # [..., R, hidden]
        emb = self.emb.expand(target)
        q = e_q[..., None, :].expand(target)
        return self.mlp(torch.cat([emb, q], dim=-1))