mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-15 05:51:52 +00:00
Maxime Ellerbach
18eee1b477
* 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
276 lines
9.6 KiB
Python
276 lines
9.6 KiB
Python
#!/usr/bin/env python
|
|
"""Shared fixtures and helpers for VLA-JEPA tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from torch import Tensor, nn
|
|
|
|
from lerobot.configs.types import FeatureType, PolicyFeature
|
|
from lerobot.policies.vla_jepa.configuration_vla_jepa import VLAJEPAConfig
|
|
from lerobot.utils.constants import ACTION, OBS_IMAGES, OBS_STATE
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Shared constants
|
|
# ---------------------------------------------------------------------------
|
|
|
|
BATCH_SIZE = 2
|
|
ACTION_DIM = 3
|
|
STATE_DIM = 4
|
|
IMAGE_SIZE = 8
|
|
ACTION_HORIZON = 4
|
|
N_ACTION_STEPS = 2
|
|
NUM_VIDEO_FRAMES = 3
|
|
QWEN_HIDDEN_SIZE = 16 # hidden size produced by _FakeQwenBackbone
|
|
|
|
EXPECTED_ACTION_CHUNK_SHAPE = (BATCH_SIZE, ACTION_HORIZON, ACTION_DIM)
|
|
EXPECTED_SELECT_ACTION_SHAPE = (BATCH_SIZE, ACTION_DIM)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def set_seed_all(seed: int) -> None:
|
|
np.random.seed(seed)
|
|
torch.manual_seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
|
|
|
|
def make_config(
|
|
action_dim: int = ACTION_DIM,
|
|
state_dim: int = STATE_DIM,
|
|
action_horizon: int = ACTION_HORIZON,
|
|
num_video_frames: int = NUM_VIDEO_FRAMES,
|
|
) -> VLAJEPAConfig:
|
|
config = VLAJEPAConfig(
|
|
input_features={
|
|
f"{OBS_IMAGES}.laptop": PolicyFeature(type=FeatureType.VISUAL, shape=(3, IMAGE_SIZE, IMAGE_SIZE)),
|
|
OBS_STATE: PolicyFeature(type=FeatureType.STATE, shape=(state_dim,)),
|
|
},
|
|
output_features={
|
|
ACTION: PolicyFeature(type=FeatureType.ACTION, shape=(action_dim,)),
|
|
},
|
|
device="cpu",
|
|
chunk_size=action_horizon,
|
|
n_action_steps=min(N_ACTION_STEPS, action_horizon),
|
|
action_dim=action_dim,
|
|
state_dim=state_dim,
|
|
num_video_frames=num_video_frames,
|
|
num_action_tokens_per_timestep=2,
|
|
num_embodied_action_tokens_per_instruction=3,
|
|
num_inference_timesteps=2,
|
|
action_hidden_size=QWEN_HIDDEN_SIZE,
|
|
action_model_type="DiT-test",
|
|
action_num_layers=1,
|
|
predictor_depth=1,
|
|
predictor_num_heads=2,
|
|
predictor_mlp_ratio=2.0,
|
|
jepa_tubelet_size=1,
|
|
)
|
|
config.validate_features()
|
|
return config
|
|
|
|
|
|
def make_train_batch(
|
|
batch_size: int = BATCH_SIZE,
|
|
action_dim: int = ACTION_DIM,
|
|
state_dim: int = STATE_DIM,
|
|
action_horizon: int = ACTION_HORIZON,
|
|
num_video_frames: int = NUM_VIDEO_FRAMES,
|
|
) -> dict[str, Tensor | list[str]]:
|
|
return {
|
|
f"{OBS_IMAGES}.laptop": torch.rand(batch_size, num_video_frames, 3, IMAGE_SIZE, IMAGE_SIZE),
|
|
OBS_STATE: torch.randn(batch_size, 1, state_dim),
|
|
ACTION: torch.randn(batch_size, action_horizon, action_dim),
|
|
"task": ["pick up the cube"] * batch_size,
|
|
}
|
|
|
|
|
|
def make_inference_batch(
|
|
batch_size: int = BATCH_SIZE,
|
|
state_dim: int = STATE_DIM,
|
|
) -> dict[str, Tensor | list[str]]:
|
|
return {
|
|
f"{OBS_IMAGES}.laptop": torch.rand(batch_size, 3, IMAGE_SIZE, IMAGE_SIZE),
|
|
OBS_STATE: torch.randn(batch_size, state_dim),
|
|
"task": ["pick up the cube"] * batch_size,
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fake external models (replace Qwen3-VL and V-JEPA at test time)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeLanguageLayer(nn.Module):
|
|
"""Leaf module whose forward hook is captured by _qwen_last_decoder_hidden."""
|
|
|
|
def __init__(self, hidden_size: int) -> None:
|
|
super().__init__()
|
|
self._hidden_size = hidden_size
|
|
|
|
def forward(self, hidden: Tensor, **_: object) -> tuple[Tensor, ...]:
|
|
return (hidden,)
|
|
|
|
|
|
class _FakeLanguageModel(nn.Module):
|
|
def __init__(self, hidden_size: int) -> None:
|
|
super().__init__()
|
|
self._hidden_size = hidden_size
|
|
self.layers = nn.ModuleList([_FakeLanguageLayer(hidden_size)])
|
|
|
|
def forward(self, input_ids: Tensor, **_: object) -> SimpleNamespace:
|
|
batch_size, seq_len = input_ids.shape
|
|
hidden = torch.zeros(batch_size, seq_len, self._hidden_size, device=input_ids.device)
|
|
self.layers[-1](hidden)
|
|
return SimpleNamespace()
|
|
|
|
|
|
class _FakeQwenInnerModel(nn.Module):
|
|
"""Mimics the `.model.model` level that _qwen_last_decoder_hidden walks into."""
|
|
|
|
def __init__(self, hidden_size: int) -> None:
|
|
super().__init__()
|
|
self.language_model = _FakeLanguageModel(hidden_size)
|
|
|
|
def forward(self, input_ids: Tensor, **kwargs: object) -> SimpleNamespace:
|
|
return self.language_model(input_ids)
|
|
|
|
|
|
class _FakeQwenBackbone(nn.Module):
|
|
def __init__(self, hidden_size: int) -> None:
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(1))
|
|
self.config = SimpleNamespace(
|
|
hidden_size=hidden_size,
|
|
text_config=SimpleNamespace(hidden_size=hidden_size),
|
|
)
|
|
self.model = _FakeQwenInnerModel(hidden_size)
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
return self.weight.device
|
|
|
|
def forward(self, input_ids: Tensor, **_: object) -> SimpleNamespace:
|
|
batch_size, seq_len = input_ids.shape
|
|
hidden_size = self.config.hidden_size
|
|
values = torch.arange(
|
|
batch_size * seq_len * hidden_size,
|
|
device=input_ids.device,
|
|
dtype=torch.float32,
|
|
).view(batch_size, seq_len, hidden_size)
|
|
hidden = values / values.numel() + self.weight
|
|
self.model(input_ids) # call through so the forward hook on layers[-1] fires
|
|
return SimpleNamespace(hidden_states=[hidden])
|
|
|
|
|
|
class _FakeQwenInterface(nn.Module):
|
|
def __init__(self, config: VLAJEPAConfig) -> None:
|
|
super().__init__()
|
|
self.config = config
|
|
self.model = _FakeQwenBackbone(hidden_size=QWEN_HIDDEN_SIZE)
|
|
|
|
@staticmethod
|
|
def _get_torch_dtype(dtype_name: str) -> torch.dtype:
|
|
return torch.float32 if dtype_name == "float32" else torch.bfloat16
|
|
|
|
def expand_tokenizer(self) -> tuple[list[str], list[int], int]:
|
|
max_action_tokens = self.config.chunk_size * self.config.num_action_tokens_per_timestep
|
|
action_tokens = [self.config.special_action_token.format(idx) for idx in range(max_action_tokens)]
|
|
action_token_ids = list(range(1000, 1000 + max_action_tokens))
|
|
return action_tokens, action_token_ids, 2000
|
|
|
|
def build_inputs(
|
|
self,
|
|
images: list[list[Tensor]],
|
|
instructions: list[str],
|
|
action_prompt: str,
|
|
embodied_prompt: str,
|
|
) -> dict[str, Tensor]:
|
|
batch_size = len(images)
|
|
del images, instructions, action_prompt, embodied_prompt
|
|
action_count = (self.config.num_video_frames - 1) * self.config.num_action_tokens_per_timestep
|
|
token_ids = (
|
|
[10]
|
|
+ list(range(1000, 1000 + action_count))
|
|
+ [2000] * self.config.num_embodied_action_tokens_per_instruction
|
|
+ [11]
|
|
)
|
|
return {
|
|
"input_ids": torch.tensor(
|
|
[token_ids] * batch_size,
|
|
device=self.model.device,
|
|
dtype=torch.long,
|
|
)
|
|
}
|
|
|
|
@staticmethod
|
|
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):
|
|
def __init__(self, hidden_size: int = 8, tubelet_size: int = 1) -> None:
|
|
super().__init__()
|
|
self.weight = nn.Parameter(torch.ones(1))
|
|
# image_size must be >= patch_size (16) so the predictor grid is non-zero.
|
|
# Setting image_size=16 gives a 1x1 grid (1 patch per frame).
|
|
self.config = SimpleNamespace(hidden_size=hidden_size, tubelet_size=tubelet_size, image_size=16)
|
|
|
|
@property
|
|
def device(self) -> torch.device:
|
|
return self.weight.device
|
|
|
|
def get_vision_features(self, pixel_values_videos: Tensor) -> Tensor:
|
|
batch_size, num_frames = pixel_values_videos.shape[:2]
|
|
hidden_size = self.config.hidden_size
|
|
frame_values = pixel_values_videos.float().mean(dim=(2, 3, 4), keepdim=False)
|
|
return frame_values[:, :, None].expand(batch_size, num_frames, hidden_size)
|
|
|
|
|
|
class _FakeVideoProcessor:
|
|
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}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture
|
|
def patch_vla_jepa_external_models(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
from lerobot.policies.vla_jepa import modeling_vla_jepa
|
|
|
|
monkeypatch.setattr(modeling_vla_jepa, "Qwen3VLInterface", _FakeQwenInterface)
|
|
monkeypatch.setattr(
|
|
modeling_vla_jepa.AutoModel,
|
|
"from_pretrained",
|
|
lambda *args, **kwargs: _FakeVideoEncoder(),
|
|
)
|
|
monkeypatch.setattr(
|
|
modeling_vla_jepa.AutoVideoProcessor,
|
|
"from_pretrained",
|
|
lambda *args, **kwargs: _FakeVideoProcessor(),
|
|
)
|