rtc/controller.py

"""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)