mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-11 20:11:48 +00:00
Maxime Ellerbach
2e9cd87bbd
* first commit * feat(policies): add VLA-JEPA * feat(policies): add VLA-JEPA * support vla_jepa * (feat)policies: add VLA-JEPA * linting * adding deps to pyproject.toml * updating uv lock * adding guards to avoid needing transformers and diffusers for type checking and basic tests * fixing action and state dim * fix warnings with qwen processor kwargs * fixing wm_loss not propagating * adjusting obs steps, tublets size to match original implementation * some more fixes to be closer to the original implem * adding more tests to ensure good coverage * align VLA-JEPA architecture with original checkpoint - Remove stale `action_num_heads` / `action_attention_head_dim` config fields; DiT head dimensions are now always derived from the preset (DiT-B/L/test). - Add `num_target_vision_tokens` and `action_max_seq_len` config fields required by the action head's future-token embedding and positional embedding tables. - Fix default `qwen_model_name` to 2B (matches all released checkpoints). - Rename `ActionEncoder` attrs w1/w2/w3 → layer1/layer2/layer3 to match checkpoint key names; replace `nn.Sequential` decoder/state-encoder with `_MLP2` (layer1/layer2 naming). - Fix `VLAJEPAActionHead` to size ActionEncoder and StateEncoder at `inner_dim` (DiT input width) rather than `action_hidden_size` (DiT output width). - Rename `DiT.blocks` → `transformer_blocks` and `attn` → `attn1` to match checkpoint; add alternating cross/self attention (even blocks cross-attend to Qwen context, odd blocks self-attend). - Add `DiT-test` preset for unit tests. - Rewrite `ActionConditionedVideoPredictor` with explicit ViT-style blocks (`_PredictorBlock` with fused qkv) to match checkpoint structure; rename `encoder`/`norm`/`proj` → `predictor_blocks`/`predictor_norm`/`predictor_proj`. * propagate action_is_pad masking through VLA-JEPA policy pipeline Pass the `action_is_pad` tensor from the batch through to the action head so padded timesteps are excluded from the flow-matching loss. * update VLA-JEPA tests for arch changes and action_is_pad - Switch conftest to use `action_model_type="DiT-test"` now that `action_num_heads` / `action_attention_head_dim` have been removed. - Add action_head tests covering fully-padded loss (zero) and equivalence of action_is_pad=None vs all-zeros mask. - Remove obsolete `test_native_to_lerobot_wm_only` test. * add VLA-JEPA documentation Covers architecture overview, pretrained checkpoints, config reference, training/eval commands for LIBERO-10, and guidance on fine-tuning for single-camera datasets. * add one-shot script to convert ginwind/VLA-JEPA checkpoints to safetensors (will remove once migrated) * make default params more aligned with paper and pretrained models - adding possibility of freezing qwen backbone and world model - added tests for weight loading * trying out to re-init the action head to avoid pretraining dimension mismatch * allow different state dim and action dim * removing missleading future_action_window_size to just use chunk_size * lots of changes to make existing weights work, need to massively refactor the pre and post processing * refactoring into using pre and post processor * pre-commit cleanup * fixing doc defaults args Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net> * adressing dtype zeros issue * adding guard for diffusers * fixing training and exal examples * trying to close success rate gap * fix qwen norm layer output libero eval is now as expected * adding instructions for different embodiement + fixing some tests * smol fix to avoid having default CPU device when training * fixing misconception about multiview / singleview handling * removing conversion script * adding licences * adding .mdx docs and shortening polivy_vla_jepa_README.md * removing useless pre-processor * cleanup * removing swish in favor of silu * adding configuration gripper index and threshold * fixing simlink --------- Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net> Co-authored-by: ginwind <ginwind@mail.ustc.edu.cn>
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
#!/usr/bin/env python
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from lerobot.policies.vla_jepa.world_model import (
|
|
ActionConditionedVideoPredictor,
|
|
)
|
|
|
|
_ACTION_EMBED_DIM = 8
|
|
|
|
|
|
def _make_predictor(
|
|
embed_dim: int = 8,
|
|
action_embed_dim: int = _ACTION_EMBED_DIM,
|
|
predictor_embed_dim: int = 24,
|
|
num_action_tokens: int = 2,
|
|
tokens_per_frame: int = 1,
|
|
) -> ActionConditionedVideoPredictor:
|
|
return ActionConditionedVideoPredictor(
|
|
num_frames=1,
|
|
img_size=(1, tokens_per_frame),
|
|
patch_size=1,
|
|
tubelet_size=1,
|
|
embed_dim=embed_dim,
|
|
action_embed_dim=action_embed_dim,
|
|
predictor_embed_dim=predictor_embed_dim,
|
|
depth=1,
|
|
num_heads=2,
|
|
mlp_ratio=2.0,
|
|
num_action_tokens_per_step=num_action_tokens,
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"batch,num_steps,tokens_per_frame,embed_dim",
|
|
[
|
|
(1, 2, 1, 8),
|
|
(2, 3, 4, 8),
|
|
(4, 5, 2, 16),
|
|
],
|
|
)
|
|
def test_predictor_output_shape(batch: int, num_steps: int, tokens_per_frame: int, embed_dim: int) -> None:
|
|
predictor = _make_predictor(
|
|
embed_dim=embed_dim, action_embed_dim=_ACTION_EMBED_DIM, tokens_per_frame=tokens_per_frame
|
|
)
|
|
frame_tokens = torch.randn(batch, num_steps * tokens_per_frame, embed_dim)
|
|
action_tokens = torch.randn(batch, num_steps * 2, _ACTION_EMBED_DIM)
|
|
out = predictor(frame_tokens, action_tokens)
|
|
assert tuple(out.shape) == (batch, num_steps * tokens_per_frame, embed_dim)
|
|
assert torch.isfinite(out).all()
|
|
|
|
|
|
def test_predictor_step_mismatch_raises() -> None:
|
|
predictor = _make_predictor(tokens_per_frame=4)
|
|
frame_tokens = torch.randn(2, 3 * 4, 8) # 3 steps, 4 tokens each
|
|
with pytest.raises(RuntimeError):
|
|
predictor(frame_tokens, torch.randn(2, 2 * 2, 8)) # 2 steps → mismatch
|