# VLA-JEPA 이것은 **VLA-JEPA**의 LeRobot 포트입니다. VLA-JEPA는 Qwen3-VL 언어 백본, 자기지도 학습 비디오 월드 모델(V-JEPA2), flow-matching DiT 액션 헤드를 결합한 Vision-Language-Action 모델입니다. --- ## 아키텍처 개요 VLA-JEPA에는 세 가지 주요 구성 요소가 있습니다: | 컴포넌트 | 모듈 | 역할 | | ----------------------- | --------------------------------- | ------------------------------------------------------- | | **Qwen3-VL 백본** | `Qwen3VLInterface` | 이미지와 언어 지시문을 컨텍스트 토큰으로 융합 | | **DiT-B 액션 헤드** | `VLAJEPAActionHead` | 액션 청크에 대한 Flow-matching diffusion | | **V-JEPA2 월드 모델** | `ActionConditionedVideoPredictor` | 자기지도형 비디오 예측 손실(학습 전용) | ### 데이터 흐름 **학습:** 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`를 통해 사용할 수 있는 프리셋: | 프리셋 | 헤드 | Head dim | | ------- | ----- | -------- | | `DiT-B` | 12 | 64 | | `DiT-L` | 32 | 48 | 프리셋은 주의(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`](https://huggingface.co/collections/lerobot/vla-jepa), [ginwind/VLA-JEPA](https://huggingface.co/ginwind/VLA-JEPA)에서 변환되었습니다: | 체크포인트 | 데이터셋 | 카메라 | 월드 모델 | 액션 차원 | | ----------------------------- | ----------------- | ----------------------- | ----------- | ---------- | | `lerobot/VLA-JEPA-LIBERO` | LIBERO-10 | 2 (agentview + wrist) | 활성화됨 | 7 | | `lerobot/VLA-JEPA-Pretrain` | DROID 1.0.1 | 2 (외부 왼쪽 뷰) | 활성화됨 | 7 | | `lerobot/VLA-JEPA-SimplerEnv` | OXE Bridge / RT-1 | 1 (2회 중복 표시) | 활성화됨 | 7 | 모든 체크포인트는 언어 백본으로 `Qwen/Qwen3-VL-2B-Instruct`를 사용합니다. --- ## 구성 `VLAJEPAConfig`의 주요 매개변수: | 매개변수 | 기본값 | 설명 | | ---------------------------- | ------------- | ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | `chunk_size` | 7 | 추론 호출당 예측되는 액션 수 | | `n_action_steps` | 7 | 재계획 전에 예측된 청크에서 실행된 단계 | | `num_video_frames` | 8 | 월드 모델에 입력되는 비디오 클립 길이 | | `enable_world_model` | `True` | V-JEPA2 예측기를 로드하고 학습할지 여부 | | `world_model_loss_weight` | 0.1 | JEPA 예측 손실의 액션 손실에 대한 가중치 | | `causal_world_model_context` | `False` | 세계 모델 컨텍스트를 인과적으로 인코딩하여(컨텍스트 위치마다 prefix-only V-JEPA2 패스 1회) 양방향 어텐션이 예측기 입력에 미래 프레임을 누출하지 못하도록 합니다. `t_enc_ctx`만큼의 추가 인코더 호출 비용이 발생합니다 | | `num_inference_timesteps` | 4 | 행동 노이즈 제거를 위한 Euler 적분 단계 | | `freeze_qwen` | `False` | Qwen3-VL 백본을 동결하고 액션 헤드만 학습합니다 | | `reinit_modules` | `None` | 로드 시 무작위로 재초기화할 수 있는 주요 접두사(크로스 임보디먼트 전이에 대해 [다른 임보디먼트에서 파인튜닝](#fine-tuning-on-a-different-embodiment)을 참조) | | `resize_images_to` | `None` | `(height, width)`는 Qwen3-VL 비전 타워에 입력되기 전에 모든 카메라 프레임이 리사이즈되는 크기입니다. `None`을 사용하면 네이티브 해상도가 유지되며, Qwen3-VL의 패치 수가 해상도에 따라 증가하므로 720x1280 카메라의 경우 GPU 메모리가 소진될 수 있습니다. 공개된 체크포인트는 `[224, 224]`를 사용합니다. | | `gripper_dim` | 6 | 액션 벡터에서 그리퍼 차원의 인덱스입니다. `gripper_joint_names`가 데이터셋의 액션 이름과 일치하는 경우 무시됩니다. | | `gripper_joint_names` | `["gripper"]` | 그리퍼를 식별하는 액션 차원 이름; 일치하는 인덱스가 `gripper_dim`보다 우선합니다 | | `gripper_threshold` | 0.5 | `pre_snap_gripper_action`와 `binarize_gripper_action`에서 사용되는 임계값입니다. `binarize`는 비정규화 _이후_에 실행되므로, 이 값은 그리퍼의 물리적 값과 비교됩니다. | | `pre_snap_gripper_action` | `False` | unnormalization 전에 gripper dim을 {0, 1}로 snap합니다. LIBERO 전용이며, 아래를 참조하세요 | | `binarize_gripper_action` | `False` | unnormalization 후 gripper dim을 {-1, 1}로 이진화합니다. LIBERO 전용이며, 아래를 참조하세요 | | `clip_normalized_actions` | `True` | `ACTION`이 `MIN_MAX`를 사용하는 경우에만 적용되며, `MEAN_STD`에서는 1 sigma에서 잘리게 되므로 무시됩니다(경고와 함께). 비정규화하기 전에 정규화된 동작을 [-1, 1]로 클리핑합니다. | | `world_model_num_views` | `None` | Camera world-model predictor가 구축된 대상입니다. 체크포인트 형식에 고정되어 있습니다. `None`은 공개된 체크포인트에 인코딩된 `jepa_tubelet_size`로 대체됩니다. | `pre_snap_gripper_action`와 `binarize_gripper_action`은 starVLA의 LIBERO eval 루프를 이식한 것으로, LIBERO의 action 규약에 대해서만 정확합니다. `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 ```bash lerobot-train \ --policy.type=vla_jepa \ --policy.repo_id=your_org/your_repo \ --dataset.repo_id=your_org/your_dataset ``` ### 사전 학습된 체크포인트에서 파인튜닝 ```bash 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`를 설정하세요: ```bash 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`에 의존하는 레이어는 다음과 같습니다: | 레이어 | 주요 접두사 | | ----------------------------------------- | ----------------------------------- | | 액션 인코더 (action_dim → inner_dim) | `model.action_model.action_encoder` | | 액션 디코더 (hidden_size → action_dim) | `model.action_model.action_decoder` | | 상태 인코더 (state_dim → inner_dim) | `model.action_model.state_encoder` | ```bash 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을 의미합니다. ```bash lerobot-train \ --policy.path=lerobot/VLA-JEPA-Pretrain \ --policy.repo_id=your_org/your_repo \ --dataset.repo_id=HuggingFaceVLA/libero \ --steps=30000 ``` **학습된 LIBERO-10 체크포인트 평가:** ```bash 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 ``` 일부 작업만 평가하려면: ```bash 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 ``` **예상 결과:** | Suite | 에피소드 | 성공 사례 | 성공률 | | -------------- | -------- | --------- | ------------ | | libero_spatial | 100 | 93 | **95.0%** | | libero_object | 100 | 100 | **100.0%** | | libero_goal | 100 | 98 | **98.0%** | | libero_10 | 100 | 96 | **93.0%** | | **전체** | **400** | **387** | **96.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 백본과 액션 헤드만 로드되고 학습됩니다. 이는 좋은 액션 성능을 위해 충분합니다. ```bash 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)는 카메라 수와 무관하므로 사전 학습된 체크포인트에서 재사용할 수 있습니다. --- ## 인용 ```bibtex @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](https://huggingface.co/ginwind/VLA-JEPA) 저장소의 라이선스 조항(**Apache 2.0 License**)에 따라 배포됩니다. LeRobot 통합 코드는 **Apache 2.0 License**를 따릅니다.