Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility

FLEX-π: 컴퓨팅 유연성을 갖춘 멀티 스트림 월드-액션 모델(Multi-Stream World-Action Model)에 대한 기술 요약

문제 정의

월드-액션 모델(WAM)은 미래의 관측값과 행동을 공동으로 예측함으로써 강력한 시공간적 표현을 학습하고 대규모 비디오 생성 모델의 사전 지식을 상속받는, 비전-언어-액션 모델(VLA)의 강력한 대안으로 부상했습니다. 그러나 현재의 범용 WAM들은 거의 예외 없이 미래의 RGB 이미지 잠재 공간(latents)만을 예측합니다. 순수하게 픽셀 재구성을 위해 학습된 이러한 잠재 공간은 외관의 세부 사항은 포착하지만, 정밀한 로봇 조작에 필요한 명시적인 3D 기하학 및 객체 수준의 의미론(semantics)과 정렬되는 데는 종종 실패합니다.

기하학 및 의미론에 대해 WAM을 직접 감독(supervising)하는 것은 일반적으로 다음과 같은 세 가지 중요한 비용을 발생시킵니다:

1. 추가 센서: 새로운 하드웨어 모달리티(예: 깊이 센서)를 요구함.
2. 새로운 사전 학습: 새로운 시각적 사전 지식을 처음부터 다시 학습시켜야 함.
3. 추론 지연: 복잡한 다중 모달 프로세싱으로 인해 배포 속도를 늦춤.

본 논문은 다음과 같은 질문을 던집니다: 어떻게 하면 이러한 희생 없이 WAM의 강점(결합 예측 및 강력한 시각적 사전 지식)을 3D 기하학 및 의미론까지 확장할 수 있을까?

방법론: FLEX-π

저자들은 미래의 RGB 관측값, 3D 포인트맵(pointmaps), 그리고 객체 중심의 DINO 의미론을 동시에 예측하는 6B 파라미터 규모의 WAM인 FLEX-π를 소개합니다. 핵심 혁신은 새로운 센서나 사전 학습 없이도 이러한 "공짜 점심(free lunch)" 식의 다중 모달 그라운딩(grounding)을 달성하는 것입니다.

1. 고정된 VAE를 통한 공유 잠재 공간

FLEX-π는 사전 학습된 비디오 생성 모델(Wan-2.2)의 고정된 변이형 오토인코더(VAE)를 활용합니다.

• RGB 인코딩: VAE는 일반적인 방식대로 RGB 이미지를 인코딩합니다.
• 3D 포인트맵 인코딩: 놀랍게도 저자들은 오직 RGB 픽셀로만 학습된 이 고정된 VAE가 3D 포인트맵을 거의 손실 없이 인코딩하고 디코딩할 수 있다는 것을 발견했습니다. 이를 통해 3D 기하학을 새로운 인코더 학습 없이 RGB와 동일한 잠재 공간으로 투영할 수 있습니다.
• 의미론적 인코딩: 객체 중심의 의미론은 고정된 DINOv3 인코더를 사용하여 도출됩니다.

2. 다중 스트림 혼합 트랜스포머(MoT) 백본

FLEX-π는 세 가지 시각적 스트림(RGB, 포인트맵, DINO)을 공유된 토큰 스트림 잠재 공간에 임베딩합니다.

• 아키텍처: 혼합 트랜스포머(Mixture-of-Transformers) 백본(5B 파라미터)이 이러한 스트림을 처리합니다. 이는 공유된 비전 트랜스포머 백본을 사용하되, 각 모달리티에 대해 독립적인 피드포워드 투영 네트워크를 유지합니다.
• 액션 전문가(Action Expert): 별도의 더 좁은 액션 전문가 트랜스포머(1B 파라미터)가 시각적 스트림에 크로스 어텐션(cross-attend)하여 액션 청크(action chunks)를 생성합니다.
• 인과적 결합 생성: 액션 토큰은 현재의 관측값에 어텐션하고, 활성화된 미래의 시각적 토큰에 크로스 어텐션합니다. 이를 통해 액션 생성이 모델의 진화하는 미래 표현으로부터 혜택을 받을 수 있게 합니다.

3. 스트림 드롭아웃 및 교차 모달리티 강제 학습을 통한 유연한 학습

단일 체크포인트가 추론 시 모든 입력/출력 조합을 지원할 수 있도록, FLEX-π는 두 가지 핵심 학습 전략을 채택합니다.

• 시각적 스트림 드롭아웃(Visual Stream Dropout): 학습 중에 입력 스트림을 무작위로 드롭(확률 0.5)하여, 모델이 부분적인 관측(예: RGB만 있는 경우)만으로도 작동할 수 있도록 보장합니다.
• 교차 모달리티 강제(Cross-Modality Forcing): 결정적으로, 특정 스트림이 입력에서 드롭되더라도, 모델은 여전히 해당 스트림을 디노이징하고 예측하도록 강제됩니다. 예를 들어, 3D 포인트맵이 입력에서 누락되더라도 모델은 RGB와 DINO 입력을 바탕으로 미래의 포인트맵을 합성해야 합니다. 이는 모델이 외관, 기하학, 의미론이 서로를 예측할 수 있는 통합된 표현을 내재화하도록 강제합니다.

4. 추론 유연성

배포 시 사용자는 입력/출력 스트림의 선택을 통해 지연 시간과 성능 사이의 균형을 맞출 수 있습니다:

• 액션 전용 모드(Action-Only Mode): 모델이 액션만을 예측합니다 (가장 빠름).
• 전체 결합 생성 모드(Full Joint Generation): 액션과 함께 미래의 RGB, 포인트맵, DINO 특징을 모두 예측합니다 (최고 성능).
• 부분 생성 모드(Partial Generation): 스트림의 임의의 하위 집합을 생성할 수 있습니다.

주요 기여

1. 무료 기하학 및 의미론적 그라운딩: 고정된 RGB 학습 비디오 VAE가 3D 포인트맵을 손실 없이 인코딩할 수 있음을 입증하여, 새로운 센서나 사전 학습 없이도 WAM이 3D 기하학 및 DINO 의미론에 대해 감독될 수 있음을 보여주었습니다.
2. 컴퓨팅 유연성: 단일 학습된 체크포인트가 모든 입력/출력 스트림의 하위 집합을 지원하여, 실무자가 배포 시점에 속도와 성능의 경계선에서 운영 지점을 선택할 수 있게 합니다.
3. 교차 모달리티 강제 학습: 누락된 모달리티를 관측된 모달리티로부터 합성하도록 강제하는 학습 체계를 통해, 강력한 공유 표현을 학습하고 정책 성능을 크게 향상시켰습니다.

실험 결과

저자들은 시뮬레이션 벤치마크(RoboTwin, LIBERO, LIBERO-Plus)와 YAM 로봇을 이용한 실제 양팔 조작 작업에서 FLEX-π를 평가했습니다.

• 실제 환경 성능: 다섯 가지 정교하고 정밀한 실제 작업(예: "Self-Repair Gripper" 및 "Soft-Bag Zipping")에서 FLEX-π는 강력한 베이스라인들(π0.5, ManiFlow, Fast-WAM)보다 성공률 측면에서 2~7배 높은 성능을 보였습니다.
  • 특히, 미래의 시각적 정보를 생성하지 않는 FLEX-π의 액션 전용(action-only) 변체조차도 모든 베이스라인보다 우수한 성능을 보였으며 π0.5보다 빠르게 작동했습니다.
  • 전체 결합 생성(Full joint generation)은 성능을 더욱 높여, 어려운 작업에서 가장 강력한 베이스라인 대비 최대 42.7%의 절대적 성공률 향상을 달성했습니다.
• 데이터 효율성: FLEX-π는 데이터 효율성이 매우 높습니다. RoboTwin에서 적은 데이터 환경(데모 50개)일 때 베이스라인보다 1.9~4.5배 높은 성능을 보였습니다. 실제 데이터를 50%만 사용하여 학습했을 때도 FLEX-π는 전체 데이터를 사용해 학습한 베이스라인보다 뛰어난 성능을 보였습니다.
• 일반화 (OOD): FLEX-π는 분포 외 변화(보지 못한 물체, 장애물, 방해 요소)에 대한 강건함을 보여줍니다. ManiFlow나 π0.5와 같은 베이스라인들이 분포 변화 시 상당한 성능 저하(25~26%)를 겪는 반면, FLEX-π의 저하는 5% 미만이었습니다.
• 절제 실험(Ablations):
  • 학습 중 DINO 또는 포인트맵 입력을 제거하면 성공률이 각각 6.8%와 20% 감소했습니다.
  • **교차 모달리티 강제(Cross-Modality Forcing)**를 제거하면 성공률이 21% 감소하여, 누락된 모달리티를 합성하도록 강제하는 것이 강력한 공유 표현을 학습하는 데 결정적임을 확인했습니다.

의의 및 주장

본 논문은 FLEX-π가 사실상 공짜로(essentially for free) 기하학적 및 의미론적 그 grounding을 얻을 수 있음을 입증한다고 주장합니다. 사전 학습된 비디오 모델의 잠재 공간을 활용하고 교차 모달리티 강제 기법을 적용함으로써, 모델은 다음을 달성했습니다:

• 더 나은 정책 성능: 복잡하고 접촉이 많으며 장기적인(long-horizon) 작업에서 훨씬 높은 성공률을 기록했습니다.
• 하드웨어 오버헤드 없음: RGB 이상의 추가 센서가 필요하지 않습니다.
• 사전 학습 오버헤드 없음: 새로운 시각적 사전 지식을 학습할 필요가 없습니다.
• 배포 유연성: 동일한 체크포인트를 사용하여 빠른 액션 전용 모드 또는 고성능 전체 결합 모드로 실행할 수 있습니다.

저자들은 FLEX-π가 매우 효과적이지만, 더 단순한 모델들보다 수렴하는 데 더 많은 학습 에포크(epochs)가 필요하며, 전체 결합 모드가 파라미터 규모가 비슷한 VLA보다 느리다는 점을 인정합니다. 이들은 향후 연구가 더 강력한 VLM 백본과의 결래 및 더 큰 로봇 데이터셋에 대한 접근을 통해 개선될 수 있다고 제안합니다.