FastWAM

FastWAM은 로봇 제어를 위한 World Action Model 정책입니다. LeRobot 통합은 FastWAM을 표준 정책 API를 통해 노출하여 policy.type=fastwam으로 구성하고, lerobot-train으로 학습하며, LeRobot 사전 학습 정책 인터페이스를 통해 로드할 수 있도록 합니다.

모델 개요

FastWAM은 학습 중에 비디오 모델링을 유지하지만, 추론 시에는 반복적으로 미래 관측을 생성하는 대신 직접 동작을 예측합니다. 이 LeRobot 정책은 FastWAM 동작 모델을 래핑하고, LeRobot 배치에서 FastWAM 학습 샘플로 변환하며, 정규화 및 동작 후처리를 위한 표준 프로세서 파이프라인을 제공합니다.

구현은 기본적으로 Wan-AI/Wan2.2-TI2V-5B에서 시각적 월드 모델 구성 요소를 초기화하고, [batch, action_horizon, action_dim] 형태의 액션 청크를 예측합니다.

LeRobot 통합의 범위

• LeRobot를 통한 표준 policy.type=fastwam 구성
• 이미지, 상태, 행동 및 언어 작업 배치 적응
• select_action와 predict_action_chunk를 통한 액션 청크 추론
• LeRobot 정책 API를 통한 Checkpoint 저장/로드
• 구성 가능한 LIBERO 그리퍼 액션 후처리

설치 요구 사항

LeRobot을 소스에서 설치한 후, FastWAM의 종속성을 설치합니다:

pip install -e ".[fastwam]"

pyproject.toml에서 FastWAM 정책 추가 항목을 설치합니다: transformers, diffusers, ftfy, regex 및 LeRobot의 기본 의존성.

LIBERO 평가의 경우, 벤치마크 의존성도 설치해야 합니다:

pip install -e ".[fastwam,libero]"

이 명령은 두 가지 extras를 모두 설치합니다. 위의 FastWAM 의존성 외에도, libero extra는 LeRobot dataset 의존성, Linux 환경의 hf-libero, 그리고 scipy를 설치합니다.

FastWAM은 Wan2.2 TI2V 백본을 사용합니다. 기본 모델 id는 다음과 같습니다:

policy.model_id=Wan-AI/Wan2.2-TI2V-5B

데이터 요구 사항

FastWAM은 다음을 포함하는 LeRobot 데이터셋을 기대합니다:

• policy.image_size[1]로 이어 붙이면 하나의 폭이 되는 하나 이상의 시각적 관찰
• policy.proprio_dim이 None이 아닌 경우 observation.state
• action
• 데이터셋의 task 필드를 통한 언어 작업 지시 또는 사전 계산된 context 및 context_mask 텐서

기본 시각 설정은 (3, 224, 448) 형상의 observation.images.image라는 이름의 이미지 특징 하나입니다. 데이터셋에 두 개의 카메라가 사용되는 경우, policy.input_features를 구성하여 두 이미지의 높이가 224와 일치하고 너비의 합이 448가 되도록 설정합니다.

사용

다음으로 새 FastWAM 정책을 생성합니다:

lerobot-train \
  --dataset.repo_id=your-org/your-dataset \
  --policy.type=fastwam \
  --policy.action_dim=7 \
  --policy.proprio_dim=8 \
  --policy.action_horizon=32 \
  --policy.n_action_steps=10 \
  --policy.image_size='[224,448]' \
  --output_dir=./outputs/fastwam_training \
  --job_name=fastwam_training \
  --steps=300000 \
  --batch_size=8 \
  --policy.device=cuda

다음으로 기존 LeRobot 형식 체크포인트를 LIBERO-10에서 평가합니다:

lerobot-eval \
  --policy.path=lerobot/fastwam_libero_uncond_2cam224 \
  --policy.device=cuda \
  --policy.dtype=float32 \
  --policy.n_action_steps=10 \
  --env.type=libero \
  --env.task=libero_10 \
  --env.observation_height=224 \
  --env.observation_width=224 \
  --eval.batch_size=1 \
  --eval.n_episodes=50 \
  --seed=0 \
  --env.episode_length=600

libero_goal, libero_spatial, libero_object의 경우 --env.episode_length=300을(를) 사용하세요.

이 체크포인트는 VISUAL: IDENTITY 전처리를 사용합니다. FastWAM은 모델 내부에서 이미지를 [-1, 1]로 다시 스케일링하므로, 전처리기에 이미지 정규화를 적용하면 두 번 정규화됩니다.

실제 로봇 롤아웃을 위해 동일한 체크포인트 경로를 사용합니다:

lerobot-rollout \
  --robot.type=so101_follower \
  --robot.port=/dev/ttyACM0 \
  --policy.path=your-org/fastwam-real-robot

구성 참고 사항

이미지 기능

policy.image_size는 연결된 FastWAM 이미지 텐서의 크기를 (height, width)로 나타냅니다. 각 구성된 이미지 특징은 (3, height, camera_width) 형태를 가져야 하며, 모든 camera_width의 합은 구성된 width와 같아야 합니다.

액션 청킹

policy.action_horizon은 학습 중에 감독되는 미래 동작의 수와 추론 중에 예측되는 동작의 수를 제어합니다. policy.n_action_steps은 정책이 새로운 청크를 예측하기 전에 소비되는 동작의 수를 제어합니다. policy.n_action_steps는 policy.action_horizon보다 작거나 같아야 합니다.

상대적 동작

FastWAM은 기본적으로 절대 동작을 예측합니다. policy.use_relative_actions=true를 사용하면 대신 현재 observation.state에 대한 오프셋을 예측하므로 policy.proprio_dim를 설정해야 합니다. policy.relative_exclude_joints에 지정된 관절(기본값 ["gripper"])은 절대값으로 유지됩니다. 먼저 --operation.chunk_size를 policy.action_horizon으로 설정하고 상대 공간에서 데이터셋 통계를 다시 계산하십시오(π₀.₅ 상대 동작 참고):

lerobot-train \
  --dataset.repo_id=your-org/your-dataset \
  --policy.type=fastwam \
  --policy.use_relative_actions=true \
  --policy.relative_exclude_joints='["gripper"]' \
  ...

lerobot/fastwam_base와 같이 상대 동작(relative actions)이 존재하기 전에 저장된 체크포인트는 이 옵션을 사용하여 해당 체크포인트에서 파인튜닝할 때 상대 동작 프로세서 단계가 추가됩니다. 이는 파인튜닝 시에만 적용되며, 평가 및 롤아웃은 체크포인트와 함께 저장된 프로세서를 그대로 사용합니다.

텍스트 인코더 장치

동결된 UMT5-XXL 텍스트 인코더는 약 11 GB를 차지합니다. policy.text_encoder_device=cpu를 사용하면 GPU에서 이를 제외할 수 있습니다. 프롬프트는 CPU에서 인코딩되고 임베딩은 policy.device로 이동합니다. 추론 중에는 인코딩된 작업이 policy.reset()가 호출될 때까지 캐시되므로, 느린 CPU 인코딩은 에피소드당 한 번씩만 실행되고 청크당 한 번씩 실행되지 않습니다. 학습에서는 모든 배치에 대해 인코딩이 수행되므로, CPU 배치로 인해 각 단계가 느려집니다.

Wan 구성 요소

FastWAM은 구성된 Wan 모델 디렉터리 또는 Hugging Face Hub 모델 ID에서 Wan VAE, video DiT, text encoder, tokenizer를 로드합니다. save_pretrained로 저장된 LeRobot 형식의 FastWAM 체크포인트는 from_pretrained에 필요한 로컬 Wan 구성 요소 파일도 복사합니다.

Attention 백엔드

FastWAM의 DiT는 모든 attention에 PyTorch의 scaled_dot_product_attention(SDPA)를 사용합니다. FlashAttention은 사용하지 않습니다. Mixture-of-Transformers(MoT) 라우팅에는 임의의 boolean [query, key] attention mask가 필요하며, FlashAttention의 varlen API로는 이를 표현할 수 없습니다. 따라서 flash-attn 패키지를 설치해도 FastWAM 경로에는 영향을 미치지 않습니다. (참고로 SDPA 자체는 내부적으로 PyTorch의 flash / memory-efficient / math 커널을 선택할 수 있습니다. 이는 flash-attn 패키지와는 무관합니다.)

컴파일된 액션 추론

FastWAM은 반복적인 추론 단계에서 Python 및 CUDA 커널 실행 오버헤드를 줄이기 위해 비디오 프리필 및 캐시된 액션 디노이징 그래프를 컴파일할 수 있습니다:

--policy.compile_action_infer=true

이 옵션은 기본적으로 비활성화되어 있습니다. 첫 번째 동작 예측은 현재 입력 형식에 대해 그래프를 컴파일하므로 더 느리며, 동일한 형식을 가진 후속 예측에서는 이를 재사용합니다. 형식이 변경되면 재컴파일이 트리거될 수 있습니다. 지원되지 않는 PyTorch 또는 가속기 구성에서는 eager 경로가 여전히 사용 가능합니다.

LIBERO 동작 토글

FastWAM LIBERO 체크포인트는 기본적으로 policy.toggle_action_dimensions=[-1]을 사용하여 원래 FastWAM 평가 파이프라인에서 사용된 그리퍼 액션 규약과 일치하도록 설정됩니다:

--policy.toggle_action_dimensions='[-1]'

결과

LIBERO에서 lerobot/fastwam_libero_uncond_2cam224로 평가:

재현: lerobot-eval --policy.path=lerobot/fastwam_libero_uncond_2cam224 --policy.device=cuda --policy.dtype=float32 --policy.n_action_steps=10 --env.type=libero --env.task=libero_spatial --env.observation_height=256 --env.observation_width=256 --eval.batch_size=1 --eval.n_episodes=50 --seed=0 --env.episode_length=300 (1x H20 140 GB).

참고

• Fast-WAM 논문
• Fast-WAM 프로젝트 페이지
• Fast-WAM 코드
• 업스트림에 릴리스된 체크포인트
• Wan2.2 TI2V 5B

인용

@article{yuan2026fastwam,
  title = {Fast-WAM: Do World Action Models Need Test-time Future Imagination?},
  author = {Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao},
  journal = {arXiv preprint arXiv:2603.16666},
  year = {2026},
  url = {https://arxiv.org/abs/2603.16666}
}