VLA-JEPA

이것은 VLA-JEPA의 LeRobot 포트입니다. VLA-JEPA는 Qwen3-VL 언어 백본, 자기지도 학습 비디오 월드 모델(V-JEPA2), flow-matching DiT 액션 헤드를 결합한 Vision-Language-Action 모델입니다.

아키텍처 개요

VLA-JEPA에는 세 가지 주요 구성 요소가 있습니다:

데이터 흐름

학습:

1. num_video_frames 프레임의 비디오 클립은 V-JEPA2에 의해 프레임별 패치 토큰으로 인코딩됩니다.
2. Qwen3-VL 백본은 다중 뷰 이미지와 작업 지시문을 처리하여, 월드 모델 조건화에 사용되는 특수 동작 토큰과 embodied 토큰을 포함하는 컨텍스트 토큰 시퀀스를 생성합니다.
3. 액션 헤드는 해당 컨텍스트 토큰을 크로스 어텐션의 키/값으로 받아 flow matching을 통해 노이즈가 제거된 액션 청크를 예측합니다.
4. world model predictor는 Qwen에서 추출된 action token을 사용하여 미래의 V-JEPA2 프레임 임베딩을 예측합니다. 해당 예측에 대한 회귀 손실이 action 손실에 추가됩니다.

추론: Qwen과 액션 헤드만 사용됩니다. 추론 시에는 월드 모델이 필요하지 않습니다.

액션 헤드 세부 정보

action_model_type를 통해 사용할 수 있는 프리셋:

프리셋은 주의(attention) 기하 구조만 설정하며, 각 항목은 action_num_heads / action_attention_head_dim으로 재정의할 수 있습니다. 여기서 두 가지 폭이 도출됩니다:

• DiT의 내부 폭은 heads x head_dim이며(DiT-B의 경우 768), 설정값이 아닌 파생 값입니다;
• DiT의 출력 폭과 action-decoder 및 state-encoder MLP의 폭은 action_hidden_size(기본값 1024)입니다.

따라서 DiT-B는 768 폭의 transformer를 실행하여 1024로 투영합니다. 이 두 가지는 독립적입니다.

월드 모델 세부 정보

비디오 예측기는 ViT 스타일의 트랜스포머(ActionConditionedVideoPredictor)로 다음을 입력으로 받습니다:

• 프레임 토큰: predictor_embed_dim으로 투영된 V-JEPA2 패치 임베딩
• 액션 토큰: predictor_embed_dim으로 투영된 Qwen 액션 토큰 임베딩

블록 인과적(block-causal) 어텐션을 사용하여 각 시간 단계가 모든 이전 단계에 어텐션을 기울일 수 있습니다. 예측기의 입력 embed_dim은 num_views × video_encoder_hidden_size와 같습니다(예: 사전 학습된 체크포인트의 경우 2 views × 1024 = 2048).

사전학습된 체크포인트

LeRobot 조직에서 다음 세 개의 체크포인트를 직접 사용할 수 있습니다: lerobot/VLA-JEPA, ginwind/VLA-JEPA에서 변환되었습니다:

모든 체크포인트는 언어 백본으로 Qwen/Qwen3-VL-2B-Instruct를 사용합니다.

구성

VLAJEPAConfig의 주요 매개변수:

pre_snap_gripper_action와 binarize_gripper_action은 starVLA LIBERO eval 루프의 이식된 코드이며, LIBERO의 액션 규약에 대해서만 정확합니다. pre_snap은 normalized 공간에 {0, 1}을 기록하고, unnormalizer는 해당 값을 중간점과 최대값으로 매핑하며, binarize는 그런 다음 해당 physical 값을 gripper_threshold(0.5)와 비교합니다. 그립퍼가 도, mm 또는 [0, 100] 단위로 측정되는 경우, 두 값 모두 임계값을 초과하여 명령된 그립퍼가 상수가 됩니다. 이 때문에 기본값은 False입니다. LIBERO 스타일의 설정에서만 활성화하고, 활성화할 경우 gripper_threshold를 그립퍼 자체의 단위로 설정하십시오. 데이터셋 통계에서 범위가 작동할 수 없음을 나타내면 프로세서 팩토리가 경고를 표시합니다.

학습

학습 단계 수는 데이터셋 크기와 컴퓨트 예산에 따라 달라질 수 있습니다. 원 논문에서는 ssv2 + droid를 결합하여 50k 단계로 사전 학습한 후, LIBERO를 위해 추가 30k 단계를 수행했지만, 제공된 사전 학습 체크포인트에서 파인튜닝을 수행할 경우 더 적은 단계로도 좋은 성능을 얻을 수 있습니다.

Full training from scratch

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

사전 학습된 체크포인트에서 파인튜닝

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

Qwen 백본을 동결하고 액션 헤드만 학습하려면 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

다른 embodiment에 대한 Fine-tuning

대상 로봇의 동작 또는 상태 차원이 사전 학습된 체크포인트와 다르면, 동작 헤드의 입력/출력 투영 레이어의 형상이 일치하지 않아 직접 로드할 수 없습니다. reinit_modules는 형상 불일치가 허용되는 키 접두사를 나열할 수 있게 하며, 해당 레이어는 무작위로 재초기화되고 나머지 모든 가중치는 체크포인트에서 재사용됩니다. 나열된 접두사 범위를 벗어난 형상 불일치는 오류를 발생시킵니다.

action_dim과 state_dim에 의존하는 레이어는 다음과 같습니다:

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

로봇에 고유 감각 상태(proprioceptive state)가 없으면 목록에서 model.action_model.state_encoder를 생략하세요.

LIBERO 결과 재현

LIBERO에서 학습: Pretrain 체크포인트에서 학습을 시작하고 LIBERO 데이터셋에서 30k 스텝 동안 학습합니다. 원 논문에서는 8개의 GPU에서 배치 크기가 32로 학습한다고 언급되어 있으며, 이는 전역 배치 크기가 256임을 의미합니다.

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

학습된 LIBERO-10 체크포인트 평가:

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

일부 작업만 평가하려면:

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

예상 결과:

카메라 수가 다른 데이터셋으로 Fine-tuning

사전 학습된 월드 모델 예측기는 embed_dim = world_model_num_views × 1024, 즉 두 개의 카메라 뷰로 학습되었습니다.

이 뷰 카운트는 temporal tubelet size를 지정하는 JEPA 인코더의 jepa_tubelet_size에서 읽혔습니다. 이제 해당 필드는 world_model_num_views입니다. None으로 두면 jepa_tubelet_size로 폴백되어 공개된 체크포인트가 변경 없이 로드됩니다.

기본 동작 — 패딩/트림 보기 (조치 불필요)

VLA-JEPA-Pretrain에서 파인튜닝할 때 모델은 world_model_num_views에 맞춰 월드 모델에 입력되는 뷰의 수를 자동으로 조정합니다:

• 단일 뷰 데이터셋(예: BridgeV2): 단일 뷰 잠재 변수를 복제하여 두 뷰 월드 모델 입력을 생성합니다. 이를 통해 가중치 불일치 없이 JEPA 자기 감독 신호를 유지합니다.
• >2개의 뷰를 가진 데이터셋(예: 3개의 뷰를 가진 DROID): 모든 뷰는 더 풍부한 컨텍스트를 위해 Qwen 백본에 전달되지만, 월드 모델에는 구성된 뷰 순서에 따라 첫 번째 world_model_num_views개의 뷰(손목 뷰 1개 + 서드 퍼슨 뷰 1개)만 사용됩니다.

옵션 1 — 월드 모델 비활성화

enable_world_model=False를 설정하면 JEPA 손실을 완전히 건너뜁니다. Qwen 백본과 액션 헤드만 로드되고 학습됩니다. 이는 좋은 액션 성능을 위해 충분합니다.

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

옵션 2 — 예측기 입력 투영 재초기화

world_model_num_views를 2가 아닌 다른 값으로 변경하려면, strict=False로 체크포인트를 로드하고 새로운 embed_dim에 맞춰 model.video_predictor.predictor_embed를 다시 초기화해야 합니다. 나머지 예측기 블록 가중치(attention, MLP, norm, output projection)는 카메라 수와 무관하므로 사전 학습된 체크포인트에서 재사용할 수 있습니다.

인용

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

라이선스

가중치는 원래 ginwind/VLA-JEPA 저장소의 라이선스 조항(Apache 2.0 License)에 따라 배포됩니다. LeRobot 통합 코드는 Apache 2.0 License를 따릅니다.