| """Soft prefix-attention masks for RTC guided inpainting (paper Sec. 3.2 / App. A). |
| |
| Black, Galliker, Levine. "Real-Time Execution of Action Chunking Flow |
| Policies." NeurIPS 2025. arXiv:2506.07339. |
| |
| RECONSTRUCTED MODULE. The original was lost; this restores it from the paper and |
| from how :mod:`rtc.sampler` calls it (``start=inference_delay``, |
| ``end=horizon-execution_horizon``, ``total=horizon``). The three properties that |
| make RTC *correct* are asserted by ``test_masking.py``: |
| |
| 1. ``w[i] == 1`` for ``i < start`` -- the first ``d`` actions are already |
| committed to the robot before the new chunk lands, so they are a HARD |
| constraint, not a preference. Getting this wrong breaks the guarantee. |
| 2. ``w`` is monotonically non-increasing -- confidence in the old chunk decays |
| with distance into the future. |
| 3. ``w[i] == 0`` for ``i >= end`` -- past the overlap there is no previous |
| chunk to agree with, so the new one is unconstrained. |
| |
| The exact interior curve between ``start`` and ``end`` is a smoothness |
| hyperparameter, not a correctness property. ``linear`` reproduces the reference |
| implementation's documented example exactly:: |
| |
| get_prefix_weights(start=2, end=6, total=10, schedule="linear") |
| -> [1, 1, 4/5, 3/5, 2/5, 1/5, 0, 0, 0, 0] |
| |
| ``exp`` is a reconstruction: same endpoints and monotonicity, but a convex decay |
| that holds the old chunk more tightly just after the frozen prefix and releases |
| it faster near ``end``. If you need bit-exact parity with the authors' ``exp``, |
| compare against Physical-Intelligence/real-time-chunking-kinetix ``src/model.py`` |
| and adjust :func:`_exp_curve` only -- nothing else depends on its shape. |
| |
| ``zeros`` is the naive hard-masking baseline the paper compares against: freeze |
| the prefix, ignore the overlap entirely. |
| """ |
|
|
| from __future__ import annotations |
|
|
| from typing import Literal |
|
|
| import numpy as np |
|
|
| PrefixAttentionSchedule = Literal["ones", "zeros", "linear", "exp"] |
|
|
| _SCHEDULES = ("ones", "zeros", "linear", "exp") |
|
|
|
|
| def _exp_curve(u: np.ndarray) -> np.ndarray: |
| """Convex remap of the linear ramp ``u`` in [0, 1], fixing u=0 -> 0, u=1 -> 1. |
| |
| ``u * expm1(u) / (e - 1)``: monotone increasing, so weights stay ordered. |
| """ |
| return u * np.expm1(u) / (np.e - 1.0) |
|
|
|
|
| def get_prefix_weights_np( |
| start: int, |
| end: int, |
| total: int, |
| schedule: PrefixAttentionSchedule = "exp", |
| ) -> np.ndarray: |
| """Soft mask ``W`` of shape ``(total,)``, float64 numpy. |
| |
| Args: |
| start: ``d``, the inference delay. Indices ``< start`` are frozen (w = 1). |
| end: ``H - s``, the prefix attention horizon. Indices ``>= end`` get w = 0. |
| total: ``H``, the chunk length. |
| schedule: see the module docstring. |
| """ |
| if schedule not in _SCHEDULES: |
| raise ValueError(f"unknown schedule {schedule!r}; pick one of {_SCHEDULES}") |
| if total <= 0: |
| raise ValueError(f"total must be positive, got {total}") |
| |
| |
| start = int(min(start, end)) |
| end = int(min(max(end, 0), total)) |
| i = np.arange(total) |
|
|
| if schedule == "ones": |
| w = np.ones(total, dtype=np.float64) |
| elif schedule == "zeros": |
| w = (i < start).astype(np.float64) |
| else: |
| span = end - start + 1 |
| w = np.clip((end - i) / max(span, 1), 0.0, 1.0) |
| if schedule == "exp": |
| w = _exp_curve(w) |
| w = np.where(i < start, 1.0, w) |
|
|
| return np.where(i >= end, 0.0, w) |
|
|
|
|
| def get_prefix_weights( |
| start: int, |
| end: int, |
| total: int, |
| schedule: PrefixAttentionSchedule = "exp", |
| device=None, |
| dtype=None, |
| ): |
| """Torch flavour of :func:`get_prefix_weights_np` (what ``rtc.sampler`` uses).""" |
| import torch |
|
|
| w = get_prefix_weights_np(start, end, total, schedule) |
| return torch.as_tensor(w, device=device, dtype=dtype or torch.float32) |
|
|
|
|
| def get_prefix_weights_jax( |
| start: int, |
| end: int, |
| total: int, |
| schedule: PrefixAttentionSchedule = "exp", |
| dtype=None, |
| ): |
| """JAX flavour, for the JAX/orbax dual-expert backend (``rtc.sampler_jax``).""" |
| import jax.numpy as jnp |
|
|
| w = get_prefix_weights_np(start, end, total, schedule) |
| return jnp.asarray(w, dtype=dtype or jnp.float32) |
|
|
|
|
| def guidance_weight(tau: float, beta: float) -> float: |
| """PiGDM guidance weight ``min(beta, (1 - tau) / (tau * r2_tau))`` (paper Eqs. 2/4). |
| |
| Lives here rather than in a backend sampler so the torch and JAX paths cannot |
| drift apart:: |
| |
| inv_r2 = (tau**2 + (1 - tau)**2) / (1 - tau)**2 # 1 / r2_tau (Eq. 4) |
| c = (1 - tau) / tau # -> inf at tau == 0 |
| weight = min(c * inv_r2, beta) |
| |
| At ``tau == 0`` the unclipped weight is infinite (the score is undefined at pure |
| noise), so ``beta`` is what makes it finite -- exactly its role in App. A.2. |
| """ |
| if tau <= 0.0: |
| return beta |
| inv_r2 = (tau**2 + (1.0 - tau) ** 2) / ((1.0 - tau) ** 2) |
| c = (1.0 - tau) / tau |
| return min(c * inv_r2, beta) |
|
|
