From 4ae2fbca36a4a9ac5efef44ec2865a488c3972fb Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Fri, 10 Jul 2026 16:29:58 +0000 Subject: [PATCH] draft for unifying prediction visualization --- src/lerobot/configs/default.py | 3 + .../lingbot_va/configuration_lingbot_va.py | 3 - .../lingbot_va/modeling_lingbot_va.py | 62 ++++++---- src/lerobot/policies/pretrained.py | 11 ++ src/lerobot/rollout/configs.py | 19 +++ src/lerobot/rollout/context.py | 1 + src/lerobot/rollout/inference/base.py | 9 ++ src/lerobot/rollout/inference/factory.py | 2 + src/lerobot/rollout/inference/sync.py | 53 ++++++++- src/lerobot/rollout/strategies/core.py | 8 +- src/lerobot/scripts/lerobot_eval.py | 110 +++++++++-------- src/lerobot/utils/constants.py | 3 + src/lerobot/utils/foxglove_visualization.py | 31 ++++- src/lerobot/utils/rerun_visualization.py | 111 ++++++++++++------ src/lerobot/utils/visualization_utils.py | 11 +- 15 files changed, 313 insertions(+), 124 deletions(-) diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 38991a665..280f7cff9 100644 --- a/src/lerobot/configs/default.py +++ b/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: diff --git a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py b/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py index 424ea7c63..699bd4c9b 100644 --- a/src/lerobot/policies/lingbot_va/configuration_lingbot_va.py +++ b/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( diff --git a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py b/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py index 0f70ad290..e1e3bd9ec 100644 --- a/src/lerobot/policies/lingbot_va/modeling_lingbot_va.py +++ b/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.""" diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 112c46170..dd731a4e0 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -96,6 +96,17 @@ class ActionSelectKwargs(TypedDict, total=False): 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): """ Base class for policy models. diff --git a/src/lerobot/rollout/configs.py b/src/lerobot/rollout/configs.py index 639e2ba29..83cea98fb 100644 --- a/src/lerobot/rollout/configs.py +++ b/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") diff --git a/src/lerobot/rollout/context.py b/src/lerobot/rollout/context.py index 20a7d715a..2ae667219 100644 --- a/src/lerobot/rollout/context.py +++ b/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 --------------------------------------------------- diff --git a/src/lerobot/rollout/inference/base.py b/src/lerobot/rollout/inference/base.py index f269aa5fe..ecfd6e05d 100644 --- a/src/lerobot/rollout/inference/base.py +++ b/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 ``"."`` (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.""" diff --git a/src/lerobot/rollout/inference/factory.py b/src/lerobot/rollout/inference/factory.py index e600bed63..efb5d61d5 100644 --- a/src/lerobot/rollout/inference/factory.py +++ b/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( diff --git a/src/lerobot/rollout/inference/sync.py b/src/lerobot/rollout/inference/sync.py index 2bb05b6ab..e2d35ac84 100644 --- a/src/lerobot/rollout/inference/sync.py +++ b/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() diff --git a/src/lerobot/rollout/strategies/core.py b/src/lerobot/rollout/strategies/core.py index 460ad12e5..8fb77a638 100644 --- a/src/lerobot/rollout/strategies/core.py +++ b/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, ) diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 722763d6e..45e8f5d16 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/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: diff --git a/src/lerobot/utils/constants.py b/src/lerobot/utils/constants.py index 8f735fe6d..941dc6339 100644 --- a/src/lerobot/utils/constants.py +++ b/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" diff --git a/src/lerobot/utils/foxglove_visualization.py b/src/lerobot/utils/foxglove_visualization.py index fc4136e12..b8fc43131 100644 --- a/src/lerobot/utils/foxglove_visualization.py +++ b/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 ".", 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 "." (e.g. "images.predicted"); route images to + # /prediction/images/ 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(): diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py index 46f2c0b4b..bd1149859 100644 --- a/src/lerobot/utils/rerun_visualization.py +++ b/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 ".", 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) diff --git a/src/lerobot/utils/visualization_utils.py b/src/lerobot/utils/visualization_utils.py index 09a89b20a..e24154efb 100644 --- a/src/lerobot/utils/visualization_utils.py +++ b/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}.")