refactor(vla-jepa): removing gpu roundtrip (#3750)

* refactor(vla-jepa): removing gpu roundtrip for the preprocessing part

* major refactor of the forward pass and model input conversion

* linting

* adressing suggestions from reviews
* removing redundant state dtype conversion
* avoiding recreating the same tensor each foward pass
* api simplification of `_encode_qwen`
* avoiding useless video assembly during inference
* guard against video=None for the wm loss

src/lerobot/policies/vla_jepa/modeling_vla_jepa.py

@@ -17,12 +17,10 @@ from __future__ import annotations
 import logging
 from collections import deque
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
-import numpy as np
 import torch
 import torch.nn.functional as F  # noqa: N812
-from PIL import Image
 from torch import Tensor, nn
 
 from lerobot.policies.pretrained import PreTrainedPolicy, T
@@ -55,12 +53,13 @@ class VLAJEPAModel(nn.Module):
       - DiT-B: flow-matching action head for future action prediction
       - V-JEPA: world model for video frame prediction
 
-    Input: List[dict] native format (same as original starVLA)
+    Inputs are batched tensors kept on the model device
-      - "image": List[PIL.Image] (multi-view images)
+      - images: List[List[Tensor [C, H, W]]] (float [0,1]) — per sample, per view (Qwen messages)
-      - "video": np.ndarray [V, T, H, W, 3]
+      - instructions: List[str]
-      - "lang": str (task instruction)
+      - videos: Tensor [B, V, T, C, H, W] (float [0,1], world model only)
-      - "action": np.ndarray [T, action_dim] (optional, training only)
+      - actions: Tensor [B, T, action_dim] (optional, training only)
-      - "state": np.ndarray [1, state_dim] (optional)
+      - state: Tensor [B, 1, state_dim] (optional)
+      - action_is_pad: Tensor [B, T] (optional)
     """
 
     def __init__(self, config: VLAJEPAConfig) -> None:
@@ -75,6 +74,11 @@ class VLAJEPAModel(nn.Module):
         self.action_tokens, self.action_token_ids, self.embodied_action_token_id = (
             self.qwen.expand_tokenizer()
         )
+        self.register_buffer(
+            "_action_token_ids_t",
+            torch.tensor(self.action_token_ids, dtype=torch.long),
+            persistent=False,
+        )
 
         # Action head (flow-matching DiT)
         self.action_model = VLAJEPAActionHead(config, cross_attention_dim=self.qwen.model.config.hidden_size)
@@ -161,166 +165,123 @@ class VLAJEPAModel(nn.Module):
 
     # ---- Native VLA-JEPA forward (follows original VLA_JEPA.py) ----
 
-    def forward(self, examples: list[dict]) -> dict[str, Tensor]:
+    def _encode_qwen(
-        """
+        self, images: list[list[Tensor]], instructions: list[str], *, need_action_tokens: bool
-        Native forward pass following original starVLA VLA_JEPA.forward.
+    ) -> tuple[Tensor, Tensor, Tensor | None]:
-
+        """Run Qwen and gather the embodied-action (and optionally action) token hidden states."""
-        Args:
-            examples: List of per-sample dicts with keys:
-                "image"  : List[PIL.Image]  — multi-view images
-                "video"  : np.ndarray [V, T, H, W, 3]
-                "lang"   : str — task instruction
-                "action" : np.ndarray [T, action_dim] (optional)
-                "state"  : np.ndarray [1, state_dim] (optional)
-
-        Returns:
-            dict with "action_loss" and "wm_loss" keys (scalar Tensors).
-        """
-        # Unpack native format (same pattern as original VLA_JEPA.py)
-        batch_images = [ex["image"] for ex in examples]  # List[List[PIL.Image]]
-        batch_videos = [ex["video"] for ex in examples]  # List[np.ndarray]
-        instructions = [ex["lang"] for ex in examples]  # List[str]
-        has_action = "action" in examples[0] and examples[0]["action"] is not None
-        actions = [ex["action"] for ex in examples] if has_action else None
-        has_state = "state" in examples[0] and examples[0]["state"] is not None
-        state = [ex["state"] for ex in examples] if has_state else None
-        action_is_pad = (
-            [ex["action_is_pad"] for ex in examples]
-            if has_action and "action_is_pad" in examples[0] and examples[0]["action_is_pad"] is not None
-            else None
-        )
-
-        # Stack videos: [B, V, T, H, W, 3] -> [B, V, T, 3, H, W]
-        batch_videos = np.stack(batch_videos)
-        batch_videos = batch_videos.transpose(0, 1, 2, 5, 3, 4)  # [B, V, T, 3, H, W]
-
-        # Adjust number of views for the world model:
-        # - fewer views than expected: duplicate the first view to fill up
-        # - more views than expected: keep only the first num_views_world_model views
-        num_views_world_model = self.config.jepa_tubelet_size
-        if batch_videos.shape[1] < num_views_world_model:
-            num_missing_views = num_views_world_model - batch_videos.shape[1]
-            first_view = np.repeat(batch_videos[:, :1], num_missing_views, axis=1)
-            batch_videos = np.concatenate([batch_videos, first_view], axis=1)
-        elif batch_videos.shape[1] > num_views_world_model:
-            batch_videos = batch_videos[:, :num_views_world_model]
-
-        # ---- Step 1: QwenVL encode (same as original) ----
         qwen_inputs = self.qwen.build_inputs(
-            images=batch_images,
+            images=images,
             instructions=instructions,
             action_prompt=self.replace_prompt,
             embodied_prompt=self.embodied_replace_prompt,
         )
-
+        input_ids = qwen_inputs["input_ids"]
-        # Locate embodied-action tokens (always needed for action head)
+        embodied_idx = (input_ids == self.embodied_action_token_id).nonzero(as_tuple=True)
-        embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
+        action_idx = None
-        embodied_indices = embodied_mask.nonzero(as_tuple=True)
+        if need_action_tokens:
-
+            action_mask = torch.isin(input_ids, self._action_token_ids_t)
-        # Locate action tokens (only needed for world model predictor)
+            action_idx = action_mask.nonzero(as_tuple=True)
-        if self.config.enable_world_model:
-            action_mask = torch.isin(
-                qwen_inputs["input_ids"],
-                torch.tensor(self.action_token_ids, device=qwen_inputs["input_ids"].device),
-            )
-            action_indices = action_mask.nonzero(as_tuple=True)
 
         device_type = next(self.parameters()).device.type
-
         with torch.autocast(device_type=device_type, dtype=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)
+            action_tokens = (
+                last_hidden[action_idx[0], action_idx[1], :].view(b, -1, h)
+                if action_idx is not None
+                else None
+            )
+        return embodied_action_tokens, action_tokens
 
-            if self.config.enable_world_model:
+    def _world_model_loss(self, videos: Tensor, action_tokens: Tensor) -> Tensor:
-                action_tokens = last_hidden[action_indices[0], action_indices[1], :].view(b, -1, h)
+        """JEPA encode + predictor L1 loss. `videos` is [B, V, T, C, H, W] float in [0, 1]."""
+        # 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:
+            missing = num_views - videos.shape[1]
+            videos = torch.cat([videos, videos[:, :1].repeat(1, missing, 1, 1, 1, 1)], dim=1)
+        elif videos.shape[1] > num_views:
+            videos = videos[:, :num_views]
 
-            embodied_action_tokens = last_hidden[embodied_indices[0], embodied_indices[1], :].view(b, -1, h)
+        b, v, t_frames, c, h_img, w_img = videos.shape
+        flat = videos.reshape(b * v, t_frames, c, h_img, w_img)
+        # Fast (torchvision) video processor on-device, do_rescale=False (frames already in [0, 1]).
+        video_pixels = self.video_processor(
+            videos=list(flat),
+            return_tensors="pt",
+            device=self.video_encoder.device,
+            do_rescale=False,
+        )["pixel_values_videos"]  # [B*V, T, C, H, W]
 
-        # ---- Step 2+3: JEPA Encoder + Predictor ----
+        with torch.no_grad():
-        device_wm = last_hidden.device
+            video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
-        if not self.config.enable_world_model:
+            # Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
-            wm_loss = torch.tensor(0.0, device=device_wm)
+            video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
+
+        tubelet_size = self.video_encoder.config.tubelet_size
+        # 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)
+
+        # Shift-by-one JEPA split: input_states = positions 0..T-2, gt_states = positions 1..T-1
+        t_enc_ctx = t_enc_total - 1
+        tokens_per_frame = video_embeddings.shape[1] // t_enc_total
+        input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
+        gt_states = video_embeddings[:, tokens_per_frame:, :]
+
+        expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
+        if action_tokens.shape[1] < expected_actions:
+            pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
+            action_tokens = torch.cat([action_tokens, pad], dim=1)
+
+        predicted_states = self.video_predictor(
+            input_states.float(), action_tokens[:, :expected_actions].float()
+        )
+        return F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
+
+    def _action_loss(
+        self,
+        embodied_action_tokens: Tensor,
+        actions: Tensor,
+        state: Tensor | None,
+        action_is_pad: Tensor | None,
+    ) -> Tensor:
+        """Flow-matching action-head loss, repeated over `repeated_diffusion_steps`."""
+        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
+            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)
+
+    def forward(
+        self,
+        images: list[list[Tensor]],
+        instructions: list[str],
+        videos: Tensor | None = None,
+        actions: Tensor | None = None,
+        state: Tensor | None = None,
+        action_is_pad: Tensor | None = None,
+    ) -> dict[str, Tensor]:
+        """Native forward: Qwen encode → optional world-model loss → optional action-head loss."""
+        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)
         else:
-            b, v, t_frames, c, h_img, w_img = batch_videos.shape
+            wm_loss = torch.zeros((), device=embodied_action_tokens.device)
-            batch_videos_flat = batch_videos.reshape(b * v, t_frames, c, h_img, w_img)
 
-            video_pixels = self.video_processor(videos=list(batch_videos_flat), return_tensors="pt")[
+        if actions is None:
-                "pixel_values_videos"
-            ].to(self.video_encoder.device)  # [B*V, T, C, H, W]
-
-            with torch.no_grad():
-                video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels)
-                # Merge views: [B*V, ...] -> [B, ..., V*embed_dim]
-                video_embeddings = torch.cat(torch.chunk(video_embeddings, chunks=v, dim=0), dim=2)
-
-            tubelet_size = self.video_encoder.config.tubelet_size
-            device_wm = video_embeddings.device
-            # 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:
-                wm_loss = torch.tensor(0.0, device=device_wm)
-            else:
-                # Shift-by-one JEPA split (matches original VLA_JEPA.py lines 231-232):
-                # input_states: positions 0..T-2, gt_states: positions 1..T-1
-                t_enc_ctx = t_enc_total - 1
-                tokens_per_frame = video_embeddings.shape[1] // t_enc_total
-
-                input_states = video_embeddings[:, : tokens_per_frame * t_enc_ctx, :]
-                gt_states = video_embeddings[:, tokens_per_frame:, :]
-
-                expected_actions = t_enc_ctx * self.config.num_action_tokens_per_timestep
-                if action_tokens.shape[1] < expected_actions:
-                    pad = action_tokens[:, -1:].repeat(1, expected_actions - action_tokens.shape[1], 1)
-                    action_tokens = torch.cat([action_tokens, pad], dim=1)
-
-                predicted_states = self.video_predictor(
-                    input_states.float(),
-                    action_tokens[:, :expected_actions].float(),
-                )
-
-                wm_loss = F.l1_loss(predicted_states, gt_states.float(), reduction="mean")
-
-        if not has_action:
             return {"wm_loss": wm_loss}
 
-        # ---- Step 4: Action Head ----
+        action_loss = self._action_loss(embodied_action_tokens, actions, state, action_is_pad)
-        with torch.autocast(device_type=device_type, dtype=torch.float32):
-            actions_tensor = torch.tensor(
-                np.array(actions), device=last_hidden.device, dtype=torch.float32
-            )  # [B, T_full, action_dim]
-            action_horizon = self.config.chunk_size
-            actions_target = actions_tensor[:, -action_horizon:, :]
-
-            state_tensor = None
-            if state is not None:
-                state_tensor = torch.tensor(
-                    np.array(state), device=last_hidden.device, dtype=last_hidden.dtype
-                )  # [B, 1, state_dim]
-
-            repeated_diffusion_steps = self.config.repeated_diffusion_steps
-            actions_target = actions_target.repeat(repeated_diffusion_steps, 1, 1)
-            embodied_action_tokens = embodied_action_tokens.repeat(repeated_diffusion_steps, 1, 1)
-            if state_tensor is not None:
-                state_tensor = state_tensor.repeat(repeated_diffusion_steps, 1, 1)
-
-            action_is_pad_rep = None
-            if action_is_pad is not None:
-                pad_tensor = torch.stack(
-                    [
-                        p.to(actions_target.device)
-                        if isinstance(p, Tensor)
-                        else torch.tensor(p, device=actions_target.device)
-                        for p in action_is_pad
-                    ]
-                )  # [B, T_full]
-                pad_tensor = pad_tensor[:, -action_horizon:]  # [B, action_horizon]
-                action_is_pad_rep = pad_tensor.repeat(repeated_diffusion_steps, 1)  # [B*R, action_horizon]
-
-            action_loss = self.action_model(
-                embodied_action_tokens, actions_target, state_tensor, action_is_pad_rep
-            )
-
         return {"action_loss": action_loss, "wm_loss": wm_loss * self.config.world_model_loss_weight}
 
     # ---- Native predict_action (follows original VLA_JEPA.predict_action) ----
@@ -328,58 +289,23 @@ class VLAJEPAModel(nn.Module):
     @torch.no_grad()
     def predict_action(
         self,
-        batch_images: list[list[Image.Image]],
+        images: list[list[Tensor]],
         instructions: list[str],
-        state: np.ndarray | None = None,
+        state: Tensor | None = None,
-    ) -> np.ndarray:
+    ) -> Tensor:
-        """
+        """Predict an action chunk. `images` is per-sample, per-view float [0,1] [C, H, W] tensors."""
-        Native action prediction following original VLA_JEPA.predict_action.
-
-        Args:
-            batch_images: List of samples; each is List[PIL.Image] (multi-view).
-            instructions: Task instructions, one per sample.
-            state: Optional [B, state_dim] numpy array.
-
-        Returns:
-            np.ndarray [B, action_horizon, action_dim] — predicted actions.
-        """
         if self.config.resize_images_to is not None:
             height, width = self.config.resize_images_to
-            resampling = getattr(Image, "Resampling", Image).BOX
+            images = [
-            batch_images = [
+                [F.interpolate(img[None], size=(height, width), mode="area")[0] for img in views]
-                [image.resize((width, height), resample=resampling) for image in sample_images]
+                for views in images
-                for sample_images in batch_images
             ]
 
-        qwen_inputs = self.qwen.build_inputs(
+        embodied_action_tokens, _ = self._encode_qwen(images, instructions, need_action_tokens=False)
-            images=batch_images,
+        return self.action_model.predict_action(
-            instructions=instructions,
+            embodied_action_tokens.float(), state.float() if state is not None else None
-            action_prompt=self.replace_prompt,
-            embodied_prompt=self.embodied_replace_prompt,
         )
 
-        embodied_mask = qwen_inputs["input_ids"] == self.embodied_action_token_id
-        embodied_indices = embodied_mask.nonzero(as_tuple=True)
-
-        device_type = next(self.parameters()).device.type
-
-        with torch.autocast(device_type=device_type, dtype=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_indices[0], embodied_indices[1], :].view(b, -1, h)
-
-        state_tensor = None
-        if state is not None:
-            state_tensor = torch.from_numpy(np.array(state)).to(
-                device=last_hidden.device, dtype=last_hidden.dtype
-            )
-
-        pred_actions = self.action_model.predict_action(
-            embodied_action_tokens.float(), state_tensor.float() if state_tensor is not None else None
-        )  # [B, action_horizon, action_dim]
-
-        return pred_actions.detach().cpu().numpy()
-
 
 # ============================================================================
 # LeRobot Adapter Layer - converts between LeRobot batch format and native VLA-JEPA format
@@ -390,9 +316,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
     """
     LeRobot adapter for VLA-JEPA.
 
-    Converts LeRobot's standard batch format (dict[str, Tensor]) to the native
+    Converts LeRobot's standard batch format (dict[str, Tensor]) to the batched tensors
-    VLA-JEPA format (List[dict]), calls the native model, and converts outputs
+    the native model expects (keeping everything on-device), calls the native model, and
-    back to LeRobot format.
+    converts outputs back to LeRobot format.
     """
 
     config_class = VLAJEPAConfig
@@ -419,9 +345,8 @@ class VLAJEPAPolicy(PreTrainedPolicy):
 
     # ---- Format Conversion: LeRobot → Native ----
 
-    def _prepare_model_inputs(self, batch: dict[str, Tensor]) -> list[dict]:
+    def _prepare_model_inputs(self, batch: dict[str, Tensor], training=True) -> dict[str, Any]:
-        """
+        """Convert a LeRobot batch to the model's batched, on-device inputs.
-        Convert LeRobot batch format to native VLA-JEPA examples format.
 
         LeRobot format:
             batch = {
@@ -431,65 +356,25 @@ class VLAJEPAPolicy(PreTrainedPolicy):
                 "task": str | List[str],  (optional instruction)
             }
 
-        Native format (List[dict]):
+        Returns the kwargs for `VLAJEPAModel.forward` / `.predict_action` (everything stays
-            {
+        on the batch device; no per-sample shredding): `images` (per-sample, per-view list for
-                "image": List[PIL.Image],       # multi-view images per sample
+        Qwen messages), `instructions`, and the batched `videos` / `actions` / `state` /
-                "video": np.ndarray [V, T, H, W, 3],
+        `action_is_pad` when present.
-                "lang": str,                     # task instruction
-                "action": np.ndarray [T, action_dim],  # optional
-                "state": np.ndarray [1, state_dim],    # optional
-            }
         """
-        # Determine batch size from the first image feature
         image_keys = list(self.config.image_features.keys())
         if not image_keys:
             raise ValueError("VLAJEPA requires at least one image feature.")
-        first_key = image_keys[0]
+        batch_size = batch[image_keys[0]].shape[0]
-        first_tensor = batch[first_key]
-        batch_size = first_tensor.shape[0]
 
-        # ---- Collect images per sample ----
+        # Current-frame image per view ([B, C, H, W]); regroup per sample for Qwen messages.
-        # images_per_sample[b][v] = PIL.Image for view v
+        frames = []
-        images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)]
         for key in image_keys:
-            tensor = batch[key]  # [B, C, H, W] or [B, T, C, H, W]
+            t = batch[key]
-            if tensor.ndim == 5:
+            if t.ndim == 5:  # [B, T, C, H, W] -> current observation (delta=0)
-                # observation_delta_indices = [0, 1, ..., num_video_frames-1]
+                t = t[:, 0]
-                # index 0 is the current observation (delta=0)
+            frames.append(self.model.qwen.to_pixel_values(t))
-                tensor = tensor[:, 0]
+        images = [[frame[b] for frame in frames] for b in range(batch_size)]
-            for b in range(batch_size):
-                images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b]))
 
-        # ---- Collect videos per sample ----
-        # Build video arrays: for each sample, stack views as [V, T, H, W, 3]
-        # Check whether any image feature has a time dimension
-        video_source = None
-        for k in image_keys:
-            if k in batch:
-                video_source = batch[k]  # Use first available for shape inspection
-                break
-
-        if video_source is None:
-            raise ValueError("No image data found in batch for video construction.")
-
-        videos_per_sample = []
-        for b in range(batch_size):
-            sample_views = []
-            for k in image_keys:
-                t = batch[k][b]  # [C, H, W] or [T, C, H, W]
-                if t.ndim == 3:
-                    t = t.unsqueeze(0)  # [1, C, H, W]
-                # Convert to [T, H, W, 3] numpy
-                t_np = t.permute(0, 2, 3, 1).detach().cpu().float().numpy()
-                # Clamp to [0, 255]
-                if t_np.max() <= 1.0:
-                    t_np = t_np * 255.0
-                t_np = np.rint(t_np.clip(0, 255)).astype(np.uint8)
-                sample_views.append(t_np)
-            # Stack views: [V, T, H, W, 3]
-            videos_per_sample.append(np.stack(sample_views, axis=0))
-
-        # ---- Collect instructions ----
         tasks = batch.get("task")
         if tasks is None:
             instructions = ["Execute the robot action."] * batch_size
@@ -498,52 +383,32 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         else:
             instructions = list(tasks)
 
-        # ---- Collect actions (training only) ----
+        inputs: dict[str, Any] = {"images": images, "instructions": instructions}
-        actions_list = None
-        action_is_pad_list = None
-        actions_tensor = batch.get(ACTION)
-        if actions_tensor is not None:
-            if actions_tensor.ndim == 2:
-                actions_tensor = actions_tensor.unsqueeze(1)
-            actions_list = [actions_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
-            action_is_pad_tensor = batch.get("action_is_pad")
-            if action_is_pad_tensor is not None:
-                action_is_pad_list = [action_is_pad_tensor[b].detach().cpu() for b in range(batch_size)]
 
-        # ---- Collect state ----
+        # Videos [B, V, T, C, H, W] - only assembled during training when the world model consumes them.
-        state_list = None
+        if self.model.config.enable_world_model and training:
-        state_tensor = batch.get(OBS_STATE)
+            views = [batch[k].unsqueeze(1) if batch[k].ndim == 4 else batch[k] for k in image_keys]
-        if state_tensor is not None:
+            inputs["videos"] = self.model.qwen.to_pixel_values(torch.stack(views, dim=1))
-            if state_tensor.ndim > 2:
-                state_tensor = state_tensor[:, -1, :]
-            if state_tensor.ndim == 2:
-                state_tensor = state_tensor.unsqueeze(1)  # [B, 1, state_dim]
-            state_list = [state_tensor[b].detach().cpu().float().numpy() for b in range(batch_size)]
 
-        # ---- Assemble native examples ----
+        actions = batch.get(ACTION)
-        examples = []
+        if actions is not None:
-        for b in range(batch_size):
+            inputs["actions"] = (actions.unsqueeze(1) if actions.ndim == 2 else actions).float()
-            example = {
+            if (pad := batch.get("action_is_pad")) is not None:
-                "image": images_per_sample[b],
+                inputs["action_is_pad"] = pad
-                "video": videos_per_sample[b],
-                "lang": instructions[b],
-            }
-            if actions_list is not None:
-                example["action"] = actions_list[b]
-            if action_is_pad_list is not None:
-                example["action_is_pad"] = action_is_pad_list[b]
-            if state_list is not None:
-                example["state"] = state_list[b]
-            examples.append(example)
 
-        return examples
+        state = batch.get(OBS_STATE)
+        if state is not None:
+            if state.ndim > 2:
+                state = state[:, -1, :]
+            inputs["state"] = (state.unsqueeze(1) if state.ndim == 2 else state).float()  # [B, 1, dim]
+
+        return inputs
 
     # ---- LeRobot Policy Interface ----
 
     def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]:
         """LeRobot train forward: convert → native forward → aggregate losses."""
-        examples = self._prepare_model_inputs(batch)
+        native_output = self.model.forward(**self._prepare_model_inputs(batch, training=True))
-        native_output = self.model.forward(examples)
 
         ref = next(iter(native_output.values()))
         zero = torch.zeros((), device=ref.device, dtype=ref.dtype)
@@ -561,16 +426,9 @@ class VLAJEPAPolicy(PreTrainedPolicy):
         self.eval()
         self._queues = populate_queues(self._queues, batch, exclude_keys=[ACTION])
 
-        examples = self._prepare_model_inputs(batch)
+        inputs = self._prepare_model_inputs(batch, training=False)
-        batch_images = [ex["image"] for ex in examples]
+        actions = self.model.predict_action(inputs["images"], inputs["instructions"], inputs.get("state"))
-        instructions = [ex["lang"] for ex in examples]
+        return actions.to(device=self.config.device, dtype=torch.float32)
-
-        state_np = None
-        if "state" in examples[0] and examples[0]["state"] is not None:
-            state_np = np.stack([ex["state"] for ex in examples])
-
-        actions_np = self.model.predict_action(batch_images, instructions, state_np)
-        return torch.from_numpy(actions_np).to(device=self.config.device, dtype=torch.float32)
 
     @torch.no_grad()
     def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor:

src/lerobot/policies/vla_jepa/qwen_interface.py

@@ -17,9 +17,7 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-import numpy as np
 import torch
-from PIL import Image
 
 from lerobot.utils.import_utils import _transformers_available
 
@@ -78,7 +76,7 @@ class Qwen3VLInterface(torch.nn.Module):
 
     def build_inputs(
         self,
-        images: Sequence[Sequence[Image.Image]],
+        images: Sequence[Sequence[torch.Tensor]],
         instructions: Sequence[str],
         action_prompt: str,
         embodied_prompt: str,
@@ -94,24 +92,42 @@ class Qwen3VLInterface(torch.nn.Module):
             content.append({"type": "text", "text": prompt})
             messages.append([{"role": "user", "content": content}])
 
+        # The Qwen image processor is a torchvision-backed fast processor: passing the
+        # images as GPU tensors (with `device`) keeps the whole vision pipeline on-device
+        # and avoids a GPU->CPU->GPU roundtrip. The image tensors are forwarded through
+        # apply_chat_template untouched into Qwen3VLProcessor.__call__.
+        # do_rescale=False: images already arrive as float in [0, 1] (the dataset decoder
+        # yields float32/255 and VISUAL normalization is IDENTITY), so we skip the
+        # processor's /255 rescale instead of round-tripping through uint8.
         batch_inputs = self.processor.apply_chat_template(
             messages,
             tokenize=True,
             add_generation_prompt=True,
             return_dict=True,
-            processor_kwargs={"padding": True, "return_tensors": "pt"},
+            processor_kwargs={
+                "padding": True,
+                "return_tensors": "pt",
+                "device": self.model.device,
+                "do_rescale": False,
+            },
         )
         return batch_inputs.to(self.model.device)
 
     @staticmethod
-    def tensor_to_pil(image_tensor: torch.Tensor) -> Image.Image:
+    def to_pixel_values(image_tensor: torch.Tensor) -> torch.Tensor:
-        image = image_tensor.detach().cpu()
+        """Prepare an image/video tensor for the fast processors (used with do_rescale=False).
-        if image.ndim == 3 and image.shape[0] in (1, 3):
+
-            image = image.permute(1, 2, 0)
+        The dataset decoder yields float32 in [0, 1] (channels-first) and VISUAL
-        image = image.float()
+        normalization is IDENTITY, so the tensor already arrives in [0, 1]; we pass it
-        if image.max() <= 1.0:
+        through as float and let the processors normalize (no rescale, no uint8
-            image = image * 255.0
+        quantization). A single channel is expanded to 3 to match the RGB processors.
-        image = image.clamp(0, 255).round().to(torch.uint8).numpy()
+
-        if image.shape[-1] == 1:
+        Works for any channels-first layout (channel dim is -3): [C, H, W], [B, C, H, W],
-            image = np.repeat(image, 3, axis=-1)
+        [T, C, H, W], [B, V, T, C, H, W], ...
-        return Image.fromarray(image)
+        """
+        image = image_tensor.detach().float()
+        if image.shape[-3] == 1:
+            repeats = [1] * image.ndim
+            repeats[-3] = 3
+            image = image.repeat(*repeats)
+        return image

tests/policies/vla_jepa/conftest.py

@@ -8,7 +8,6 @@ from types import SimpleNamespace
 import numpy as np
 import pytest
 import torch
-from PIL import Image
 from torch import Tensor, nn
 
 from lerobot.configs.types import FeatureType, PolicyFeature
@@ -191,7 +190,7 @@ class _FakeQwenInterface(nn.Module):
 
     def build_inputs(
         self,
-        images: list[list[Image.Image]],
+        images: list[list[Tensor]],
         instructions: list[str],
         action_prompt: str,
         embodied_prompt: str,
@@ -214,12 +213,13 @@ class _FakeQwenInterface(nn.Module):
         }
 
     @staticmethod
-    def tensor_to_pil(image_tensor: Tensor) -> Image.Image:
+    def to_pixel_values(image_tensor: Tensor) -> Tensor:
-        image = image_tensor.detach().cpu()
+        image = image_tensor.detach().float()
-        if image.ndim == 3 and image.shape[0] in (1, 3):
+        if image.shape[-3] == 1:
-            image = image.permute(1, 2, 0)
+            repeats = [1] * image.ndim
-        image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy()
+            repeats[-3] = 3
-        return Image.fromarray(image)
+            image = image.repeat(*repeats)
+        return image
 
 
 class _FakeVideoEncoder(nn.Module):
@@ -242,12 +242,14 @@ class _FakeVideoEncoder(nn.Module):
 
 
 class _FakeVideoProcessor:
-    def __call__(self, videos, return_tensors: str) -> dict[str, Tensor]:
+    def __call__(self, videos, return_tensors: str, device=None, **kwargs) -> dict[str, Tensor]:
         assert return_tensors == "pt"
         if isinstance(videos, list):
             pixel_values = torch.stack([torch.as_tensor(v) for v in videos])
         else:
             pixel_values = torch.as_tensor(videos).unsqueeze(0)
+        if device is not None:
+            pixel_values = pixel_values.to(device)
         return {"pixel_values_videos": pixel_values}
 
 

tests/policies/vla_jepa/test_vla_jepa.py

@@ -211,40 +211,42 @@ def test_reset_clears_action_queue(patch_vla_jepa_external_models: None) -> None
 
 
 def test_prepare_model_inputs_training_format(patch_vla_jepa_external_models: None) -> None:
-    from PIL import Image
-
     policy = VLAJEPAPolicy(make_config())
-    examples = policy._prepare_model_inputs(make_train_batch())
+    inputs = policy._prepare_model_inputs(make_train_batch())
 
-    assert len(examples) == BATCH_SIZE
+    assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"}
-    for ex in examples:
+    # images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages)
-        assert set(ex) >= {"image", "video", "lang", "action", "state"}
+    assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1
-        assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image)
+    img = inputs["images"][0][0]
-        assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8  # [V,T,H,W,C]
+    assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3
-        assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM)
+    assert len(inputs["instructions"]) == BATCH_SIZE
-        assert ex["state"].shape == (1, STATE_DIM)
+    # videos: batched [B, V, T, C, H, W] float
+    assert inputs["videos"].ndim == 6 and inputs["videos"].shape[0] == BATCH_SIZE
+    assert inputs["videos"].dtype == torch.float32
+    assert inputs["actions"].shape == (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM)
+    assert inputs["state"].shape == (BATCH_SIZE, 1, STATE_DIM)
 
 
 def test_prepare_model_inputs_inference_omits_action(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
-    for ex in policy._prepare_model_inputs(make_inference_batch()):
+    inputs = policy._prepare_model_inputs(make_inference_batch())
-        assert "action" not in ex
+    assert "actions" not in inputs and "action_is_pad" not in inputs
-        assert "image" in ex and "video" in ex and "lang" in ex
+    assert {"images", "instructions", "state"} <= set(inputs)
 
 
 def test_prepare_model_inputs_missing_task_uses_default(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     del batch["task"]
-    examples = policy._prepare_model_inputs(batch)
+    instructions = policy._prepare_model_inputs(batch)["instructions"]
-    assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples)
+    assert all(isinstance(s, str) and len(s) > 0 for s in instructions)
 
 
 def test_prepare_model_inputs_string_task_broadcast(patch_vla_jepa_external_models: None) -> None:
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     batch["task"] = "open the drawer"
-    assert all(ex["lang"] == "open the drawer" for ex in policy._prepare_model_inputs(batch))
+    assert policy._prepare_model_inputs(batch)["instructions"] == ["open the drawer"] * BATCH_SIZE
 
 
 def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: None) -> None:
@@ -253,7 +255,7 @@ def test_prepare_model_inputs_no_state_omitted(patch_vla_jepa_external_models: N
     policy = VLAJEPAPolicy(make_config())
     batch = make_inference_batch()
     del batch[OBS_STATE]
-    assert all("state" not in ex for ex in policy._prepare_model_inputs(batch))
+    assert "state" not in policy._prepare_model_inputs(batch)
 
 
 # ---------------------------------------------------------------------------
@@ -446,14 +448,14 @@ def test_postprocessor_applied_after_predict_action_chunk(
     """
     from lerobot.policies.vla_jepa.processor_vla_jepa import make_vla_jepa_pre_post_processors
 
-    raw_actions = np.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=np.float32)
+    raw_actions = torch.zeros((BATCH_SIZE, ACTION_HORIZON, ACTION_DIM), dtype=torch.float32)
 
     cfg = make_config()
     cfg.clip_normalized_actions = False
     cfg.binarize_gripper_action = False
     policy = VLAJEPAPolicy(cfg)
     policy.eval()
-    monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.copy())
+    monkeypatch.setattr(policy.model, "predict_action", lambda *a, **kw: raw_actions.clone())
 
     dataset_stats = _make_dataset_stats()
     _, postprocessor = make_vla_jepa_pre_post_processors(cfg, dataset_stats)
@@ -564,9 +566,9 @@ def test_single_view_is_duplicated_for_world_model(patch_vla_jepa_external_model
     original_processor = policy.model.video_processor
 
     class _CapturingProcessor:
-        def __call__(self, videos: list, return_tensors: str) -> dict:
+        def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
             captured_videos.extend(videos)
-            return original_processor(videos=videos, return_tensors=return_tensors)
+            return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
 
     policy.model.video_processor = _CapturingProcessor()
     policy.forward(_make_multiview_train_batch(num_views=1))
@@ -587,9 +589,9 @@ def test_excess_views_trimmed_for_world_model(patch_vla_jepa_external_models: No
     original_processor = policy.model.video_processor
 
     class _CapturingProcessor:
-        def __call__(self, videos: list, return_tensors: str) -> dict:
+        def __call__(self, videos: list, return_tensors: str, **kwargs) -> dict:
             captured_videos.extend(videos)
-            return original_processor(videos=videos, return_tensors=return_tensors)
+            return original_processor(videos=videos, return_tensors=return_tensors, **kwargs)
 
     policy.model.video_processor = _CapturingProcessor()
     policy.forward(_make_multiview_train_batch(num_views=3))