admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-10 03:21:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
* 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
This commit is contained in:
Maxime Ellerbach
2026-06-29 18:50:04 +02:00
committed by GitHub
parent 5ac3b49a5f
commit 18eee1b477
4 changed files with 228 additions and 350 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/policies/vla_jepa/conftest.py
+11 -9
View File
@@ -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:
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.shape[-3] == 1:
repeats = [1] * image.ndim
repeats[-3] = 3
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
+25 -23
View File
@@ -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
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 ex["action"].shape == (ACTION_HORIZON, ACTION_DIM)
assert ex["state"].shape == (1, STATE_DIM)
assert set(inputs) >= {"images", "instructions", "videos", "actions", "state"}
# images: per-sample, per-view [C, H, W] float tensors (kept as a list for Qwen messages)
assert len(inputs["images"]) == BATCH_SIZE and len(inputs["images"][0]) == 1
img = inputs["images"][0][0]
assert isinstance(img, torch.Tensor) and img.dtype == torch.float32 and img.ndim == 3
assert len(inputs["instructions"]) == BATCH_SIZE
# 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()):
assert "action" not in ex
assert "image" in ex and "video" in ex and "lang" in ex
inputs = policy._prepare_model_inputs(make_inference_batch())
assert "actions" not in inputs and "action_is_pad" not in inputs
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)
assert all(isinstance(ex["lang"], str) and len(ex["lang"]) > 0 for ex in examples)
instructions = policy._prepare_model_inputs(batch)["instructions"]
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))