Drift Q-Learning

기술 요약: 드리프트 Q-러닝 (Drift Q-Learning, DriftQL)

문제 정의
오프라인 강화 학습(Offline Reinforcement Learning, RL)은 추가적인 환경 상호작용 없이 고정된 데이터셋으로부터 최적의 정책을 학습하는 것을 목표로 합니다. 이 설정의 핵심 과제는 "분포 외(out-of-distribution, OOD)" 문제입니다. 정적인 데이터로 학습된 가치 함수는 데이터셋에 존재하지 않는 행동에 대해 신뢰할 수 없는 추정치를 생성합니다. 이를 완화하기 위해서는 학습된 정책이 데이터 분포를 향하도록 제약하면서도, 동시에 높은 가치를 가진 행동을 향해 개선될 수 있도록 허용하는 행동 정규화(behavior regularization)가 필요합니다.

디퓨전(diffusion) 및 플로우 기반(flow-based) 정책과 같은 표현력이 풍부한 생성 모델을 활용하는 기존 솔루션들은 다봉형(multimodal) 행동 분포의 구조를 효과적으로 모델링합니다. 그러나 이들은 실질적인 한계점을 가집니다. 즉, 추론 시 반복적인 디노이징(denoising)이나 무거운 솔버(solver) 통합이 필요하여 실행 속도가 느려집니다. 추론 속도를 높이기 위한 증류(distillation) 기반의 대안들도 존재하지만, 이는 추가적인 네트워크와 2단계 학습 파이프라인을 요구하여 복잡성을 초래합니다.

방법론
저자들은 드리프트 기반의 행동 정규화와 크리틱(critic) 중심의 정책 개선을 결합한 단일 단계(one-step) 생성 방법인 **Drift Q-Learning (DriftQL)**을 제안합니다. DriftQL은 사전 노이즈(prior noise)와 조건부 신호로부터 데이터로 직접 가는 단일 단계 푸시포워드 맵(pushforward map)을 학습하는 "드리프팅 모델(Drifting Models)" 프레임워크를 기반으로 구축되었습니다.

DriftQL의 핵심은 상태 $s$와 가우시안 노이즈 $\epsilon$을 입력받아 행동 $a$를 매핑하는 확률적 액터 네트워크 $f_\theta(s, \epsilon)$입니다. 학습 목적 함수는 두 가지 주요 메커니즘을 통해 분포 모델링과 보상 탐색을 통합합니다:

1. 조건부 드리프트 필드 (Conditional Drift Field): 정책은 두 가지 상반된 힘으로 구성된 상태 조건부 드리프트 필드에 의해 감독됩니다.

   • 인력 (Attraction, $V^+$): 생성된 행동을 관측된 데이터셋의 행동 쪽으로 끌어당깁니다(행동 복제, behavior cloning). 일반적으로 상태당 하나의 관측된 행동이 존재하는 오프라인 RL 설정에서, 이는 생성된 행동과 데이터셋 행동 사이의 직접적인 벡터 차이로 단순화됩니다.
   • 척력 (Repulsion, $V^-$): 생성된 행동들이 서로 밀어내도록 합니다. 이는 정책이 단일 모드(데이터셋의 행동)로 붕괴되는 것을 방지하고, 높은 가치를 가진 영역 내에서 다양한 후보군을 유지하게 합니다. 척력은 단일 상태에 대해 생성된 $N$개의 후보 행동들에 대한 커널 가중 평균 이동(kernel-weighted mean-shift)을 사용하여 계산됩니다.

2. 크리틱 중심의 개선 (Critic-Driven Improvement): 일반적인 비조건부 생성을 위한 드리프팅 모델과 달리, DriftQL은 가치 기반 목적 함수를 포함합니다. 클립된 더블 Q-크리틱(clipped double-Q critic)이 행동의 가치를 추정합니다. 액터는 다음의 결합 손실을 최소화합니다:
   $$L_{actor}(\theta) = \alpha L_{drift}(\theta) + L_Q(\theta)$$
   여기서 $L_{drift}$는 드리프트 타겟을 향해 행동을 회귀시키는 드리프트 손실이며, $L_Q$는 음의 기대 Q-값입니다. 가치 그래디언트는 정책이 데이터 지지(support) 내의 고가치 영역을 향하도록 편향시켜, 데이터에 여러 고가치 모드가 존재할 경우 정책이 이를 집중적으로 학습할 수 있게 합니다.

주요 기여

• 단일 단계 생성 정책: DriftQL은 단 한 번의 순방향 패스(forward pass)로 행동을 생성하므로, 디퓨전 및 플ow 기반 방식에서 요구되는 반복적인 디노이징 체인, ODE 솔버 또는 증류 단계가 필요 없습니다.
• 통합된 학습 목적 함수: 드리프트 필드(행동 정규화)와 Q-크리틱(정책 개선)을 단일 네트워크의 통합된 손실 함수로 통합하여, 다단계 학습 파이프라인의 복잡성을 피했습니다.
• 모드 선택 및 다양성: 인력(데이터로의 유인)과 척력(후보 간의 밀어냄)의 결합을 통해, 정책은 학습 중에 후보들 사이에 질량을 유지하면서도 고가치 모드로 수렴할 수 있게 하여, 결정론적인 단일 행동으로 붕에되거나 지원되지 않는 영역으로 흩어지는 것을 방지합니다.

실험 결과
DriftQL은 연속 제어, 내비게이션 및 조작 작업을 포함하는 D4RL 및 OGBench 벤치마크에서 평가되었습니다.

• 성능: DriftQL은 78개의 오프라인 학습 태스크에서 최신 디퓨전 및 플로우 기반 베이스라인(예: Diffusion-QL, FQL, IDQL)과 대등하거나 이를 능가하는 성능을 일관되게 보여주었습니다. 특히 장기 호라이즌 내비게이션 태스크(antmaze-large-st, antmaze-giant-st)와 복잡한 조작 태스크에서 강력한 이점을 보였습니다.
• 강건성: 데이터 품질이 저하된 상황(일부 행동을 무작위 노이즈로 교체하여 시뮬레이션)에서도, DriftQL은 경쟁 방법론들이 노이즈 증가에 따라 눈에 띄게 고전하는 것과 달리, 깨끗한 데이터 베이스라인에 근접한 성능을 유지했습니다.
• 효율성: 추론 지연 시간 측면에서 DriftQL은 FQL보다 약 2배, Diffusion-QL보다 3배, IDQL/IFQL보다 4배 빠릅니다. 이는 별도의 증류 단계를 거치지 않고도 달성되었습니다. 학습 비용은 쌍 기반 커널 계산으로 인해 생성된 샘플 수($N$)에 따라 스케일링되지만, 디퓨전 모델의 반복적인 디노이징 오버헤드에 비하면 여전히 완만한 수준입니다.

의의 및 주장
본 논문은 DriftQL을 디퓨전 및 플로우 기반 오프라인 RL 방식의 유망한 대안으로 제시합니다. 저자들은 DriftQL이 생성 모델의 표현력(다봉 분포 처리)을 유지하면서도 결정론적 접근 방식의 단순함과 효율성(단일 순방향 패스)을 갖추고 있다고 주장합니다.

저자들은 DriftQL이 오프라인 RL의 상충하는 목표를 성공적으로 조화시킨다고 주장합니다. 즉, 신뢰할 수 있는 가치 추정을 위해 정책을 데이터 지지 영역 내로 제한하는 동시에, 높은 보상을 주는 영역을 향한 개선을 허용합니다. 반복적인 샘플링의 계산 오버헤드를 제거함으로써, DriftQL은 추론 속도가 중요한 실제 환경 배포에 더 실용적인 솔루션을 제공하며, 데이터 오염에 대한 강건성 또한 유지합니다. 이 연구는 드리프팅(drifting)이 다단계 생성 정책과 관련된 테스트 타임 비용 없이도 확률적이고 분포적인 학습 신호를 제공하는, 오프라인 RL을 위한 유용한 "중간 지점"임을 시사합니다.