# LingBot-VA LingBot-VA는 **Wan2.2** 비디오-확산 스택을 기반으로 구축된 **자기회귀 비디오-액션 월드-모델 정책**입니다. 하나의 자기회귀 시퀀스 내에서 미래 **비디오 라텐트**와 **로봇 액션**의 예측을 번갈아 수행합니다("VA" = Video-Action). LeRobot 통합은 LingBot-VA를 표준 학습, 평가 및 프로세서 인터페이스에 연결합니다. ## 모델 개요 LingBot-VA is a **dual-stream "mixture-of-transformers"**: a video/latent stream (`patch_embedding_mlp → blocks → proj_out`) and an action stream (`action_embedder → blocks → action_proj_out`) share the same 30 transformer blocks and text conditioning. | 컴포넌트 | 클래스 | 역할 | | ------------------------ | ----------------------- | ----------------------------------------------------------- | | DiT 백본(학습 가능) | `WanTransformer3DModel` | 약 5B 파라미터의 듀얼 스트림 트랜스포머. | | VAE (동결) | `AutoencoderKLWan` | Wan2.2 VAE, `z_dim=48`. 소스 저장소에서 지연 로드됩니다. | | 텍스트 인코더(동결) | `UMT5EncoderModel` | UMT5-XXL, `d_model=4096`. 소스 저장소에서 지연 로드됩니다. | 추론 시 정책은 청크마다 자기회귀 루프를 실행합니다. 두 개의 독립적인 flow-matching 스케줄러를 사용하여 비디오-잠재 변수 스트림(CFG, 약 20단계)과 액션 스트림(약 50단계)을 노이즈 제거하며, 청크 간에 KV 캐시를 유지합니다. 청크가 실행되는 동안 실제 관찰된 키프레임이 KV 캐시로 피드백됩니다(클로즈드 루프 월드 모델링). ### LeRobot 통합의 범위 - LeRobot를 통해 표준 `policy.type=lingbot_va` 구성을 수행합니다. - Hub에 있는 LeRobot 형식의 바로 사용 가능한 체크포인트(공개된 업스트림 버전에서 변환됨). - 표준 `select_action` 인터페이스 뒤에 있는 자기회귀 이중 스트림 추론(단일 환경 평가, `--eval.batch_size=1`). - 평가/학습 중 정책의 **예측(상상)된 비디오**를 저장하도록 선택적으로 활성화합니다. - LIBERO와 RoboTwin에서 `lerobot-eval`을 사용한 평가. - 이중 스트림 플로우 매칭 손실(`policy.forward`)을 통한 학습/파인튜닝, 아래를 참조하세요. ## 설치 1. [설치 가이드](./installation)를 따라 LeRobot을 설치하세요. 2. LingBot-VA 추가 항목을 설치합니다: ```bash pip install -e ".[lingbot_va]" ``` ## 체크포인트 공개된 업스트림 체크포인트는 LeRobot 형식으로 변환되어 Hub에 업로드되었습니다: | 변형 | LeRobot 체크포인트 | | ---------------------- | -------------------------------- | | LIBERO-Long 사후 학습 | `lerobot/lingbot_va_libero_long` | | RoboTwin 사후 학습 | `lerobot/lingbot_va_robotwin` | | 사전학습된 베이스 | `lerobot/lingbot_va_base` | 학습 가능한 ~5B transformer만 LeRobot `model.safetensors`에 저장됩니다. 동결된 VAE + UMT5 + tokenizer(~20 GB)는 로드 시점에 `config.wan_pretrained_path`에서 가져오며(기본값은 소스 `robbyant/*` 저장소입니다), UMT5-XXL 텍스트 인코더는 기본적으로 CPU에서 실행됩니다(`config.text_encoder_device`). 이를 통해 5B transformer + VAE가 단일 24–32 GB GPU에 적합합니다. ## Evaluation (LIBERO) ```bash lerobot-eval \ --policy.path=lerobot/lingbot_va_libero_long \ --policy.device=cuda \ --env.type=libero --env.task=libero_10 \ --env.observation_height=128 --env.observation_width=128 \ --eval.n_episodes=50 --eval.batch_size=1 \ --output_dir=outputs/eval/lingbot_va_libero ``` LingBot-VA의 스트리밍 추론(KV 캐시 + 관찰된 키프레임 피드백)은 단일 환경 평가용으로 구현되어 있으며, `--eval.batch_size=1`을 사용해야 합니다. ## Evaluation (RoboTwin) RoboTwin 2.0은 SAPIEN + CuRobo 시뮬레이터 스택이 필요합니다. `docker/Dockerfile.sim`의 `robotwin_native` 단계는 참고용 설치입니다(LingBot-VA도 `warp-lang==1.3.1`과 고정된 `pytorch3d`가 필요하며, CuRobo는 `TORCH_CUDA_ARCH_LIST`에 GPU의 compute capability를 지정하여 빌드해야 합니다). 시뮬레이터 서버의 RoboTwin 설정은 joint control을 사용하므로, 아래 end-effector 모드는 로컬에서 실행됩니다. RoboTwin은 **end-effector-pose control**을 사용하므로 `--env.action_mode=ee`로 실행하세요: 정책은 각 팔의 `xyz+quaternion+gripper` 델타(`robotwin_tshape` latent layout)를 예측하며, 이는 에피소드의 초기 eef pose에 합성되어 CuRobo IK를 통해 실행됩니다. ```bash lerobot-eval \ --policy.path=lerobot/lingbot_va_robotwin \ --policy.device=cuda \ --env.type=robotwin --env.task=beat_block_hammer --env.action_mode=ee \ --eval.n_episodes=10 --eval.batch_size=1 \ --output_dir=outputs/eval/lingbot_va_robotwin ``` ### 예측(상상)된 비디오 저장 `--policy.save_predicted_video=true`를 설정하면 예측된 비디오 latents를 VAE-decode하여 env-rendered `eval_episode_*.mp4` 비디오 옆에 `pred_episode_*.mp4`를 추가로 작성합니다. 동일한 플래그는 `lerobot-train` 중 주기적 eval에도 적용됩니다. ### 상상된 비디오를 실시간으로 시청 `lerobot-rollout --display_data=true` 실행 시 `--policy.return_predicted_video=true`를 설정하면 각 청크의 예측된 비디오를 디코딩하여, 각 액션이 실행될 때 해당 액션에 속하는 상상된 카메라 프레임을 표시합니다(`prediction.images.`, 각 실제 카메라 옆에). 이 작업은 청크당 VAE 디코딩 1회의 비용이 발생합니다. ## 학습 / 파인튜닝 `LingBotVAPolicy.forward(batch)`은 논문에서 제시한 이중 스트림 **flow-matching** 손실(`latent_loss + action_loss`, 타임스텝 가중, 액션 마스크 적용)을 구현합니다. 이 메서드는 카메라 클립을 VAE로 인코딩하여 비디오 라텐트를 생성하고, UMT5로 태스크를 인코딩한 뒤, 두 스트림 모두에 노이즈를 추가하고, 트랜스포머의 블록 인과 학습 패스를 실행한 후 `(loss, metrics)`를 반환합니다. 옵티마이저 프리셋은 AdamW이며, 선형 웜업 후 일정하게 유지되는 스케줄을 사용합니다(업스트림과 일치). 요구 사항: - 블록 인과 마스크는 PyTorch **flex-attention**를 사용하므로, 학습 시에는 `--policy.attn_mode=flex`로 정책을 구성해야 합니다(기본 `torch` SDPA는 추론 전용입니다). - AdamW를 사용하는 경우, 전체 5B DiT는 단일 24–32 GB GPU에 맞지 않습니다. **LoRA**(`--peft.method_type=LORA`) 및/또는 옵티마이저 오프로드로 파인튜닝하세요. `get_optim_params`는 학습 가능한(예: 어댑터) 파라미터만 반환하며, VAE + UMT5 텍스트 인코더는 고정(frozen)된 상태로 유지됩니다. PEFT 지원을 활성화하려면 `lerobot[peft]` 추가 항목을 설치하세요([PEFT 학습 가이드](./peft_training) 참조). ```bash lerobot-train \ --policy.path=lerobot/lingbot_va_libero_long --policy.attn_mode=flex \ --peft.method_type=LORA --peft.r=32 --peft.lora_alpha=32 \ --peft.target_modules='transformer\.blocks\.\d+\.attn[12]\.(to_q|to_v)' \ --dataset.repo_id= \ --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va ``` SmolVLA / π₀와 달리 LingBot-VA에는 내장된 기본 LoRA 대상이 제공되지 않으므로 `--peft.target_modules`를 명시적으로 전달해야 합니다. 학습 가능한 모듈은 `self.transformer`(dual-stream Wan transformer)뿐입니다. 위 예제는 표준 LoRA 대상 집합인 해당 transformer의 self-attention(`attn1`)과 cross-attention(`attn2`) 블록의 query/value 투영을 적응합니다. 더 높은 용량의 어댑터가 필요하면 범위를 넓히세요(예: `to_k`/`to_out` 또는 `ffn` 레이어 추가). `--peft.method_type`를 전달하면 PEFT가 적용되므로 `--policy.use_peft=true`는 필수 사항이 아닙니다. 데이터셋은 카메라 클립(카메라별 시간 창, `frame_chunk_size`개의 잠재 프레임으로 VAE 인코딩됨)와 항목당 `frame_chunk_size * action_per_frame`개의 액션 스텝을 제공해야 합니다. ## 데이터 형식(액션 채널 및 카메라 순서) LingBot-VA는 **엔드 이펙터(카르테지안) 자세** 정책으로, 관절 위치가 아니라 EEF 자세와 그립퍼를 예측합니다. 동작은 고정된 멀티 엠프디먼트 **30차원** 레이아웃에 존재합니다. 로봇의 동작 차원을 해당 채널에 매핑하고 나머지 채널은 `0`으로 패딩하세요(`used_action_channel_ids`는 특정 체크포인트가 실제로 사용하는 채널을 선택합니다): | 채널 | 의미 | | -------- | ----------------------------------------------------- | | 0–6 | 왼쪽 팔의 엔드 이펙터 자세 | | 7–13 | 오른쪽 팔의 엔드 이펙터 자세 | | 14–20 | 왼쪽 팔 관절(발행된 체크포인트에서 미사용) | | 21–27 | 오른쪽 팔 관절(발행된 체크포인트에서 미사용) | | 28 | 왼쪽 그립퍼 | | 29 | 오른쪽 그립퍼 | - **LIBERO**는 `0–6` 채널을 사용합니다: 6-DoF EEF 델타(xyz + 회전) + 그리퍼(단일 암). - **RoboTwin**은 채널 `[0–6, 28, 7–13, 29]`을 사용합니다: 왼쪽 EEF(xyz + quaternion) + 왼쪽 gripper + 오른쪽 EEF + 오른쪽 gripper(16차원). env는 CuRobo IK를 통해 이러한 포즈를 관절 궤적으로 변환합니다 — 관절은 절대 예측되지 않습니다. Joint-space 데이터셋(또는 다른 EEF 규약)은 이러한 체크포인트를 파인튜닝하기 전에 이 스키마로 재매핑되어야 합니다. **카메라 순서는 고정되어 있으며 순서에 민감합니다**. 카메라별 latents는 `obs_cam_keys` 순서대로 공간적으로 연결되므로, 물리적 카메라→슬롯 매핑은 학습 시와 일치해야 합니다: | 벤치마크 | `obs_cam_keys` (순서대로) | `camera_layout` | | --------- | ----------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------- | | LIBERO | `observation.images.image` (agentview / 3인칭), `observation.images.image2` (eye-in-hand 손목) | `width_concat` (너비 방향으로 연결된 latents) | | RoboTwin | `observation.images.head_camera`, `observation.images.left_camera`, `observation.images.right_camera` | `robotwin_tshape` (아래에 풀 해상도 헤드, 위에 반 해상도 손목 2개) | 첫 번째 카메라는 외부/헤드 뷰이며, 나머지는 손목 뷰입니다. ## 추론 하이퍼파라미터 (LIBERO) | 키 | 값 | | -------------------------------------- | --------------------------------------------------------------------------------- | | 높이 × 너비 | 128 × 128 | | 카메라 | `observation.images.image`(agentview), `observation.images.image2`(eye-in-hand) | | 사용된 action 채널 | 0–6 (7-DoF 암 + 그리퍼) | | action_per_frame / frame_chunk_size | 4 / 4 | | attn_window | 30 | | video / action 노이즈 제거 단계 | 20 / 50 | | guidance_scale / action_guidance_scale | 5 / 1 | | snr_shift / action_snr_shift | 5.0 / 0.05 | 다음은 `LingBotVAConfig`의 기본값입니다. `--policy.=...`을 통해 원하는 값을 덮어쓸 수 있습니다. ## The input is a single word "Notes". 참고 - **Attention 백엔드:** 추론에는 `torch` SDPA 백엔드(항상 사용 가능)가 사용됩니다. `flashattn`와 `flex` 백엔드는 선택 사항이며, `flex`는 학습에만 필요합니다. - **모델 크기:** DiT는 약 5B 파라미터이며, 동결된 VAE+UMT5는 약 20 GB를 추가합니다. 추론에는 대략 18–24 GB의 VRAM이 필요합니다. ## 라이선스 LingBot-VA는 Apache-2.0 라이선스로 공개됩니다. [업스트림 저장소](https://github.com/Robbyant/lingbot-va)를 참조하세요.