기술 요약: Discrete-WAM

문제 정의

자율 주행은 단순히 순간적인 관측에 반응하는 것을 넘어, 자차(ego)의 행동이 주변 세계의 진화에 어떻게 영향을 미치는지에 대한 추론을 근본적으로 요구한다. 현재의 엔드투엔드(E2E) 시스템은 일반적으로 행동 복제(behavior cloning)를 통한 상태-행동 직접 매핑에 의존하며, 이는 명시적인 행동 조건부 역학(action-conditioned dynamics)을 모델링하기보다는 통계적 상관관계를 포착하는 데 그친다. 연속적 잠재 세계 모델(continuous-latent world models)은 강력한 연속성을 제공하지만, 반사실적 미래(counterfactual futures)에 대한 인과적 추론에 필요한 구성적 구조가 부족한 경우가 많다. 또한, 기존 방식들은 세계 모델링, 정책 학습, 의사 결정(decision-making)을 별개의 작업이나 보조 목표로 취급하여, 통합된 생성 과정 내에서 관측, 행동, 그리고 유도된 미래 상태 사이의 일관된 인과적 대응 관계를 구축하지 못한다.

방법론: Discrete-WAM

저자들은 미래의 시각적 상태와 자차 행동을 정렬된 이산 토큰(discrete tokens)으로 표현하는 통합 잠재 비전-액션 세계 정책인 Discrete-WAM을 제안한다. 이 프레임워크는 공유된 이산 확산 아키텍처를 통해 대안적인 미래에 대한 구성적 인과 추론을 가능하게 한다.

1. 통합 이산 표현 (Unified Discrete Representation)

Discrete-WAM은 연속적인 관측과 행동을 공유된 이산 토큰 공간으로 변환한다:

• 비전 토큰화 (Vision Tokenization): 연속적인 카메라 관측값은 사전 학습된 VQ-VAE 토크나이저를 사용하여 이산 세만틱 토큰으로 인코딩된다.
• 액션 토크나이징 (Action Tokenization): 연속적인 미래 움직임은 가속도 어휘집(acceleration vocabulary) 상의 이산 토큰 호환 가능한 표현으로 변환된다. 저자들은 결정론적인 하드 양자화 오류를 피하기 위해 2D 가속도 그리드 상의 소프트 라벨 보간 전략을 채택하였으며, 이를 통해 연속적인 제어 값이 인접한 프로토타입들의 가중 합을 통해 보존될 수 있도록 한다.
• 의사 결정 토큰 (Decision Tokens): 고수준 의사 결정 조건(예: 기동 의도, 목표 차선)은 희소하고 저주파수인 의사 결정 토큰($D_t$)으로 표현된다.

2. 계층적 정책 모델링 (Hierarchical Policy Modeling)

Dense한 미래 행동을 장면 문맥으로부터 직접 생성하는 대신, Discrete-WAM은 정책 생성을 계층적 과정으로 분해한다:

• 고수준 (High-Level): 모델은 거친 기동 의도와 다중 모드(multi-modal) 주행 선택을 포착하는 희소한 의사 결정 스켈레톤($D_t$)을 예측한다.
• 저수준 (Low-Level): 이산 확산 기반 액션 토큰 플래너는 예측된 의사 결정에 조건화되어 매끄럽고 시간적으로 일관된 궤적을 생성한다.
  이러한 분리는 모델이 고수준에서 다중 모드 의사 결정을 선택하는 동시에 저수준에서 매끄러운 궤적을 실현할 수 있게 하여, 대안적인 주행 모드의 붕괴를 방지한다.

3. 통합 생성 프레임워크 (Unified Generative Framework)

Discrete-WAM은 자율 주행을 공유된 이산 토큰 공간에서의 시퀀스 모델링으로 정식화하며, 단일 아키텍처 하에서 세 가지 상호 보완적인 학습 모드를 지원한다:

• 세계 모델링 (World Modeling): 장면 문맥($C_t$)과 미래 액션 토큰($A_{t+1:t+H}$)을 조건으로 미래 시각 토큰($V_{t+1:t+H}$)을 예측한다.
• 세계-정책 모델링 (World-Policy Modeling): 인터리브드(interleaved) 토큰 시퀀스 내에서 비전과 액션 예측을 공동 학습하여, 정책 생성과 세계의 진화 사이의 결합을 촉진한다.
• 정책 모델링 (Policy Modeling): 의사 결정($D_t$)에 이어지는 액션 시퀀스($A_{t+1:t+H}$)를 예측한다.

4. 학습 및 사후 학습 (Training and Post-Training)

• 통합 사전 학습 (Unified Pretraining): 모델은 토큰 편집(token-editing) 목적 함수를 사용하여 다단계 사전 학습을 거친다. 모델은 역사적 문맥과 조건부 입력을 바탕으로 손상된 토큰(시각 또는 액션)을 복구하는 법을 배운다. 손실 함수는 토큰 분류를 위한 교차 엔트로피와 물리적 일관성을 보장하기 위한 연속 운동 재구성 손실(가속도, 궤적, 위치)을 결합한다.
• 사후 학습 (RL): 희귀하거나 안전이 중요한 행동을 다루기 위해, 저자들은 그룹 상대적 정책 최적화(GRPO)를 적용한다. 모델은 후보 궤적을 샘포링하고, 온라인 보상 함수(안전 및 규칙 임계값 포함)로 이를 평가하며, 의사 결정-궤적 조합을 정교화하기 위해 정책을 업데이트한다.

5. 추론 및 편집 (Inference and Editing)

추론 과정에서 Discrete-WAM은 미래 상태에 대한 반복적인 계획 생성 및 편집을 활용한다. replace_confidence와 같은 스케줄링 전략은 예측 신뢰도와 엔트로피에 따라 토큰을 선택적으로 업데이트하며, 이를 통해 단일 순방향 롤아웃(forward rollout)에 구애받지 않고 액션 조건부 미래에 대한 비인과적이고 전역적으로 일관된 추론을 수행한다.

주요 기여 (Key Contributions)

1. 통합 비전-액션 세계 정책: 관측과 자차 행동을 공유된 이산 잠재 공간에서 표현함으로써, 정렬된 세만틱 모델링과 액션 조건부 미래에 대한 구조적 추론을 가능하게 하는 프레임워크를 도입하였다.
2. 이산 확산 프레임워크: 세계 모델링과 액션 모델링을 위한 통합된 사전 학습 접근 방식을 통해, 단일 생성 과정 내에서 액션 조건부 미래 역학을 공동으로 포착한다.
3. 계층적 의사 결정 조건부 계획: 정책을 고수 수준의 의사 결정 스켈레톤과 저수준의 액션 정교화로 분해함으로써, 계획 성능을 향상시키는 동시에 제어 가능한 미래 생성, 반사실적 추론, 안전 인식 평가를 가능하게 함을 입증하였다.

실험 결과

실험은 대규모 자율 주행 벤치마크인 NAVSIM-v1 및 NAVSIM-v2에서 수행되었다.

• 계획 성능: Discrete-WAM은 강력한 엔드투엔드 베이스라인들과 비교하여 경쟁력 있거나 우수한 계획 성능을 달랬다. NAVSIM-v2에서 EPDMS 점수 90.4를 기록하여 WAM-Flow(87.0) 및 SparseDriveV2(90.1)와 같은 방법들을 앞질렀다. NAVSIM-v1에서는 PDMS 92.2를 달성하여 최신 SOTA 플래너들을 능가했다.
• 세계 생성: 생성 품질(FID 및 FVD) 측면에서, Discrete-WAM은 짧은 호즌(horizon) 계획에 최적화되었음에도 불구하고 기존의 생성형 드라이빙 세계 모델(예: Vista, DrivingWorld)보다 뛰어난 성능을 보이며 FID 6.6, FVD 80.0을 달성했다.
• 절제 연구 (Ablation Studies):
  • 통합 사전 학습: 통합 사전 학습으로 초기화되고 LoRA(LoRA-SFT)를 통해 미세 조정된 모델이 처음부터 학습하거나 전체 미세 조정을 거친 모델보다 우수한 성능을 보였으며, 이는 사전 학습된 세계-정책 지식을 보존하는 것의 가치를 나타낸다.
  • 의사 결정 모델링: 의사 결정 토큰($D_t$)의 포함은 계획 지표를 유의미하게 개선하였으며, 이는 의사 결정 수준의 가이드가 효과적인 행동 사전 확률(prior)을 제공함을 확인시켜 준다.
  • 사후 학습 (Post-Training): RL 사후 학습은 EPDMS를 90.0에서 90.4로 더욱 향상시켜, 의사 결정-궤적 조합을 탐색하는 것의 이점을 입증했다.
  • 반사실적 추론: 모델은 "놀라움(surprise)"(사실적 롤아웃과 반사실적 롤아웃 사이의 KL 발산)과 계획 안전 지표(PDMS) 사이의 강한 음의 상관관계를 보여주었으며, 이는 안전이 중요한 장면 역학을 포착하는 능력을 나타낸다.

의의 및 주장

본 논문은 Discrete-WAM이 반응적 모방을 넘어 구조화된 세계-정책 추론으로 나아감으로써, 자율 주행에서 더 신뢰할 수 있는 의사 결정에 대한 원칙적인 경로를 제공한다고 주장한다. 세계 모델링, 정책 학습, 의사 결정을 이산 토큰 공간 내에서 통합함으로써, 이 프레임워크는 다음을 지원한다:

• 제어 가능한 생성: 미래 관측 및 궤적의 특정 측면을 편집할 수 있는 능력.
• 반사실적 추론: 대안적인 미래를 추론하고 실행 전 안전 위험을 평가할 수 있는 능력.
• 안전 인식 예측: 충돌이나 규칙 위반에 해당하는 고-놀라움(high-surprise) 시나리오를 식별하는 능력.

저자들은 이 접근 방식이 표현 형성 및 생성 모델링의 근본적인 병목 현상을 해결하며, 다양한 주행 시나리오 전반에 걸쳐 일반화될 수 있는 체화된 지능(embodied intelligence)의 확장 가능한 토대를 제공한다고 상정한다.