From fb8c6ac8b9013febfb23b471b7e94b9495d4e162 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Fri, 17 Jul 2026 15:42:03 +0000 Subject: [PATCH] cleaned up, added docs and a few tests --- docs/source/inference.mdx | 11 ++- .../rollout/inference/prediction_buffer.py | 79 +++++++++++++++++++ src/lerobot/rollout/inference/sync.py | 49 ++++-------- tests/test_rollout.py | 60 ++++++++++++++ 4 files changed, 160 insertions(+), 39 deletions(-) create mode 100644 src/lerobot/rollout/inference/prediction_buffer.py diff --git a/docs/source/inference.mdx b/docs/source/inference.mdx index 31405b5de..7a7af7ab1 100644 --- a/docs/source/inference.mdx +++ b/docs/source/inference.mdx @@ -43,7 +43,7 @@ lerobot-rollout \ | ---------------- | ------------------------------------------------------ | | `--duration` | Run time in seconds (0 = infinite) | | `--task` | Task description passed to the policy | -| `--display_data` | Stream observations/actions to Rerun for visualization | +| `--display_data` | Stream observations/actions to the visualization backend (`--display_mode`; add `--display_extra_data` for a policy's imagined predictions) | ### Sentry (`--strategy.type=sentry`) @@ -251,9 +251,12 @@ See the [Real-Time Chunking](./rtc) guide for details on tuning RTC parameters. | `--duration` | Run time in seconds (0 = infinite) | 0 | | `--device` | Torch device (`cpu`, `cuda`, `mps`) | auto | | `--task` | Task description (used when no dataset is provided) | -- | -| `--display_data` | Stream telemetry to Rerun visualization | false | -| `--display_ip` / `--display_port` | Remote Rerun server address | -- | -| `--interpolation_multiplier` | Action interpolation factor | 1 | +| `--display_data` | Stream telemetry to the visualization backend | false | +| `--display_mode` | Visualization backend: `rerun` or `foxglove` | rerun | +| `--display_extra_data` | Also stream a policy's intermediate predictions (e.g. a world model's imagined video) on a dedicated channel; implies `--display_data`, sync inference only | false | +| `--display_compressed_images` | JPEG-compress images before streaming (less bandwidth, more CPU) | false | +| `--display_ip` / `--display_port` | Remote Rerun server address (or bind interface/port for foxglove) | -- | +| `--interpolation_multiplier` | Action interpolation factor (upsamples the control rate) | 1 | | `--use_torch_compile` | Enable `torch.compile` for inference | false | | `--resume` | Resume a previous recording session | false | | `--play_sounds` | Vocal synthesis for events | true | diff --git a/src/lerobot/rollout/inference/prediction_buffer.py b/src/lerobot/rollout/inference/prediction_buffer.py new file mode 100644 index 000000000..08c1c1954 --- /dev/null +++ b/src/lerobot/rollout/inference/prediction_buffer.py @@ -0,0 +1,79 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Paced playback of a policy's temporal predictions for visualization. + +Intermediate predictions (e.g. a world model's imagined clip) arrive as a ``{key: sequence}`` stack +with time on axis 0, one fresh stack per predicted chunk. Advancement is kept separate from reading +so the loop's two clocks don't get conflated: :meth:`advance` steps the playhead at policy cadence +(one new prediction per chunk), while :meth:`current` is a side-effect-free read for the display, +which may run faster (e.g. action-interpolated). Between policy ticks the viewer just holds the last +frame. Modality-agnostic: a step may be an image, vector, or scalar — only axis 0 is assumed to be +time. +""" + +from __future__ import annotations + + +class PacedPredictionBuffer: + """Maps a temporal prediction stack onto a chunk's policy ticks for paced visualization. + + Drive from policy inference (:meth:`load` a fresh chunk, :meth:`advance` on its later ticks); + read from the display loop (:meth:`current`). + """ + + def __init__(self, ticks_per_chunk: int | None = None) -> None: + # Policy ticks per predicted chunk; the T steps of each stack are spread across this span. + # ``None`` falls back to one step per policy tick (clamped to the stack length). + self._ticks_per_chunk = ticks_per_chunk + self._stacks: dict = {} # key -> sequence with time on axis 0, for the current chunk + self._cursor = 0 # policy ticks elapsed since the current chunk's stacks were loaded + + def reset(self) -> None: + """Drop the current chunk's stacks and rewind the playhead.""" + self._stacks = {} + self._cursor = 0 + + def load(self, stacks: dict) -> None: + """Store a freshly predicted chunk's stacks and rewind the playhead to its first step.""" + self._stacks = stacks + self._cursor = 0 + + def advance(self) -> None: + """Move the playhead forward one policy tick (no-op until a chunk is loaded).""" + if self._stacks: + self._cursor += 1 + + def current(self) -> dict | None: + """Current playhead step per key (``None`` until a chunk is loaded); no side effects. + + Maps each stack's ``T`` steps onto the ``ticks_per_chunk`` span (one step/tick, clamped, when + the span is unknown). + """ + if not self._stacks: + return None + tick = self._cursor + span = self._ticks_per_chunk + out: dict = {} + for key, stack in self._stacks.items(): + n = len(stack) + if n == 0: + continue + idx = round(tick / (span - 1) * (n - 1)) if span and span > 1 else tick + idx = min(max(idx, 0), n - 1) + step = stack[idx] + if hasattr(step, "detach"): # torch tensor -> numpy for the visualization backend + step = step.detach().cpu().numpy() + out[key] = step + return out or None diff --git a/src/lerobot/rollout/inference/sync.py b/src/lerobot/rollout/inference/sync.py index e2d35ac84..b4135ef73 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/src/lerobot/rollout/inference/sync.py @@ -27,6 +27,7 @@ from lerobot.policies.utils import make_robot_action, prepare_observation_for_in from lerobot.processor import PolicyProcessorPipeline from .base import InferenceEngine +from .prediction_buffer import PacedPredictionBuffer logger = logging.getLogger(__name__) @@ -75,14 +76,12 @@ class SyncInferenceEngine(InferenceEngine): self._device = torch.device(device or "cpu") self._robot_type = robot_type - # Intermediate-prediction visualization (e.g. a world model's imagined video). When on, - # ``get_action`` requests predictions and keeps the current chunk's frame stacks; a playhead - # (``get_intermediate_predictions``) advances one step per tick, paced across the chunk's tick - # span so the imagined clip stays wall-clock aligned with execution. + # Intermediate-prediction visualization (e.g. a world model's imagined future). ``get_action`` + # (once per policy tick) loads/advances the buffer; ``get_intermediate_predictions`` reads it + # for display. Pacing on policy ticks keeps playback aligned with execution regardless of + # action-interpolation upsampling. self._visualize_predictions = visualize_predictions - self._pred_stacks: dict = {} # key -> [T, H, W, 3] frame stack for the current chunk - self._pred_cursor = 0 # ticks elapsed since the current chunk's frames arrived - self._ticks_per_chunk = getattr(getattr(policy, "config", None), "chunk_size", None) + self._predictions = PacedPredictionBuffer(getattr(policy.config, "chunk_size", None)) logger.info( "SyncInferenceEngine initialized (device=%s, action_keys=%d, visualize_predictions=%s)", self._device, @@ -104,33 +103,11 @@ class SyncInferenceEngine(InferenceEngine): self._policy.reset() self._preprocessor.reset() self._postprocessor.reset() - self._pred_stacks = {} - self._pred_cursor = 0 + self._predictions.reset() def get_intermediate_predictions(self) -> dict | None: - """Serve one imagined frame per key for this tick, advancing the playhead. - - Maps the current chunk's ``T`` decoded frames onto its ``ticks_per_chunk`` control ticks so - the imagined video plays back in step with execution (falls back to one frame/tick, clamped, - when the chunk's tick span is unknown). Returns ``None`` until a chunk with frames arrives. - """ - if not self._pred_stacks: - return None - tick = self._pred_cursor - span = self._ticks_per_chunk - out: dict = {} - for key, stack in self._pred_stacks.items(): - n = len(stack) - if n == 0: - continue - idx = round(tick / (span - 1) * (n - 1)) if span and span > 1 else tick - idx = min(max(idx, 0), n - 1) - frame = stack[idx] - if hasattr(frame, "detach"): - frame = frame.detach().cpu().numpy() - out[key] = frame - self._pred_cursor += 1 - return out or None + """Read the current chunk's paced prediction frame for display (``None`` if none).""" + return self._predictions.current() def get_action(self, obs_frame: dict | None) -> torch.Tensor | None: """Run the full inference pipeline on ``obs_frame`` and return an action tensor.""" @@ -155,9 +132,11 @@ class SyncInferenceEngine(InferenceEngine): self._policy.select_action(observation, return_intermediate_predictions=True) ) if predictions: - # A fresh chunk was predicted this tick — store its frame stacks and restart the playhead. - self._pred_stacks = predictions - self._pred_cursor = 0 + # A fresh chunk was predicted this tick — load its stacks and rewind the playhead. + self._predictions.load(predictions) + else: + # Same chunk, later policy tick — step the playhead through the imagined frames. + self._predictions.advance() else: action = self._policy.select_action(observation) action = self._postprocessor(action) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 85a29ff4c..d271220cb 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -19,6 +19,7 @@ from __future__ import annotations import dataclasses from unittest.mock import MagicMock +import numpy as np import pytest import torch @@ -348,3 +349,62 @@ def test_rollout_context_fields(): field_names = {f.name for f in dataclasses.fields(RolloutContext)} assert field_names == {"runtime", "hardware", "policy", "processors", "data"} + + +# --------------------------------------------------------------------------- +# Paced prediction buffer +# --------------------------------------------------------------------------- + + +def _walk_policy_ticks(buffer, key, n_ticks): + """Read the buffer over ``n_ticks`` policy ticks (load counts as the first).""" + seen = [buffer.current()[key]] + for _ in range(n_ticks - 1): + buffer.advance() + seen.append(buffer.current()[key]) + return seen + + +def test_prediction_buffer_paces_over_chunk(): + from lerobot.rollout.inference.prediction_buffer import PacedPredictionBuffer + + buffer = PacedPredictionBuffer(ticks_per_chunk=4) + assert buffer.current() is None # nothing until a chunk is loaded + + # 2 steps spread over a 4-policy-tick chunk: first half -> step 0, second half -> step 1. + buffer.load({"x": ["a", "b"]}) + assert _walk_policy_ticks(buffer, "x", 4) == ["a", "a", "b", "b"] + + buffer.reset() + assert buffer.current() is None + + # Unknown span -> one step per policy tick, clamped at the last step. + unpaced = PacedPredictionBuffer(ticks_per_chunk=None) + unpaced.load({"x": ["a", "b"]}) + assert _walk_policy_ticks(unpaced, "x", 3) == ["a", "b", "b"] + + +def test_prediction_buffer_read_is_independent_of_display_rate(): + from lerobot.rollout.inference.prediction_buffer import PacedPredictionBuffer + + # Pacing fix: the playhead advances per policy tick; the display may read it many times per tick + # (e.g. interpolation multiplier=3), each read returning the same frame with no side effects — so + # playback never races to the end and freezes. + buffer = PacedPredictionBuffer(ticks_per_chunk=2) + buffer.load({"x": ["a", "b"]}) + assert [buffer.current()["x"] for _ in range(3)] == ["a", "a", "a"] # policy tick 0 + buffer.advance() + assert [buffer.current()["x"] for _ in range(3)] == ["b", "b", "b"] # policy tick 1 + + +def test_prediction_buffer_is_modality_agnostic(): + from lerobot.rollout.inference.prediction_buffer import PacedPredictionBuffer + + # Non-image temporal predictions pass through; torch tensors are detached to numpy on the way out. + buffer = PacedPredictionBuffer(ticks_per_chunk=2) + states = torch.arange(6).reshape(2, 3) # [T=2, D=3] + buffer.load({"observation.state": states}) + out = buffer.current()["observation.state"] + assert isinstance(out, np.ndarray) and np.array_equal(out, states[0].numpy()) + buffer.advance() + assert np.array_equal(buffer.current()["observation.state"], states[1].numpy())