기술 요약: 접두사 최적 생성 정책 (Prefix-Optimal Generative Policies, POGP)

문제 정의

확산 정책(Diffusion policies)은 복잡하고 다봉적인(multimodal) 행동 분포를 표현할 수 있는 강력한 클래스의 모델로 연속 제어 분야에서 부상했습니다. 그러나 이들의 실질적인 채택은 상당한 계산 병목 현상에 의해 저해됩니다. 즉, 일반적으로 $K$번의 순차적인 신경망 평가를 요구하는 반복적인 디노이징(denoising) 과정은 작업의 현재 난이도와 관계없이 고정된 횟수를 수행해야 합니다.

표준 확산 학습은 오직 최종 출력($a_0$)만을 최적화하며, 중간 디노이징 단계($a_1, \dots, a_{K-1}$)는 비지도 상태로 남겨둡니다. 결과적으로 이러한 중간 상태들은 고품질의 실행 가능한 결정이 되도록 학습되지 않습니다. 이는 훈련-테스트 간의 불일치(train-test mismatch)를 유발합니다. 만약 계산량을 절약하기 위해(실시간 또는 에지 제약 조건에 따라) 추론을 조기에 중단(truncation)한다면, 생성된 행동은 종종 최적이 아니게 됩니다. D3P와 같은 기존의 동적 확산 방식은 정지 시간을 추정하기 위해 별도의 어댑터를 학습하려고 시도하지만, 중간 단계 자체를 최적화하지는 않으므로 정지 기준과 부분 출력의 실제 품질 사이에 괴리가 발생합니다.

방법론

본 논문은 **접두사 최적 생성 정책(POGP)**을 제안합니다. 이는 확산 정책과 **접두사 가치 함수(Prefix Value Function, PVF)**를 공동으로 학습하여, 적응형 계산을 가능하게 함과 동시에 디노이징 체인의 모든 단계에서 행동 품질을 개선하는 프레임워크입니다.

1. 접두사 가치 함수 (Prefix Value Function, PVF)

POGP는 $K$단계 디노이징 체인($a_K \to \dots \to a_0$)의 모든 중간 단계 $t$에 대한 가치 함수 $V_t(s, a_t)$를 도입합니다. 환경의 상태-행동 쌍을 평가하는 표준 가치 함수와 달리, PVF는 현재의 중간 행동 $a_t$로부터 디노이징 체인을 완료했을 때의 기대 수익을 추정합니다.

• 정의: $V_t^\pi(s, a_t) := Q^\pi(s, \pi_{\text{full}}(s, a_t))$이며, 여기서 $\pi_{\text{full}}$은 단계 $t$부터 0까지 디노이저를 실행한 결과입니다.
• 학습 메커니즘: $O(K^2)$의 비용이 드는 롤아웃을 피하기 위해, POGP는 디노이징 체인에 대한 **벨만 스타일 재귀(Bellman-style recursion)**를 도출합니다. $V_t$의 타겟은 최종 행동의 Q-값과 다음 단계의 부트스트랩된 가치의 가중 결합입니다:
  $$V_t(s, a_t) = h \cdot Q(s, a_0) + (1-h) \cdot V_{t-1}(s, \pi_\theta(s, a_t, t))$$
  여기서 $h$는 균등 위험률(uniform hazard rate, $1/K$로 설정)입니다. 이를 통해 네트워크는 단일 공유 체인을 사용하여 중간 단계의 가치를 학습할 수 있습니다.

2. 게이트 오브젝티브를 통한 공동 최적화

POGP는 세 가지 구성 요소를 공동으로 최적화합니다:

• 환경 크리틱 (Environment Critic, $Q$): 환경 전이에 대한 표준 더블 Q-러닝(double Q-learning).
• 접두사 크리틱 (Prefix Critic, $V$): 위에서 설명한 Hazard Bellman 재귀를 통해 학습됨.
• 액터 (Actor, $\pi$): 최종 행동의 가치를 극대화하는 것과 높은 접두사 값을 유지하는 것 사이의 균형을 맞추는 **게이트 손실(gated loss)**을 통해 최적화됩니다. 게이팅 스케줄 $w(n)$은 PVF가 안정화될 때까지 접두사 감독을 지연시켜, 초기 단계의 노이즈가 정책을 불안정하게 만드는 것을 방지합니다.

3. 동적 정지 및 완성

테스트 시, 학습된 PVF는 원칙적인 정지 기준을 제공합니다. 알고리즘은 연속적인 $m$개 단계 동안 한계 이득 $\Delta_t = V_{t-1} - V_t$가 임계값 $\epsilon$ 미만으로 떨어질 때 디노이징을 중단합니다.

• 완성 맵 (Completion Map): 중단된 행동($a_T$)에는 잔여 노이즈가 포함될 수 있으므로, POGP는 경량 MLP($g_\omega$)를 학습하여 이러한 중단된 행동을 완전히 디노이징된 대응물($a_0$)로 투영함으로써, 조기에 중단되더라도 실행 가능한 행동이 고품질을 유지하도록 보장합니다.

주요 기여

1. 접두사 가치 함수: 디노이징 체인에 대한 TD 부트스트래핑을 통해 학습되는 새로운 가치 함수입니다. 이는 중간 출력을 고품질의 실행 가능한 행동으로 형성하는 보조 목적 함수 역할을 하며, 표준 확산 모델의 "비지도 중간 단계" 문제를 해결합니다.
2. 학습된 조기 정지: 연속적인 접두사 값의 차이에서 직접 도출된 정지 규칙입니다. 이는 별도의 분리된 정지 네트워크를 필요로 하지 않으며, 정책이 추가적인 정교화의 한계 효용을 예측하여 계산량을 적응적으로 할당할 수 있게 합니다.
3. 경험적 검증: 본 방법론은 중간 단계의 감독이 전체 체인의 성능을 향상시킬 뿐만 아니라 보조적인 목적 함수로서도 유익하다는 것을 입증하며, 이를 통해 실험적으로 검증되었습니다.

실험 결과

저자들은 네 가지 MuJoCo 환경(HalfCheetah, Walker2d, Ant, Hopper)에서 12개의 베이스라인(고전적 RL 방법인 TD3, SAC, PPO 및 최첨단 확산 정책인 Diff-QL, D3P, SDAC 포함)을 대상으로 POGP를 평가했습니다.

• 전체 체인 성능: $K=20$ 단계를 모두 실행하도록 강제했을 때, POGP는 모든 베이스라인을 능가하여 0.942 (95% CI: [0.91, 0.97])의 사분위 평균(IQM)을 달형성했습니다. 이는 최고 베이스라인인 SDAC보다 약 4.6%, 동적 베이스라인인 D3P보다 3.5% 향상된 수치입니다.
• 적응형 계산 효율성: 적응형 정지 규칙 하에서 POGP는 전체 체인 성능의 97.8% 이상을 유지하면서도 평균 디노이징 반복 횟수를 약 2.7배 줄였습니다 (예: HalfCheetah에서 20단계에서 약 6.8단계로 감소).
• D3P와의 비교: POGP는 D3P보다 수익(return) 측면에서 3.5% 더 높았으며, 반복 횟수는 18.2% 더 적게 사용했습니다. 일상적인 상태에서도 높은 반복 횟수를 유지하는 D3P와 달리, POGP는 안정적인 보행 중에는 계산량을 35단계로 동적으로 줄이고, 섭동(perturbation)으로부터 회복할 때만 1520단계로 늘립니다.

의의

본 논문은 가변적인 확산 깊이를 정책 최적화의 일급 객체(first-class object)로 취급함으로써, 적응형 계산과 강화 학습 사이의 간극을 메운다고 주장합니다. 그 의의는 다음과 같은 이중적 이점에 있습니다:

1. 효율성: 성능 저하 없이 상당한 계산 절감(최대 2.7배)을 가능하게 하여, 확산 정책을 실시간 및 에지 배포에 적합하게 만듭니다.
2. 품질: 접두사 학습이 최종 터미널 행동의 품질을 개선하는 정규화 보조 목적 함수로 작용하며, 이는 중간 단계의 감독이 조기 정지 메커니즘과 독립적으로도 견고한 정책을 학습하는 데 본질적으로 유익함을 시사합니다.

이 연구는 중간 디노이징 단계가 실행 가능한 결정이 되도록 최적화될 수 있음을 확립하여, 중단된 확산 추론에 내재된 훈련-테스트 불일치 문제를 해결했습니다.