기술 요약: 세계 모델 강화 비전 - 언어 - 행동 모델을 위한 듀얼 스트림 확산

1. 문제 제기

비전 - 언어 - 행동 (VLA) 모델은 인터넷 규모의 비전 - 언어 모델 (VLM) 의 지각적 근거를 활용하여 범용 로봇 정책 학습을 위한 강력한 패러다임으로 부상했습니다. 그러나 표준 VLA 는 종종 근본적인 물리적 과정에 대한 명시적인 이해가 부족하여, 자신의 행동이 환경을 어떻게 변화시키는지 모델링하는 데 어려움을 겪습니다.

최근 세계 모델링 목표 (미래 상태와 행동을 공동 예측) 를 통해 VLA 를 강화하여 이러한 격차를 해소하려는 시도들은 상당한 아키텍처적 도전에 직면해 있습니다:

• 통합 결합 확산 (Unified Joint Diffusion): 행동과 비전 토큰을 단일 잠재 공간에 연결하는 접근 방식은 종종 모달리티 간 격차로 인해 실패합니다. 이러한 접근 방식은 저차원이고 시간적으로 부드러운 행동 궤적의 예측과 고차원이며 공간적으로 복잡한 시각 데이터의 예측을 조화시키는 데 어려움을 겪습니다.
• 인과/단방향 설계: 모달리티를 별도의 모델로 분리하는 접근 방식 (예: 세계 모델 출력에 조건을 부여하여 행동 생성) 은 모달리티별 학습을 용이하게 하지만, 본질적으로 양방향 지식 전달을 방해하여 모델이 교차 모달 인과 관계를 학습하는 능력을 제한합니다.

핵심적인 과제는 교차 모달 통합과 인과 추론을 효과적으로 가능하게 하면서도 모달리티별 충실도를 유지하는 프레임워크를 설계하는 것입니다.

2. 방법론: 듀얼 스트림 확산 (DUST)

저자들은 교차 모달 통합과 모달리티별 충실도 사이의 트레이드오프를 해결하도록 설계된 세계 모델 강화 VLA 프레임워크인 DUST(Dual-Stream Diffusion) 를 제안합니다.

아키텍처

멀티모달 확산 트랜스포머 (MMDiT) 를 기반으로 구축된 DUST 는 행동과 관찰 (미래 시각 임베딩) 을 위한 별도의 토큰 스트림을 유지하면서 교차 모달 교환을 용이하게 합니다:

• 듀얼 스트림: 행동 및 비전 토큰은 MMDiT 블록 내에서 별도의 경로를 통해 전파됩니다.
• 교차 모달 어텐션: 스트림은 공유되는 교차 모달 어텐션 레이어 동안에만 상호작용하여, 통합된 잠재 표현을 강제하지 않고 정보 교환을 가능하게 합니다.
• 모달리티별 처리: 공유 레이어 이후, 스트림은 전문화된 디노이징을 위해 모달리티별 DiT 블록으로 라우팅됩니다. 비전 경로는 의미적으로 일관된 미래 임베딩을 재구성하는 데 중점을 두고, 행동 경로는 저수준의 모터 제어를 정제합니다.
• 조건부 (Conditioning): 모델은 현재 관찰과 작업 지시문에 기반하여 고정된 VLM 백본 (Eagle-2) 에서 추출된 의미 특징에 조건을 부여받습니다.

학습 알고리즘: 분해된 흐름 매칭

DUST 는 양방향 인과 의존성을 학습하기 위해 분해된 확산 학습 알고리즘을 사용합니다:

• 독립적 노이즈 스케줄링: 모든 토큰이 노이즈 스케줄을 공유하는 표준 확산과 달리, DUST 는 각 모달리티에 대해 독립적인 노이즈 교란 ($\tau_A$는 행동용, $\tau_o$는 비전용) 을 적용합니다.
• 양방향 인과 학습: 모델이 다양한 노이즈 구성 (예: 완전한 노이즈가 적용된 행동과 함께 깨끗한 미래 상태, 또는 그 반대) 에 노출됨으로써, 모델은 순방향 동역학 ("이 행동의 결과는 무엇인가?") 과 역방향 동역학 ("이 상태에 이르게 하는 행동은 무엇인가?") 을 모두 학습합니다.
• 손실 함수: 모델은 모달리티별 흐름 매칭 손실로 구성된 공동 목적 함수를 통해 최적화됩니다:
  $$L_{Joint}(\theta) = L_A(\theta) + \lambda_{WM} L_{WM}(\theta)$$
  여기서 $L_A$는 행동 손실이고 $L_{WM}$은 세계 모델링 손실 (고주파 노이즈를 피하기 위해 원시 픽셀 대신 미래 관찰 임베딩을 예측) 입니다.

추론: 비동기 결합 샘플링

모달리티 간의 계산적 비대칭성 (더 많은 단계가 필요한 고차원 비전 vs. 더 빠르게 수렴하는 저차원 행동) 을 해결하기 위해 DUST 는 비동기 샘플링 전략을 도입합니다:

• 테스트 시간 확장: 비전 토큰은 모든 확산 단계에서 업데이트되는 반면, 행동 토큰은 덜 빈번하게 (매 $q$단계마다) 업데이트됩니다.
• 메커니즘: 이는 행동에 대한 과도한 디노이징 단계로 인한 성능 저하나 지연을 초래하지 않으면서 시각 표현의 더 세밀한 정제를 가능하게 합니다.

3. 주요 기여

1. 듀얼 스트림 아키텍처: 행동과 관찰을 위한 별도의 스트림을 유지하여 통합된 잠재 공간을 강제하지 않으면서 교차 모달 지식 전달을 가능하게 하는 새로운 MMDiT 기반 설계.
2. 분해된 학습: 각 모달리티에 독립적인 노이즈 스케줄을 활용하여 모델이 행동과 미래 상태 간의 강력한 양방향 인과 관계를 학습하도록 강제하는 학습 알고리즘.
3. 비동기 샘플링: 비전과 행동 토큰을 서로 다른 속도로 진화시킴으로써 추론 효율성과 정확도 사이의 균형을 맞추는 테스트 시간 확장 방법.
4. 전이 학습 능력: 로봇 - 인간 혼합 데이터의 공동 학습 및 행동이 없는 비디오 데이터셋에서의 사전 학습을 포함하여 이질적인 데이터 소스를 효과적으로 활용할 수 있음을 입증.

4. 실험 결과

저자들은 DUST 를 시뮬레이션 벤치마크 (RoboCasa, GR-1), 실제 세계 작업 (Franka Research 3), 그리고 전이 학습 시나리오에서 평가했습니다.

• 시뮬레이션 벤치마크:

  • RoboCasa: DUST 는 다양한 작업 카테고리와 데모 규모 (100~1000 개 데모) 에서 GR00T-N1.5 및 FLARE 를 포함한 최첨단 베이스라인보다 최대 6% 높은 성공률을 달성했습니다.
  • GR-1: DUST 는 픽 앤 플레이 및 관절형 작업 모두에서 베이스라인보다 5~6% 더 높은 성능을 보였습니다.
  • 테스트 시간 확장: 행동 단계를 고정하고 비전 디노이징 단계를 늘리는 것은 성공률에서 추가적인 2~5% 개선을 제공했습니다.

• 실제 세계 성능:

  • 7 가지 다양한 작업 (픽 앤 플레이, 삽입, 도구 사용) 에 걸쳐 Franka Research 3 로봇에서 DUST 는 베이스라인보다 10% 높은 성공률로 우수한 물리적 근거를 입증했습니다.

• 전이 학습:

  • 공동 학습: RoboCasa, GR-1, EgoDex(인간) 데이터의 혼합으로 학습했을 때, DUST 는 베이스라인보다 상당한 개선을 보이며 이질적인 로봇 및 인간 데이터셋 간의 격차를 효과적으로 해소했습니다.
  • 사전 학습: 행동이 없는 비디오 데이터 (BridgeV2) 로 사전 학습한 후 로봇 데이터로 미세 조정하는 것은 사전 학습 없이 학습하는 경우보다 성공률을 8.4% 향상시켰으며, 정책 학습을 위한 수동 비디오 데이터의 유용성을 강조했습니다.

• 애블레이션 연구:

  • 듀얼 스트림 구조나 분해된 노이즈 스케줄링을 제거하면 각각 8% 와 12% 의 성능 저하가 발생하여 두 구성 요소의 필요성을 확인했습니다.
  • 세계 모델링 타겟으로 시맨틱 임베딩 (SigLIP-2) 을 사용하는 것이 픽셀 수준 재구성 (Flux.1 VAE) 보다 우월함이 입증되어, 고주파 시각 세부 사항이 종종 제어와 무관함을 시사했습니다.

5. 중요성과 주장

이 논문은 DUST 가 교차 모달 교환을 가능하게 하면서도 별개의 모달리티 스트림을 유지함으로써 행동과 관찰의 공동 모델링에 내재된 갈등을 성공적으로 해결했다고 주장합니다. 노이즈 스케줄을 분해함으로써 모델은 표면적 상관관계가 아닌 의미 있는 인과 의존성을 학습합니다.

저자들은 DUST 가 로봇 및 인간을 포함한 이질적인 데이터로 VLA 모델을 확장하기 위한 통합 프레임워크를 제공하며, 로봇 데모 수집의 높은 비용을 해결하기 위해 행동이 없는 비디오를 사전 학습에 활용한다고 강조합니다. 비동기 샘플링 전략은 또한 테스트 시간 확장을 위한 실용적인 메커니즘을 제공하여, 작업 요구 사항에 따라 추론 속도와 예측 정확도 사이의 트레이드오프를 가능하게 합니다. 이러한 결과는 모달리티별 충실도를 고려하여 설계된 세계 모델 강화 VLA 가 시뮬레이션과 실제 세계 환경 모두에서 로봇 정책 학습을 크게 향상시킬 수 있음을 시사합니다.