OSWAM이 상상 토큰으로 행동 확산을 한 스텝으로 단축한다

무엇인가

자율주행과 구현체 제어(embodied control)에서 월드 액션 모델(WAM)은 미래 영상 예측을 중간 변수로 두고 행동을 뽑는다. 문제는 추론 비용이다. 미래를 상상한 뒤 실행하는 방식은 테스트 시점에 비디오 롤아웃을 돌려 지연이 크고, 이를 없앤 direct-policy WAM(Fast-WAM, SimWAM 등)조차 행동 생성에 flow matching 기반 디노이징을 수십 스텝 돌려야 경쟁력 있는 품질이 나온다. 논문은 이 지점을 정면으로 겨냥한다. 스텝을 줄이면 비선형 다단계 궤적이 뭉개져 성능이 급락하는데, 저자들이 제시한 수치는 1스텝 설정에서 NAVSIM -20%, RoboTwin -40%다.

어떻게 동작하나

저자들의 진단은 명확하다. 문제는 랜덤 가우시안 노이즈에서 출발한다는 점, 즉 초기화에 행동 구조(action prior)가 없다는 것이다. 기존 few-step 기법인 스텝 증류나 궤적 정류(rectified flow)도 여전히 가우시안 노이즈에서 시작하므로 이 결핍을 메우지 못한다고 본다. 그래서 제안하는 것이 상상 토큰(imagination tokens)이다. 비디오 DiT(Wan2.2-5B 백본)에 행동 상상 토큰(AIT)과 장면 상상 토큰(SIT) 두 묶음을 노이즈로 초기화해 넣고, AIT가 1프레임 visual 토큰·SIT·다른 AIT와 상호작용하며 물리적 맥락을 축적하게 한다. 토큰 수가 1프레임 visual 토큰보다 훨씬 적고 한 번의 forward pass로 얻어지므로 오버헤드가 작다.

무엇과 다른가

절차는 이렇다. AIT는 행동 노이즈 ε를 인코딩해 고정 타임스텝 t_max=1000으로 두고, 보조 디코더가 이를 행동 속도 v̂_A로 매핑해 구조화된 사전 ā = ε − v̂_A를 만든다. 이 사전은 그래디언트를 끊은 뒤(sg) 같은 ε로 다시 노이즈를 입혀 중간 상태 a_init = (1−σ_r)·sg(ā) + σ_r·ε (τ_r=500, σ_r=0.5)를 만든다. Action DiT는 여기서 정제 속도 v̂_act를 예측해 â = a_init − σ_r·v̂_act 한 번으로 행동 청크를 낸다. 학습 목표는 (a_init − a_gt)/σ_r로, σ_r 크기의 오일러 업데이트 한 번이 정답 행동으로 정확히 이동하도록 스케일링한다. 여기에 SIT가 장면 맥락 슬롯 역할을 하고, Role-aware 3D RoPE가 비디오 토큰·AIT·SIT에 서로 다른 공간 앵커(비디오 그리드 밖, 그리드 중심)와 공유 시간 좌표 t_n = n·s_a/(s_v·d_t)를 부여해 역할을 구분한다.

어떻게 쓰나

두 번째 축은 특권 정보 증류(privileged-information distillation)다. direct-policy WAM은 1프레임만 보므로 장기 맥락이 표현에 눌러 담긴다. 교사 Video DiT는 전체 미래 클립의 visual 토큰을 보게 하고(미래 latent는 고정 타임스텝 τ_v^T=200으로 노이즈 처리), 학생인 OSWAM은 1프레임만 보게 둔 채 AIT 특징을 블록 [8,16,24,30]에서 정렬한다(λ_dist=10). 추론 시 교사는 버려지므로 1프레임 전용 인터페이스가 유지된다. 전체 손실은 비디오 flow matching + AIT 속도 손실 + 증류 손실 + 행동 손실의 합이다.

전제와 한계

결과는 NAVSIM에서 PDMS 90.3, 지연 136ms다. SimWAM과 PDMS를 맞추면서 디노이징을 10스텝에서 1스텝으로 줄여 3.01배 빠르고, Metis보다 PDMS 1.2 높으면서 3.24배, DriveVA보다 4.36배 빠르고 PDMS는 0.2 낮은 수준이다. RoboTwin 2.0(50개 과제)에서는 평균 성공률 90.00%, 지연 142ms로 Fast-WAM 대비 2.80배 빠르면서 성공률 하락은 1.83%p다. 1스텝 Flash-WAM 변형 대비 평균 성공률 +8.59%에 4.29배 빠르고, Light-WAM은 더 빠르지만 OSWAM이 35ms 추가로 성공률을 크게 앞선다. 캐시된 언어 컨텍스트 기준으로 RoboTwin 행동 GPU 시간이 67.61ms에서 6.78ms로 줄고 상상 토큰이 visual 시간을 5.48ms 늘려 전체 3.39배(323.67→95.35ms), 피크 메모리는 거의 동일하다. torch.compile과 CUDA Graph를 적용하면 RoboTwin 40.81ms, NAVSIM 52.64ms로 10Hz 이상 추론이 가능하다.

절제 실험도 설계의 근거를 뒷받침한다. 1스텝 Fast-WAM 베이스라인에 AIT만 붙여 NAVSIM PDMS가 71.04에서 89.81로, RoboTwin 성공률이 57.38%→87.42%(Clean), 55.80%→88.16%(Randomized)로 뛴다. 증류를 더하면 89.96 / 89.42% / 89.62%가 되고, SIT까지 넣은 구성이 최고 성능이다. 증류 가중치는 0.01이면 베이스라인보다도 slightly 못하고 λ=10에서 PDMS 90.28, ADE 0.587로 최적이다. 초기화 방식은 rectified-flow 정제보다 재노이징이 전부 우세했고 τ_r은 500이 최적(400은 부족, 600은 열화, 400~600 균일 샘플링도 열등)이었다. 교사 비디오 타임스텝을 0~800으로 흔들어도 PDMS는 90.16~90.28로 거의 변하지 않았고, 미래 프레임이 순수 노이즈인 800에서도 90.16이 나왔다.

실무 관점에서 이 논문의 쓸모는 명확하다. 로봇 팔·차량처럼 폐루프 제어 주기가 짧은 곳에서 월드 모델의 비디오 사전학습 이점은 유지하면서 디노이징 예산을 1스텝으로 깎는 패턴은 그대로 이식 가능하다. 다만 확인할 것도 있다. 첫째, 이득의 대부분은 AIT 초기화에서 나오므로 자기 도메인에서 그 사전이 얼마나 잘 학습되는지가 관건이다. 둘째, 이득은 백본(Wan2.2-5B)과 LoRA rank 8, 고정 하이퍼파라미터(t_max=1000, τ_r=500, λ_dist=10)에 의존하므로 그대로 옮길 때 재튜닝이 필요하다. 셋째, 논문 수치는 H20 GPU와 컴파일 추론 기준이며 실차·실기 배포에서의 체감 지연은 별도 검증 대상이다. 실제 PiPER 로봇 팔 실험(빵 85% vs Fast-WAM 80%, 지우개 80% vs 75%, 각 50개 시연·20회 시행)에서 성공 에피소드의 작업 완료 시간이 34.2%, 21.8% 줄었다는 점은 시뮬레이션 밖 근거로 참고할 만하다.

저자들이 밝힌 한계는 범위에 관한 것이다. 이 연구는 Action DiT를 위한 구조화된 초기화 학습에 집중했고, SIT와 특권 정보 증류가 Action DiT 자체를 직접 향상시킬 잠재력은 탐구하지 않았다. 장면 맥락 조건화와 특권 지도를 행동 전문가(action expert)까지 확장하는 것이 향후 과제로 남아 있다.