rtc zero padding issue?

- stop() restores the original predict_action_chunk so the policy object is
  not left permanently patched after the engine goes away
- snapshot the chunk anchor with clone() (guarded against None on the first
  tick) so it survives any future in-place mutation of the cached state
- document the engine->step action_names lazy-fill and the torch.compile
  ordering the probe relies on
- add tests for reset() re-anchoring across episodes and stop() probe restore

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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/factory.py

@@ -29,8 +29,10 @@ from lerobot.configs import FeatureType, PreTrainedConfig
 from lerobot.envs import EnvConfig, env_to_policy_features
 from lerobot.processor import (
     AbsoluteActionsProcessorStep,
+    NormalizerProcessorStep,
     PolicyProcessorPipeline,
     RelativeActionsProcessorStep,
+    UnnormalizerProcessorStep,
     batch_to_transition,
     policy_action_to_transition,
     transition_to_batch,
@@ -84,6 +86,52 @@ def _reconnect_relative_absolute_steps(
             step.relative_step = relative_step
 
 
+def _ensure_relative_actions(
+    preprocessor: PolicyProcessorPipeline, postprocessor: PolicyProcessorPipeline, policy_cfg
+) -> None:
+    """Enable (or inject) the relative/absolute action steps in a loaded pipeline.
+
+    When loading from a pretrained checkpoint, the saved processor is authoritative. If the base
+    predates the relative-action feature (e.g. FastWAM/LingBot bases) its pipeline has no
+    RelativeActionsProcessorStep, so lerobot-train's override cannot enable one — those override
+    keys are popped before `from_pretrained` (else it raises) and we reconstruct the steps here.
+    Bases that DO ship the (disabled) steps (e.g. pi0/pi05) are simply flipped on. No-op unless
+    ``policy_cfg.use_relative_actions`` is set, so non-relative runs are untouched.
+    """
+    if not getattr(policy_cfg, "use_relative_actions", False):
+        return
+
+    exclude_joints = list(getattr(policy_cfg, "relative_exclude_joints", []) or [])
+    action_names = getattr(policy_cfg, "action_feature_names", None)
+
+    pre_steps = list(preprocessor.steps)
+    relative_step = next((s for s in pre_steps if isinstance(s, RelativeActionsProcessorStep)), None)
+    if relative_step is None:
+        relative_step = RelativeActionsProcessorStep(
+            enabled=True, exclude_joints=exclude_joints, action_names=action_names
+        )
+        # Insert right before the normalizer (raw -> relative -> normalize); fall back to the front.
+        idx = next((i for i, s in enumerate(pre_steps) if isinstance(s, NormalizerProcessorStep)), 0)
+        pre_steps.insert(idx, relative_step)
+        preprocessor.steps = pre_steps
+    else:
+        relative_step.enabled = True
+        relative_step.exclude_joints = exclude_joints
+        relative_step.action_names = action_names
+
+    post_steps = list(postprocessor.steps)
+    absolute_step = next((s for s in post_steps if isinstance(s, AbsoluteActionsProcessorStep)), None)
+    if absolute_step is None:
+        absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+        # Insert right after the unnormalizer (unnormalize -> absolute); fall back to the front.
+        idx = next((i for i, s in enumerate(post_steps) if isinstance(s, UnnormalizerProcessorStep)), -1)
+        post_steps.insert(idx + 1, absolute_step)
+        postprocessor.steps = post_steps
+    else:
+        absolute_step.enabled = True
+        absolute_step.relative_step = relative_step
+
+
 def get_policy_class(name: str) -> type[PreTrainedPolicy]:
     """
     Retrieves a policy class by its registered name.
@@ -320,12 +368,20 @@ def make_pre_post_processors(
                 ),
             )
 
+        # The relative/absolute override keys only match if the saved base already contains those
+        # steps (e.g. pi0/pi05). For bases that predate the feature (FastWAM/LingBot) they would
+        # raise "Override keys ... do not match any step". Pop them here and let
+        # _ensure_relative_actions() enable-or-inject the steps after loading (handles both cases).
+        pre_overrides = dict(kwargs.get("preprocessor_overrides") or {})
+        post_overrides = dict(kwargs.get("postprocessor_overrides") or {})
+        pre_overrides.pop("relative_actions_processor", None)
+        post_overrides.pop("absolute_actions_processor", None)
         preprocessor = PolicyProcessorPipeline.from_pretrained(
             pretrained_model_name_or_path=pretrained_path,
             config_filename=kwargs.get(
                 "preprocessor_config_filename", f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json"
             ),
-            overrides=kwargs.get("preprocessor_overrides", {}),
+            overrides=pre_overrides,
             to_transition=batch_to_transition,
             to_output=transition_to_batch,
             revision=pretrained_revision,
@@ -335,11 +391,12 @@ def make_pre_post_processors(
             config_filename=kwargs.get(
                 "postprocessor_config_filename", f"{POLICY_POSTPROCESSOR_DEFAULT_NAME}.json"
             ),
-            overrides=kwargs.get("postprocessor_overrides", {}),
+            overrides=post_overrides,
             to_transition=policy_action_to_transition,
             to_output=transition_to_policy_action,
             revision=pretrained_revision,
         )
+        _ensure_relative_actions(preprocessor, postprocessor, policy_cfg)
         _reconnect_relative_absolute_steps(preprocessor, postprocessor)
         if isinstance(policy_cfg, Evo1Config):
             from .evo1.processor_evo1 import reconcile_evo1_processors

src/lerobot/policies/fastwam/configuration_fastwam.py

@@ -27,6 +27,8 @@ from lerobot.configs import (
 from lerobot.optim import AdamWConfig
 from lerobot.utils.constants import ACTION, OBS_STATE
 
+from ..rtc.configuration_rtc import RTCConfig
+
 WAN22_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B"
 WAN22_DIFFUSERS_MODEL_ID = "Wan-AI/Wan2.2-TI2V-5B-Diffusers"
 FASTWAM_BASE_MODEL_ID = "lerobot/fastwam_base"
@@ -188,12 +190,25 @@ class FastWAMConfig(PreTrainedConfig):
     action_video_freq_ratio: int = 4
     image_size: tuple[int, int] = (224, 448)
     context_len: int = 128
+
+    # Relative actions: converts absolute actions to relative (action -= state) during
+    # preprocessing, and reverses it at postprocessing. Requires `proprio_dim` (OBS_STATE).
+    use_relative_actions: bool = False
+    # Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
+    relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
+    # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
+    action_feature_names: list[str] | None = None
     model_id: str = WAN22_MODEL_ID
     tokenizer_model_id: str = WAN_T5_TOKENIZER_ID
     text_encoder_model_id: str = WAN22_DIFFUSERS_MODEL_ID
     base_model_id: str | None = FASTWAM_BASE_MODEL_ID
     tokenizer_max_len: int = 128
     load_text_encoder: bool = True
+    # Device for the frozen ~11GB UMT5-XXL text encoder. `None` keeps it on the main
+    # policy `device` (default). Set to e.g. "cpu" to keep it off the GPU and save VRAM;
+    # prompts are then encoded on that device and the resulting embeddings moved to the
+    # policy device. Trades GPU memory for slower (CPU) text encoding.
+    text_encoder_device: str | None = None
     mot_checkpoint_mixed_attn: bool = False
     torch_dtype: str = "bfloat16"
     prompt_template: str = (
@@ -201,6 +216,11 @@ class FastWAMConfig(PreTrainedConfig):
     )
     num_inference_steps: int = 10
     inference_seed: int | None = 42
+    # Real-Time Chunking (RTC): async chunk generation with prefix guidance so a new chunk
+    # inpaints smoothly onto the still-executing tail of the previous one. `None` disables it
+    # (default synchronous inference). Consumed by `RTCInferenceEngine`, which calls
+    # `predict_action_chunk(..., inference_delay=, prev_chunk_left_over=)`.
+    rtc_config: RTCConfig | None = None
     rand_device: str = "cpu"
     text_cfg_scale: float = 1.0
     negative_prompt: str = ""
@@ -279,6 +299,10 @@ class FastWAMConfig(PreTrainedConfig):
         finally:
             self.pretrained_path = pretrained_path
 
+    @property
+    def chunk_size(self) -> int:
+        return self.action_horizon
+
     def get_optimizer_preset(self) -> AdamWConfig:
         return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay)
 

src/lerobot/policies/fastwam/modeling_fastwam.py

@@ -22,6 +22,7 @@ import torch
 from torch import Tensor
 
 from lerobot.policies.pretrained import PreTrainedPolicy
+from lerobot.policies.rtc.modeling_rtc import RTCProcessor
 from lerobot.utils.constants import OBS_STATE
 from lerobot.utils.import_utils import require_package
 
@@ -85,8 +86,29 @@ class FastWAMPolicy(PreTrainedPolicy):
                 for layer in mot.layers:
                     if "video" in layer.blocks:
                         layer.blocks["video"].requires_grad_(False)
+        self.init_rtc_processor()
         self.reset()
 
+    def init_rtc_processor(self) -> None:
+        """Attach a Real-Time Chunking processor to the core model when configured.
+
+        Mirrors the PI0/PI05 pattern: the policy owns the `RTCProcessor` and hands it to
+        the core `FastWAM` model, which consults it inside `infer_action`'s denoising loop.
+        Must stay public and named exactly `init_rtc_processor`: the rollout loader
+        (`lerobot.rollout.context`) sets `policy.config.rtc_config = cfg.inference.rtc` and
+        then calls `policy.init_rtc_processor()` to (re)build the processor after load, so
+        `--inference.type=rtc` alone is enough to enable guidance — no separate policy-side
+        `rtc_config` needed. A private/renamed method would be silently skipped (guidance
+        off), degrading RTC to unguided async chunk-swapping.
+        """
+        self.rtc_processor = None
+        if self.config.rtc_config is not None:
+            self.rtc_processor = RTCProcessor(self.config.rtc_config)
+            self.model.rtc_processor = self.rtc_processor
+
+    def _rtc_enabled(self) -> bool:
+        return self.config.rtc_config is not None and self.config.rtc_config.enabled
+
     @classmethod
     def _load_as_safetensor(cls, model, model_file: str, map_location: str, strict: bool):
         """Shape-aware load that supports cross-embodiment fine-tuning.
@@ -150,6 +172,24 @@ class FastWAMPolicy(PreTrainedPolicy):
 
     def reset(self) -> None:
         self._action_queue: deque[Tensor] = deque([], maxlen=self.config.n_action_steps)
+        # Per-episode text-embedding cache (mirrors LingBot-VA's `_prompt_embeds`). The task
+        # is fixed for an episode, so the ~11GB UMT5 encoder runs once on the first chunk and
+        # the resulting context is reused for every subsequent chunk. Cleared here on reset so
+        # a new episode's (possibly different) task is re-encoded. Proprio is still appended
+        # fresh each chunk downstream, so only the text-only context is cached.
+        self._cached_prompt: Any = None
+        self._cached_context: Tensor | None = None
+        self._cached_context_mask: Tensor | None = None
+
+    def _encode_prompt_cached(self, prompt: Any) -> tuple[Tensor, Tensor]:
+        """Encode `prompt` to `(context, context_mask)`, reusing the cache when the prompt is
+        unchanged so UMT5 runs at most once per episode (per distinct task)."""
+        if self._cached_context is None or self._cached_prompt != prompt:
+            context, context_mask = self.model.encode_prompt(prompt)
+            self._cached_prompt = prompt
+            self._cached_context = context
+            self._cached_context_mask = context_mask
+        return self._cached_context, self._cached_context_mask
 
     def _batch_to_training_sample(self, batch: dict[str, Tensor]) -> dict[str, Tensor]:
         """Adapt a standard LeRobot batch to the FastWAM-native sample that
@@ -183,7 +223,9 @@ class FastWAMPolicy(PreTrainedPolicy):
                 sample["proprio"] = state.unsqueeze(1) if state.ndim == 2 else state
         return sample
 
-    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict[str, Any]]:
+    def forward(
+        self, batch: dict[str, Tensor], reduction: str = "mean"
+    ) -> tuple[Tensor, dict[str, Any]]:
         """Compute FastWAM training loss for a LeRobot batch.
 
         Args:
@@ -191,24 +233,42 @@ class FastWAMPolicy(PreTrainedPolicy):
                 (`video`, `action`, `context`, `context_mask`) or LeRobot keys
                 that can be adapted (`observation.images.*`, `observation.state`,
                 `action`, `action_is_pad`).
+            reduction (str): "mean" returns the scalar loss (default, backward
+                compatible); "none" returns per-sample losses of shape (batch_size,)
+                for sample weighting (RA-BC).
 
         Returns:
-            tuple[Tensor, dict[str, Any]]: The scalar loss to backprop, and a dict of
-            logging metrics (e.g. `loss_video`, `loss_action`) — the `(loss, output_dict)`
-            contract the LeRobot training loop expects.
+            tuple[Tensor, dict[str, Any]]: The loss to backprop (scalar for "mean",
+            per-sample (B,) for "none"), and a dict of logging metrics (e.g.
+            `loss_video`, `loss_action`) — the `(loss, output_dict)` contract the
+            LeRobot training loop expects.
         """
 
         sample = self._batch_to_training_sample(batch)
-        loss, metrics = self.model.training_loss(sample)
+        loss, metrics = self.model.training_loss(sample, reduction=reduction)
         return loss, dict(metrics or {})
 
     @torch.no_grad()
-    def predict_action_chunk(self, batch: dict[str, Tensor], **_: Any) -> Tensor:
+    def predict_action_chunk(
+        self,
+        batch: dict[str, Tensor],
+        inference_delay: int | None = None,
+        prev_chunk_left_over: Tensor | None = None,
+        execution_horizon: int | None = None,
+        **_: Any,
+    ) -> Tensor:
         """Predict a chunk of actions from the current FastWAM observation.
 
         Args:
             batch (dict[str, Tensor]): Inference batch with `input_image` or
                 image observation keys, plus `context/context_mask` or `prompt`.
+            inference_delay (int | None): RTC — number of prefix steps assumed already
+                executed by the time this chunk lands (from measured inference latency).
+            prev_chunk_left_over (Tensor | None): RTC — the previous chunk's unexecuted
+                action tail `[T_prev, action_dim]` in model space; guides denoising so the
+                new chunk inpaints onto it. `None` (default) = plain synchronous inference.
+            execution_horizon (int | None): RTC — override for the prefix-weight horizon;
+                `None` falls back to `rtc_config.execution_horizon`.
 
         Returns:
             Tensor: Action chunk with shape `[B, action_horizon, action_dim]`.
@@ -216,6 +276,20 @@ class FastWAMPolicy(PreTrainedPolicy):
 
         self.eval()
         infer_kwargs = _batch_to_infer_kwargs(batch=batch, config=self.config)
+        # Encode the task once per episode and reuse it (LingBot-VA parity): swap the raw
+        # `prompt` for the cached `context`/`context_mask` so `infer_action` skips `encode_prompt`
+        # and the text encoder isn't re-run every chunk. Skipped when the caller supplies its own
+        # precomputed `context` (the two are mutually exclusive downstream).
+        if infer_kwargs.get("context") is None and infer_kwargs.get("prompt") is not None:
+            context, context_mask = self._encode_prompt_cached(infer_kwargs["prompt"])
+            infer_kwargs["prompt"] = None
+            infer_kwargs["context"] = context
+            infer_kwargs["context_mask"] = context_mask
+        # RTC guidance args flow straight to `infer_action`; they are inert unless an
+        # RTCProcessor is attached, enabled, and `prev_chunk_left_over` is provided.
+        infer_kwargs["inference_delay"] = inference_delay
+        infer_kwargs["prev_chunk_left_over"] = prev_chunk_left_over
+        infer_kwargs["execution_horizon"] = execution_horizon
         batch_size = _infer_kwargs_batch_size(infer_kwargs)
         if batch_size == 1:
             action = _action_from_model_output(self.model.infer_action(**infer_kwargs))
@@ -260,9 +334,10 @@ class FastWAMPolicy(PreTrainedPolicy):
             mixtures={"video": video_expert, "action": action_expert},
             mot_checkpoint_mixed_attn=config.mot_checkpoint_mixed_attn,
         )
+        text_encoder_device = config.text_encoder_device or device
         text_encoder = (
             load_pretrained_wan_text_encoder(
-                model_id=config.text_encoder_model_id, torch_dtype=dtype, device=device
+                model_id=config.text_encoder_model_id, torch_dtype=dtype, device=text_encoder_device
             )
             if config.load_text_encoder
             else None
@@ -273,6 +348,7 @@ class FastWAMPolicy(PreTrainedPolicy):
             mot=mot,
             vae=load_pretrained_wan_vae(torch_dtype=dtype, device=device),
             text_encoder=text_encoder,
+            text_encoder_device=config.text_encoder_device,
             tokenizer=build_wan_tokenizer(
                 model_id=config.tokenizer_model_id, tokenizer_max_len=config.tokenizer_max_len
             ),

src/lerobot/policies/fastwam/processor_fastwam.py

@@ -21,13 +21,16 @@ import torch
 
 from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.processor import (
+    AbsoluteActionsProcessorStep,
     ActionProcessorStep,
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
     NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
+    ProcessorStep,
     ProcessorStepRegistry,
+    RelativeActionsProcessorStep,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
     policy_action_to_transition,
@@ -105,10 +108,20 @@ def make_fastwam_pre_post_processors(
     # anyway) and unsafe across fine-tuning: its `resize_size` would be inherited from the base
     # checkpoint's camera geometry, not this dataset's, making the concatenation N_cameras x too wide.
 
-    input_steps = [
+    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
+    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
+    # below so its cached raw state (set during preprocessing) flows to postprocessing.
+    relative_step = RelativeActionsProcessorStep(
+        enabled=config.use_relative_actions,
+        exclude_joints=getattr(config, "relative_exclude_joints", []),
+        action_names=getattr(config, "action_feature_names", None),
+    )
+
+    input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
         DeviceProcessorStep(device=config.device),
+        relative_step,
         NormalizerProcessorStep(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
@@ -116,12 +129,13 @@ def make_fastwam_pre_post_processors(
             device=config.device,
         ),
     ]
-    output_steps = [
+    output_steps: list[ProcessorStep] = [
         UnnormalizerProcessorStep(
             features=config.output_features,
             norm_map=config.normalization_mapping,
             stats=normalization_stats,
         ),
+        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
     ]
     if config.toggle_action_dimensions:
         output_steps.append(

src/lerobot/policies/fastwam/wan/modular.py

@@ -839,6 +839,7 @@ class FastWAM(torch.nn.Module):
         text_dim: int | None = None,
         proprio_dim: int | None = None,
         device: str = "cpu",
+        text_encoder_device: str | torch.device | None = None,
         torch_dtype: torch.dtype = torch.float32,
         video_train_shift: float = 5.0,
         video_infer_shift: float = 5.0,
@@ -907,12 +908,27 @@ class FastWAM(torch.nn.Module):
         self.train_scheduler = self.train_video_scheduler
         self.infer_scheduler = self.infer_video_scheduler
 
+        # Optional Real-Time Chunking processor (set by the policy wrapper). When present and
+        # enabled it guides the action denoising loop in `infer_action` so a freshly generated
+        # chunk inpaints onto the previous chunk's unexecuted tail. Plain attribute (not an
+        # nn.Module) — carries no parameters and stays out of state_dict / device moves.
+        self.rtc_processor = None
         self.device = torch.device(device)
+        # When pinned (e.g. "cpu"), the frozen text encoder stays on this device instead
+        # of following the model onto the GPU — `_apply` skips it and `encode_prompt` runs
+        # it here, moving embeddings back to `self.device`. `None` = follow `self.device`.
+        self._text_encoder_device = (
+            torch.device(text_encoder_device) if text_encoder_device is not None else None
+        )
         self.torch_dtype = torch_dtype
         self.loss_lambda_video = float(loss_lambda_video)
         self.loss_lambda_action = float(loss_lambda_action)
 
         self.to(self.device)
+        # `self.to` above (via `_apply`) skips a pinned text encoder; make sure it actually
+        # sits on the pinned device (it was loaded there, but this is a cheap safety net).
+        if self.text_encoder is not None and self._text_encoder_device is not None:
+            self.text_encoder._apply(lambda t: t.to(self._text_encoder_device))
 
     @classmethod
     def from_wan22_pretrained(
@@ -1003,7 +1019,8 @@ class FastWAM(torch.nn.Module):
         # while staying out of `state_dict()` / `parameters()`.
         super()._apply(fn, *args, **kwargs)
         self.vae._apply(fn)
-        if self.text_encoder is not None:
+        # A pinned text encoder (e.g. on CPU) must NOT follow device moves — leave it put.
+        if self.text_encoder is not None and self._text_encoder_device is None:
             self.text_encoder._apply(fn)
         return self
 
@@ -1024,9 +1041,12 @@ class FastWAM(torch.nn.Module):
                 "Prompt encoding requires loaded text encoder/tokenizer. "
                 "Set `load_text_encoder=true` or provide precomputed `context/context_mask`."
             )
+        # Run the encoder on its own device (may be pinned to CPU to save VRAM), then
+        # move the resulting embeddings/mask to the model device for the DiT.
+        te_device = self._text_encoder_device or self.device
         ids, mask = self.tokenizer(prompt, return_mask=True, add_special_tokens=True)
-        ids = ids.to(self.device)
-        mask = mask.to(self.device, dtype=torch.bool)
+        ids = ids.to(te_device)
+        mask = mask.to(te_device, dtype=torch.bool)
         prompt_emb = self.text_encoder(ids, mask)
         seq_lens = mask.gt(0).sum(dim=1).long()
         for i, v in enumerate(seq_lens):
@@ -1034,7 +1054,7 @@ class FastWAM(torch.nn.Module):
         # Match FastWAM/Wan2.2 context semantics: padding embeddings are zeroed,
         # while cross-attention still sees a fixed-length context.
         mask = torch.ones_like(mask)
-        return prompt_emb.to(device=self.device), mask
+        return prompt_emb.to(device=self.device), mask.to(device=self.device)
 
     def _append_proprio_to_context(
         self,
@@ -1359,7 +1379,9 @@ class FastWAM(torch.nn.Module):
         pred_action = self.action_expert.post_dit(tokens_out["action"], action_pre)
         return pred_video, pred_action
 
-    def _compute_training_video_loss(self, inputs, pred_video, target_video, timestep_video):
+    def _compute_training_video_loss(
+        self, inputs, pred_video, target_video, timestep_video, reduction: str = "mean"
+    ):
         include_initial_video_step = inputs["first_frame_latents"] is None
         if inputs["first_frame_latents"] is not None:
             pred_video = pred_video[:, :, 1:]
@@ -1374,9 +1396,13 @@ class FastWAM(torch.nn.Module):
             loss_video_per_sample.device,
             dtype=loss_video_per_sample.dtype,
         )
-        return (loss_video_per_sample * video_weight).mean()
+        weighted = loss_video_per_sample * video_weight
+        # reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC).
+        return weighted if reduction == "none" else weighted.mean()
 
-    def _compute_training_action_loss(self, inputs, pred_action, target_action, timestep_action):
+    def _compute_training_action_loss(
+        self, inputs, pred_action, target_action, timestep_action, reduction: str = "mean"
+    ):
         action_loss_token = functional.mse_loss(
             pred_action.float(), target_action.float(), reduction="none"
         ).mean(dim=2)
@@ -1393,9 +1419,11 @@ class FastWAM(torch.nn.Module):
             action_loss_per_sample.device,
             dtype=action_loss_per_sample.dtype,
         )
-        return (action_loss_per_sample * action_weight).mean()
+        weighted = action_loss_per_sample * action_weight
+        # reduction="none" returns the per-sample vector (B,) for sample weighting (RA-BC).
+        return weighted if reduction == "none" else weighted.mean()
 
-    def training_loss(self, sample, tiled: bool = False):
+    def training_loss(self, sample, tiled: bool = False, reduction: str = "mean"):
         inputs = self.build_inputs(sample, tiled=tiled)
         targets = self._sample_training_targets(inputs)
         pred_video, pred_action = self._run_training_mot(inputs=inputs, targets=targets)
@@ -1404,17 +1432,20 @@ class FastWAM(torch.nn.Module):
             pred_video=pred_video,
             target_video=targets["target_video"],
             timestep_video=targets["timestep_video"],
+            reduction=reduction,
         )
         loss_action = self._compute_training_action_loss(
             inputs=inputs,
             pred_action=pred_action,
             target_action=targets["target_action"],
             timestep_action=targets["timestep_action"],
+            reduction=reduction,
         )
+        # With reduction="none" both terms are (B,), so loss_total is the per-sample loss (B,).
         loss_total = self.loss_lambda_video * loss_video + self.loss_lambda_action * loss_action
         loss_dict = {
-            "loss_video": self.loss_lambda_video * float(loss_video.detach().item()),
-            "loss_action": self.loss_lambda_action * float(loss_action.detach().item()),
+            "loss_video": self.loss_lambda_video * float(loss_video.detach().mean().item()),
+            "loss_action": self.loss_lambda_action * float(loss_action.detach().mean().item()),
         }
         return loss_total, loss_dict
 
@@ -1799,6 +1830,9 @@ class FastWAM(torch.nn.Module):
         seed: int | None = None,
         rand_device: str = "cpu",
         tiled: bool = False,
+        inference_delay: int | None = None,
+        prev_chunk_left_over: torch.Tensor | None = None,
+        execution_horizon: int | None = None,
     ) -> dict[str, Any]:
         self.eval()
         if str(getattr(self.video_expert, "video_attention_mask_mode", "")) != "first_frame_causal":
@@ -1851,18 +1885,43 @@ class FastWAM(torch.nn.Module):
             dtype=latents_action.dtype,
             shift_override=sigma_shift,
         )
+        rtc_active = (
+            self.rtc_processor is not None
+            and getattr(self.rtc_processor.rtc_config, "enabled", False)
+            and prev_chunk_left_over is not None
+        )
+        num_train_timesteps = float(self.infer_action_scheduler.num_train_timesteps)
         for step_t_action, step_delta_action in zip(infer_timesteps_action, infer_deltas_action, strict=True):
             timestep_action = step_t_action.unsqueeze(0).to(dtype=latents_action.dtype, device=self.device)
 
-            pred_action = self._predict_action_noise_with_cache(
-                latents_action=latents_action,
-                timestep_action=timestep_action,
-                context=context,
-                context_mask=context_mask,
-                video_kv_cache=video_kv_cache,
-                attention_mask=attention_mask,
-                video_seq_len=video_seq_len,
-            )
+            def denoise(x_t, ts=timestep_action):
+                return self._predict_action_noise_with_cache(
+                    latents_action=x_t,
+                    timestep_action=ts,
+                    context=context,
+                    context_mask=context_mask,
+                    video_kv_cache=video_kv_cache,
+                    attention_mask=attention_mask,
+                    video_seq_len=video_seq_len,
+                )
+
+            if rtc_active:
+                # `time` is the flow-matching noise level sigma in [0, 1]: FastWAM's model
+                # predicts velocity v = noise - clean, so the clean-action estimate is
+                # x1 = x_t - sigma * v — exactly RTC's `x1_t = x_t - time * v_t`.
+                sigma = float(step_t_action.item()) / num_train_timesteps
+                pred_action = self.rtc_processor.denoise_step(
+                    x_t=latents_action,
+                    prev_chunk_left_over=prev_chunk_left_over.to(
+                        device=latents_action.device, dtype=latents_action.dtype
+                    ),
+                    inference_delay=inference_delay or 0,
+                    time=sigma,
+                    original_denoise_step_partial=denoise,
+                    execution_horizon=execution_horizon,
+                )
+            else:
+                pred_action = denoise(latents_action)
 
             latents_action = self.infer_action_scheduler.step(pred_action, step_delta_action, latents_action)
 

src/lerobot/policies/lingbot_va/configuration_lingbot_va.py

@@ -28,7 +28,11 @@ from dataclasses import dataclass, field
 from lerobot.configs.policies import PreTrainedConfig
 from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
 from lerobot.optim.optimizers import AdamWConfig
-from lerobot.optim.schedulers import ConstantWithWarmupSchedulerConfig, LRSchedulerConfig
+from lerobot.optim.schedulers import (
+    ConstantWithWarmupSchedulerConfig,
+    CosineAnnealingWithWarmupSchedulerConfig,
+    LRSchedulerConfig,
+)
 from lerobot.utils.constants import ACTION
 
 
@@ -92,6 +96,15 @@ 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)))
 
+    # Relative actions: converts absolute actions to relative (action -= state) during
+    # preprocessing, and reverses it at postprocessing. Requires the dataset to provide
+    # observation.state whose leading dims align 1:1 with the used action channels.
+    use_relative_actions: bool = False
+    # Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
+    relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
+    # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
+    action_feature_names: list[str] | None = None
+
     # Opt-in: VAE-decode predicted video latents to ``self.last_predicted_frames`` for saving MP4s.
     save_predicted_video: bool = False
 
@@ -112,6 +125,11 @@ class LingBotVAConfig(PreTrainedConfig):
     optimizer_weight_decay: float = 1e-4
     optimizer_grad_clip_norm: float = 1.0
     scheduler_warmup_steps: int = 1000
+    # Scheduler after warmup. "constant_with_warmup" (upstream default: warmup then flat peak LR)
+    # or "cosine_annealing_with_warmup" (warmup then cosine anneal peak->0 over the remaining steps).
+    # Cosine tightens the loss tail and often nudges final loss down; it does NOT reduce the
+    # flow-matching estimator's step-to-step noise (that's metric variance, LR-independent).
+    scheduler_type: str = "constant_with_warmup"
 
     def __post_init__(self):
         super().__post_init__()
@@ -150,7 +168,10 @@ class LingBotVAConfig(PreTrainedConfig):
         )
 
     def get_scheduler_preset(self) -> LRSchedulerConfig | None:
-        # Upstream uses a linear warmup followed by a constant LR (warmup_constant_lambda).
+        # Default (upstream): linear warmup then constant LR (warmup_constant_lambda).
+        # Optionally cosine-anneal peak->0 over the remaining steps via scheduler_type.
+        if self.scheduler_type == "cosine_annealing_with_warmup":
+            return CosineAnnealingWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
         return ConstantWithWarmupSchedulerConfig(num_warmup_steps=self.scheduler_warmup_steps)
 
     @property

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)
@@ -257,8 +253,12 @@ class LingBotVAPolicy(PreTrainedPolicy):
             "grid_id": grid_id,
         }
 
-    def _flow_matching_loss(self, input_dict, pred):
-        """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``)."""
+    def _flow_matching_loss(self, input_dict, pred, reduction: str = "mean"):
+        """Dual-stream flow-matching loss (port of upstream ``Trainer.compute_loss``).
+
+        ``reduction="mean"`` returns scalar (latent_loss, action_loss); ``"none"`` returns
+        per-sample vectors of shape ``(B,)`` each (averaged over latent frames) for RA-BC.
+        """
         latent_pred, action_pred = pred
         ld, ad = input_dict["latent_dict"], input_dict["action_dict"]
         action_pred = rearrange(action_pred, "b (f n) c -> b c f n 1", f=ad["targets"].shape[-3])
@@ -278,7 +278,8 @@ class LingBotVAPolicy(PreTrainedPolicy):
         latent_loss = (
             (latent_loss * lw[:, None, :, None, None]).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
         )
-        latent_loss = (latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)).mean()
+        # per (batch*frame) mean over spatial/channel -> (B*F,)
+        latent_loss = latent_loss.sum(dim=1) / (torch.ones_like(latent_loss).sum(dim=1) + 1e-6)
 
         amask = ad["actions_mask"].float()
         action_loss = F.mse_loss(action_pred.float(), ad["targets"].float().detach(), reduction="none")
@@ -286,10 +287,14 @@ class LingBotVAPolicy(PreTrainedPolicy):
             (action_loss * aw[:, None, :, None, None] * amask).permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
         )
         amask_f = amask.permute(0, 2, 3, 4, 1).flatten(0, 1).flatten(1)
-        action_loss = (action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)).mean()
-        return latent_loss, action_loss
+        action_loss = action_loss.sum(dim=1) / (amask_f.sum(dim=1) + 1e-6)
 
-    def training_loss_from_streams(self, latents, actions, actions_mask, text_emb):
+        if reduction == "none":
+            # (B*F,) -> (B, F) -> (B,): per-sample losses for RA-BC weighting.
+            return latent_loss.reshape(bn, fn).mean(dim=1), action_loss.reshape(bn, fn).mean(dim=1)
+        return latent_loss.mean(), action_loss.mean()
+
+    def training_loss_from_streams(self, latents, actions, actions_mask, text_emb, reduction: str = "mean"):
         """Core dual-stream training loss given prepared latents / actions / text embeddings.
 
         ``latents``: ``[B, in_channels, F, h, w]`` (normalized video latents).
@@ -318,20 +323,24 @@ class LingBotVAPolicy(PreTrainedPolicy):
             "window_size": int(torch.randint(4, 65, (1,)).item()),
         }
         pred = self.transformer(input_dict, train_mode=True)
-        latent_loss, action_loss = self._flow_matching_loss(input_dict, pred)
+        latent_loss, action_loss = self._flow_matching_loss(input_dict, pred, reduction)
+        # reduction="none": latent_loss/action_loss are (B,) -> loss is per-sample (B,).
         loss = latent_loss + action_loss
-        return loss, {"latent_loss": latent_loss.item(), "action_loss": action_loss.item()}
+        return loss, {"latent_loss": latent_loss.mean().item(), "action_loss": action_loss.mean().item()}
 
-    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
+    def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict | None]:
         """Training forward: dual-stream flow-matching loss.
 
         Builds the (video-latent, action, text) training streams from a LeRobot batch
         (VAE-encoding the camera frames and UMT5-encoding the task), then runs the flow-matching
         dual-stream loss. Requires the policy to be built with ``attn_mode='flex'``.
+
+        ``reduction="mean"`` returns the scalar loss (default); ``"none"`` returns per-sample
+        losses of shape ``(B,)`` for sample weighting (RA-BC).
         """
         self._ensure_frozen_modules()
         latents, actions, actions_mask, text_emb = self._build_training_streams(batch)
-        return self.training_loss_from_streams(latents, actions, actions_mask, text_emb)
+        return self.training_loss_from_streams(latents, actions, actions_mask, text_emb, reduction=reduction)
 
     @torch.no_grad()
     def _build_training_streams(self, batch):
@@ -400,22 +409,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 +445,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 +491,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 +499,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 +868,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/lingbot_va/processor_lingbot_va.py

@@ -25,12 +25,14 @@ import torch
 
 from lerobot.configs.types import FeatureType, NormalizationMode
 from lerobot.processor import (
+    AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
     NormalizerProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
+    RelativeActionsProcessorStep,
     RenameObservationsProcessorStep,
     UnnormalizerProcessorStep,
 )
@@ -52,9 +54,19 @@ def make_lingbot_va_pre_post_processors(
 ]:
     """Build the pre/post processor pipelines for LingBot-VA."""
 
+    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
+    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
+    # below so its cached raw state (set during preprocessing) flows to postprocessing.
+    relative_step = RelativeActionsProcessorStep(
+        enabled=config.use_relative_actions,
+        exclude_joints=getattr(config, "relative_exclude_joints", []),
+        action_names=getattr(config, "action_feature_names", None),
+    )
+
     input_steps: list[ProcessorStep] = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
+        relative_step,
         NormalizerProcessorStep(
             features={**config.input_features, **config.output_features},
             norm_map=config.normalization_mapping,
@@ -63,13 +75,16 @@ def make_lingbot_va_pre_post_processors(
         DeviceProcessorStep(device=config.device),
     ]
 
-    # Unnormalize actions from [-1, 1] to physical units (QUANTILES) using q01/q99 restored from the checkpoint.
+    # Unnormalize actions back to physical units. Config-driven norm_map (was hardcoded QUANTILES)
+    # so it stays symmetric with the preprocessor's NormalizerProcessorStep — required for
+    # use_relative_actions with ACTION=IDENTITY (and unchanged for QUANTILES runs).
     output_steps: list[ProcessorStep] = [
         UnnormalizerProcessorStep(
             features=config.output_features,
-            norm_map={FeatureType.ACTION: NormalizationMode.QUANTILES},
+            norm_map=config.normalization_mapping,
             stats=dataset_stats,
         ),
+        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step),
         DeviceProcessorStep(device="cpu"),
     ]
 

src/lerobot/policies/pretrained.py

@@ -23,12 +23,13 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack
 
-import packaging
-import safetensors
 from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
+import packaging.version
+import safetensors
 from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor
+import torch
 from torch import Tensor, nn
 
 from lerobot.__version__ import __version__
@@ -93,6 +94,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):
@@ -221,6 +234,14 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
+        # safetensors' load_file maps the bare string "cuda" to cuda:0 regardless of the current
+        # device (unlike torch's .to("cuda"), which honors torch.cuda.current_device()). Under
+        # multi-GPU accelerate/FSDP every rank would then load its weights onto GPU 0, OOMing it
+        # before sharding. Resolve "cuda" to the concrete current-device index so each rank loads
+        # onto its own GPU.
+        if map_location == "cuda" and torch.cuda.is_available():
+            map_location = f"cuda:{torch.cuda.current_device()}"
+
         # Create base kwargs
         kwargs = {"strict": strict}
 
@@ -231,16 +252,6 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC):
         # Load the model with appropriate kwargs
         missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
         log_model_loading_keys(missing_keys, unexpected_keys)
-
-        # For older versions, manually move to device if needed
-        if "device" not in kwargs and map_location != "cpu":
-            logging.warning(
-                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
-                " This means that the model is loaded on 'cpu' first and then copied to the device."
-                " This leads to a slower loading time."
-                " Please update safetensors to version 0.4.3 or above for improved performance."
-            )
-            model.to(map_location)
         return model
 
     @abc.abstractmethod
@@ -273,20 +284,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/policies/rtc/action_queue.py

@@ -236,11 +236,21 @@ class ActionQueue:
         if action_index_before_inference is not None:
             indexes_diff = max(0, self.last_index - action_index_before_inference)
             if indexes_diff != real_delay:
+                # The latency estimate (`real_delay`) and the number of actions the robot
+                # actually consumed during inference (`indexes_diff`) disagree. This happens
+                # when the queue starved (robot idle) or on the first chunk (nothing consumed
+                # yet). Discarding `real_delay` here would drop actions the arm never executed
+                # and splice the queue `real_delay` steps ahead of the physical pose — a hard
+                # jump/slam, worst on slow policies where `real_delay` is large. Never discard
+                # more than was actually consumed.
+                resolved = min(real_delay, indexes_diff)
                 logger.warning(
-                    "Indexes diff is not equal to real delay. indexes_diff=%d, real_delay=%d",
+                    "Indexes diff != real delay (indexes_diff=%d, real_delay=%d); "
+                    "clamping discard to %d to avoid a queue-splice jump.",
                     indexes_diff,
                     real_delay,
+                    resolved,
                 )
-                return real_delay
+                return resolved
 
         return effective_delay

src/lerobot/policies/vla_jepa/action_head.py

@@ -282,6 +282,7 @@ class VLAJEPAActionHead(nn.Module):
         actions: torch.Tensor,
         state: torch.Tensor | None = None,
         action_is_pad: torch.Tensor | None = None,
+        reduction: str = "mean",
     ) -> torch.Tensor:
         noise = torch.randn_like(actions)
         t = self.sample_time(actions.shape[0], actions.device, actions.dtype)
@@ -302,6 +303,10 @@ class VLAJEPAActionHead(nn.Module):
 
         loss = F.mse_loss(pred_actions, velocity, reduction="none")  # [B, T, action_dim]
         valid_mask = ~action_is_pad.unsqueeze(-1)  # [B, T, 1]
+        if reduction == "none":
+            # Per-sample loss (B,) for sample weighting (RA-BC): mask-average over T and action_dim.
+            per_sample_valid = valid_mask.sum(dim=(1, 2)) * loss.shape[-1]  # [B]
+            return (loss * valid_mask).sum(dim=(1, 2)) / per_sample_valid.clamp_min(1)
         num_valid = valid_mask.sum() * loss.shape[-1]
         return (loss * valid_mask).sum() / num_valid.clamp_min(1)
 

src/lerobot/policies/vla_jepa/configuration_vla_jepa.py

@@ -56,6 +56,14 @@ class VLAJEPAConfig(PreTrainedConfig):
     action_dim: int = 7
     state_dim: int = 8
 
+    # Relative actions: converts absolute actions to relative (action -= state) during
+    # preprocessing, and reverses it at postprocessing. Requires `state_dim` (OBS_STATE).
+    use_relative_actions: bool = False
+    # Joint names to keep absolute (not converted to relative). Empty list = all dims relative.
+    relative_exclude_joints: list[str] = field(default_factory=lambda: ["gripper"])
+    # Populated at runtime from dataset metadata by make_policy (used to build the exclude mask).
+    action_feature_names: list[str] | None = None
+
     num_action_tokens_per_timestep: int = 8
     num_embodied_action_tokens_per_instruction: int = 32
     num_inference_timesteps: int = 4

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -194,8 +194,12 @@ class VLAJEPAModel(nn.Module):
             )
         return embodied_action_tokens, action_tokens
 
-    def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor:
-        """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1]."""
+    def _world_model_loss(self, videos: Tensor, action_tokens: Tensor, reduction: str = "mean") -> Tensor:
+        """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1].
+
+        `reduction="none"` returns a per-sample loss (B,) for sample weighting (RA-BC);
+        "mean" returns the scalar loss.
+        """
         # Match the world model's expected view count: pad with the first view, or trim extras.
         num_views = self.config.jepa_tubelet_size
         if videos.shape[1] < num_views:
@@ -223,7 +227,8 @@ class VLAJEPAModel(nn.Module):
         # num_video_frames raw frames → t_enc_total temporal positions after tubelet compression
         t_enc_total = self.config.num_video_frames // tubelet_size
         if t_enc_total < 2:
-            return torch.zeros((), device=video_embeddings.device)
+            zero_shape = (video_embeddings.shape[0],) if reduction == "none" else ()
+            return torch.zeros(zero_shape, device=video_embeddings.device)
 
         # Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1
         t_enc_ctx = t_enc_total - 1
@@ -239,6 +244,10 @@ class VLAJEPAModel(nn.Module):
         predicted_states = self.video_predictor(
             input_states.float(), action_tokens[:, :expected_actions].float()
         )
+        if reduction == "none":
+            # Per-sample loss (B,): mean over all non-batch dims (tokens, feature).
+            l = F.l1_loss(predicted_states, gt_states.float(), reduction="none")
+            return l.mean(dim=tuple(range(1, l.ndim)))
         return F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
 
     def _action_loss(
@@ -247,17 +256,27 @@ class VLAJEPAModel(nn.Module):
         actions: Tensor,
         state: Tensor | None,
         action_is_pad: Tensor | None,
+        reduction: str = "mean",
     ) -> Tensor:
-        """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
+        """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`.
+
+        `reduction="none"` returns a per-sample loss (B,) — the `repeated_diffusion_steps`
+        independent noise draws are averaged back per original sample — for RA-BC weighting.
+        """
         device_type = next(self.parameters()).device.type
         with torch.autocast(device_type=device_type, dtype=torch.float32):
             r = self.config.repeated_diffusion_steps
             horizon = self.config.chunk_size
+            b = embodied_action_tokens.shape[0]
             actions_target = actions[:, -horizon:, :].to(torch.float32).repeat(r, 1, 1)
             embodied = embodied_action_tokens.repeat(r, 1, 1)
             state_rep = state.to(embodied_action_tokens.dtype).repeat(r, 1, 1) if state is not None else None
             pad_rep = action_is_pad[:, -horizon:].repeat(r, 1) if action_is_pad is not None else None
-            return self.action_model(embodied, actions_target, state_rep, pad_rep)
+            loss = self.action_model(embodied, actions_target, state_rep, pad_rep, reduction=reduction)
+            if reduction == "none":
+                # `.repeat(r, 1, 1)` tiles as [rep0(b0..b_{B-1}), rep1(...), ...] → (r, B); mean over reps.
+                return loss.view(r, b).mean(dim=0)
+            return loss
 
     def forward(
         self,
@@ -267,21 +286,29 @@ class VLAJEPAModel(nn.Module):
         actions: Tensor | None = None,
         state: Tensor | None = None,
         action_is_pad: Tensor | None = None,
+        reduction: str = "mean",
     ) -> dict[str, Tensor]:
-        """Native forward: Qwen encode → optional world-model loss → optional action-head loss."""
+        """Native forward: Qwen encode → optional world-model loss → optional action-head loss.
+
+        `reduction="none"` makes both loss terms per-sample (B,) for RA-BC weighting; "mean"
+        returns scalar losses.
+        """
         embodied_action_tokens, action_tokens = self._encode_qwen(
             images, instructions, need_action_tokens=self.config.enable_world_model
         )
 
         if self.config.enable_world_model and videos is not None:
-            wm_loss = self._world_model_loss(videos, action_tokens)
+            wm_loss = self._world_model_loss(videos, action_tokens, reduction=reduction)
         else:
-            wm_loss = torch.zeros((), device=embodied_action_tokens.device)
+            zero_shape = (embodied_action_tokens.shape[0],) if reduction == "none" else ()
+            wm_loss = torch.zeros(zero_shape, device=embodied_action_tokens.device)
 
         if actions is None:
             return {"wm_loss": wm_loss}
 
-        action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad)
+        action_loss = self._action_loss(
+            embodied_action_tokens, actions, state, action_is_pad, reduction=reduction
+        )
         return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
 
     # ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
@@ -367,12 +394,19 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         batch_size = batch[image_keys[0]].shape[0]
 
         # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages.
+        # Resize to config.resize_images_to (as predict_action does) so training and inference feed
+        # Qwen the same resolution. Critical for memory: native camera frames (e.g. 720x1280) would
+        # otherwise blow up the Qwen3-VL vision-tower attention (patch count grows with resolution).
+        resize_hw = tuple(self.config.resize_images_to) if self.config.resize_images_to else None
         frames = []
         for key in image_keys:
             t = batch[key]
             if t.ndim == 5:  # [B, T, C, H, W] -> current observation (delta=0)
                 t = t[:, 0]
-            frames.append(self.model.qwen.to_pixel_values(t))
+            px = self.model.qwen.to_pixel_values(t)  # [B, C, H, W]
+            if resize_hw is not None and tuple(px.shape[-2:]) != resize_hw:
+                px = F.interpolate(px.float(), size=resize_hw, mode="area")
+            frames.append(px)
         images = [[frame[b] for frame in frames] for b in range(batch_size)]
 
         tasks = batch.get("task")
@@ -388,7 +422,26 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         # Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them.
         if self.model.config.enable_world_model and training:
             views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys]
-            inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1))
+            # The world model consumes a SINGLE stacked [B, V, T, C, H, W] tensor, so all camera
+            # views must share a spatial size. Cameras can differ (e.g. base 480x640 vs wrist
+            # 720x1280), so resize each view to a common size before stacking — config.resize_images_to
+            # if set (same target predict_action uses), else the first view's size (a no-op when all
+            # views already match, preserving behavior for single-resolution datasets). The vjepa video
+            # processor does the final resize to the encoder resolution downstream.
+            cfg = self.model.config
+            target_hw = tuple(cfg.resize_images_to) if cfg.resize_images_to else tuple(views[0].shape[-2:])
+            resized = []
+            for v in views:
+                if tuple(v.shape[-2:]) != target_hw:
+                    b, t, c = v.shape[0], v.shape[1], v.shape[2]
+                    v = F.interpolate(
+                        v.reshape(b * t, c, v.shape[3], v.shape[4]).float(),
+                        size=target_hw,
+                        mode="bilinear",
+                        align_corners=False,
+                    ).reshape(b, t, c, target_hw[0], target_hw[1])
+                resized.append(v)
+            inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(resized, dim=1))
 
         actions = batch.get(ACTION)
         if actions is not None:
@@ -406,15 +459,17 @@ class VLAJEPAPolicy(PreTrainedPolicy):
 
     # ---- LeRobot Policy Interface ----
 
-    def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
+    def forward(self, batch: dict[str, Tensor], reduction: str = "mean") -> tuple[Tensor, dict]:
         """LeRobot train forward: convert → native forward → aggregate losses."""
-        native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True))
+        native_output = self.model.forward(
+            **self._prepare_model_inputs(batch, training=True), reduction=reduction
+        )
 
         ref = next(iter(native_output.values()))
-        zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
+        zero = torch.zeros_like(ref)
         total_loss = native_output.get("action_loss", zero) + native_output.get("wm_loss", zero)
-        logs = {k: v.detach().item() for k, v in native_output.items()}
-        logs["loss"] = total_loss.detach().item()
+        logs = {k: v.detach().mean().item() for k, v in native_output.items()}
+        logs["loss"] = total_loss.detach().mean().item()
         return total_loss, logs
 
     def get_optim_params(self) -> dict:

src/lerobot/policies/vla_jepa/processor_vla_jepa.py

@@ -20,6 +20,7 @@ import torch
 
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.processor import (
+    AbsoluteActionsProcessorStep,
     AddBatchDimensionProcessorStep,
     DeviceProcessorStep,
     EnvTransition,
@@ -28,6 +29,7 @@ from lerobot.processor import (
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
+    RelativeActionsProcessorStep,
     RenameObservationsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
@@ -112,10 +114,21 @@ def make_vla_jepa_pre_post_processors(
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
 ]:
     features = {**config.input_features, **config.output_features}
+
+    # Shared relative-action step (OpenPI order: raw -> relative -> normalize -> model ->
+    # unnormalize -> absolute). The SAME instance is passed to AbsoluteActionsProcessorStep
+    # below so its cached raw state (set during preprocessing) flows to postprocessing.
+    relative_step = RelativeActionsProcessorStep(
+        enabled=config.use_relative_actions,
+        exclude_joints=getattr(config, "relative_exclude_joints", []),
+        action_names=getattr(config, "action_feature_names", None),
+    )
+
     input_steps = [
         RenameObservationsProcessorStep(rename_map={}),
         AddBatchDimensionProcessorStep(),
         DeviceProcessorStep(device=config.device),
+        relative_step,
         NormalizerProcessorStep(
             features=features,
             norm_map=config.normalization_mapping,
@@ -136,6 +149,11 @@ def make_vla_jepa_pre_post_processors(
             stats=dataset_stats,
         )
     )
+    # Reverse the relative conversion on the unnormalized action, before gripper binarization.
+    # gripper is kept absolute by relative_exclude_joints, so the two steps touch disjoint dims.
+    output_steps.append(
+        AbsoluteActionsProcessorStep(enabled=config.use_relative_actions, relative_step=relative_step)
+    )
     if config.binarize_gripper_action:
         output_steps.append(
             BinarizeGripperProcessorStep(gripper_dim=config.gripper_dim, threshold=config.gripper_threshold)

src/lerobot/processor/relative_action_processor.py

@@ -51,6 +51,12 @@ def to_relative_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) ->
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
     if state.device != actions.device or state.dtype != actions.dtype:
         state = state.to(device=actions.device, dtype=actions.dtype)
+    # When the observation is temporally stacked (e.g. LingBot-VA loads several obs steps via
+    # observation_delta_indices, giving state shape (B, T_obs, state_dim)), the relative reference
+    # is the CURRENT frame (delta == 0, i.e. index 0). Collapse to it so the offset broadcasts over
+    # the action horizon. pi0/pi05 pass a 2D (B, state_dim) state and are unaffected.
+    if state.ndim == 3:
+        state = state[:, 0]
     state_offset = state[..., :dims] * mask_t
     if actions.ndim == 3:
         state_offset = state_offset.unsqueeze(-2)
@@ -73,6 +79,10 @@ def to_absolute_actions(actions: Tensor, state: Tensor, mask: Sequence[bool]) ->
     # DeviceProcessorStep moves the transition, so it can be on CPU while actions are on CUDA.
     if state.device != actions.device or state.dtype != actions.dtype:
         state = state.to(device=actions.device, dtype=actions.dtype)
+    # Mirror to_relative_actions: collapse a temporally-stacked (B, T_obs, state_dim) state to the
+    # current frame (index 0) so the round-trip stays symmetric with the relative conversion.
+    if state.ndim == 3:
+        state = state[:, 0]
     state_offset = state[..., :dims] * mask_t
     if actions.ndim == 3:
         state_offset = state_offset.unsqueeze(-2)
@@ -126,7 +136,7 @@ class RelativeActionsProcessorStep(ProcessorStep):
         observation = transition.get(TransitionKey.OBSERVATION, {})
         state = observation.get(OBS_STATE) if observation else None
 
-        # Always cache state for the paired AbsoluteActionsProcessorStep
+        # Always cache state for the paired AbsoluteActionsProcessorStep.
         if state is not None:
             self._last_state = state
 
@@ -146,6 +156,11 @@ class RelativeActionsProcessorStep(ProcessorStep):
         """Return the cached ``observation.state`` used as the reference point for relative/absolute action conversions."""
         return self._last_state
 
+    def set_cached_state(self, state: torch.Tensor | None) -> None:
+        """Override the cached anchor state, e.g. to re-pin a chunk's anchor after the
+        per-tick pipeline overwrote it (see ``SyncInferenceEngine``)."""
+        self._last_state = state
+
     def get_config(self) -> dict[str, Any]:
         return {
             "enabled": self.enabled,

src/lerobot/rewards/pretrained.py

@@ -21,8 +21,6 @@ from pathlib import Path
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING, Any, TypeVar
 
-import packaging
-import safetensors
 from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download
 from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE
 from huggingface_hub.errors import HfHubHTTPError
@@ -129,29 +127,13 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC):
 
     @classmethod
     def _load_as_safetensor(cls, model: T, model_file: str, map_location: str, strict: bool) -> T:
-        # Create base kwargs
-        kwargs = {"strict": strict}
-
-        # Add device parameter for newer versions that support it
-        if packaging.version.parse(safetensors.__version__) >= packaging.version.parse("0.4.3"):
-            kwargs["device"] = map_location
-
-        # Load the model with appropriate kwargs
-        missing_keys, unexpected_keys = load_model_as_safetensor(model, model_file, **kwargs)
+        missing_keys, unexpected_keys = load_model_as_safetensor(
+            model, model_file, strict=strict, device=map_location
+        )
         if missing_keys:
             logging.warning(f"Missing key(s) when loading model: {missing_keys}")
         if unexpected_keys:
             logging.warning(f"Unexpected key(s) when loading model: {unexpected_keys}")
-
-        # For older versions, manually move to device if needed
-        if "device" not in kwargs and map_location != "cpu":
-            logging.warning(
-                "Loading model weights on other devices than 'cpu' is not supported natively in your version of safetensors."
-                " This means that the model is loaded on 'cpu' first and then copied to the device."
-                " This leads to a slower loading time."
-                " Please update safetensors to version 0.4.3 or above for improved performance."
-            )
-            model.to(map_location)
         return model
 
     def get_optim_params(self):

src/lerobot/rollout/configs.py

@@ -223,9 +223,22 @@ class RolloutConfig:
     fps: float = 30.0
     duration: float = 0.0  # 0 = infinite (24/7 mode)
     interpolation_multiplier: int = 1
+    # Safety net (opt-in): if any commanded joint's target differs from the robot's
+    # currently-measured position by more than this many units in a single control
+    # step, the action is treated as an unsafe jump — the robot is NOT commanded and
+    # the rollout stops. Units are the robot's raw `.pos` values (degrees for the
+    # SO-arms). Gripper joints are excluded (different unit/range). `None` disables it.
+    # Guards against chunk-splice / bad-chunk slams; a normal per-tick move at fps=30
+    # is only a few degrees, so a threshold like 20-30 catches slams without tripping
+    # on legitimate fast motion.
+    max_action_jump_deg: float | None = None
     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 +268,26 @@ 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,
+            )
+
+        if self.max_action_jump_deg is not None and self.max_action_jump_deg <= 0:
+            raise ValueError(
+                f"max_action_jump_deg must be positive when set, got {self.max_action_jump_deg}"
+            )
+
         # --- 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

@@ -43,7 +43,6 @@ from lerobot.processor import (
     make_default_processors,
     rename_stats,
 )
-from lerobot.processor.relative_action_processor import RelativeActionsProcessorStep
 from lerobot.robots import make_robot_from_config
 from lerobot.teleoperators import Teleoperator, make_teleoperator_from_config
 from lerobot.utils.feature_utils import combine_feature_dicts, hw_to_dataset_features
@@ -52,7 +51,6 @@ from .configs import BaseStrategyConfig, DAggerStrategyConfig, RolloutConfig
 from .inference import (
     InferenceEngine,
     RTCInferenceConfig,
-    SyncInferenceConfig,
     create_inference_engine,
 )
 from .robot_wrapper import ThreadSafeRobot
@@ -399,15 +397,6 @@ def build_rollout_context(
         },
     )
 
-    if isinstance(cfg.inference, SyncInferenceConfig) and any(
-        isinstance(step, RelativeActionsProcessorStep) and step.enabled
-        for step in getattr(preprocessor, "steps", ())
-    ):
-        raise NotImplementedError(
-            "SyncInferenceEngine does not support policies with relative actions for now."
-            "Use --inference.type=rtc or remove relative action processor steps from the policy pipeline."
-        )
-
     # --- 7. Inference strategy (needs policy + pre/post + hardware) --
     logger.info(
         "Creating inference engine (type=%s)...",
@@ -429,6 +418,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(
@@ -116,7 +118,8 @@ def create_inference_engine(
             postprocessor=postprocessor,
             robot_wrapper=robot_wrapper,
             rtc_config=config.rtc,
-            hw_features=hw_features,
+            dataset_features=dataset_features,
+            ordered_action_keys=ordered_action_keys,
             task=task,
             fps=fps,
             device=device,

src/lerobot/rollout/inference/rtc.py

@@ -34,12 +34,13 @@ import torch
 from lerobot.policies.pretrained import PreTrainedPolicy
 from lerobot.policies.rtc import ActionQueue, LatencyTracker, reanchor_relative_rtc_prefix
 from lerobot.policies.rtc.configuration_rtc import RTCConfig
-from lerobot.policies.utils import prepare_observation_for_inference
+from lerobot.policies.utils import make_robot_action, prepare_observation_for_inference
 from lerobot.processor import (
     NormalizerProcessorStep,
     PolicyProcessorPipeline,
     RelativeActionsProcessorStep,
 )
+from lerobot.utils.constants import ACTION
 from lerobot.utils.feature_utils import build_dataset_frame
 
 from ..robot_wrapper import ThreadSafeRobot
@@ -63,17 +64,28 @@ _RTC_JOIN_TIMEOUT_S: float = 3.0
 
 
 def _normalize_prev_actions_length(prev_actions: torch.Tensor, target_steps: int) -> torch.Tensor:
-    """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference."""
+    """Pad or truncate RTC prefix actions to a fixed length for stable compiled inference.
+
+    Padding repeats the last real action ("hold") rather than filling with zeros. The RTC
+    guidance pulls the new chunk toward this prefix at the padded indices (they fall inside
+    the weighted region when the real leftover is shorter than ``target_steps``). A zero in
+    the model's normalized action space decodes to the dataset *mean* action — a nonzero
+    offset that yanks the spliced action toward a mean/neutral pose for one step, producing
+    an intermittent seam (e.g. 95 -> 103 -> 95). Holding the last real action keeps the
+    padded targets continuous with the prefix, so no fake target enters the guided region.
+    The fixed output length is preserved so ``torch.compile`` policies keep stable shapes.
+    """
     if prev_actions.ndim != 2:
         raise ValueError(f"Expected 2D [T, A] tensor, got shape={tuple(prev_actions.shape)}")
-    steps, action_dim = prev_actions.shape
+    steps, _ = prev_actions.shape
     if steps == target_steps:
         return prev_actions
     if steps > target_steps:
         return prev_actions[:target_steps]
-    padded = torch.zeros((target_steps, action_dim), dtype=prev_actions.dtype, device=prev_actions.device)
-    padded[:steps] = prev_actions
-    return padded
+    if steps == 0:
+        raise ValueError("Cannot pad an empty prefix: no last action to hold.")
+    hold = prev_actions[-1:].expand(target_steps - steps, -1)
+    return torch.cat([prev_actions, hold], dim=0)
 
 
 # ---------------------------------------------------------------------------
@@ -97,7 +109,8 @@ class RTCInferenceEngine(InferenceEngine):
         postprocessor: PolicyProcessorPipeline,
         robot_wrapper: ThreadSafeRobot,
         rtc_config: RTCConfig,
-        hw_features: dict,
+        dataset_features: dict,
+        ordered_action_keys: list[str],
         task: str,
         fps: float,
         device: str | None,
@@ -111,7 +124,31 @@ class RTCInferenceEngine(InferenceEngine):
         self._postprocessor = postprocessor
         self._robot = robot_wrapper
         self._rtc_config = rtc_config
-        self._hw_features = hw_features
+        # Build observations with the SAME feature spec sync uses (post
+        # `robot_observation_processor`), not the raw-hardware spec. `build_dataset_frame`
+        # orders `observation.state` by this spec's `names`; using the raw-hardware order
+        # here (as before) desynced the state vector from sync whenever the observation
+        # processor reorders/renames state keys, corrupting both normalization and the
+        # relative-action anchor. The `prefix="observation"` filter ignores the action
+        # entries in the combined dict.
+        self._obs_features = dataset_features
+        # The model emits actions in `dataset_features[ACTION]` order (the order it was
+        # trained on); the robot expects them in `ordered_action_keys` order. Sync remaps
+        # by NAME via `make_robot_action` + reindex (sync.py) before returning; RTC must do
+        # the SAME, otherwise the engine-agnostic strategy (`send_next_action`) maps the raw
+        # model-order tensor onto `ordered_action_keys` positionally and mis-assigns joints
+        # whenever the two orders differ — a per-joint permutation that drives the arm wrong.
+        self._ordered_action_keys = ordered_action_keys
+        state_ft = dataset_features.get("observation.state")
+        if state_ft is not None:
+            logger.info("RTC observation.state layout: %s", state_ft.get("names"))
+        action_ft = dataset_features.get(ACTION)
+        if action_ft is not None:
+            logger.info(
+                "RTC action layout: model/dataset=%s -> robot=%s",
+                action_ft.get("names"),
+                self._ordered_action_keys,
+            )
         self._task = task
         self._fps = fps
         self._device = device or "cpu"
@@ -230,10 +267,19 @@ class RTCInferenceEngine(InferenceEngine):
     # ------------------------------------------------------------------
 
     def get_action(self, obs_frame: dict | None) -> torch.Tensor | None:
-        """Pop the next action from the RTC queue (ignores ``obs_frame``)."""
+        """Pop the next action from the RTC queue (ignores ``obs_frame``).
+
+        The queued action is in the model's ``dataset_features[ACTION]`` order; remap it
+        by NAME into ``ordered_action_keys`` order before returning, so the engine-agnostic
+        strategy maps values onto the correct joints. Mirrors ``SyncInferenceEngine.get_action``.
+        """
         if self._action_queue is None:
             return None
-        return self._action_queue.get()
+        action = self._action_queue.get()
+        if action is None:
+            return None
+        action_dict = make_robot_action(action, self._obs_features)
+        return torch.tensor([action_dict[k] for k in self._ordered_action_keys])
 
     def notify_observation(self, obs: dict) -> None:
         """Publish the latest observation for the RTC thread to consume."""
@@ -252,6 +298,8 @@ class RTCInferenceEngine(InferenceEngine):
             policy_device = torch.device(self._device)
 
             warmup_required = max(1, self._compile_warmup_inferences) if self._use_torch_compile else 0
+            # exclude the first N inferences from the latency tracker to avoid cold-start spikes
+            latency_warmup_required = max(1, warmup_required)
             inference_count = 0
             consecutive_errors = 0
 
@@ -273,10 +321,10 @@ class RTCInferenceEngine(InferenceEngine):
                         idx_before = queue.get_action_index()
                         prev_actions = queue.get_left_over()
 
-                        latency = latency_tracker.max()
+                        latency = latency_tracker.p95()
                         delay = math.ceil(latency / time_per_chunk) if latency else 0
 
-                        obs_batch = build_dataset_frame(self._hw_features, obs, prefix="observation")
+                        obs_batch = build_dataset_frame(self._obs_features, obs, prefix="observation")
                         obs_batch = prepare_observation_for_inference(
                             obs_batch, policy_device, self._task, self._robot.robot_type
                         )
@@ -316,7 +364,8 @@ class RTCInferenceEngine(InferenceEngine):
                         inference_count += 1
                         consecutive_errors = 0
                         is_warmup = self._use_torch_compile and inference_count <= warmup_required
-                        if is_warmup:
+                        # Ignore the first N inferences for latency tracking to avoid cold-start spikes
+                        if inference_count <= latency_warmup_required:
                             latency_tracker.reset()
                         else:
                             latency_tracker.add(new_latency)

src/lerobot/rollout/inference/sync.py

@@ -22,28 +22,23 @@ 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
+from lerobot.processor import PolicyProcessorPipeline, RelativeActionsProcessorStep
 
 from .base import InferenceEngine
 
 logger = logging.getLogger(__name__)
 
 
-# TODO(Steven): support relative-action policies.  The per-tick flow refreshes
-# ``RelativeActionsProcessorStep._last_state`` every call, so cached chunk
-# actions popped on later ticks get reanchored to the *current* robot state and
-# absolute targets drift through the chunk.  Relative-action policies are
-# rejected at context-build time today; RTC postprocesses the whole chunk and
-# is unaffected.
-#
-# Candidate fix: drive the policy via ``predict_action_chunk`` and serve a
-# local FIFO of postprocessed actions.  Eliminates drift by construction and
-# saves per-tick pre/post work, but bypasses ``select_action`` — needs
-# fallbacks for SAC (raises), ACT temporal ensembling (ensembler lives in
-# ``select_action``), and Diffusion-family (obs-history queues populated as a
-# side effect of ``select_action``).
+# Relative-action support: a predicted chunk of offsets is anchored to the robot
+# state at prediction time, but the sync engine reruns the pre/post pipeline every
+# tick, so ``RelativeActionsProcessorStep`` would re-anchor cached actions to the
+# current (moved) state and drift through the chunk. We pin the anchor per chunk:
+# a probe on the policy's public ``predict_action_chunk`` flags the ticks that
+# predict a fresh chunk; on the others the engine restores the anchor the relative
+# step overwrote. ``select_action`` stays on the hot path, so per-tick side effects
+# (e.g. LingBot-VA keyframe feedback) are preserved.
 
 
 class SyncInferenceEngine(InferenceEngine):
@@ -64,6 +59,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 +69,45 @@ class SyncInferenceEngine(InferenceEngine):
         self._task = task
         self._device = torch.device(device or "cpu")
         self._robot_type = robot_type
+
+        # Find an enabled RelativeActionsProcessorStep to pin its anchor per chunk
+        # (see module comment), mirroring the RTC engine.
+        self._relative_step = next(
+            (
+                s
+                for s in getattr(preprocessor, "steps", ())
+                if isinstance(s, RelativeActionsProcessorStep) and s.enabled
+            ),
+            None,
+        )
+        # Set by the probe for the current tick / ever, respectively.
+        self._chunk_predicted = False
+        self._ever_predicted_chunk = False
+        self._original_predict_action_chunk = None  # set while the probe is installed
+        if self._relative_step is not None:
+            # ``action_names`` is optional on the step; fill it lazily from the
+            # policy/dataset so the relative<->absolute mask is built correctly. This is
+            # a deliberate engine->step side effect (the step is configured by its consumer).
+            if self._relative_step.action_names is None:
+                cfg_names = getattr(policy.config, "action_feature_names", None)
+                self._relative_step.action_names = list(cfg_names) if cfg_names else list(ordered_action_keys)
+            self._install_chunk_probe()
+            logger.info("Relative actions enabled: chunk anchor pinned per predicted chunk")
+
+        # 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:
@@ -85,6 +116,11 @@ class SyncInferenceEngine(InferenceEngine):
 
     def stop(self) -> None:
         """No background resources to stop."""
+        # Undo the probe so the policy object isn't left permanently patched
+        # (it may outlive this engine or be reused by another).
+        if self._original_predict_action_chunk is not None:
+            self._policy.predict_action_chunk = self._original_predict_action_chunk
+            self._original_predict_action_chunk = None
         logger.info("SyncInferenceEngine stopped")
 
     def reset(self) -> None:
@@ -93,6 +129,54 @@ class SyncInferenceEngine(InferenceEngine):
         self._policy.reset()
         self._preprocessor.reset()
         self._postprocessor.reset()
+        # New episode: the next tick predicts a fresh chunk and re-anchors.
+        self._chunk_predicted = False
+        self._ever_predicted_chunk = False
+        self._pred_stacks = {}
+        self._pred_cursor = 0
+
+    def _install_chunk_probe(self) -> None:
+        """Wrap the policy's public ``predict_action_chunk`` so we learn which ticks
+        predict a fresh chunk (when the anchor must advance) without introspecting any
+        private action queue. Chunking policies call it from ``select_action``.
+
+        Wraps whatever callable is currently bound (e.g. an already-``torch.compile``d
+        one, since ``build_rollout_context`` compiles before building the engine); undone
+        in ``stop()``."""
+        self._original_predict_action_chunk = self._policy.predict_action_chunk
+        inner = self._original_predict_action_chunk
+
+        def probe(*args, **kwargs):
+            self._chunk_predicted = True
+            self._ever_predicted_chunk = True
+            return inner(*args, **kwargs)
+
+        self._policy.predict_action_chunk = probe
+
+    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."""
@@ -107,12 +191,34 @@ class SyncInferenceEngine(InferenceEngine):
             if self._device.type == "cuda" and self._policy.config.use_amp
             else nullcontext()
         )
+        # Snapshot the chunk anchor before the preprocessor overwrites it with this
+        # tick's state; restore it below if this tick only served a cached action.
+        # ``clone`` so the snapshot survives even if the cached tensor is ever mutated
+        # in place (today it is only rebound, but the copy is cheap for a state vector).
+        anchor_before = None
+        if self._relative_step is not None:
+            cached = self._relative_step.get_cached_state()
+            anchor_before = cached.clone() if cached is not None else None
+        self._chunk_predicted = False
         with torch.inference_mode(), autocast_ctx:
             observation = prepare_observation_for_inference(
                 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)
+            # Hold the anchor only for a chunking policy serving a cached action this
+            # tick; policies that never chunk or that recomputed keep refreshing.
+            if self._relative_step is not None and self._ever_predicted_chunk and not self._chunk_predicted:
+                self._relative_step.set_cached_state(anchor_before)
             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,
         )
 
@@ -267,6 +273,36 @@ def estimate_max_episode_seconds(
 # ---------------------------------------------------------------------------
 
 
+def detect_unsafe_action_jumps(
+    action_dict: dict,
+    obs_raw: dict,
+    max_jump: float,
+    exclude_substrings: tuple[str, ...] = ("gripper",),
+) -> dict[str, float]:
+    """Return ``{joint: jump}`` for joints whose commanded target exceeds the current
+    measured position by more than ``max_jump`` in a single control step; empty if safe.
+
+    Only keys present in both ``action_dict`` and ``obs_raw`` and not matching an
+    excluded substring are checked (grippers use a different unit/range). A large
+    single-step jump is the signature of a chunk-splice slam: a well-behaved trajectory
+    advances a joint by only a few units per tick.
+    """
+    violations: dict[str, float] = {}
+    for key, target in action_dict.items():
+        if any(token in key for token in exclude_substrings):
+            continue
+        current = obs_raw.get(key)
+        if current is None:
+            continue
+        try:
+            jump = abs(float(target) - float(current))
+        except (TypeError, ValueError):
+            continue
+        if jump > max_jump:
+            violations[key] = jump
+    return violations
+
+
 def send_next_action(
     obs_processed: dict,
     obs_raw: dict,
@@ -300,6 +336,25 @@ def send_next_action(
     if len(interp) != len(ordered_keys):
         raise ValueError(f"Interpolated tensor length ({len(interp)}) != action keys ({len(ordered_keys)})")
     action_dict = {k: interp[i].item() for i, k in enumerate(ordered_keys)}
+
+    # Safety net: refuse to command a single-step joint jump larger than the configured
+    # limit and stop the rollout. Catches chunk-splice / bad-chunk slams before they reach
+    # the motors. Compares the commanded target against the freshly-measured pose in
+    # `obs_raw` (same `.pos` units); no-op when the limit is unset.
+    max_jump = ctx.runtime.cfg.max_action_jump_deg
+    if max_jump is not None:
+        unsafe = detect_unsafe_action_jumps(action_dict, obs_raw, max_jump)
+        if unsafe:
+            detail = ", ".join(f"{k}={v:.1f}" for k, v in sorted(unsafe.items(), key=lambda kv: -kv[1]))
+            logger.error(
+                "SAFETY STOP: commanded joint jump exceeds %.1f in one step (%s); "
+                "refusing to send action and signalling shutdown.",
+                max_jump,
+                detail,
+            )
+            ctx.runtime.shutdown_event.set()
+            return None
+
     processed = ctx.processors.robot_action_processor((action_dict, obs_raw))
     ctx.hardware.robot_wrapper.send_action(processed)
     return action_dict

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/scripts/lerobot_train.py

@@ -20,6 +20,7 @@ Requires: pip install 'lerobot[training]'  (includes dataset + accelerate + wand
 
 import dataclasses
 import logging
+import os
 import sys
 import time
 from contextlib import nullcontext
@@ -171,6 +172,9 @@ def update_policy(
     train_metrics.update_s = time.perf_counter() - start_time
     if torch.cuda.is_available():
         train_metrics.gpu_mem_gb = torch.cuda.max_memory_allocated() / (1024**3)
+    # Aggregate the policy's scalar outputs for logging and rank-reduction across the log window.
+    if output_dict:
+        train_metrics.update_metrics(output_dict)
     return train_metrics, output_dict
 
 
@@ -461,6 +465,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
     # declares language columns; otherwise stay on PyTorch's default
     # collate so non-language training runs are unaffected.
     collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None
+    # Multi-node fork-OOM mitigation: on EFA nodes (vm.overcommit_memory=0, no swap),
+    # forking dataloader workers from a multi-GB rank reserve-charges the rank's full virtual
+    # footprint, so 8 ranks x num_workers forking at once trips OSError(ENOMEM) despite free
+    # RAM. Honor LEROBOT_DATALOADER_MP_CONTEXT (e.g. "forkserver"/"spawn") to spawn workers
+    # from a clean context instead of fork(). Only meaningful when workers are used.
+    dataloader_mp_context = os.environ.get("LEROBOT_DATALOADER_MP_CONTEXT") or None
+    if cfg.num_workers == 0:
+        dataloader_mp_context = None
     dataloader = torch.utils.data.DataLoader(
         dataset,
         num_workers=cfg.num_workers,
@@ -472,6 +484,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         collate_fn=collate_fn,
         prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
         persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
+        multiprocessing_context=dataloader_mp_context,
     )
 
     # Build eval dataloader if a held-out split exists
@@ -499,6 +512,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             collate_fn=eval_collate_fn,
             prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None,
             persistent_workers=cfg.persistent_workers and cfg.num_workers > 0,
+            multiprocessing_context=dataloader_mp_context,
         )
 
     # Prepare everything with accelerator
@@ -572,7 +586,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         batch = preprocessor(batch)
         train_tracker.dataloading_s = time.perf_counter() - start_time
 
-        train_tracker, output_dict = update_policy(
+        train_tracker, _ = update_policy(
             train_tracker,
             policy,
             batch,
@@ -605,9 +619,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
                     train_tracker.samples_per_s = effective_batch_size / step_time
                 logging.info(train_tracker)
                 if wandb_logger:
+                    # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the
+                    # tracker by update_policy, so to_dict() already carries their windowed,
+                    # rank-reduced averages — no per-step output_dict passthrough needed.
                     wandb_log_dict = train_tracker.to_dict()
-                    if output_dict:
-                        wandb_log_dict.update(output_dict)
                     # Log sample weighting statistics if enabled
                     if sample_weighter is not None:
                         weighter_stats = sample_weighter.get_stats()

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/logging_utils.py

@@ -104,6 +104,7 @@ class MetricsTracker:
         "episodes",
         "epochs",
         "accelerator",
+        "_caller_metrics",
     ]
 
     def __init__(
@@ -129,6 +130,9 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
         self.accelerator = accelerator
+        # Meter names the caller registered up front. update_metrics() leaves these untouched, so a
+        # policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter.
+        self._caller_metrics: set[str] = set(self.metrics)
 
     def __getattr__(self, name: str) -> int | dict[str, AverageMeter] | AverageMeter | Any:
         if name in self.__dict__:
@@ -156,6 +160,21 @@ class MetricsTracker:
         self.episodes = self.samples / self._avg_samples_per_ep
         self.epochs = self.samples / self._num_frames
 
+    def update_metrics(self, values: dict[str, Any]) -> None:
+        """Accumulate a dict of scalar metrics, auto-registering a meter for each new key.
+
+        Non-numeric values and bools are ignored.
+        Caller-registered metrics (those passed to the constructor) are never overridden.
+        """
+        for name, value in values.items():
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                continue
+            if name in self._caller_metrics:
+                continue
+            if name not in self.metrics:
+                self.metrics[name] = AverageMeter(name, ":.3f", reduction="mean")
+            self.metrics[name].update(float(value))
+
     def reduce_across_ranks(self) -> None:
         """
         Synchronises the running averages of every metric whose ``reduction`` is not ``"none"``

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}.")
 

tests/policies/test_relative_actions.py

@@ -346,3 +346,10 @@ def test_state_not_modified_by_relative_processor(dataset, action_dim):
 
     result_state = result[TransitionKey.OBSERVATION][OBS_STATE]
     torch.testing.assert_close(result_state, original_state)
+
+
+def test_cached_anchor_not_in_config():
+    """The cached anchor is ephemeral runtime state and must not leak into the config."""
+    step = RelativeActionsProcessorStep(enabled=True)
+    step.set_cached_state(torch.tensor([[1.0, 2.0, 3.0, 4.0]]))
+    assert set(step.get_config()) == {"enabled", "exclude_joints", "action_names"}

tests/test_rollout.py

@@ -348,3 +348,200 @@ def test_rollout_context_fields():
 
     field_names = {f.name for f in dataclasses.fields(RolloutContext)}
     assert field_names == {"runtime", "hardware", "policy", "processors", "data"}
+
+
+# ---------------------------------------------------------------------------
+# Sync engine: relative-action anchoring (drift-free chunk execution)
+# ---------------------------------------------------------------------------
+
+_REL_ACTION_NAMES = ["j0.pos", "j1.pos", "j2.pos", "gripper.pos"]
+_REL_ACTION_DIM = len(_REL_ACTION_NAMES)
+
+
+def _relative_pre_post(exclude_joints=None):
+    """Pre/post processors wrapping the real relative (caches anchor) and absolute
+    (relative + cached state) steps, mirroring what the sync engine feeds them."""
+    from lerobot.processor import (
+        AbsoluteActionsProcessorStep,
+        RelativeActionsProcessorStep,
+        TransitionKey,
+        create_transition,
+    )
+    from lerobot.utils.constants import OBS_STATE
+
+    relative_step = RelativeActionsProcessorStep(
+        enabled=True, exclude_joints=exclude_joints or [], action_names=list(_REL_ACTION_NAMES)
+    )
+    absolute_step = AbsoluteActionsProcessorStep(enabled=True, relative_step=relative_step)
+
+    class _Pre:
+        steps = [relative_step]
+
+        def __call__(self, observation):
+            # Run the relative step so it caches the anchor, then pass the batch through.
+            transition = create_transition(observation={OBS_STATE: observation[OBS_STATE]})
+            relative_step(transition)
+            return observation
+
+        def reset(self):
+            pass
+
+    class _Post:
+        def __call__(self, action):
+            transition = create_transition(action=action)
+            return absolute_step(transition)[TransitionKey.ACTION]
+
+        def reset(self):
+            pass
+
+    return _Pre(), _Post(), relative_step
+
+
+def _fake_relative_policy(chunk_rel, n_action_steps, chunking=True):
+    """Fake relative-action policy for the sync engine.
+
+    ``chunking=True`` buffers a chunk and serves it one action per tick, calling the
+    public ``predict_action_chunk`` only on refill (pi0/fastwam/lingbot). ``False``
+    returns an action directly and never calls it. The engine's anchor probe keys off
+    that public call, so the fake routes through it rather than any private queue.
+    """
+    from collections import deque
+
+    policy = MagicMock()
+    policy.config.use_amp = False
+    policy.config.action_feature_names = list(_REL_ACTION_NAMES)
+    state = {"predict_calls": 0}
+    queue = deque(maxlen=n_action_steps)
+
+    def predict_action_chunk(_batch=None, **_kwargs):
+        state["predict_calls"] += 1
+        return chunk_rel.unsqueeze(0)  # [B=1, n, dim]
+
+    def select_action(_observation):
+        if not chunking:
+            return chunk_rel[0].unsqueeze(0)
+        if len(queue) == 0:
+            actions = policy.predict_action_chunk(_observation)
+            queue.extend(actions.transpose(0, 1))  # [n, 1, dim]
+        return queue.popleft()
+
+    policy.predict_action_chunk.side_effect = predict_action_chunk
+    policy.select_action.side_effect = select_action
+    policy.reset.side_effect = queue.clear
+    policy._predict_state = state
+    return policy
+
+
+def _build_sync_engine(policy, pre, post):
+    from lerobot.rollout import SyncInferenceEngine
+
+    return SyncInferenceEngine(
+        policy=policy,
+        preprocessor=pre,
+        postprocessor=post,
+        dataset_features={"action": {"names": list(_REL_ACTION_NAMES)}},
+        ordered_action_keys=list(_REL_ACTION_NAMES),
+        task="test",
+        device="cpu",
+        robot_type="mock",
+    )
+
+
+def _obs_frame(state_values):
+    import numpy as np
+
+    return {"observation.state": np.asarray(state_values, dtype=np.float32)}
+
+
+def test_sync_relative_holds_anchor_across_chunk():
+    """Every action popped within a chunk must anchor to the tick-0 state (no drift)."""
+    n = 4
+    # A distinct relative offset per chunk step so a wrong anchor would be visible.
+    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.1 * (i + 1)) for i in range(n)])
+    pre, post, relative_step = _relative_pre_post()
+    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
+    engine = _build_sync_engine(policy, pre, post)
+
+    assert engine._relative_step is relative_step  # introspection wired the step
+
+    s0 = [1.0, 2.0, 3.0, 4.0]
+    outputs = []
+    for tick in range(n):
+        # Feed a *different* state each tick; a drifting anchor would use it.
+        state = [v + tick for v in s0]
+        outputs.append(engine.get_action(_obs_frame(state)))
+
+    # Exactly one chunk was predicted across the n ticks.
+    assert policy._predict_state["predict_calls"] == 1
+    for tick in range(n):
+        expected = torch.tensor(s0) + chunk_rel[tick]
+        torch.testing.assert_close(outputs[tick], expected)
+
+    # Next tick empties the queue -> fresh chunk -> anchor advances to the new state.
+    s_next = [10.0, 20.0, 30.0, 40.0]
+    out = engine.get_action(_obs_frame(s_next))
+    assert policy._predict_state["predict_calls"] == 2
+    torch.testing.assert_close(out, torch.tensor(s_next) + chunk_rel[0])
+    # The anchor now reflects the fresh-chunk state, not the held one.
+    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_next]))
+
+
+def test_sync_relative_reset_reanchors_new_episode():
+    """After ``reset()`` the first tick of the next episode anchors to the new state."""
+    n = 3
+    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
+    pre, post, relative_step = _relative_pre_post()
+    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
+    engine = _build_sync_engine(policy, pre, post)
+
+    # Episode 1: one tick anchors to s0 and leaves cached actions in the queue.
+    engine.get_action(_obs_frame([1.0, 1.0, 1.0, 1.0]))
+    assert policy._predict_state["predict_calls"] == 1
+
+    engine.reset()  # clears the queue and the per-episode chunk flags
+
+    # Episode 2: a fresh state must produce a fresh chunk anchored to that state,
+    # not carry over the previous episode's anchor.
+    s_new = [7.0, 8.0, 9.0, 10.0]
+    out = engine.get_action(_obs_frame(s_new))
+    assert policy._predict_state["predict_calls"] == 2
+    torch.testing.assert_close(out, torch.tensor(s_new) + chunk_rel[0])
+    torch.testing.assert_close(relative_step.get_cached_state(), torch.tensor([s_new]))
+
+
+def test_sync_relative_non_chunking_policy_refreshes_every_tick():
+    """A policy that never calls ``predict_action_chunk`` must not freeze the anchor."""
+    n = 3
+    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.5) for _ in range(n)])
+    pre, post, _ = _relative_pre_post()
+    policy = _fake_relative_policy(chunk_rel, n_action_steps=n, chunking=False)
+    engine = _build_sync_engine(policy, pre, post)
+
+    s0 = [1.0, 1.0, 1.0, 1.0]
+    for tick in range(3):
+        state = [v + tick for v in s0]
+        out = engine.get_action(_obs_frame(state))
+        # Anchor must track the current state every tick (no chunk => no hold).
+        torch.testing.assert_close(out, torch.tensor(state) + chunk_rel[0])
+    assert policy._predict_state["predict_calls"] == 0
+
+
+def test_sync_engine_no_relative_step_is_none():
+    """Without an enabled relative step, the engine takes the plain select_action path."""
+    policy = MagicMock()
+    policy.config.use_amp = False
+    engine = _build_sync_engine(policy, MagicMock(steps=[]), MagicMock())
+    assert engine._relative_step is None
+
+
+def test_sync_relative_stop_restores_policy_method():
+    """``stop()`` un-patches the probe so the policy object isn't permanently modified."""
+    n = 3
+    chunk_rel = torch.stack([torch.full((_REL_ACTION_DIM,), 0.2) for _ in range(n)])
+    pre, post, _ = _relative_pre_post()
+    policy = _fake_relative_policy(chunk_rel, n_action_steps=n)
+    original = policy.predict_action_chunk
+    engine = _build_sync_engine(policy, pre, post)
+    assert policy.predict_action_chunk is not original  # probe installed
+    engine.stop()
+    assert policy.predict_action_chunk is original  # restored

tests/utils/test_logging_utils.py

@@ -233,3 +233,37 @@ def test_metrics_tracker_reduce_across_ranks_invokes_reduce():
     # accumulate against the cluster view rather than the stale per-rank sum.
     meter = tracker.update_s
     assert meter.sum / meter.count == pytest.approx(meter.avg)
+
+
+def test_metrics_tracker_update_metrics_registers_and_averages():
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
+    tracker.update_metrics({"latent_loss": 0.2, "action_loss": 0.4})
+    tracker.update_metrics({"latent_loss": 0.4, "action_loss": 0.6})
+
+    # New keys are auto-registered as mean-reduced meters and averaged over the window.
+    assert tracker.metrics["latent_loss"].reduction == "mean"
+    assert tracker.metrics["latent_loss"].avg == pytest.approx(0.3)
+    assert tracker.metrics["action_loss"].avg == pytest.approx(0.5)
+    assert tracker.to_dict()["latent_loss"] == pytest.approx(0.3)
+
+
+def test_metrics_tracker_update_metrics_skips_non_numeric():
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics={})
+    tracker.update_metrics({"loss": 0.5, "head_mode": "sparse", "enabled": True})
+
+    # strings and bools ignored
+    assert "loss" in tracker.metrics
+    assert "head_mode" not in tracker.metrics
+    assert "enabled" not in tracker.metrics
+
+
+def test_metrics_tracker_update_metrics_does_not_override_caller_meter():
+    # A policy that echoes "loss" in its output dict must not overwrite the caller-owned,
+    # already-aggregated loss meter.
+    metrics = {"loss": AverageMeter("loss", ":.3f", reduction="mean")}
+    tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics)
+    tracker.loss = 1.0  # caller-set optimized loss
+    tracker.update_metrics({"loss": 99.0, "latent_loss": 0.2})
+
+    assert tracker.metrics["loss"].avg == pytest.approx(1.0)  # snapshot ignored
+    assert tracker.metrics["latent_loss"].avg == pytest.approx(0.2)