RTC: wire real-time chunking into /act

rtc.py

@@ -0,0 +1,237 @@
+"""Real-Time Chunking (RTC) for MolmoAct2 — inference-time, no retraining.
+
+RTC (arXiv 2506.07339) removes the discontinuity you get when a freshly generated
+action chunk replaces the one currently executing. It treats the overlap as an
+INPAINTING problem during flow sampling: the actions that will inevitably execute
+while we're computing are pinned to the previous chunk, a middle band is softly
+guided toward it, and the tail beyond the old chunk is generated freely.
+
+    idx <  d          "frozen"  weight 1      (executes during the inference delay)
+    d <= idx < s      "guided"  weight 1->0   (EXP schedule)
+    idx >= s          "free"    weight 0
+
+Guidance is NOT an overwrite. Following the paper (and LeRobot's reference
+implementation) it is a pseudoinverse-guidance (PiGDM) correction injected into the
+velocity field at EVERY denoising step, via a vector-Jacobian product:
+
+    x1     = x_t + (1 - tau) * v            # predicted clean sample
+    err    = (prev_chunk - x1) * W          # weighted target error
+    corr   = VJP(x1, x_t, err)              # d x1/d x_t ^T @ err
+    v_rtc  = v + w(tau) * corr
+    w(tau) = min(beta, ((1-tau)^2 + tau^2) / (tau * (1-tau)))
+
+SIGN NOTE (the highest-risk part of this port): LeRobot integrates time 1->0 with a
+velocity pointing toward NOISE and writes `v - w*corr`. MolmoAct2 integrates 0->1
+with a velocity pointing toward DATA (`trajectory = trajectory + dt * velocity`), so
+the correction ADDS here. Both end up moving the trajectory by +corr; only the
+velocity convention differs. `selftest_guidance_direction()` asserts this
+empirically rather than trusting the derivation.
+
+COSTS (measured/reported, worth knowing before enabling):
+  * needs autograd -> the model's @torch.no_grad() must be lifted, and the
+    CUDA-graph fast path must be disabled (you cannot backprop a captured graph).
+  * ~20% extra latency on top of that. We are already NETWORK-bound, so RTC is
+    only worth enabling once the round-trip is short.
+  * requires d <= s <= H - d. With H=30, s=10: d <= 10. At 10fps that's ~1.0s of
+    tolerable delay; at 15fps only ~0.66s.
+"""
+
+import math
+from typing import Optional
+
+import torch
+
+
+def prefix_weights(
+    delay: int, execution_horizon: int, total: int, schedule: str = "exp"
+) -> torch.Tensor:
+    """Per-timestep guidance weights (port of LeRobot's get_prefix_weights).
+
+    `delay` (d) actions are pinned at 1.0; weights decay to 0 by `execution_horizon`
+    (s); everything at/after s is free (0.0)."""
+    start = min(delay, execution_horizon)
+    end = execution_horizon
+    if schedule == "zeros":
+        w = torch.zeros(total)
+        w[:start] = 1.0
+        return w
+    if schedule == "ones":
+        w = torch.ones(total)
+        w[end:] = 0.0
+        return w
+
+    # linear ramp over the guided band, exclusive of the 1.0 and 0.0 endpoints
+    skip = max(total - end, 0)
+    steps = total - skip - start
+    lin = torch.linspace(1, 0, steps + 2)[1:-1] if (end > start and steps > 0) else torch.tensor([])
+    if schedule == "exp":
+        # decay harder than linear: w * expm1(w) / (e - 1)
+        lin = lin * torch.expm1(lin).div(math.e - 1)
+    if total - end > 0:
+        lin = torch.cat([lin, torch.zeros(total - end)])
+    if min(start, total) > 0:
+        lin = torch.cat([torch.ones(min(start, total)), lin])
+    return lin
+
+
+def guidance_weight(tau: float, max_weight: float) -> float:
+    """w(tau) = min(beta, ((1-tau)^2 + tau^2) / (tau*(1-tau))), clamped at both ends."""
+    one_minus = 1.0 - tau
+    if tau <= 0.0 or one_minus <= 0.0:
+        return float(max_weight)
+    w = ((one_minus ** 2) + (tau ** 2)) / (tau * one_minus)
+    return float(min(w, max_weight))
+
+
+def feasible(delay: int, horizon: int, execution_horizon: int) -> bool:
+    """RTC needs d <= s <= H - d. Past that the frozen prefix and the free tail
+    overlap and the guidance is not well defined — better to skip RTC for that
+    request than to emit a silently wrong chunk."""
+    return 0 <= delay <= execution_horizon <= horizon - delay
+
+
+def pick_execution_horizon(delay: int, horizon: int) -> Optional[int]:
+    """Smallest feasible s that leaves the guided band some room, or None if the
+    delay is too large for this horizon (d > H/2)."""
+    s = min(max(delay + 4, delay), horizon - delay)
+    return s if feasible(delay, horizon, s) else None
+
+
+class RTCState:
+    """Per-session guidance target, kept in the model's NORMALIZED action space.
+
+    We cache the previous chunk as the raw flow output rather than asking the client
+    to send actions back: the flow operates on normalized actions, so a robot-scale
+    chunk from the client would have to be re-normalized (and our client also applies
+    a joint calibration). Caching the model's own output sidesteps both."""
+
+    def __init__(self) -> None:
+        self.prev: Optional[torch.Tensor] = None  # (B, H, A) normalized
+        self.enabled = False
+        self.consumed = 0        # actions executed since `prev` was produced (= alignment shift)
+        self.delay = 0           # d: actions that will execute during THIS inference
+        self.execution_horizon = 10
+        self.max_guidance_weight = 10.0
+        self.schedule = "exp"
+        self.applied = 0         # count of guided steps, for observability
+
+    def target(self, like: torch.Tensor) -> Optional[torch.Tensor]:
+        """Previous chunk aligned to the new chunk's timeline: drop the `consumed`
+        actions that already executed, then zero-pad to the new chunk's shape."""
+        if not self.enabled or self.prev is None:
+            return None
+        left = self.prev[:, self.consumed:, :]
+        if left.shape[1] == 0:
+            return None
+        out = torch.zeros_like(like)
+        n = min(left.shape[1], out.shape[1])
+        a = min(left.shape[2], out.shape[2])
+        out[:, :n, :a] = left[:, :n, :a].to(out.device, out.dtype)
+        return out
+
+
+def install_rtc(model, state: RTCState):
+    """Monkeypatch the model's Euler flow loop to apply RTC guidance.
+
+    Returns the original bound method so the caller can restore it. The patched loop
+    is a no-op (bit-identical to upstream) whenever `state` has no target, so leaving
+    it installed costs nothing when RTC is off."""
+    original = model._run_action_flow_loop
+
+    def guided_loop(inputs, steps: int) -> torch.Tensor:
+        trajectory = inputs.trajectory
+        target = state.target(trajectory)
+        if target is None:
+            # No usable prefix (first call of a session, or the previous chunk is
+            # fully consumed). Run upstream verbatim — but still CAPTURE the output,
+            # or the next call has nothing to guide toward.
+            out = original(inputs, steps)
+            state.prev = out.detach()
+            return out
+
+        action_expert = model._require_action_expert()
+        dt = 1.0 / steps
+        pad = inputs.action_dim_is_pad
+        mask_enabled = model.config.mask_action_dim_padding
+        W = prefix_weights(state.delay, state.execution_horizon,
+                           trajectory.shape[1], state.schedule)
+        W = W.to(trajectory.device, trajectory.dtype).view(1, -1, 1)
+
+        for idx in range(steps):
+            tau = idx / steps  # 0 = noise, 1 = data (MolmoAct2 integrates forward)
+            x_t = trajectory.detach().requires_grad_(True)
+            with torch.enable_grad():
+                velocity = action_expert.forward_with_context(
+                    x_t,
+                    inputs.modulations[idx].conditioning,
+                    context=inputs.context,
+                    modulation=inputs.modulations[idx],
+                )
+                velocity = model._mask_action_dim_tensor(
+                    velocity, action_dim_is_pad=pad, enabled=mask_enabled
+                )
+                x1 = x_t + (1.0 - tau) * velocity          # predicted clean sample
+                err = ((target - x1) * W).detach()          # weighted pull toward prev
+                corr = torch.autograd.grad(x1, x_t, err, retain_graph=False)[0]
+
+            w = guidance_weight(tau, state.max_guidance_weight)
+            # + (not -): our velocity points toward DATA -- see SIGN NOTE above.
+            velocity = (velocity + w * corr).detach()
+            trajectory = model._mask_action_dim_tensor(
+                trajectory.detach() + dt * velocity, action_dim_is_pad=pad, enabled=mask_enabled
+            )
+            state.applied += 1
+        state.prev = trajectory.detach()
+        return trajectory
+
+    model._run_action_flow_loop = guided_loop
+    return original
+
+
+# --------------------------------------------------------------------------- tests
+def selftest_guidance_direction() -> dict:
+    """Assert the SIGN empirically on a toy linear flow, with no model involved.
+
+    Toy: velocity = (goal - x). Euler-integrating it drives x -> goal. With RTC
+    guidance toward `prev`, the FROZEN prefix must end up closer to `prev` than the
+    unguided run does, and the free tail must be left alone."""
+    H, A, steps = 30, 4, 10
+    goal = torch.zeros(1, H, A)
+    prev = torch.ones(1, H, A) * 5.0
+    W = prefix_weights(4, 10, H, "exp").view(1, -1, 1)
+
+    def run(guided: bool) -> torch.Tensor:
+        x = torch.full((1, H, A), -5.0)
+        for idx in range(steps):
+            tau = idx / steps
+            xt = x.detach().requires_grad_(True)
+            with torch.enable_grad():
+                v = goal - xt
+                x1 = xt + (1.0 - tau) * v
+                if guided:
+                    err = ((prev - x1) * W).detach()
+                    corr = torch.autograd.grad(x1, xt, err, retain_graph=False)[0]
+                    v = v + guidance_weight(tau, 10.0) * corr
+            x = (x.detach() + (1.0 / steps) * v.detach())
+        return x
+
+    plain, rtc = run(False), run(True)
+    d_plain = (plain[0, 0] - prev[0, 0]).abs().mean().item()
+    d_rtc = (rtc[0, 0] - prev[0, 0]).abs().mean().item()
+    tail_shift = (rtc[0, -1] - plain[0, -1]).abs().mean().item()
+    return {
+        "prefix_dist_unguided": round(d_plain, 4),
+        "prefix_dist_rtc": round(d_rtc, 4),
+        "prefix_pulled_toward_prev": d_rtc < d_plain,
+        "free_tail_unchanged": tail_shift < 1e-5,
+    }
+
+
+if __name__ == "__main__":
+    w = prefix_weights(4, 10, 30, "exp")
+    print("weights[:12]:", [round(float(x), 3) for x in w[:12]])
+    print("frozen prefix all 1.0 :", bool((w[:4] == 1.0).all()))
+    print("free tail all 0.0     :", bool((w[10:] == 0.0).all()))
+    print("monotonic in guided band:", bool((w[4:10].diff() <= 0).all()))
+    print("w(tau) mid/edges      :", [round(guidance_weight(t, 10.0), 3) for t in (0.0, 0.1, 0.5, 0.9, 1.0)])
+    print("direction selftest    :", selftest_guidance_direction())