VLA-JEPA

Ini adalah port LeRobot dari VLA-JEPA, sebuah model Vision-Language-Action yang menggabungkan backbone bahasa Qwen3-VL dengan model dunia video self-supervised (V-JEPA2) dan kepala aksi flow-matching DiT.

Ikhtisar Arsitektur

VLA-JEPA memiliki tiga komponen utama:

Alur data

Pelatihan:

1. Klip video berisi num_video_frames frame dikodekan oleh V-JEPA2 menjadi token patch per frame.
2. Backbone Qwen3-VL memproses gambar multi-view + instruksi tugas dan menghasilkan urutan token konteks yang mencakup token aksi khusus (untuk kondisioning model dunia) dan token embodied.
3. Kepala aksi menerima token konteks tersebut sebagai kunci/nilai cross-attention dan memprediksi chunk aksi yang telah didenoise melalui flow matching.
4. Prediktor model dunia menggunakan token aksi yang diekstrak dari Qwen untuk memprediksi embedding frame V-JEPA2 masa depan; kerugian regresi pada prediksi tersebut ditambahkan ke kerugian aksi.

Inferensi: Hanya Qwen + kepala aksi yang digunakan. Model dunia tidak diperlukan pada saat inferensi.

Detail kepala aksi

Preset yang tersedia melalui action_model_type:

Preset hanya mengatur geometri perhatian, dan setiap entri dapat ditimpa oleh action_num_heads / action_attention_head_dim. Dua lebar diturunkan darinya:

• lebar internal DiT adalah heads x head_dim (768 untuk DiT-B), yang diturunkan, bukan dikonfigurasi;
• lebar output DiT, serta lebar MLP action-decoder dan state-encoder, adalah action_hidden_size (default 1024).

Jadi DiT-B menjalankan transformer dengan lebar 768 yang diproyeksikan ke 1024. Keduanya independen.

Detail model dunia

Prediktor video adalah transformer bergaya ViT (ActionConditionedVideoPredictor) yang menerima:

• Token frame: embedding patch V-JEPA2 yang diproyeksikan ke predictor_embed_dim
• Token aksi: embedding token aksi Qwen yang diproyeksikan ke predictor_embed_dim

Itu menggunakan perhatian kausal blok sehingga setiap langkah temporal dapat memperhatikan semua langkah sebelumnya. Input embed_dim dari prediktor sama dengan num_views × video_encoder_hidden_size (misalnya 2 views × 1024 = 2048 untuk checkpoint yang telah dilatih sebelumnya).

Checkpoint Pra-terlatih

Tiga checkpoint tersedia langsung di dalam org LeRobot di sini: lerobot/VLA-JEPA, dikonversi dari ginwind/VLA-JEPA:

Semua checkpoint menggunakan Qwen/Qwen3-VL-2B-Instruct sebagai tulang punggung bahasa.

Konfigurasi

Parameter utama dalam VLAJEPAConfig:

pre_snap_gripper_action dan binarize_gripper_action adalah port dari loop eval starVLA LIBERO dan hanya benar untuk konvensi aksi LIBERO. pre_snap menulis {0, 1} ke ruang normalized, unnormalizer memetakan nilai-nilai tersebut ke titik tengah dan maksimum, lalu binarize membandingkan nilai physical tersebut dengan gripper_threshold (0.5). Untuk gripper yang diukur dalam derajat, mm, atau [0, 100], kedua nilai tersebut berada di atas ambang batas dan gripper yang diperintahkan menjadi konstan. Karena alasan itu, nilai default-nya adalah False; aktifkan hanya untuk setup bergaya LIBERO, dan atur gripper_threshold dalam satuan gripper itu sendiri jika Anda melakukannya. Pabrik processor memberikan peringatan ketika statistik dataset menunjukkan bahwa rentang tersebut tidak dapat berfungsi.

Pelatihan

Jumlah langkah pelatihan dapat bervariasi berdasarkan ukuran dataset dan anggaran komputasi. Paper asli melakukan pra-pelatihan selama 50k langkah pada ssv2 + droid secara bersama-sama, kemudian 30k langkah tambahan untuk LIBERO, tetapi langkah yang lebih sedikit mungkin masih menghasilkan kinerja yang baik saat melakukan fine-tuning dari checkpoint pra-terlatih yang disediakan.

Pelatihan penuh dari awal

lerobot-train \
  --policy.type=vla_jepa \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Fine-tuning dari checkpoint pretrained

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/your_dataset

Jika Anda ingin membekukan backbone Qwen dan hanya melatih action head, atur policy.freeze_qwen=True:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --dataset.repo_id=your_org/your_dataset

Fine-tuning pada embodiment yang berbeda

Ketika robot target memiliki dimensi aksi atau keadaan yang berbeda dari checkpoint pra-terlatih, lapisan proyeksi input/output dari kepala aksi akan memiliki bentuk yang tidak cocok dan tidak dapat dimuat secara langsung. reinit_modules memungkinkan Anda untuk mencantumkan prefiks kunci yang diizinkan untuk tidak cocok — lapisan-lapisan tersebut akan diinisialisasi ulang secara acak sementara setiap bobot lainnya digunakan kembali dari checkpoint. Setiap ketidakcocokan bentuk di luar prefiks yang dicantumkan akan memicu kesalahan.

Lapisan yang bergantung pada action_dim dan state_dim adalah:

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --policy.freeze_qwen=true \
  --policy.reinit_modules='["model.action_model.action_encoder", "model.action_model.action_decoder", "model.action_model.state_encoder"]' \
  --dataset.repo_id=your_org/your_dataset

Jika robot Anda tidak memiliki keadaan proprioseptif, hilangkan model.action_model.state_encoder dari daftar.

Mereproduksi hasil LIBERO

Latihan pada LIBERO: memulai pelatihan dari checkpoint Pretrain, melatih selama 30k langkah pada dataset LIBERO. Paper asli menyebutkan pelatihan di 8 GPU dengan batch size 32, yang berarti global batch size 256.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=HuggingFaceVLA/libero \
  --steps=30000

Mengevaluasi checkpoint LIBERO-10 yang telah dilatih sebelumnya:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_spatial,libero_object,libero_goal,libero_10 \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Untuk mengevaluasi hanya sebagian tugas:

lerobot-eval \
  --policy.path=lerobot/VLA-JEPA-LIBERO \
  --env.type=libero \
  --env.task=libero_10 \
  --env.task_ids='[0,1,2]' \
  --eval.n_episodes=10 \
  --eval.batch_size=5

Hasil yang diharapkan:

Fine-tuning pada dataset dengan jumlah kamera yang berbeda

Prediktor model dunia pra-terlatih dilatih dengan embed_dim = world_model_num_views × 1024, yaitu dua tampilan kamera.

Jumlah view ini sebelumnya dibaca dari jepa_tubelet_size, yang juga menamai ukuran tubelet temporal dari encoder JEPA. world_model_num_views adalah field untuknya sekarang; jika dibiarkan pada None, nilai akan kembali ke jepa_tubelet_size sehingga checkpoint yang dipublikasikan tetap dapat dimuat tanpa perubahan.

Perilaku default — melihat padding / pemotongan (tidak perlu tindakan)

Saat melakukan fine-tuning dari VLA-JEPA-Pretrain, model secara otomatis menyesuaikan jumlah view yang diberikan ke world model agar sesuai dengan world_model_num_views:

• Dataset satu-pandangan (misalnya BridgeV2): laten satu-pandangan diduplikasi untuk menghasilkan input model dunia dua-pandangan, sehingga mempertahankan sinyal self-supervised JEPA tanpa ketidakcocokan bobot.
• >Dataset 2-view (misalnya DROID dengan 3 view): semua view diteruskan ke backbone Qwen (untuk konteks yang lebih kaya), tetapi hanya world_model_num_views view pertama (satu wrist + satu third-person, mengikuti urutan view yang dikonfigurasi) yang digunakan untuk world model.

Opsi 1 — Nonaktifkan model dunia

Setel enable_world_model=False untuk sepenuhnya melewati loss JEPA. Hanya backbone Qwen dan action head yang dimuat dan dilatih. Ini sudah cukup untuk kinerja aksi yang baik.

lerobot-train \
  --policy.path=lerobot/VLA-JEPA-Pretrain \
  --policy.enable_world_model=false \
  --policy.repo_id=your_org/your_repo \
  --dataset.repo_id=your_org/single_camera_dataset

Opsi 2 — Inisialisasi ulang proyeksi input predictor

Jika Anda ingin mengubah world_model_num_views ke nilai selain 2, muat checkpoint dengan strict=False dan inisialisasi ulang model.video_predictor.predictor_embed untuk embed_dim baru. Semua bobot blok predictor lainnya (attention, MLP, norm, output projection) tidak bergantung pada jumlah kamera dan dapat digunakan kembali dari checkpoint pretrained.

Sitasi

@misc{sun2026vlajepaenhancingvisionlanguageactionmodel,
  title         = {VLA-JEPA: Enhancing Vision-Language-Action Model with Latent World Model},
  author        = {Jingwen Sun and Wenyao Zhang and Zekun Qi and Shaojie Ren and Zezhi Liu and Hanxin Zhu and Guangzhong Sun and Xin Jin and Zhibo Chen},
  year          = {2026},
  eprint        = {2602.10098},
  archivePrefix = {arXiv},
  primaryClass  = {cs.RO},
  url           = {https://arxiv.org/abs/2602.10098},
}

Lisensi

Bobot didistribusikan di bawah ketentuan lisensi repositori asli ginwind/VLA-JEPA (Apache 2.0 License). Kode integrasi LeRobot mengikuti Apache 2.0 License.