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.

추론 시 정책은 청크마다 자기회귀 루프를 실행합니다. 두 개의 독립적인 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. 설치 가이드를 따라 LeRobot을 설치하세요.
2. LingBot-VA 추가 항목을 설치합니다:

pip install -e ".[lingbot_va]"

체크포인트

공개된 업스트림 체크포인트는 LeRobot 형식으로 변환되어 Hub에 업로드되었습니다:

학습 가능한 ~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에 적합합니다.

평가 (LIBERO)

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을 사용해야 합니다.

평가 (RoboTwin)

RoboTwin 2.0에는 SAPIEN + CuRobo 시뮬레이터 스택이 필요합니다. 벤치마크 Docker 이미지(docker/Dockerfile.benchmark.robotwin, warp-lang==1.3.1 및 TORCH_CUDA_ARCH_LIST에 GPU의 compute capability를 지정하여 빌드한 CuRobo도 필요함)를 사용할 수 있습니다. RoboTwin은 end-effector-pose control을 사용하므로 --env.action_mode=ee로 실행하세요: 정책은 팔당 xyz+quaternion+gripper 델타(robotwin_tshape latent layout)를 예측하며, 이는 에피소드의 초기 eef pose에 합성되어 CuRobo IK를 통해 실행됩니다.

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.<camera>, 각 실제 카메라 옆에). 이 작업은 청크당 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 학습 가이드 참조).

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=<your LeRobot-format dataset> \
  --batch_size=1 --steps=... --output_dir=outputs/train/lingbot_va

SmolVLA / π₀와 달리 LingBot-VA에는 기본 LoRA 대상이 내장되어 있지 않으므로 --peft.target_modules를 명시적으로 전달해야 합니다. 학습 가능한 모듈은 self.transformer(dual-stream Wan transformer)뿐입니다. 위 예제는 표준 LoRA 대상 집합인 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는 특정 체크포인트가 실제로 사용하는 채널을 선택합니다):

• 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 순서대로 공간적으로 연결되므로, 물리적 카메라→슬롯 매핑은 학습 시와 일치해야 합니다:

첫 번째 카메라는 외부/헤드 뷰이며, 나머지는 손목 뷰입니다.

추론 하이퍼파라미터 (LIBERO)

다음은 LingBotVAConfig의 기본값입니다. --policy.<name>=...을 통해 원하는 값을 덮어쓸 수 있습니다.

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 라이선스로 공개됩니다. 업스트림 저장소를 참조하세요.