File size: 8,769 Bytes
bdca518 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 | """Real-time asynchronous controller for RTC (paper Algorithm 1).
Faithful reimplementation of Algorithm 1 ("Real-Time Chunking") from
Black, Galliker, Levine. "Real-Time Execution of Action Chunking Flow
Policies." NeurIPS 2025. arXiv:2506.07339.
The controller decouples *execution* from *inference*:
* ``get_action(obs)`` is called by the robot every control period ``dt`` to
consume one action and provide the latest observation. It always returns
immediately (the real-time guarantee).
* a background ``inference_loop`` thread continuously generates the *next*
chunk while the current one executes, freezing the first ``d`` actions
(guaranteed to be executed before the new chunk is ready) and inpainting
the rest via RTC guided sampling.
The inference function itself is injected, so this controller is independent of
the policy/backend. For pi0.5 pass a closure around
:func:`rtc.pi0.rtc_sample_actions` (see ``examples/pi05_realtime.py``).
Notation (matching the paper)
------------------------------
H prediction horizon (chunk length)
s execution horizon; ``s = max(d, smin)`` for each chunk
smin minimum desired execution horizon (user hyperparameter)
d inference delay = #control steps between receiving o_t and A_t ready,
estimated *conservatively* as ``max`` of a buffer of recent delays
b delay-buffer size
t #actions consumed from the current chunk since the last inference start
"""
from __future__ import annotations
import threading
from collections import deque
from typing import Callable, Generic, TypeVar
# We keep the controller numerics-library agnostic: a "chunk" is anything
# indexable along its first axis (e.g. a numpy array or torch tensor of shape
# (H, A)). `InferenceFn(obs, prev_chunk, inference_delay, execution_horizon)`
# returns the next chunk in the *same* layout, aligned so that index 0
# corresponds to the same controller timestep as ``prev_chunk[s]``.
Obs = TypeVar("Obs")
Chunk = TypeVar("Chunk")
InferenceFn = Callable[[Obs, "Chunk", int, int], "Chunk"]
class RealTimeChunkingController(Generic[Obs, Chunk]):
"""Asynchronous RTC execution system (paper Algorithm 1).
Args:
inference_fn: ``f(obs, prev_chunk, inference_delay, execution_horizon)
-> next_chunk``. For pi0.5 wrap :func:`rtc.pi0.rtc_sample_actions`
with a single (un-batched along the controller axis) chunk.
horizon: ``H``, the chunk length (number of actions per chunk).
initial_chunk: ``A_init``, the first chunk to execute (already generated,
e.g. by a plain ``sample_actions`` call). Length must be ``H``.
smin: minimum execution horizon (paper real-world default 25).
initial_delay: ``d_init``, seed for the delay buffer.
delay_buffer_size: ``b``, number of recent delays kept for the (max)
delay estimate (paper real-world default 10).
"""
def __init__(
self,
inference_fn: InferenceFn,
horizon: int,
initial_chunk: Chunk,
*,
smin: int = 1,
initial_delay: int = 0,
delay_buffer_size: int = 10,
) -> None:
if len(initial_chunk) != horizon: # type: ignore[arg-type]
raise ValueError(f"initial_chunk length {len(initial_chunk)} != horizon {horizon}") # type: ignore[arg-type]
self._inference_fn = inference_fn
self._H = horizon
self._smin = smin
self._b = delay_buffer_size
# ---- mutex-protected shared state (paper INITIALIZESHAREDSTATE) ------
self._lock = threading.Lock()
self._cond = threading.Condition(self._lock) # condition variable C on M
self._t = 0 # #actions consumed from current chunk since last inference start
self._cur: Chunk = initial_chunk # A_cur
self._obs: Obs | None = None # o_cur
self._initial_delay = initial_delay
self._thread: threading.Thread | None = None
self._stop = False
self._last_observed_delay = initial_delay
# ------------------------------------------------------------------ API --
def start(self) -> None:
"""Launch the background inference loop thread."""
if self._thread is not None:
raise RuntimeError("controller already started")
self._stop = False
self._thread = threading.Thread(target=self._inference_loop, name="rtc-inference", daemon=True)
self._thread.start()
def stop(self) -> None:
"""Signal the inference loop to exit and join the thread."""
with self._lock:
self._stop = True
self._cond.notify_all()
if self._thread is not None:
self._thread.join()
self._thread = None
def get_action(self, next_obs: Obs, future_k: int = 0):
"""Consume one action and provide the latest observation (paper GETACTION).
Called by the controller every ``dt``. Returns immediately with the next
action to execute. Raises if the chunk is exhausted before a new one is
ready (the ``d <= H - s`` real-time constraint was violated).
When ``future_k > 0`` also returns ``(action, future)`` where ``future`` is a
copy of the current chunk from the returned action onward, up to ``future_k``
rows (fewer near the chunk end). Taken under the same lock as the action, so
the two are always from the same chunk even if the background loop swaps it.
Used by forward-looking waypoint interpolation (e.g. --bspline).
"""
with self._lock:
self._t += 1
self._obs = next_obs
self._cond.notify_all() # wake the inference loop (paper: notify C)
idx = self._t - 1
if idx >= self._H:
raise RuntimeError(
f"chunk exhausted (t={self._t} > H={self._H}): inference delay exceeded "
f"H - s; the real-time constraint d <= H - s was violated."
)
if future_k and future_k > 0:
import numpy as _np
return self._cur[idx], _np.asarray(self._cur[idx:idx + future_k]).copy()
return self._cur[idx]
@property
def last_observed_delay(self) -> int:
"""The most recently *measured* inference delay (in control steps)."""
return self._last_observed_delay
# ------------------------------------------------------------- internals --
def _inference_loop(self) -> None:
"""Background thread continuously generating the next chunk (paper INFERENCELOOP)."""
self._lock.acquire()
try:
# Buffer of recent inference delays; estimate the next delay as the
# buffer max (conservative), per Algorithm 1 line 17.
queue: deque[int] = deque([self._initial_delay], maxlen=self._b)
while True:
# Conservative delay estimate for the *upcoming* inference.
d = max(queue)
# Wait until enough actions have been consumed to start the next
# chunk. Algorithm 1 waits for t >= smin; the prose fixes the
# execution horizon at s = max(d, smin), so we wait for
# t >= max(smin, d) to guarantee the freeze prefix d <= s.
threshold = max(self._smin, d)
while not self._stop and (self._t < threshold or self._obs is None):
self._cond.wait()
if self._stop:
return
s = self._t # actions executed since last inference start
prev_chunk = self._cur # full current chunk (sliced inside inference_fn)
obs = self._obs
# Run inference with the lock released so get_action stays live.
self._lock.release()
try:
new_chunk = self._inference_fn(obs, prev_chunk, d, s)
finally:
self._lock.acquire()
# Swap in the new chunk as soon as it is available.
self._cur = new_chunk
# Reset t so it indexes into the new chunk. The number of actions
# consumed *during* inference is exactly the observed delay.
self._t -= s
observed_delay = self._t
self._last_observed_delay = observed_delay
queue.append(observed_delay)
finally:
if self._lock.locked():
self._lock.release()
def execution_horizon_for(delay_estimate: int, smin: int) -> int:
"""``s = max(d, smin)`` -- the per-chunk execution horizon (paper Sec. 3.3)."""
return max(delay_estimate, smin)
|