cleaned up, added docs and a few tests

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_data`                  | Stream telemetry to the visualization backend                     | false   |
-| `--display_ip` / `--display_port` | Remote Rerun server address                                       | --      |
+| `--display_mode`                  | Visualization backend: `rerun` or `foxglove`                      | rerun   |
-| `--interpolation_multiplier`      | Action interpolation factor                                       | 1       |
+| `--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

@@ -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

@@ -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,
+        # Intermediate-prediction visualization (e.g. a world model's imagined future). ``get_action``
-        # ``get_action`` requests predictions and keeps the current chunk's frame stacks; a playhead
+        # (once per policy tick) loads/advances the buffer; ``get_intermediate_predictions`` reads it
-        # (``get_intermediate_predictions``) advances one step per tick, paced across the chunk's tick
+        # for display. Pacing on policy ticks keeps playback aligned with execution regardless of
-        # span so the imagined clip stays wall-clock aligned with execution.
+        # action-interpolation upsampling.
         self._visualize_predictions = visualize_predictions
-        self._pred_stacks: dict = {}  # key -> [T, H, W, 3] frame stack for the current chunk
+        self._predictions = PacedPredictionBuffer(getattr(policy.config, "chunk_size", None))
-        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)
         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._predictions.reset()
-        self._pred_cursor = 0
 
     def get_intermediate_predictions(self) -> dict | None:
-        """Serve one imagined frame per key for this tick, advancing the playhead.
+        """Read the current chunk's paced prediction frame for display (``None`` if none)."""
-
+        return self._predictions.current()
-        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
 
     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.
+                    # A fresh chunk was predicted this tick — load its stacks and rewind the playhead.
-                    self._pred_stacks = predictions
+                    self._predictions.load(predictions)
-                    self._pred_cursor = 0
+                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

@@ -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())