refactoring processors

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

@@ -26,6 +26,7 @@ from torch import Tensor, nn
 from lerobot.policies.pretrained import PreTrainedPolicy, T
 from lerobot.policies.utils import populate_queues
 from lerobot.utils.constants import ACTION, OBS_STATE
+from lerobot.utils.device_utils import get_autocast_context
 from lerobot.utils.import_utils import _transformers_available, require_package
 
 if TYPE_CHECKING or _transformers_available:
@@ -183,7 +184,7 @@ class VLAJEPAModel(nn.Module):
             action_idx = action_mask.nonzero(as_tuple=True)
 
         device_type = next(self.parameters()).device.type
-        with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
+        with get_autocast_context(device_type, torch.bfloat16):
             last_hidden = self._qwen_last_decoder_hidden(qwen_inputs)  # [B, seq_len, H]
             b, _, h = last_hidden.shape
             embodied_action_tokens = last_hidden[embodied_idx[0], embodied_idx[1], :].view(b, -1, h)
@@ -194,8 +195,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 +228,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 +245,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).
+            elementwise = F.l1_loss(predicted_states, gt_states.float(), reduction="none")
+            return elementwise.mean(dim=tuple(range(1, elementwise.ndim)))
         return F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
 
     def _action_loss(
@@ -247,17 +257,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):
+        with get_autocast_context(device_type, 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 +287,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 +395,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 +423,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 +460,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

@@ -17,14 +17,19 @@ from __future__ import annotations
 from typing import Any
 
 import torch
+import torch.nn.functional as F  # noqa: N812
 
+from lerobot.configs import PipelineFeatureType, PolicyFeature
 from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
 from lerobot.processor import (
+    AbsoluteActionsProcessorStep,
     EnvTransition,
+    ObservationProcessorStep,
     PolicyAction,
     PolicyProcessorPipeline,
     ProcessorStep,
     ProcessorStepRegistry,
+    RelativeActionsProcessorStep,
     TransitionKey,
     UnnormalizerProcessorStep,
     make_default_policy_processor_steps,
@@ -32,6 +37,69 @@ from lerobot.processor import (
 )
 
 
+@ProcessorStepRegistry.register(name="vla_jepa_image_prep")
+class ImagePrepProcessorStep(ObservationProcessorStep):
+    """Prepares image observations for the VLA-JEPA model: float cast, 1->3 channel expand, resize.
+
+    This makes explicit (in the serialized pipeline) the image prep the model used to do
+    internally. The model keeps the same operations as idempotent guards, so:
+      - checkpoints saved WITHOUT this step (older uploads) are unaffected — the model still
+        does the prep;
+      - checkpoints saved WITH this step get it done here, and the model-side guards no-op.
+
+    Mirrors `Qwen3VLInterface.to_pixel_values` + the `F.interpolate(mode="area")` resize in
+    `VLAJEPAPolicy._prepare_model_inputs`/`predict_action`. Deliberately does NOT clamp (the
+    model path doesn't), so values stay bit-identical. Handles [C,H,W], [B,C,H,W]/[T,C,H,W]
+    and [B,T,C,H,W] image tensors.
+    """
+
+    def __init__(self, resize_to: tuple[int, int] | None = None, expand_channels: bool = True):
+        self.resize_to = tuple(resize_to) if resize_to is not None else None
+        self.expand_channels = expand_channels
+
+    def observation(self, observation: dict) -> dict:
+        new_observation = dict(observation)
+        for key in observation:
+            if "image" not in key:
+                continue
+            image = observation[key].float()
+            if self.expand_channels and image.shape[-3] == 1:
+                repeats = [1] * image.ndim
+                repeats[-3] = 3
+                image = image.repeat(*repeats)
+            if self.resize_to is not None and tuple(image.shape[-2:]) != self.resize_to:
+                device = image.device
+                # NOTE: no "area" kernel on mps; resize on cpu then move back.
+                if device.type == "mps":
+                    image = image.cpu()
+                lead = image.shape[:-3]
+                c, h, w = image.shape[-3:]
+                flat = image.reshape(-1, c, h, w)
+                flat = F.interpolate(flat, size=self.resize_to, mode="area")
+                image = flat.reshape(*lead, c, *self.resize_to).to(device)
+            new_observation[key] = image
+        return new_observation
+
+    def get_config(self) -> dict[str, Any]:
+        return {
+            "resize_to": list(self.resize_to) if self.resize_to is not None else None,
+            "expand_channels": self.expand_channels,
+        }
+
+    def transform_features(self, features):
+        for key in features[PipelineFeatureType.OBSERVATION]:
+            if "image" not in key:
+                continue
+            feat = features[PipelineFeatureType.OBSERVATION][key]
+            # Match `to_pixel_values`: only a single channel is expanded to 3.
+            nb_channel = 3 if (self.expand_channels and feat.shape[0] == 1) else feat.shape[0]
+            spatial = self.resize_to if self.resize_to is not None else tuple(feat.shape[1:])
+            features[PipelineFeatureType.OBSERVATION][key] = PolicyFeature(
+                type=feat.type, shape=(nb_channel, *spatial)
+            )
+        return features
+
+
 @ProcessorStepRegistry.register(name="vla_jepa_clip_actions")
 class ClipActionsProcessorStep(ProcessorStep):
     """Clips action tensor to [-1, 1] before unnormalization."""
@@ -109,10 +177,24 @@ def make_vla_jepa_pre_post_processors(
 ]:
     features = {**config.input_features, **config.output_features}
     steps = make_default_policy_processor_steps(config, dataset_stats)
+
+    # 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 = [
         steps.rename_observations,
         steps.add_batch_dim,
         steps.to_device,
+        ImagePrepProcessorStep(
+            resize_to=tuple(config.resize_images_to) if config.resize_images_to else None,
+        ),
+        relative_step,
         steps.normalize,
     ]
     output_steps: list[ProcessorStep] = []
@@ -131,6 +213,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/utils/device_utils.py

@@ -15,6 +15,7 @@
 # limitations under the License.
 
 import logging
+from contextlib import nullcontext
 
 import torch
 
@@ -121,3 +122,25 @@ def is_amp_available(device: str):
         return False
     else:
         raise ValueError(f"Unknown device '{device}.")
+
+
+def get_autocast_context(device_type: str, dtype: torch.dtype = torch.bfloat16):
+    """Return a device-safe autocast context manager.
+
+    Hardcoding `torch.autocast(dtype=torch.bfloat16)` breaks on backends without AMP
+    (MPS) and silently misbehaves on pre-Ampere CUDA GPUs that lack bf16 support. This
+    picks a safe context per device:
+      - no AMP support (e.g. mps): `nullcontext()` (run in the tensors' native dtype)
+      - CUDA requesting bf16 on compute capability < 8.0 (pre-Ampere): fall back to fp16
+      - otherwise: `torch.autocast(device_type, dtype)`
+    """
+    if not is_amp_available(device_type):
+        return nullcontext()
+    if (
+        device_type == "cuda"
+        and dtype == torch.bfloat16
+        and torch.cuda.is_available()
+        and torch.cuda.get_device_capability()[0] < 8
+    ):
+        dtype = torch.float16
+    return torch.autocast(device_type=device_type, dtype=dtype)

tests/policies/vla_jepa/test_processor_vla_jepa.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the VLA-JEPA image-prep processor step and its back-compat contract with the model.
+
+The step moves image resize + 1->3 channel-expand out of the model into the (serialized)
+preprocessor. The model keeps the same ops as idempotent guards, so:
+  - old checkpoints (JSON without the step) are unaffected — the model still does the prep;
+  - new checkpoints (JSON with the step) get it done in the step, and the model guards no-op.
+These tests pin the step's numerics (bit-identical to the model's F.interpolate(area)) and the
+equivalence of the two paths on the Qwen image path.
+"""
+
+from __future__ import annotations
+
+from copy import deepcopy
+
+import pytest
+import torch
+import torch.nn.functional as F  # noqa: N812
+
+pytest.importorskip("transformers")
+pytest.importorskip("diffusers")
+
+from conftest import (  # noqa: E402
+    BATCH_SIZE,
+    IMAGE_SIZE,
+    make_config,
+    make_inference_batch,
+    make_train_batch,
+)
+
+from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature  # noqa: E402
+from lerobot.policies.vla_jepa.modeling_vla_jepa import VLAJEPAPolicy  # noqa: E402
+from lerobot.policies.vla_jepa.processor_vla_jepa import (  # noqa: E402
+    ImagePrepProcessorStep,
+    make_vla_jepa_pre_post_processors,
+)
+from lerobot.processor import ProcessorStepRegistry  # noqa: E402
+from lerobot.utils.constants import OBS_IMAGES, OBS_STATE  # noqa: E402
+
+RESIZE = (IMAGE_SIZE // 2, IMAGE_SIZE // 2)  # (4, 4)
+IMG_KEY = f"{OBS_IMAGES}.laptop"
+
+
+# ---------------------------------------------------------------------------
+# Step numerics / shape handling
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "shape",
+    [
+        (3, IMAGE_SIZE, IMAGE_SIZE),  # [C, H, W] (raw single-sample inference)
+        (BATCH_SIZE, 3, IMAGE_SIZE, IMAGE_SIZE),  # [B, C, H, W]
+        (BATCH_SIZE, 2, 3, IMAGE_SIZE, IMAGE_SIZE),  # [B, T, C, H, W] (video stack)
+    ],
+)
+def test_image_prep_resize_shapes_and_area_numerics(shape: tuple[int, ...]) -> None:
+    step = ImagePrepProcessorStep(resize_to=RESIZE)
+    x = torch.rand(*shape)
+    out = step.observation({IMG_KEY: x})[IMG_KEY]
+
+    assert out.shape[:-2] == x.shape[:-2]  # leading + channel dims unchanged
+    assert tuple(out.shape[-2:]) == RESIZE
+    assert out.dtype == torch.float32
+
+    # bit-identical to the model-side F.interpolate(mode="area"), no clamp
+    ref = F.interpolate(x.float().reshape(-1, *x.shape[-3:]), size=RESIZE, mode="area").reshape(
+        *x.shape[:-2], *RESIZE
+    )
+    assert torch.equal(out, ref)
+
+
+def test_image_prep_channel_expand() -> None:
+    step = ImagePrepProcessorStep(resize_to=None, expand_channels=True)
+    x = torch.rand(BATCH_SIZE, 1, IMAGE_SIZE, IMAGE_SIZE)
+    out = step.observation({IMG_KEY: x})[IMG_KEY]
+    assert out.shape[1] == 3
+    # all three channels are copies of the single input channel
+    assert torch.equal(out[:, 0], x[:, 0]) and torch.equal(out[:, 1], x[:, 0])
+
+
+def test_image_prep_resize_skip_when_already_target_size() -> None:
+    step = ImagePrepProcessorStep(resize_to=RESIZE)
+    x = torch.rand(BATCH_SIZE, 3, *RESIZE)
+    out = step.observation({IMG_KEY: x})[IMG_KEY]
+    # size already matches -> only the float cast happens, values preserved exactly
+    assert torch.equal(out, x)
+
+
+def test_image_prep_leaves_non_image_keys_untouched() -> None:
+    step = ImagePrepProcessorStep(resize_to=RESIZE)
+    state = torch.randn(BATCH_SIZE, 4)
+    out = step.observation({IMG_KEY: torch.rand(BATCH_SIZE, 3, IMAGE_SIZE, IMAGE_SIZE), OBS_STATE: state})
+    assert torch.equal(out[OBS_STATE], state)
+
+
+def test_image_prep_config_roundtrip_via_registry() -> None:
+    step = ImagePrepProcessorStep(resize_to=RESIZE, expand_channels=True)
+    cfg = step.get_config()
+    assert cfg == {"resize_to": [RESIZE[0], RESIZE[1]], "expand_channels": True}
+    rebuilt = ProcessorStepRegistry.get("vla_jepa_image_prep")(**cfg)
+    assert rebuilt.resize_to == RESIZE
+    assert rebuilt.expand_channels is True
+
+
+def test_image_prep_transform_features() -> None:
+    step = ImagePrepProcessorStep(resize_to=RESIZE, expand_channels=True)
+    features = {
+        PipelineFeatureType.OBSERVATION: {
+            IMG_KEY: PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_SIZE, IMAGE_SIZE)),
+            "observation.images.depth": PolicyFeature(
+                type=FeatureType.VISUAL, shape=(1, IMAGE_SIZE, IMAGE_SIZE)
+            ),
+            OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(4,)),
+        }
+    }
+    out = step.transform_features(features)[PipelineFeatureType.OBSERVATION]
+    assert out[IMG_KEY].shape == (3, *RESIZE)  # already 3-channel, only resized
+    assert out["observation.images.depth"].shape == (3, *RESIZE)  # 1->3 expanded
+    assert out[OBS_STATE].shape == (4,)  # non-image untouched
+
+
+# ---------------------------------------------------------------------------
+# Pipeline wiring + back-compat with the model
+# ---------------------------------------------------------------------------
+
+
+def test_image_prep_step_wired_into_preprocessor() -> None:
+    cfg = make_config()
+    cfg.resize_images_to = RESIZE
+    preprocessor, _ = make_vla_jepa_pre_post_processors(cfg, dataset_stats=None)
+    prep_steps = [s for s in preprocessor.steps if isinstance(s, ImagePrepProcessorStep)]
+    assert len(prep_steps) == 1
+    assert prep_steps[0].resize_to == RESIZE
+
+
+@torch.no_grad()
+@pytest.mark.parametrize("batch_fn", [make_inference_batch, make_train_batch])
+def test_image_prep_matches_model_qwen_path(patch_vla_jepa_external_models: None, batch_fn) -> None:
+    """The Qwen image path is identical whether the step resized (new ckpt) or the model does (old ckpt).
+
+    Both use F.interpolate(mode="area"), so pre-resizing in the step then letting the model's
+    size guard no-op yields byte-identical Qwen inputs to the pure model path. This is the
+    contract that keeps already-uploaded checkpoints correct.
+    """
+    cfg = make_config()
+    cfg.resize_images_to = RESIZE
+    policy = VLAJEPAPolicy(cfg)
+    policy.eval()
+    training = batch_fn is make_train_batch
+
+    batch = batch_fn()
+
+    # Path A (old checkpoint, no processor step): the model resizes internally.
+    imgs_a = policy._prepare_model_inputs(deepcopy(batch), training=training)["images"]
+
+    # Path B (new checkpoint): the step resizes first; the model's guard becomes a no-op.
+    step = ImagePrepProcessorStep(resize_to=RESIZE)
+    resized = step.observation({IMG_KEY: batch[IMG_KEY]})
+    batch_b = deepcopy(batch)
+    batch_b[IMG_KEY] = resized[IMG_KEY]
+    imgs_b = policy._prepare_model_inputs(batch_b, training=training)["images"]
+
+    assert len(imgs_a) == len(imgs_b) == BATCH_SIZE
+    for views_a, views_b in zip(imgs_a, imgs_b, strict=True):
+        for a, b in zip(views_a, views_b, strict=True):
+            assert torch.equal(a, b)