ReflectDrive-2: 이산 확산 주행을 위한 강화학습 정렬 자기 수정 기법

기존의 모방 학습 기반 자율주행 모델은 전문가 데이터를 단순히 따라 하느라 예기치 못한 상황에서 발생하는 오류를 스스로 수정하지 못하는 한계가 있었다. ReflectDrive-2는 별도의 네트워크 없이도 주행 계획을 스스로 검토하고 수정하는 AutoEdit 메커니즘을 도입하여 안전성과 주행 효율을 동시에 개선했다.

자율주행에서 Transformer 기반의 VLA(Vision-Language-Action) 모델은 주로 다음 토큰을 예측하는 방식으로 경로를 생성한다. 하지만 이 방식은 한 번 생성된 경로에 오류가 있더라도 이를 되돌리기 어렵고, 전체 시퀀스를 다시 생성해야 하는 비효율성이 존재한다. ReflectDrive-2는 이를 해결하기 위해 경로를 이산적인 토큰 시퀀스로 표현하고, Masked Discrete Diffusion 방식을 채택하여 경로의 특정 부분만 다시 마스킹하고 채워 넣는 '제자리 수정'이 가능하게 했다.

단순히 수정 기능만 추가하는 것으로는 부족하다. 모델이 스스로 만든 초안이 나중에 수정될 것임을 미리 알고 초안을 작성해야 하며, 수정 단계에서도 최종 주행 결과가 좋아지는 방향으로 고쳐야 한다. 이를 위해 연구진은 초안 작성과 수정 단계를 하나의 연속된 과정으로 보고, 최종 결과에 대해서만 보상을 주는 강화학습(RL)을 적용했다. 이 과정을 통해 초안 작성기는 '수정하기 좋은 초안'을 내놓고, 수정기는 '실제 주행 점수를 높이는 수정'을 수행하도록 상호 최적화된다.

결과적으로 모델은 전문가의 주행을 흉내 내는 수준을 넘어, 스스로의 실수를 인지하고 교정하는 능력을 갖추게 된다. 이는 마치 사람이 운전할 때 순간적인 판단 착오를 즉각적으로 핸들이나 페달 조작으로 바로잡는 것과 유사한 메커니즘을 딥러닝 구조 내에서 구현한 것이다.

ReflectDrive-2는 Decision-Draft-Reflect의 3단계 파이프라인으로 작동한다. 먼저 다중 모달 컨텍스트를 입력받아 목표 지점(Goal Token)을 제안하고, 이를 조건으로 Masked Discrete Diffusion을 통해 4초 분량의 궤적 초안을 병렬 디코딩한다. 이후 AutoEdit 단계에서 동일한 모델 가중치를 재사용하여 신뢰도가 낮은 토큰을 선택적으로 재작성한다.

학습은 2단계로 진행된다. 1단계 지도 학습(SFT)에서는 전문가 궤적에 종방향 및 횡방향 섭동을 가한 SAP(Structure-Aware Perturbation) 데이터를 사용하여 복구 손실 함수를 최적화한다. SAP Loss는 [오류가 포함된 궤적 토큰 입력 → Transformer 인코딩 → 원본 전문가 토큰 예측 → Cross-Entropy 계산] 과정을 거쳐 모델이 구조적 오류를 교정하는 법을 배우게 한다.

2단계 강화학습(RL)에서는 전체 롤아웃(초안+수정)에 대해 단일 보상을 부여한다. [초안 및 수정 단계의 모든 토큰 전이 확률 입력 → 최종 궤적 생성 → NAVSIM PDMS 기반 보상 계산 → Policy Gradient 업데이트]를 수행한다. 이 때 보상 Ag는 [현재 궤적 점수 R(τg) - 샘플링된 궤적들의 평균 점수]로 계산되어, 모델이 평균보다 더 나은 주행을 하도록 가중치를 갱신한다.

NAVSIM 벤치마크에서 카메라 입력만으로 91.0 PDMS를 기록하며 기존 SOTA 모델인 AutoVLA(89.1)를 능가했다. 특히 강화학습을 적용했을 때 AutoEdit의 기여도가 지도 학습만 했을 때(+0.3)보다 훨씬 높은 +1.9 PDMS로 나타나, 초안 작성과 수정 단계의 결합이 성능 향상의 핵심임을 증명했다.

추론 효율성 측면에서는 NVIDIA Thor 플랫폼에서 최적화된 스택을 통해 지연 시간을 45.0ms에서 31.8ms로 약 1.42배 단축했다. Shared-prefix KV reuse를 통해 어텐션 연산 지연 시간을 0.28ms에서 0.08ms로 줄였으며, Fused CUDA unmasking 커널을 통해 토큰 업데이트 속도를 7.5배 향상시켰다.

모델 아키텍처는 0.7B 파라미터의 Masked-diffusion 언어 백본과 0.1B ViT 비전 인코더로 구성된다. 궤적은 8개의 웨이포인트를 각각 종방향/횡방향 좌표 토큰으로 이산화하여 총 16개의 토큰으로 표현한다. 기존의 연속적 확산 모델과 달리 가우시안 노이즈 대신 [MASK] 토큰을 사용하는 이산 확산 방식을 사용하여 텍스트 생성 모델과의 호환성을 높였다.

추론 시에는 Alternating Step Decode(ASD) 기법을 도입했다. 이는 이전 프레임에서 생성된 궤적을 현재 프레임의 좌표계로 변환(Rigid-body transform)한 뒤, 이를 AutoEdit의 입력으로 사용하여 처음부터 다시 생성하는 비용을 절감하는 시간적 정제(Temporal Refiner) 방식이다. 또한 Action-expert FFN을 도입하여 행동 관련 연산의 차원을 4096에서 1024로 축소함으로써 연산 효율을 극대화했다.

현재 ReflectDrive-2는 고정된 해상도의 BEV 좌표 토큰을 사용하기 때문에 생성된 웨이포인트의 공간적 정밀도가 토큰의 빈(Bin) 크기에 의해 제한된다. 또한 현재의 RL 보상은 가벼운 폐쇄 루프 주행 점수에 의존하고 있어, 실제 복잡한 도로 환경의 모든 안전 요구사항을 완벽하게 대변하기에는 한계가 있을 수 있다.