RealtimeWAM, 세계 행동 모델을 한 단계 비동기 추론으로 전환

개요

세계 행동 모델(World Action Models, WAMs)은 비디오 생성 백본이 학습한 시각 표현을 로봇 행동 예측에 활용한다. 현재 관측만 사용하는 정책과 달리, 환경의 변화와 가능한 미래 움직임에 관한 시각적 정보를 행동 생성에 반영할 수 있다는 장점이 있다. 그러나 비디오 전문가와 행동 전문가의 계산량이 크고, 행동 예측에 여러 번의 디노이징이 필요하면 실시간 제어에 부담이 된다.

논문 ‘RealtimeWAM: One-Step Asynchronous World Action Models’는 이 문제를 모델 내부의 반복 계산과 전문가 사이의 실행 순서라는 두 측면에서 다룬다. 제안 방식은 행동을 한 단계로 생성하면서, 비디오와 행동 전문가가 가능한 구간에서 동시에 실행되도록 추론 구조를 재배치한다.

핵심 기술

• TACD로 행동 생성을 한 단계로 축소. 확산 기반 행동 모델은 일반적으로 노이즈가 섞인 행동을 여러 단계에 걸쳐 정제한다. Teacher-Anchored Consistency Distillation(TACD)은 여러 단계를 수행하는 동결 교사 모델의 롤아웃 최종 지점을 학생 모델에 직접 학습시켜 한 단계 생성을 가능하게 한다.
• 국소 일관성과 최종 결과의 차이 해결. 인접한 추론 상태 사이의 일관성만 낮추면 국소 오차는 작아질 수 있지만 최종 행동이 정확하다는 보장은 없다. TACD는 국소 일관성 손실에 교사의 다단계 롤아웃 종점에 대한 명시적 감독을 추가해 전체적인 목표 결과를 보존하려 한다.
• CEWP로 전문가 대기 제거. 기존 파이프라인에서는 비디오 전문가가 시각 표현을 모두 계산한 뒤에야 행동 전문가가 작업을 시작할 수 있다. Cross-Expert Wavefront Pipelining(CEWP)은 비디오 전문가의 KV 캐시를 블록 단위로 공유하고, 해당 블록이 행동 어텐션에 필요해지는 시점 직전에만 동기화한다. 이 방식으로 두 전문가의 계산을 겹쳐 실행한다.
• 다양한 로봇 벤치마크 평가. 실험은 LIBERO, LIBERO-Plus, RoboTwin과 Fast-WAM, Faster-WAM 등 여러 모델 변형을 대상으로 진행됐다. 논문 요약에 따르면 RealtimeWAM은 정확도 저하를 1% 미만으로 유지하면서 H100에서 최대 25배의 엔드투엔드 속도 향상을 보였다.

의미와 한계

RealtimeWAM의 핵심은 모델을 단순히 축소하는 데 있지 않다. TACD는 행동 전문가 내부의 반복 추론을 줄이고, CEWP는 비디오 전문가가 끝날 때까지 행동 전문가가 기다리는 시간을 줄인다. 즉, 실시간성을 높이기 위해 개별 연산의 최적화와 전체 추론 그래프의 스케줄링을 함께 다룬다.

비디오 기반 시각 사전지식의 이점을 유지하면서 지연을 줄일 수 있다면, 세계 모델을 활용한 폐루프 로봇 제어의 실용성이 높아질 수 있다. 다만 제공된 자료만으로는 서로 다른 로봇 형태, 센서 지연, 장기 작업에서의 안정성까지 판단하기 어렵다. 공개된 코드와 체크포인트를 이용한 재현 실험이 방법의 범용성을 검증하는 다음 단계가 될 것이다.

출처: Hugging Face Daily Papers