draft for unifying prediction visualization

src/lerobot/configs/default.py

@@ -93,6 +93,9 @@ class EvalConfig:
     recording_repo_id: str | None = None
     # Whether the pushed recording repositories should be private.
     recording_private: bool = False
+    # Whether to save the policy's imagined/predicted video (world-model policies only) as mp4s.
+    # Requests intermediate predictions from the policy each step; policies that produce none are unaffected.
+    save_predicted_video: bool = False
 
     def __post_init__(self) -> None:
         if self.recording_repo_id is not None and not self.recording:

src/lerobot/policies/lingbot_va/configuration_lingbot_va.py

@@ -92,9 +92,6 @@ class LingBotVAConfig(PreTrainedConfig):
     # (un)normalization quantiles live in the checkpoint's ``policy_postprocessor.json``, not here.
     used_action_channel_ids: list[int] = field(default_factory=lambda: list(range(7)))
 
-    # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s.
-    save_predicted_video: bool = False
-
     # Normalization: IDENTITY here; images are scaled + VAE-encoded and actions are
     # quantile-(un)normalized inside the policy / dedicated processor steps.
     normalization_mapping: dict[str, NormalizationMode] = field(

src/lerobot/policies/lingbot_va/modeling_lingbot_va.py

@@ -38,7 +38,7 @@ import torch.nn.functional as F  # noqa: N812
 from einops import rearrange
 from torch import Tensor
 
-from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.pretrained import PreTrainedPolicy, unpack_action_output
 from lerobot.utils.constants import ACTION
 from lerobot.utils.import_utils import require_package
 
@@ -99,8 +99,6 @@ class LingBotVAPolicy(PreTrainedPolicy):
         # from ``config.wan_pretrained_path`` the first time inference runs.
         self._frozen: dict = {}
 
-        self.last_predicted_frames: Tensor | None = None
-        self.last_predicted_latents: Tensor | None = None
         self.reset()
 
     # Frozen-module lazy loading (VAE + UMT5 + tokenizer)
@@ -170,8 +168,6 @@ class LingBotVAPolicy(PreTrainedPolicy):
         self._prompt: str | None = None
         self._prompt_embeds = None
         self._negative_prompt_embeds = None
-        self.last_predicted_frames = None
-        self.last_predicted_latents = None
         self._use_cfg = (cfg.guidance_scale > 1) or (cfg.action_guidance_scale > 1)
         # Two independent flow-matching schedulers (video latent + action streams).
         self._scheduler = FlowMatchScheduler(shift=cfg.snr_shift, sigma_min=0.0, extra_one_step=True)
@@ -400,22 +396,31 @@ class LingBotVAPolicy(PreTrainedPolicy):
         return torch.cat(per_cam, dim=-1).to(self.config.device)
 
     @torch.no_grad()
-    def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
+    def select_action(
+        self, batch: dict[str, Tensor], return_intermediate_predictions: bool = False, **kwargs
+    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
         """Return one action, refilling the chunk (and feeding back observed keyframes) as needed.
 
         Mirrors the upstream LIBERO client loop (``evaluation/libero/client.py``): the first obs is
         the conditioning frame; every observation produced afterwards is buffered as a keyframe and,
         once the chunk's actions are exhausted, the buffered frames + executed actions are fed back
         into the KV cache before the next chunk is predicted.
+
+        When ``return_intermediate_predictions=True`` returns ``(action, predictions)``. Predictions
+        are produced only on the ticks that predict a fresh chunk (first tick and each chunk refill);
+        on the intermediate ticks that just pop a cached action, ``predictions`` is an empty dict.
         """
         self.eval()
         self._ensure_frozen_modules()
         self._maybe_init_prompt(batch)
 
+        predictions: dict[str, Tensor] = {}
         if not self._started:
             # First call: this observation conditions the first chunk (it is *not* a keyframe).
             self._started = True
-            actions = self.predict_action_chunk(batch)  # [B, chunk_size, n_used]
+            actions, predictions = unpack_action_output(
+                self.predict_action_chunk(batch, return_intermediate_predictions=return_intermediate_predictions)
+            )  # [B, chunk_size, n_used]
             self._action_queue.extend(actions.transpose(0, 1))  # [chunk_size, B, n_used]
             self._obs_buffer = []
             self._exec_step = 0
@@ -427,17 +432,31 @@ class LingBotVAPolicy(PreTrainedPolicy):
             if len(self._action_queue) == 0:
                 # All actions for the current chunk have been executed; feed the observed
                 # keyframes + executed actions back and predict the next chunk.
-                actions = self.predict_action_chunk(None)
+                actions, predictions = unpack_action_output(
+                    self.predict_action_chunk(
+                        None, return_intermediate_predictions=return_intermediate_predictions
+                    )
+                )
                 self._action_queue.extend(actions.transpose(0, 1))
                 self._exec_step = 0
 
         self._prev_j = self._exec_step % self.config.action_per_frame
         self._exec_step += 1
-        return self._action_queue.popleft()
+        action = self._action_queue.popleft()
+        if return_intermediate_predictions:
+            return action, predictions
+        return action
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
-        """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized)."""
+    def predict_action_chunk(
+        self, batch: dict[str, Tensor], return_intermediate_predictions: bool = False, **kwargs
+    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
+        """Run one autoregressive chunk and return actions ``[B, chunk_size, n_used]`` (normalized).
+
+        When ``return_intermediate_predictions=True`` returns ``(actions, predictions)`` where
+        ``predictions`` holds this chunk's VAE-decoded imagined video under ``"images.predicted"``
+        (``[T, H, W, 3]`` uint8 on CPU).
+        """
         self.eval()
         self._ensure_frozen_modules()
         self._maybe_init_prompt(batch)
@@ -459,12 +478,6 @@ class LingBotVAPolicy(PreTrainedPolicy):
         # actions: [B, action_dim, F, action_per_frame, 1] (model-normalized). Keep for KV feedback.
         self._executed_actions = actions
 
-        if self.config.save_predicted_video:
-            # Match upstream LingBot-VA visualization: collect chunk latents and decode the
-            # concatenated latent sequence once after the rollout finishes.
-            self.last_predicted_frames = None
-            self.last_predicted_latents = latents.detach().to("cpu")
-
         # On the first chunk, frame 0 is the conditioning frame (already "known"): the upstream
         # LIBERO client skips it (start_idx=1), so we drop the first frame's actions here.
         used = self.config.used_action_channel_ids
@@ -473,7 +486,15 @@ class LingBotVAPolicy(PreTrainedPolicy):
             a = a[:, :, 1:]  # drop frame 0 -> (F-1) frames of actions
         a = a.squeeze(-1).flatten(2)  # [B, n_used, n_steps]
         a = a.transpose(1, 2).contiguous()  # [B, n_steps, n_used]
-        return a.to(torch.float32)
+        a = a.to(torch.float32)
+
+        if return_intermediate_predictions:
+            # Decode this chunk's imagined video for visualization / eval. Per-chunk decode (the VAE
+            # has no streaming decoder) may differ slightly at chunk boundaries from a single decode
+            # over the whole concatenated latent sequence; acceptable for monitoring/inspection.
+            frames = self._decode_predicted_video(latents)  # [T, H, W, 3] uint8, CPU
+            return a, {"images.predicted": frames}
+        return a
 
     # Prompt / text encoding
     def _maybe_init_prompt(self, batch):
@@ -834,11 +855,6 @@ class LingBotVAPolicy(PreTrainedPolicy):
         return actions, latents
 
     # Predicted-video decoding (opt-in)
-    @torch.no_grad()
-    def decode_predicted_latents(self, latents) -> Tensor:
-        """Decode a concatenated predicted-latent sequence into ``[T, H, W, 3]`` uint8 frames."""
-        return self._decode_predicted_video(latents)
-
     @torch.no_grad()
     def _decode_predicted_video(self, latents) -> Tensor:
         """VAE-decode predicted latents into a uint8 frame stack ``[T, H, W, 3]`` on CPU."""

src/lerobot/policies/pretrained.py

@@ -93,6 +93,18 @@ def _build_card_context(
 
 class ActionSelectKwargs(TypedDict, total=False):
     noise: Tensor | None
+    return_intermediate_predictions: bool
+
+
+def unpack_action_output(out: Tensor | tuple[Tensor, dict[str, Tensor]]) -> tuple[Tensor, dict[str, Tensor]]:
+    """Normalize a ``select_action`` / ``predict_action_chunk`` return to ``(action, predictions)``.
+
+    These methods return a bare action ``Tensor`` by default, or a ``(action, predictions)`` tuple when
+    called with ``return_intermediate_predictions=True``. A bare tensor becomes ``(tensor, {})``.
+    """
+    if isinstance(out, tuple):
+        return out[0], out[1]
+    return out, {}
 
 
 class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
@@ -273,20 +285,34 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
+    def predict_action_chunk(
+        self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]
+    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
         """Returns the action chunk (for action chunking policies) for a given observation, potentially in batch mode.
 
         Child classes using action chunking should use this method within `select_action` to form the action chunk
         cached for selection.
+
+        By default returns just the action `Tensor`. If `return_intermediate_predictions=True`,
+        returns `(action, predictions)` where `predictions` is a (possibly empty) `dict[str, Tensor]`
+        of additional model predictions a policy may expose (e.g. world-model predicted frames).
+        Policies that produce nothing extra may ignore the kwarg.
         """
         raise NotImplementedError
 
     @abc.abstractmethod
-    def select_action(self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]) -> Tensor:
+    def select_action(
+        self, batch: dict[str, Tensor], **kwargs: Unpack[ActionSelectKwargs]
+    ) -> Tensor | tuple[Tensor, dict[str, Tensor]]:
         """Return one action to run in the environment (potentially in batch mode).
 
         When the model uses a history of observations, or outputs a sequence of actions, this method deals
         with caching.
+
+        By default returns just the action `Tensor`. If `return_intermediate_predictions=True`,
+        returns `(action, predictions)` where `predictions` is a (possibly empty) `dict[str, Tensor]`
+        of additional model predictions a policy may expose (e.g. world-model predicted frames).
+        Policies that produce nothing extra may ignore the kwarg.
         """
         raise NotImplementedError
 

src/lerobot/rollout/configs.py

@@ -226,6 +226,10 @@ class RolloutConfig:
     device: str | None = None
     task: str = ""
     display_data: bool = False
+    # Also visualize model "extras" (e.g. a world model's imagined video) alongside observations.
+    # Off by default: requesting predictions forces per-chunk decoding on the control thread and only
+    # world-model policies produce anything. Implies display_data. Sync inference only.
+    display_extra_data: bool = False
     # Visualization backend used when display_data is True: "rerun" or "foxglove".
     display_mode: str = "rerun"
     # For "rerun": IP of a remote server to send to. For "foxglove": interface to bind the WebSocket
@@ -255,6 +259,21 @@ class RolloutConfig:
 
     def __post_init__(self):
         """Validate config invariants and load the policy config from ``--policy.path``."""
+        # --- Visualization validation ---
+        # Extra-data visualization piggybacks on the display_data path (backend init + telemetry
+        # logging are both gated on display_data), so enabling it implies display_data.
+        if self.display_extra_data and not self.display_data:
+            logger.info("display_extra_data=True implies display_data=True; enabling display_data")
+            self.display_data = True
+        # Only the sync engine surfaces intermediate predictions (RTC runs the policy in a background
+        # thread); warn and let it be ignored rather than fail.
+        if self.display_extra_data and not isinstance(self.inference, SyncInferenceConfig):
+            logger.warning(
+                "display_extra_data is only supported with sync inference (--inference.type=sync); "
+                "it will be ignored for inference type '%s'",
+                self.inference.type,
+            )
+
         # --- Strategy-specific validation ---
         if isinstance(self.strategy, DAggerStrategyConfig) and self.teleop is None:
             raise ValueError("DAgger strategy requires --teleop.type to be set")

src/lerobot/rollout/context.py

@@ -429,6 +429,7 @@ def build_rollout_context(
         use_torch_compile=cfg.use_torch_compile,
         compile_warmup_inferences=cfg.compile_warmup_inferences,
         shutdown_event=shutdown_event,
+        visualize_predictions=cfg.display_extra_data,
     )
 
     # --- 8. Assemble ---------------------------------------------------

src/lerobot/rollout/inference/base.py

@@ -69,6 +69,15 @@ class InferenceEngine(abc.ABC):
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Return the next action tensor, or ``None`` if unavailable."""
 
+    def get_intermediate_predictions(self) -> dict | None:
+        """Extra display-ready model outputs to visualize this tick, or ``None``.
+
+        Lets a backend surface a world model's intermediate predictions (e.g. imagined video
+        frames) into the rollout visualization path, keyed by ``"<datatype>.<name>"`` (mirroring
+        observation feature keys). Default: nothing extra.
+        """
+        return None
+
     def notify_observation(self, obs: dict) -> None:  # noqa: B027
         """Publish the latest processed observation.  Default: no-op."""
 

src/lerobot/rollout/inference/factory.py

@@ -95,6 +95,7 @@ def create_inference_engine(
     use_torch_compile: bool = False,
     compile_warmup_inferences: int = 2,
     shutdown_event: Event | None = None,
+    visualize_predictions: bool = False,
 ) -> InferenceEngine:
     """Instantiate the appropriate inference engine from a config object."""
     logger.info("Creating inference engine: %s", config.type)
@@ -108,6 +109,7 @@ def create_inference_engine(
             task=task,
             device=device,
             robot_type=robot_wrapper.robot_type,
+            visualize_predictions=visualize_predictions,
         )
     if isinstance(config, RTCInferenceConfig):
         return RTCInferenceEngine(

src/lerobot/rollout/inference/sync.py

@@ -22,7 +22,7 @@ from copy import copy
 
 import torch
 
-from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.pretrained import PreTrainedPolicy, unpack_action_output
 from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
 from lerobot.processor import PolicyProcessorPipeline
 
@@ -64,6 +64,7 @@ class SyncInferenceEngine(InferenceEngine):
         task: str,
         device: str | None,
         robot_type: str,
+        visualize_predictions: bool = False,
     ) -> None:
         self._policy = policy
         self._preprocessor = preprocessor
@@ -73,10 +74,20 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         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.
+        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)
         logger.info(
-            "SyncInferenceEngine initialized (device=%s, action_keys=%d)",
+            "SyncInferenceEngine initialized (device=%s, action_keys=%d, visualize_predictions=%s)",
             self._device,
             len(ordered_action_keys),
+            self._visualize_predictions,
         )
 
     def start(self) -> None:
@@ -93,6 +104,33 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
+        self._pred_stacks = {}
+        self._pred_cursor = 0
+
+    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
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
         """Run the full inference pipeline on ``obs_frame`` and return an action tensor."""
@@ -112,7 +150,16 @@ class SyncInferenceEngine(InferenceEngine):
                 observation, self._device, self._task, self._robot_type
             )
             observation = self._preprocessor(observation)
-            action = self._policy.select_action(observation)
+            if self._visualize_predictions:
+                action, predictions = unpack_action_output(
+                    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
+            else:
+                action = self._policy.select_action(observation)
             action = self._postprocessor(action)
         action_tensor = action.squeeze(0).cpu()
 

src/lerobot/rollout/strategies/core.py

@@ -156,8 +156,8 @@ class RolloutStrategy(abc.ABC):
         except Exception as e:
             logger.warning("Could not return to initial position: %s", e)
 
-    @staticmethod
     def _log_telemetry(
+        self,
         obs_processed: dict | None,
         action_dict: dict | None,
         runtime_ctx: RuntimeContext,
@@ -166,10 +166,16 @@ class RolloutStrategy(abc.ABC):
         cfg = runtime_ctx.cfg
         if not cfg.display_data:
             return
+        # When extra-data visualization is on, pull any display-ready model predictions from the
+        # engine (e.g. a world model's imagined video) and log them on the dedicated prediction channel.
+        prediction = None
+        if cfg.display_extra_data and self._engine is not None:
+            prediction = self._engine.get_intermediate_predictions()
         log_visualization_data(
             cfg.display_mode,
             observation=obs_processed,
             action=action_dict,
+            prediction=prediction,
             compress_images=cfg.display_compressed_images,
         )
 

src/lerobot/scripts/lerobot_eval.py

@@ -83,6 +83,7 @@ from lerobot.envs import (
     preprocess_observation,
 )
 from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors
+from lerobot.policies.pretrained import unpack_action_output
 from lerobot.processor import PolicyProcessorPipeline
 from lerobot.types import PolicyAction
 from lerobot.utils.constants import ACTION, DONE, OBS_IMAGE, OBS_IMAGES, OBS_STR, REWARD
@@ -169,7 +170,7 @@ def rollout(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
-    predicted_latents_callback: Callable[[PreTrainedPolicy], None] | None = None,
+    save_predicted_video: bool = False,
 ) -> dict:
     """Run a batched policy rollout once through a batch of environments.
 
@@ -199,9 +200,10 @@ def rollout(
             are returned optionally because they typically take more memory to cache. Defaults to False.
         render_callback: Optional rendering callback to be used after the environments are reset, and after
             every step.
-        predicted_latents_callback: Optional callback invoked after every ``select_action`` with the policy
-            itself. World-model policies (e.g. LingBot-VA) stash predicted video latents on
-            ``policy.last_predicted_latents``; this lets the caller concatenate chunks and decode once.
+        save_predicted_video: When True, request intermediate predictions from the policy each step
+            (``select_action(..., return_intermediate_predictions=True)``) and collect any imagined
+            video frames a world-model policy returns. Collected per image key in
+            ``ret["predicted_frames"]`` as a list of ``[T, H, W, 3]`` uint8 chunk stacks.
     Returns:
         The dictionary described above.
     """
@@ -245,6 +247,9 @@ def rollout(
     all_rewards = []
     all_successes = []
     all_dones = []
+    # Imagined-video frames returned by world-model policies, collected per image key. Each entry is
+    # a chunk stack [T, H, W, 3] uint8; concatenated on the time axis by the caller.
+    predicted_frames: dict[str, list] = {}
 
     step = 0
     # Keep track of which environments are done.
@@ -279,9 +284,13 @@ def rollout(
 
             observation = preprocessor(observation)
             with torch.inference_mode():
-                action = policy.select_action(observation)
-            if predicted_latents_callback is not None:
-                predicted_latents_callback(policy)
+                extra = {"return_intermediate_predictions": True} if save_predicted_video else {}
+                action, predictions = unpack_action_output(policy.select_action(observation, **extra))
+            # World-model policies return imagined frames only on chunk-boundary ticks; collect them.
+            for key, frames in predictions.items():
+                if hasattr(frames, "detach"):
+                    frames = frames.detach().to("cpu")
+                predicted_frames.setdefault(key, []).append(frames)
             action = postprocessor(action)
 
             action_transition = {ACTION: action}
@@ -394,6 +403,9 @@ def rollout(
             stacked_observations[key] = torch.stack([obs[key] for obs in all_observations], dim=1)
         ret[OBS_STR] = stacked_observations
 
+    if save_predicted_video:
+        ret["predicted_frames"] = predicted_frames
+
     if hasattr(policy, "use_original_modules"):
         policy.use_original_modules()
 
@@ -435,11 +447,6 @@ def eval_policy(
     if max_episodes_rendered > 0 and not videos_dir:
         raise ValueError("If max_episodes_rendered > 0, videos_dir must be provided.")
 
-    # World-model policies (e.g. LingBot-VA) opt into predicted-video saving via their config.
-    save_predicted_video = save_predicted_video or bool(
-        getattr(getattr(policy, "config", None), "save_predicted_video", False)
-    )
-
     if not isinstance(policy, PreTrainedPolicy):
         exc = ValueError(
             f"Policy of type 'PreTrainedPolicy' is expected, but type '{type(policy)}' was provided."
@@ -489,16 +496,6 @@ def eval_policy(
         predicted_video_paths: list[str] = []
         n_predicted_rendered = 0
 
-    # Collect predicted-video latents across a rollout (world-model policies only). The latents are
-    # concatenated and decoded once after the rollout, matching upstream LingBot-VA's visualization path.
-    def collect_predicted_latents(policy: PreTrainedPolicy):
-        latents = getattr(policy, "last_predicted_latents", None)
-        if latents is not None:
-            pred_latents.append(
-                latents.detach().to("cpu") if hasattr(latents, "detach") else torch.as_tensor(latents).cpu()
-            )
-            policy.last_predicted_latents = None
-
     if return_episode_data:
         episode_data: dict | None = None
 
@@ -510,9 +507,6 @@ def eval_policy(
         if max_episodes_rendered > 0:
             ep_frames: list[np.ndarray] = []
 
-        if save_predicted_video:
-            pred_latents: list[torch.Tensor] = []
-
         if start_seed is None:
             seeds = None
         else:
@@ -533,7 +527,7 @@ def eval_policy(
             env_features=env_features,
             recording_repo_id=recording_repo_id,
             recording_private=recording_private,
-            predicted_latents_callback=collect_predicted_latents if save_predicted_video else None,
+            save_predicted_video=save_predicted_video,
         )
 
         # Figure out where in each rollout sequence the first done condition was encountered (results after
@@ -599,33 +593,33 @@ def eval_policy(
                 threads.append(thread)
                 n_episodes_rendered += 1
 
-        # Maybe save the policy's predicted (imagined) video for this batch's rollout.
-        if save_predicted_video and len(pred_latents) > 0:
-            predicted_latent = torch.cat(pred_latents, dim=2)
-            decoder = getattr(policy, "decode_predicted_latents", None) or getattr(
-                policy, "_decode_predicted_video", None
-            )
-            if decoder is None:
-                raise AttributeError(
-                    "Policy config requested predicted-video saving, but the policy does not expose "
-                    "`decode_predicted_latents` or `_decode_predicted_video`."
-                )
-            predicted_video = decoder(predicted_latent)
-            if hasattr(predicted_video, "detach"):
-                predicted_video = predicted_video.detach().to("cpu").numpy()
+        # Maybe save the policy's predicted (imagined) video for this batch's rollout. The policy
+        # returns display-ready frame stacks per image key; concatenate them on the time axis and
+        # write one mp4 per key (no decoding here — the policy already decoded).
+        pred_frames = rollout_data.get("predicted_frames", {}) if save_predicted_video else {}
+        if save_predicted_video and any(len(stacks) > 0 for stacks in pred_frames.values()):
             videos_dir.mkdir(parents=True, exist_ok=True)
-            predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}.mp4"
-            predicted_video_paths.append(str(predicted_video_path))
-            thread = threading.Thread(
-                target=write_video,
-                args=(
-                    str(predicted_video_path),
-                    predicted_video,
-                    env.unwrapped.metadata["render_fps"],
-                ),
-            )
-            thread.start()
-            threads.append(thread)
+            multi_key = len(pred_frames) > 1
+            for key, stacks in pred_frames.items():
+                if len(stacks) == 0:
+                    continue
+                predicted_video = torch.cat(
+                    [s if hasattr(s, "dim") else torch.as_tensor(s) for s in stacks], dim=0
+                )
+                predicted_video = predicted_video.detach().to("cpu").numpy()  # [T, H, W, 3] uint8
+                suffix = f"_{key.replace('.', '_')}" if multi_key else ""
+                predicted_video_path = videos_dir / f"pred_episode_{n_predicted_rendered}{suffix}.mp4"
+                predicted_video_paths.append(str(predicted_video_path))
+                thread = threading.Thread(
+                    target=write_video,
+                    args=(
+                        str(predicted_video_path),
+                        predicted_video,
+                        env.unwrapped.metadata["render_fps"],
+                    ),
+                )
+                thread.start()
+                threads.append(thread)
             n_predicted_rendered += 1
 
         progbar.set_postfix(
@@ -771,6 +765,11 @@ def eval_main(cfg: EvalPipelineConfig):
     recording_dir = Path(cfg.output_dir) / "recordings" if cfg.eval.recording else None
     max_episodes_rendered = 0 if cfg.eval.recording else 10
     videos_dir = None if cfg.eval.recording else Path(cfg.output_dir) / "videos"
+    # Predicted-video saving needs a directory to write mp4s into; recording mode leaves videos_dir
+    # unset, so provide one explicitly.
+    save_predicted_video = cfg.eval.save_predicted_video
+    if save_predicted_video and videos_dir is None:
+        videos_dir = Path(cfg.output_dir) / "videos"
 
     with torch.no_grad(), torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext():
         info = eval_policy_all(
@@ -790,6 +789,7 @@ def eval_main(cfg: EvalPipelineConfig):
             env_features=cfg.env.features if cfg.eval.recording else None,
             recording_repo_id=cfg.eval.recording_repo_id,
             recording_private=cfg.eval.recording_private,
+            save_predicted_video=save_predicted_video,
         )
         print("Overall Aggregated Metrics:")
         print(info["overall"])
@@ -837,6 +837,7 @@ def eval_one(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
+    save_predicted_video: bool = False,
 ) -> TaskMetrics:
     """Evaluates one task_id of one suite using the provided vec env."""
 
@@ -858,6 +859,7 @@ def eval_one(
         env_features=env_features,
         recording_repo_id=recording_repo_id,
         recording_private=recording_private,
+        save_predicted_video=save_predicted_video,
     )
 
     per_episode = task_result["per_episode"]
@@ -889,6 +891,7 @@ def run_one(
     env_features: dict | None = None,
     recording_repo_id: str | None = None,
     recording_private: bool = False,
+    save_predicted_video: bool = False,
 ):
     """
     Run eval_one for a single (task_group, task_id, env).
@@ -923,6 +926,7 @@ def run_one(
         env_features=env_features,
         recording_repo_id=task_repo_id,
         recording_private=recording_private,
+        save_predicted_video=save_predicted_video,
     )
 
     if max_episodes_rendered > 0:
@@ -949,6 +953,7 @@ def eval_policy_all(
     return_episode_data: bool = False,
     start_seed: int | None = None,
     max_parallel_tasks: int = 1,
+    save_predicted_video: bool = False,
 ) -> dict:
     """
     Evaluate a nested `envs` dict: {task_group: {task_id: vec_env}}.
@@ -1008,6 +1013,7 @@ def eval_policy_all(
         env_features=env_features,
         recording_repo_id=recording_repo_id,
         recording_private=recording_private,
+        save_predicted_video=save_predicted_video,
     )
 
     if max_parallel_tasks <= 1:

src/lerobot/utils/constants.py

@@ -30,6 +30,9 @@ OBS_LANGUAGE_SUBTASK = OBS_STR + ".subtask"
 OBS_LANGUAGE_SUBTASK_TOKENS = OBS_LANGUAGE_SUBTASK + ".tokens"
 OBS_LANGUAGE_SUBTASK_ATTENTION_MASK = OBS_LANGUAGE_SUBTASK + ".attention_mask"
 
+PREDICTION_STR = "prediction"
+PREDICTION_PREFIX = PREDICTION_STR + "."
+
 ACTION = "action"
 ACTION_PREFIX = ACTION + "."
 ACTION_TOKENS = ACTION + ".tokens"

src/lerobot/utils/foxglove_visualization.py

@@ -37,6 +37,8 @@ from .constants import (
     OBS_PREFIX,
     OBS_STATE,
     OBS_STR,
+    PREDICTION_PREFIX,
+    PREDICTION_STR,
     REWARD,
     SUCCESS,
     TRUNCATED,
@@ -283,10 +285,11 @@ def _log_foxglove_image(
 def log_foxglove_data(
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
+    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
     """
-    Logs observation and action data to a Foxglove WebSocket server for real-time visualization.
+    Logs observation, action and prediction data to a Foxglove WebSocket server for real-time visualization.
 
     Mirrors ``log_rerun_data`` but emits Foxglove messages over the server started by
     :func:`init_foxglove`. Data is mapped as follows:
@@ -302,6 +305,8 @@ def log_foxglove_data(
     Args:
         observation: An optional dictionary containing observation data to log.
         action: An optional dictionary containing action data to log.
+        prediction: An optional dictionary of display-ready model outputs (e.g. a world model's
+            imagined video), keyed "<datatype>.<name>", logged on ``/prediction/...`` topics.
         compress_images: Whether to JPEG-compress images before logging to save bandwidth in exchange
             for CPU and quality.
     """
@@ -334,6 +339,30 @@ def log_foxglove_data(
                     )
         _log_foxglove_scalars(_foxglove_topic(OBS_STATE), obs_scalars, log_time=now)
 
+    if prediction:
+        # Predicted outputs are keyed "<datatype>.<name>" (e.g. "images.predicted"); route images to
+        # /prediction/images/<name> and any scalars to an aggregate /prediction/state topic.
+        pred_scalars: dict[str, float] = {}
+        for k, v in prediction.items():
+            if v is None:
+                continue
+            key = k[len(PREDICTION_PREFIX) :] if str(k).startswith(PREDICTION_PREFIX) else str(k)
+            if _is_scalar(v):
+                pred_scalars[key] = float(v)
+            elif isinstance(v, np.ndarray):
+                if v.ndim == 1:
+                    pred_scalars.update(_labeled_scalars(key, v))
+                else:
+                    name = key[len("images.") :] if key.startswith("images.") else key
+                    _log_foxglove_image(
+                        f"/{PREDICTION_STR}/images/{_foxglove_safe_name(name)}",
+                        name,
+                        v,
+                        compress_images=compress_images,
+                        log_time=now,
+                    )
+        _log_foxglove_scalars(f"/{PREDICTION_STR}/state", pred_scalars, log_time=now)
+
     if action:
         action_scalars: dict[str, float] = {}
         for k, v in action.items():

src/lerobot/utils/rerun_visualization.py

@@ -27,7 +27,7 @@ import numpy as np
 from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit
 from lerobot.types import RobotAction, RobotObservation
 
-from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR
+from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, PREDICTION_PREFIX
 from .import_utils import require_package
 
 
@@ -37,6 +37,43 @@ def _is_scalar(x):
     )
 
 
+def _log_scalar_or_image_mapping(rr, data, prefix, scalar_paths, image_paths, compress_images):
+    """Log a mapping of scalars/images (observation- or prediction-style) under ``prefix``.
+
+    Scalars and 1D arrays go to ``scalar_paths`` (time-series); 2D/3D arrays are treated as images
+    (CHW->HWC as needed, depth for single-channel) and go to ``image_paths`` (spatial views).
+    """
+    for k, v in data.items():
+        if v is None:
+            continue
+        key = str(k) if str(k).startswith(prefix) else f"{prefix}{k}"
+
+        if _is_scalar(v):
+            rr.log(key, rr.Scalars(float(v)))
+            scalar_paths.add(key)
+        elif isinstance(v, np.ndarray):
+            arr = v
+            # Convert CHW -> HWC when needed
+            if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
+                arr = np.transpose(arr, (1, 2, 0))
+            if arr.ndim == 1:
+                rr.log(key, rr.Scalars(arr.astype(float)))
+                scalar_paths.add(key)
+            else:
+                if arr.shape[-1] == 1:
+                    # At record time, the depth unit is inferred from the frame type.
+                    depth_unit = infer_depth_unit(arr.dtype)
+                    img_entity = rr.DepthImage(
+                        arr,
+                        meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0,
+                        colormap=rr.components.Colormap.Viridis,
+                    )
+                else:
+                    img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
+                rr.log(key, entity=img_entity, static=True)
+                image_paths.add(key)
+
+
 def init_rerun(
     session_name: str = "lerobot_control_loop", ip: str | None = None, port: int | None = None
 ) -> None:
@@ -73,10 +110,16 @@ def shutdown_rerun() -> None:
     rr.rerun_shutdown()
 
 
-def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]):
-    """Build a Rerun blueprint laying out camera images, observation and action scalars in separate views.
+def _build_blueprint(
+    observation_paths: set[str],
+    action_paths: set[str],
+    image_paths: set[str],
+    prediction_paths: set[str],
+):
+    """Build a Rerun blueprint laying out camera/predicted images and scalar series in separate views.
 
-    Camera images, observation and action scalars are arranged in a grid.
+    Images (observation and prediction) each get a spatial view; observation, action, and prediction
+    scalars each get their own time-series view. All arranged in a grid.
     """
 
     # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
@@ -88,22 +131,29 @@ def _build_blueprint(observation_paths: set[str], action_paths: set[str], image_
         views.append(rrb.TimeSeriesView(name="observation", contents=sorted(observation_paths)))
     if action_paths:
         views.append(rrb.TimeSeriesView(name="action", contents=sorted(action_paths)))
+    if prediction_paths:
+        views.append(rrb.TimeSeriesView(name="prediction", contents=sorted(prediction_paths)))
 
     return rrb.Blueprint(rrb.Grid(*views))
 
 
-def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image_paths: set[str]) -> None:
-    """Build and send the blueprint once, from the first observation and action data."""
+def _ensure_blueprint(
+    observation_paths: set[str],
+    action_paths: set[str],
+    image_paths: set[str],
+    prediction_paths: set[str],
+) -> None:
+    """Build and send the blueprint once, from the first observation/action/prediction data."""
     if getattr(log_rerun_data, "blueprint", None) is not None:
         return
 
-    if not (observation_paths or action_paths or image_paths):
+    if not (observation_paths or action_paths or image_paths or prediction_paths):
         return
 
     # Safe + zero-overhead: `log_rerun_data` already ran the `require_package` guard and imported rerun.
     import rerun as rr
 
-    blueprint = _build_blueprint(observation_paths, action_paths, image_paths)
+    blueprint = _build_blueprint(observation_paths, action_paths, image_paths, prediction_paths)
     log_rerun_data.blueprint = blueprint
     rr.send_blueprint(blueprint)
 
@@ -111,10 +161,11 @@ def _ensure_blueprint(observation_paths: set[str], action_paths: set[str], image
 def log_rerun_data(
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
+    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
     """
-    Logs observation and action data to Rerun for real-time visualization.
+    Logs observation, action and prediction data to Rerun for real-time visualization.
 
     This function iterates through the provided observation and action dictionaries and sends their contents
     to the Rerun viewer. It handles different data types appropriately:
@@ -133,6 +184,8 @@ def log_rerun_data(
     Args:
         observation: An optional dictionary containing observation data to log.
         action: An optional dictionary containing action data to log.
+        prediction: An optional dictionary of display-ready model outputs (e.g. a world model's
+            imagined video), keyed "<datatype>.<name>", logged on a dedicated "prediction." channel.
         compress_images: Whether to compress images before logging to save bandwidth & memory in exchange for cpu and quality.
     """
 
@@ -142,37 +195,19 @@ def log_rerun_data(
     observation_paths: set[str] = set()
     action_paths: set[str] = set()
     image_paths: set[str] = set()
+    prediction_paths: set[str] = set()
 
     if observation:
-        for k, v in observation.items():
-            if v is None:
-                continue
-            key = k if str(k).startswith(OBS_PREFIX) else f"{OBS_STR}.{k}"
+        _log_scalar_or_image_mapping(
+            rr, observation, OBS_PREFIX, observation_paths, image_paths, compress_images
+        )
 
-            if _is_scalar(v):
-                rr.log(key, rr.Scalars(float(v)))
-                observation_paths.add(key)
-            elif isinstance(v, np.ndarray):
-                arr = v
-                # Convert CHW -> HWC when needed
-                if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4):
-                    arr = np.transpose(arr, (1, 2, 0))
-                if arr.ndim == 1:
-                    rr.log(key, rr.Scalars(arr.astype(float)))
-                    observation_paths.add(key)
-                else:
-                    if arr.shape[-1] == 1:
-                        # At record time, the depth unit is inferred from the frame type.
-                        depth_unit = infer_depth_unit(arr.dtype)
-                        img_entity = rr.DepthImage(
-                            arr,
-                            meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0,
-                            colormap=rr.components.Colormap.Viridis,
-                        )
-                    else:
-                        img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr)
-                    rr.log(key, entity=img_entity, static=True)
-                    image_paths.add(key)
+    if prediction:
+        # Predicted images share the spatial-view set (their "prediction." names keep them distinct);
+        # predicted scalars get their own time-series view.
+        _log_scalar_or_image_mapping(
+            rr, prediction, PREDICTION_PREFIX, prediction_paths, image_paths, compress_images
+        )
 
     if action:
         for k, v in action.items():
@@ -188,4 +223,4 @@ def log_rerun_data(
                 rr.log(key, rr.Scalars(v.reshape(-1).astype(float)))
                 action_paths.add(key)
 
-    _ensure_blueprint(observation_paths, action_paths, image_paths)
+    _ensure_blueprint(observation_paths, action_paths, image_paths, prediction_paths)

src/lerobot/utils/visualization_utils.py

@@ -56,14 +56,19 @@ def log_visualization_data(
     display_mode: str,
     observation: RobotObservation | None = None,
     action: RobotAction | None = None,
+    prediction: dict | None = None,
     compress_images: bool = False,
 ) -> None:
-    """Logs observation/action data to the backend selected by ``display_mode``."""
+    """Logs observation/action/prediction data to the backend selected by ``display_mode``."""
 
     if display_mode == "rerun":
-        log_rerun_data(observation=observation, action=action, compress_images=compress_images)
+        log_rerun_data(
+            observation=observation, action=action, prediction=prediction, compress_images=compress_images
+        )
     elif display_mode == "foxglove":
-        log_foxglove_data(observation=observation, action=action, compress_images=compress_images)
+        log_foxglove_data(
+            observation=observation, action=action, prediction=prediction, compress_images=compress_images
+        )
     else:
         raise ValueError(f"Unknown display_mode '{display_mode}'. Expected one of {VISUALIZATION_MODES}.")