"""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 _flow_owner(model): """Find the object that actually owns the flow loop. `_run_action_flow_loop`, `_require_action_expert` and `_mask_action_dim_tensor` are all methods of MolmoAct2Model, while what you load from the Hub is a MolmoAct2ForConditionalGeneration that HOLDS one as `.model`. Patching the outer wrapper raises AttributeError at load time and takes the whole server down, so resolve the owner instead of assuming it.""" for cand in (model, getattr(model, "model", None), getattr(model, "base_model", None)): if cand is not None and hasattr(cand, "_run_action_flow_loop"): return cand return None def install_rtc(model, state: RTCState): """Monkeypatch the flow loop to apply RTC guidance. Returns the original bound method, or None if the loop could not be located — NEVER raises. A diagnostic feature must not be able to stop the model from loading; without the patch the server simply serves un-guided chunks. 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.""" owner = _flow_owner(model) if owner is None: return None original = owner._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 = owner._require_action_expert() dt = 1.0 / steps pad = inputs.action_dim_is_pad mask_enabled = owner.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 = owner._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 = owner._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 owner._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())