From 877847c90eef8b3bd6ad9a255b118f9b979c01a7 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Tue, 9 Jun 2026 09:18:46 +0000 Subject: [PATCH] refactor(vla-jepa): removing gpu roundtrip for the preprocessing part --- .../policies/vla_jepa/modeling_vla_jepa.py | 138 ++++++++---------- .../policies/vla_jepa/qwen_interface.py | 41 ++++-- tests/policies/vla_jepa/conftest.py | 18 +-- tests/policies/vla_jepa/test_vla_jepa.py | 21 +-- 4 files changed, 107 insertions(+), 111 deletions(-) diff --git a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py index 45d83e652..1bd8305fb 100644 --- a/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py +++ b/src/lerobot/policies/vla_jepa/modeling_vla_jepa.py @@ -19,10 +19,8 @@ from collections import deque from pathlib import Path from typing import TYPE_CHECKING -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 @@ -56,11 +54,11 @@ class VLAJEPAModel(nn.Module): - V-JEPA: world model for video frame prediction Input: List[dict] native format (same as original starVLA) - - "image": List[PIL.Image] (multi-view images) - - "video": np.ndarray [V, T, H, W, 3] + - "image": List[Tensor [C, H, W]] (float [0,1], multi-view images) + - "video": Tensor [V, T, C, H, W] (float [0,1]) - "lang": str (task instruction) - - "action": np.ndarray [T, action_dim] (optional, training only) - - "state": np.ndarray [1, state_dim] (optional) + - "action": Tensor [T, action_dim] (optional, training only) + - "state": Tensor [1, state_dim] (optional) """ def __init__(self, config: VLAJEPAConfig) -> None: @@ -167,18 +165,18 @@ class VLAJEPAModel(nn.Module): Args: examples: List of per-sample dicts with keys: - "image" : List[PIL.Image] — multi-view images - "video" : np.ndarray [V, T, H, W, 3] + "image" : List[Tensor [C, H, W]] (float [0,1]) — multi-view images + "video" : Tensor [V, T, C, H, W] (float [0,1]) "lang" : str — task instruction - "action" : np.ndarray [T, action_dim] (optional) - "state" : np.ndarray [1, state_dim] (optional) + "action" : Tensor [T, action_dim] (optional) + "state" : Tensor [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] + batch_images = [ex["image"] for ex in examples] # List[List[Tensor [C, H, W]]] + batch_videos = [ex["video"] for ex in examples] # List[Tensor [V, T, C, H, W]] 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 @@ -190,9 +188,8 @@ class VLAJEPAModel(nn.Module): 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] + # Stack videos: List[[V, T, 3, H, W]] -> [B, V, T, 3, H, W] (already channels-first) + batch_videos = torch.stack(batch_videos) # Adjust number of views for the world model: # - fewer views than expected: duplicate the first view to fill up @@ -200,8 +197,8 @@ class VLAJEPAModel(nn.Module): 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) + first_view = batch_videos[:, :1].repeat(1, num_missing_views, 1, 1, 1, 1) + batch_videos = torch.cat([batch_videos, first_view], dim=1) elif batch_videos.shape[1] > num_views_world_model: batch_videos = batch_videos[:, :num_views_world_model] @@ -244,9 +241,15 @@ class VLAJEPAModel(nn.Module): b, v, t_frames, c, h_img, w_img = batch_videos.shape 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")[ - "pixel_values_videos" - ].to(self.video_encoder.device) # [B*V, T, C, H, W] + # Fast (torchvision) video processor: pass GPU float [0,1] tensors + device so the + # resize/normalize stays on-device (no GPU->CPU->GPU roundtrip). do_rescale=False + # because the frames already arrive in [0, 1]. + video_pixels = self.video_processor( + videos=list(batch_videos_flat), + return_tensors="pt", + device=self.video_encoder.device, + do_rescale=False, + )["pixel_values_videos"] # [B*V, T, C, H, W] with torch.no_grad(): video_embeddings = self.video_encoder.get_vision_features(pixel_values_videos=video_pixels) @@ -286,16 +289,16 @@ class VLAJEPAModel(nn.Module): # ---- Step 4: Action Head ---- with torch.autocast(device_type=device_type, dtype=torch.float32): - actions_tensor = torch.tensor( - np.array(actions), device=last_hidden.device, dtype=torch.float32 + actions_tensor = torch.stack(actions).to( + 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 + state_tensor = torch.stack(state).to( + device=last_hidden.device, dtype=last_hidden.dtype ) # [B, 1, state_dim] repeated_diffusion_steps = self.config.repeated_diffusion_steps @@ -307,12 +310,7 @@ class VLAJEPAModel(nn.Module): 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 - ] + [p.to(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] @@ -328,26 +326,30 @@ class VLAJEPAModel(nn.Module): @torch.no_grad() def predict_action( self, - batch_images: list[list[Image.Image]], + batch_images: list[list[Tensor]], instructions: list[str], - state: np.ndarray | None = None, - ) -> np.ndarray: + state: Tensor | None = None, + ) -> Tensor: """ Native action prediction following original VLA_JEPA.predict_action. Args: - batch_images: List of samples; each is List[PIL.Image] (multi-view). + batch_images: List of samples; each is List[Tensor [C, H, W]] (float [0,1], multi-view). instructions: Task instructions, one per sample. - state: Optional [B, state_dim] numpy array. + state: Optional [B, state_dim] tensor. Returns: - np.ndarray [B, action_horizon, action_dim] — predicted actions. + Tensor [B, action_horizon, action_dim] — predicted actions (on the model device). """ if self.config.resize_images_to is not None: height, width = self.config.resize_images_to - resampling = getattr(Image, "Resampling", Image).BOX + # PIL BOX resampling ~= area-averaging downsample; F.interpolate(mode="area") + # is the on-GPU equivalent. Images stay float [0,1] (do_rescale=False downstream). batch_images = [ - [image.resize((width, height), resample=resampling) for image in sample_images] + [ + F.interpolate(image[None], size=(height, width), mode="area")[0] + for image in sample_images + ] for sample_images in batch_images ] @@ -370,15 +372,13 @@ class VLAJEPAModel(nn.Module): 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 - ) + state_tensor = 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() + return pred_actions # ============================================================================ @@ -431,13 +431,13 @@ class VLAJEPAPolicy(PreTrainedPolicy): "task": str | List[str], (optional instruction) } - Native format (List[dict]): + Native format (List[dict]), all tensors kept on the batch device: { - "image": List[PIL.Image], # multi-view images per sample - "video": np.ndarray [V, T, H, W, 3], + "image": List[Tensor [C, H, W]] (float [0,1]), # multi-view images per sample + "video": Tensor [V, T, C, H, W] (float [0,1]), "lang": str, # task instruction - "action": np.ndarray [T, action_dim], # optional - "state": np.ndarray [1, state_dim], # optional + "action": Tensor [T, action_dim], # optional + "state": Tensor [1, state_dim], # optional } """ # Determine batch size from the first image feature @@ -449,8 +449,8 @@ class VLAJEPAPolicy(PreTrainedPolicy): batch_size = first_tensor.shape[0] # ---- Collect images per sample ---- - # images_per_sample[b][v] = PIL.Image for view v - images_per_sample: list[list[Image.Image]] = [[] for _ in range(batch_size)] + # images_per_sample[b][v] = float [0,1] Tensor [C, H, W] for view v (kept on-device) + images_per_sample: list[list[Tensor]] = [[] for _ in range(batch_size)] for key in image_keys: tensor = batch[key] # [B, C, H, W] or [B, T, C, H, W] if tensor.ndim == 5: @@ -458,20 +458,10 @@ class VLAJEPAPolicy(PreTrainedPolicy): # index 0 is the current observation (delta=0) tensor = tensor[:, 0] for b in range(batch_size): - images_per_sample[b].append(self.model.qwen.tensor_to_pil(tensor[b])) + images_per_sample[b].append(self.model.qwen.to_pixel_values(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.") - + # Build video tensors: for each sample, stack views as [V, T, C, H, W] (float [0,1], on-device) videos_per_sample = [] for b in range(batch_size): sample_views = [] @@ -479,15 +469,9 @@ class VLAJEPAPolicy(PreTrainedPolicy): 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)) + sample_views.append(self.model.qwen.to_pixel_values(t)) + # Stack views: [V, T, C, H, W] + videos_per_sample.append(torch.stack(sample_views, dim=0)) # ---- Collect instructions ---- tasks = batch.get("task") @@ -505,10 +489,10 @@ class VLAJEPAPolicy(PreTrainedPolicy): 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)] + actions_list = [actions_tensor[b].detach().float() 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)] + action_is_pad_list = [action_is_pad_tensor[b].detach() for b in range(batch_size)] # ---- Collect state ---- state_list = None @@ -518,7 +502,7 @@ class VLAJEPAPolicy(PreTrainedPolicy): 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)] + state_list = [state_tensor[b].detach().float() for b in range(batch_size)] # ---- Assemble native examples ---- examples = [] @@ -565,12 +549,12 @@ class VLAJEPAPolicy(PreTrainedPolicy): batch_images = [ex["image"] for ex in examples] instructions = [ex["lang"] for ex in examples] - state_np = None + state = None if "state" in examples[0] and examples[0]["state"] is not None: - state_np = np.stack([ex["state"] for ex in examples]) + state = torch.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) + actions = self.model.predict_action(batch_images, instructions, state) + return actions.to(device=self.config.device, dtype=torch.float32) @torch.no_grad() def select_action(self, batch: dict[str, Tensor], noise: Tensor | None = None) -> Tensor: diff --git a/src/lerobot/policies/vla_jepa/qwen_interface.py b/src/lerobot/policies/vla_jepa/qwen_interface.py index 24f530efc..d24e6aaaa 100644 --- a/src/lerobot/policies/vla_jepa/qwen_interface.py +++ b/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,37 @@ 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: - image = image_tensor.detach().cpu() - if image.ndim == 3 and image.shape[0] in (1, 3): - image = image.permute(1, 2, 0) - image = image.float() - if image.max() <= 1.0: - image = image * 255.0 - image = image.clamp(0, 255).round().to(torch.uint8).numpy() - if image.shape[-1] == 1: - image = np.repeat(image, 3, axis=-1) - return Image.fromarray(image) + def to_pixel_values(image_tensor: torch.Tensor) -> torch.Tensor: + """Prepare an image/video tensor for the fast processors (used with do_rescale=False). + + The dataset decoder yields float32 in [0, 1] (channels-first) and VISUAL + normalization is IDENTITY, so the tensor already arrives in [0, 1]; we pass it + through as float and let the processors normalize (no rescale, no uint8 + quantization). A single channel is expanded to 3 to match the RGB processors. + """ + image = image_tensor.detach().float() + if image.ndim == 3 and image.shape[0] == 1: + image = image.repeat(3, 1, 1) + return image diff --git a/tests/policies/vla_jepa/conftest.py b/tests/policies/vla_jepa/conftest.py index 5301b5bc7..799802d5b 100644 --- a/tests/policies/vla_jepa/conftest.py +++ b/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,11 @@ class _FakeQwenInterface(nn.Module): } @staticmethod - def tensor_to_pil(image_tensor: Tensor) -> Image.Image: - image = image_tensor.detach().cpu() - if image.ndim == 3 and image.shape[0] in (1, 3): - image = image.permute(1, 2, 0) - image = (image.float().clamp(0, 1) * 255).to(torch.uint8).numpy() - return Image.fromarray(image) + def to_pixel_values(image_tensor: Tensor) -> Tensor: + image = image_tensor.detach().float() + if image.ndim == 3 and image.shape[0] == 1: + image = image.repeat(3, 1, 1) + return image class _FakeVideoEncoder(nn.Module): @@ -242,12 +240,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} diff --git a/tests/policies/vla_jepa/test_vla_jepa.py b/tests/policies/vla_jepa/test_vla_jepa.py index 70194dd59..b9bc398a2 100644 --- a/tests/policies/vla_jepa/test_vla_jepa.py +++ b/tests/policies/vla_jepa/test_vla_jepa.py @@ -211,16 +211,17 @@ 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()) assert len(examples) == BATCH_SIZE for ex in examples: assert set(ex) >= {"image", "video", "lang", "action", "state"} - assert len(ex["image"]) == 1 and isinstance(ex["image"][0], Image.Image) - assert ex["video"].ndim == 5 and ex["video"].dtype == np.uint8 # [V,T,H,W,C] + assert len(ex["image"]) == 1 + assert isinstance(ex["image"][0], torch.Tensor) and ex["image"][0].dtype == torch.float32 + assert ex["image"][0].ndim == 3 # [C, H, W] + assert isinstance(ex["video"], torch.Tensor) + assert ex["video"].ndim == 5 and ex["video"].dtype == torch.float32 # [V, T, C, H, W] assert ex["action"].shape == (ACTION_HORIZON, ACTION_DIM) assert ex["state"].shape == (1, STATE_DIM) @@ -446,14 +447,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 +565,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 +588,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))