훈련전용 비디오 기반 World‑Action 모델 SimWAM

왜 중요한가

자율주행 플래너가 장래 장면을 실시간으로 합성하면 추론 지연이 크게 늘어나므로 실용성에 제약이 발생한다. SimWAM은 대규모 비디오 생성기를 훈련 신호로만 활용하고 추론 시에는 경량의 액션 전문가만 남기는 구조를 사용하여 학습 단계에서 운동학·장면 동역학 정보를 전이하면서도 배포 시 높은 처리량을 유지한다. NAVSIM에서 91.5 PDMS를 달성하고 nuScenes로의 제로샷 전이에서 충돌률 0.04%를 달성한 점은 훈련 시 세계 모델을 사용하되 추론 비용을 제거하는 방법이 실제 주행 품질과 안전성 면에서 유의미하다는 근거를 제공한다.

핵심 기여

훈련 시 비디오-액션 공동 플로우 매칭

기존 WAM과 달리 SimWAM은 비디오 생성기를 추론 루프에서 배제하고 훈련 신호로만 사용하기 위해 영상(latent)과 궤적을 같은 플로우 매칭 손실로 공동 학습한다. 현재 관찰을 인코딩한 표현 z(o_t)를 공유 어텐션으로 전달하고, 손실은 $ \mathcal{L}=\mathcal{L}_{\text{FM}}^{\text{act}}+\lambda\,\mathcal{L}_{\text{FM}}^{\text{vid}} $ 형태로 결합하여 액션 전문가가 영상 기반 동작 사전을 내재화하도록 유도한다. 이 방식은 훈련 단계에서 장면 동역학을 액션 표현에 주입하면서도 추론 시에는 비디오 분기를 제거할 수 있게 해 효율성과 성능을 동시에 확보한다.

격리 어텐션으로 추론-훈련 분리 보장

공유 어텐션 스트림에 현재 프레임 토큰, 미래 프레임 토큰, 액션 토큰을 배치하되 미래 토큰과 액션 토큰 간의 상호 참조를 차단하는 격리 마스크를 도입했다. 이 구조는 학습 시 미래 영상으로부터 관찰 표현이 강화되도록 허용하면서도 액션 분기는 미래 영상에 직접 의존하지 않도록 보장하므로 훈련 후 비디오 DiT와 VAE를 제거해 추론 지연을 크게 줄일 수 있다. 실험에서는 격리 마스크가 다른 어텐션 패턴보다 PDMS와 안전성 지표를 더 잘 유지하면서도 추론 시 비디오 합성을 불필요하게 만든다는 근거가 관찰되었다.

SDE 기반 강화학습과 LoRA 적응

확정적 흐름 ODE는 탐색성과 전이 확률이 부족하여 정책 그라디언트 적용이 어렵기 때문에 SimWAM은 ODE를 주변 분포를 보존하는 SDE로 변환하여 후보 궤적 집단을 샘플링하고 Flow-GRPO 스타일의 그룹-상대 장점으로 LoRA만 업데이트하는 강화학습을 수행했다. LoRA rank=32, α=16로 액션 전문가의 어텐션 투사에만 저비용 적응을 적용하여 훈련 비용과 추론 간단성을 모두 보존했다. 이 단계는 모방학습 결과(90.3 PDMS)에서 보상 최적화를 통해 91.5 PDMS로 향상시키는 정량적 이득을 가져왔다.

핵심 아이디어 이해하기

시퀀스 수준의 미래 장면 합성은 행동 예측에 유용한 동역학 정보를 제공하지만 실제 배포에서는 비디오 생성 비용이 병목이 된다. 따라서 현재 관찰로부터 얻은 표현 z(o_t)에 훈련 시 비디오 생성을 통한 동적 사전을 내재화하고, 액션 토큰이 미래 영상 토큰을 참조하지 못하도록 격리된 어텐션을 적용하면 학습된 동작 prior는 유지되면서 추론 비용을 크게 낮출 수 있다. 플로우 매칭으로 비디오와 궤적을 공동 학습하고 ODE를 주변 보존 SDE로 바꿔 RL 최적화를 적용하면 탐색성과 보상 기반 개선이 가능해져 안전성과 진행 거리를 동시에 개선할 수 있다.

방법론

문제 정의는 전방 카메라 관찰 o_t, 자차 상태 s_t, 내비게이션 명령 l로부터 향후 H 스텝의 ego 궤적 a_{t+1:t+H}을 예측하는 것이다. 영상 전문가는 Wan2.2-5B 기반의 Video Diffusion Transformer와 VAE, T5 텍스트 인코더를 사용해 현재 프레임을 조건으로 미래 프레임 잠재 z_{t+1:t+N}을 플로우 매칭으로 학습하고, 액션 전문가는 hidden size d_a=1024의 Diffusion Transformer로 조건 c={z(o_t),s_t,l}에서 궤적 속도장 v_{θ_a}를 예측한다. 두 전문가의 상호작용은 공유 어텐션 스트림으로 제한되며 손실은 $ \mathcal{L}=\mathcal{L}{\text{FM}}^{\text{act}}+\lambda,\mathcal{L}{\text{FM}}^{\text{vid}} $로 결합된다. 강화학습 단계에서는 ODE를 SDE(식(2)와 유사)로 변환해 G=8 후보 궤적을 샘플링하고 NAVSIM PDM 보상으로 그룹 상대 장점을 계산해 클리핑된 정책 업데이트로 LoRA(rank=32, α=16)만 갱신한다. 입력 해상도는 기본 384×672, 예측은 4초 동안 8개 웨이포인트(2Hz)이며 모방학습은 AdamW lr=1e-4로 100 에폭, RL lr=5e-5로 하드 서브셋에서 수행되었다.

주요 결과

NAVSIM navtest에서 단일 전방 카메라만 사용해 SimWAM은 91.5 PDMS를 달성했으며 이는 동종 WAM 기반 플래너들보다 우수한 성능을 보였다. 모방학습 단계만으로는 90.3 PDMS였고 RL 적용으로 91.5 PDMS로 추가 향상되었다. 구성 분석에서는 액션 전용 86.6 → 영상 공동학습 90.3 → RL 91.5로 단계별 기여가 뚜렷하게 나타났고, 격리 어텐션 마스크가 다른 어텐션 패턴보다 안전성(NC, TTC)과 PDMS 균형에서 우위를 보였다. 비디오 백본 유연성 실험에서는 Wan2.2-5B으로 90.3 PDMS, Cosmos2.5으로 90.4 PDMS가 보고되어 훈련 시 사용하는 비디오 모델의 품질이 영향을 미치지만 아키텍처 자체는 백본 교체를 허용한다. 제로샷 nuScenes 평가에서는 평균 충돌률 0.04%와 평균 L2 오차 0.96m로 안전성 측면에서 강한 일반화를 보였다.

기술 상세

플로우 매칭 손실은 보간된 x_τ에서 속도장 v_θ(x_τ,τ,c)를 예측하도록 평균 제곱오차를 최소화하는 방식으로 구성되며, 이 손실을 영상과 궤적 두 모달리티에 대해 공동 적용해 관찰 표현을 정교화한다. 원래의 확정적 ODE는 탐색이 제한적이므로 주변 보존 SDE를 도입해 Euler–Maruyama 전이 확률을 얻고 그룹 샘플링(G=8)으로 후보 궤적을 평가해 Flow-GRPO 계열의 강화학습을 수행한다. 강화학습 시에는 전체 가중치를 갱신하지 않고 LoRA(rank=32, α=16) 어댑터만 업데이트해 학습 안정성과 추론 효율을 보장한다. 샘플링 단계는 5스텝이면 대부분 성능을 회복하고 10스텝에서 최고 PDMS(90.3)를 관찰했으며 20스텝은 지연만 증가시켜 실무에서는 10스텝 이하를 권장한다.

한계점

SimWAM은 훈련 단계에서 대형 비디오 생성기와의 공동학습이 요구되므로 학습 비용이 높고 적절한 비디오 백본의 품질이 결과에 영향을 미친다. 연구에서 사용한 설정은 전방 카메라 단일 입력에 최적화되어 있어 멀티센서(라이다·다중뷰) 정보가 필요한 시나리오에 대해서는 추가 검증이 필요하다. 또한 강화학습은 하드 서브셋에서 유의미한 개선을 보였으나 15k 스텝 이후 수렴성과 과적합 위험이 관찰되어 RL 단계의 데이터 선택과 조기중단 등 실무적 제어가 요구된다.

실무 활용

SimWAM은 훈련 시 대형 비디오 생성기의 동작 사전을 흡수하고 추론 시에는 경량 액션 전문가만 배포하는 설계로 실시간 주행 제어에 적합하다. 이 접근은 고성능 비디오 모델의 학습 이점을 유지하면서도 추론 지연과 계산 비용을 낮춘다. 공개된 코드와 가중치를 통해 연구·개발 환경에서 빠르게 재현하고 기존 플래너에 적용할 수 있다.

제한된 온보드 계산 리소스에서 실시간 궤적 예측이 요구되는 자율주행 차량의 플래너로 배포
대규모 비디오 생성 모델로 얻은 동적 사전을 다른 경량 플래너로 전이하는 학습 파이프라인 실험
WAM 연구에서 훈련-전용 비디오 신호가 추론 효율에 미치는 영향을 비교하는 기준선

코드 공개 여부: 공개

코드 저장소 보기

키워드

SimWAMflow matchingDiffusion Transformerisolated attention maskSDELoRAPDMSNAVSIMnuScenes

용어 해설

플로우 매칭(Flow Matching): — 데이터 샘플과 가우시안 노이즈의 선형 보간 경로에서 네트워크가 일정한 속도(노이즈−데이터)를 예측하도록 학습시키는 손실 구성이다. 입력으로 보간된 x_τ와 τ, 조건 c를 받아 속도장 v_θ(x_τ,τ,c)를 예측하고 평균 제곱오차로 학습한다. 이 구성은 ODE 통합으로 노이즈에서 데이터로 샘플링하거나 SDE로 변환해 후보 궤적을 샘플링하는 기반이 된다.
격리 어텐션 마스크(Isolated Attention Mask): — 관찰 토큰, 미래 영상 토큰, 액션 토큰을 공유된 어텐션 스트림에 배치하되 미래 영상 토큰과 액션 토큰이 서로를 보지 못하게 만드는 어텐션 패턴이다. 두 토큰 군은 모두 현재 관찰을 참조할 수 있으나 상호 참조가 차단되어 학습 시에만 영상 정보가 액션 표현에 스며들고 추론 시 영상 분기가 제거될 수 있다. 이 구조가 SimWAM의 핵심으로 추론 효율성과 훈련용 동작 사전 전달을 동시에 가능하게 한다.
마진 보존 SDE(Marginal-preserving SDE): — 원래의 확정적 흐름 ODE와 동일한 주변 분포 p_τ(x_τ)를 보존하도록 설계된 확률 미분방정식이다. 드리프트 항에 추가적인 항과 시간 의존 노이즈 σ_τ를 도입하여 Euler–Maruyama 단계마다 가우시안 전이 확률을 얻을 수 있으므로 정책 그래디언트 등 RL 최적화에 필요한 전이 확률을 제공한다. SimWAM은 후보 궤적 탐색과 Flow-GRPO 기반 강화 학습에서 이 SDE를 사용한다.