admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 09:46:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

cleaned up, added docs and a few tests

Browse Source
This commit is contained in:
Maxime Ellerbach
2026-07-17 15:42:03 +00:00
parent 4ae2fbca36
commit fb8c6ac8b9
4 changed files with 160 additions and 39 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/inference.mdx
+7 -4
View File
@@ -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 |
src/lerobot/rollout/inference/prediction_buffer.py
+79
View File
@@ -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
src/lerobot/rollout/inference/sync.py
+14 -35
View File
@@ -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)
tests/test_rollout.py
+60
View File
@@ -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())