VLA-JEPA

Đây là phiên bản LeRobot của VLA-JEPA, một mô hình Vision-Language-Action kết hợp bộ khung ngôn ngữ Qwen3-VL với mô hình thế giới video tự giám sát (V-JEPA2) và đầu hành động flow-matching DiT.

Tổng quan kiến trúc

VLA-JEPA có ba thành phần chính:

Luồng dữ liệu

Đào tạo:

1. Một đoạn video gồm num_video_frames khung hình được V-JEPA2 mã hóa thành các patch token cho từng khung hình.
2. Backbone Qwen3-VL xử lý các hình ảnh đa góc nhìn + hướng dẫn tác vụ và tạo ra một chuỗi các token ngữ cảnh bao gồm các token hành động đặc biệt (để điều kiện hóa mô hình thế giới) và các token thân thể.
3. Action head nhận các token ngữ cảnh đó làm keys/values cho cross-attention và dự đoán một chunk hành động đã khử nhiễu thông qua flow matching.
4. Bộ dự đoán mô hình thế giới sử dụng các token hành động được trích xuất từ Qwen để dự đoán các embedding khung hình V-JEPA2 trong tương lai; một tổn thất hồi quy trên các dự đoán đó được thêm vào tổn thất hành động.

Suy luận: Chỉ sử dụng Qwen + đầu ra hành động. Mô hình thế giới không cần thiết ở thời điểm suy luận.

Chi tiết đầu hành động

Các cài đặt có sẵn thông qua action_model_type:

Cài đặt trước chỉ thiết lập hình học chú ý, và mỗi mục có thể được ghi đè bởi action_num_heads / action_attention_head_dim. Hai chiều rộng được suy ra từ nó:

• chiều rộng nội bộ của DiT là heads x head_dim (768 đối với DiT-B), được suy ra chứ không được cấu hình;
• chiều rộng output của DiT, cùng với chiều rộng của các MLP action-decoder và state-encoder, là action_hidden_size (mặc định là 1024).

Vì vậy DiT-B chạy một transformer có chiều rộng 768 và chiếu lên 1024. Hai thành phần này độc lập với nhau.

Chi tiết mô hình thế giới

Bộ dự đoán video là một transformer kiểu ViT (ActionConditionedVideoPredictor) nhận vào:

• Token khung hình: Nhúng patch của V-JEPA2 được chiếu lên predictor_embed_dim
• Token hành động: các embedding token hành động của Qwen được chiếu sang predictor_embed_dim

Nó sử dụng chú ý nhân quả theo khối (block-causal attention) để mỗi bước thời gian có thể chú ý đến tất cả các bước trước đó. Đầu vào embed_dim của bộ dự đoán bằng num_views × video_encoder_hidden_size (ví dụ: 2 góc nhìn × 1024 = 2048 đối với các checkpoint được huấn luyện trước).

Điểm kiểm tra được huấn luyện trước

Ba checkpoint có sẵn trực tiếp trong tổ chức LeRobot tại đây: lerobot/VLA-JEPA, được chuyển đổi từ ginwind/VLA-JEPA:

Tất cả các checkpoint đều sử dụng Qwen/Qwen3-VL-2B-Instruct làm bộ khung ngôn ngữ.

Cấu hình

Các tham số chính trong VLAJEPAConfig:

pre_snap_gripper_action và binarize_gripper_action là bản port của vòng eval LIBERO trong starVLA và chỉ chính xác với quy ước hành động của LIBERO. pre_snap ghi {0, 1} vào không gian normalized, bộ unnormalizer ánh xạ các giá trị đó sang điểm giữa và giá trị tối đa, sau đó binarize so sánh giá trị physical đó với gripper_threshold (0.5). Đối với gripper được đo bằng độ, mm hoặc [0, 100], cả hai giá trị đều nằm trên ngưỡng và gripper được lệnh sẽ trở thành một hằng số. Vì lý do đó, mặc định của chúng là False; chỉ bật chúng cho các cấu hình kiểu LIBERO, và đặt gripper_threshold theo đơn vị riêng của gripper nếu bạn làm vậy. Nhà máy processor sẽ cảnh báo khi thống kê dữ liệu cho thấy khoảng giá trị không thể hoạt động.

Huấn luyện

Số bước huấn luyện có thể thay đổi tùy thuộc vào kích thước tập dữ liệu và ngân sách tính toán. Bài báo gốc đã tiền huấn luyện trong 50k bước trên ssv2 + droid cùng lúc, sau đó thêm 30k bước cho LIBERO, nhưng ít bước hơn vẫn có thể cho hiệu suất tốt khi tinh chỉnh từ các checkpoint tiền huấn luyện được cung cấp.

Đào tạo đầy đủ từ đầu

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

Tinh chỉnh từ điểm kiểm tra được huấn luyện trước

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

Nếu bạn muốn đóng băng phần thân Qwen và chỉ huấn luyện đầu hành động, hãy đặt 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

Tinh chỉnh trên một cấu hình khác

Khi robot đích có chiều kích thước hành động hoặc trạng thái khác với checkpoint được huấn luyện trước, các lớp chiếu đầu vào/đầu ra của action head sẽ có hình dạng không khớp và không thể được nạp trực tiếp. reinit_modules cho phép bạn liệt kê các tiền tố khóa được phép không khớp — các lớp đó sẽ được khởi tạo lại ngẫu nhiên trong khi mọi trọng số khác được tái sử dụng từ checkpoint. Bất kỳ sự không khớp hình dạng nào nằm ngoài các tiền tố được liệt kê sẽ gây ra lỗi.

Các lớp phụ thuộc vào action_dim và state_dim là:

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

Nếu robot của bạn không có trạng thái cảm nhận bản thân, hãy bỏ model.action_model.state_encoder khỏi danh sách.

Tái tạo kết quả LIBERO

Đào tạo trên LIBERO: bắt đầu đào tạo từ checkpoint Pretrain, đào tạo trong 30k bước trên bộ dữ liệu LIBERO. Bài báo gốc đề cập đến việc đào tạo trên 8 GPU với batch size là 32, tức là global batch size là 256.

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

Đánh giá checkpoint LIBERO-10 đã được huấn luyện trước:

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

Để đánh giá một tập con các tác vụ:

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

Kết quả mong đợi:

Chỉnh huấn luyện trên các bộ dữ liệu có số lượng camera khác nhau

Bộ dự đoán mô hình thế giới được huấn luyện trước được huấn luyện với embed_dim = world_model_num_views × 1024, tức là hai góc nhìn từ camera.

Số lần xem này trước đây được đọc từ jepa_tubelet_size, vốn cũng là tên của kích thước tubelet theo thời gian của bộ mã hóa JEPA. world_model_num_views là trường dành cho nó hiện tại; để nó ở None sẽ quay lại dùng jepa_tubelet_size để các checkpoint đã công bố vẫn được nạp mà không thay đổi.

Ứng xử mặc định — xem đệm / cắt bỏ (không cần thao tác)

Khi tinh chỉnh từ VLA-JEPA-Pretrain, mô hình sẽ tự động điều chỉnh số lượng góc nhìn được đưa vào mô hình thế giới để khớp với world_model_num_views:

• Các bộ dữ liệu đơn góc nhìn (ví dụ: BridgeV2): latent đơn góc nhìn được nhân đôi để tạo ra đầu vào mô hình thế giới hai góc nhìn, giữ nguyên tín hiệu tự giám sát JEPA mà không có sự lệch trọng số nào.
• >2-view datasets (ví dụ: DROID với 3 views): tất cả các views đều được truyền vào Qwen backbone (để có ngữ cảnh phong phú hơn), nhưng chỉ có world_model_num_views views đầu tiên (một wrist + một third-person, theo thứ tự view đã cấu hình) được sử dụng cho world model.

Tùy chọn 1 — Tắt mô hình thế giới

Đặt enable_world_model=False để bỏ qua hoàn toàn tổn thất JEPA. Chỉ phần backbone Qwen và đầu action được nạp và huấn luyện. Điều này là đủ để đạt hiệu suất action tốt.

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

Tùy chọn 2 — Khởi tạo lại phép chiếu đầu vào của bộ dự đoán

Nếu bạn muốn thay đổi world_model_num_views thành một giá trị khác 2, hãy tải checkpoint với strict=False và khởi tạo lại model.video_predictor.predictor_embed cho embed_dim mới. Tất cả các trọng số khối predictor khác (attention, MLP, norm, output projection) không phụ thuộc vào số lượng camera và có thể được tái sử dụng từ checkpoint đã huấn luyện trước.

Trích dẫn

@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},
}

Giấy phép

Trọng số được phân phối theo các điều khoản giấy phép của kho ginwind/VLA-JEPA gốc (Apache 2.0 License). Mã tích hợp LeRobot tuân theo Apache 2.0 License.