기술 요약: Adjoint Matching을 통한 확산 정책의 확장 가능한 최대 엔트로피 강화 학습

문제 정의

확산 정책(Diffusion policies)은 강화 학습(RL)에서 복잡하고 다봉형(multi-modal)인 행동 분포를 표현하기 위한 강력한 패러다임으로 부상했습니다. 그러나 이들의 온라인 RL 적용은 정답 데이터(ground-truth data)가 없는 상황에서의 확장 가능한 학습 문제로 인해 제한적입니다. 스코어 매칭(score matching)이나 브릿지 매칭(bridge matching)과 같은 표준 학습 목적 함수는 온라인 RL에 정적인 최적 샘플 데이터셋이 없기 때문에 직접적으로 적용하기 어렵습니다.

온라인 환경에서 확산 정책을 학습시키려는 기존 방식들은 세 가지 주요 한계에 직면해 있습니다:

1. 역 KL 발산(Reverse KL Divergence): 역 KL 발산을 최소화하는 방법들은 전체 확산 과정(diffusion process)에 대한 역전파(backpropagation)를 요구합니다. 이는 막대한 메모리 오버헤드와 계산 비효효율성을 초래하여 더 큰 규모의 문제로의 확장을 방해합니다.
2. 중요도 가중치 매칭(Importance Weighted Matching): 예측된 행동과 타겟 행동 사이의 오차를 최소화하기 위해 중요도 샘플링을 사용하는 방식은 고차원 행동 공간으로 확장하는 데 어려움을 겪습니다.
3. Q-스코어 방법(Q-Score Methods): Q-함수의 그래디언트($\nabla_a Q$)를 활용하는 기법들은 흔히 과감쇄 랑제빈 역학(overdamped Langevin dynamics)에 의존합니다. 이러한 과정은 대개 유한한 단계 이후에 절단(truncation)이 필요하여 유한 시간 영역(finite-time regime)에서 이론적 보장이 부족하거나, 실제 타겟 분포로부터 샘플을 생성한다는 보장이 없는 휴리스틱한 샘플링 기법에 의존합니다.

방법론: Adjoint Matching Diffusion Policy (AMDP)

저자들은 **Adjoint Matching Diffusion Policy (AMDP)**를 제안하며, 이는 최대 엔트로피 강화 학습(Max Ent-RL)을 위한 확산 정책 최적화를 확률적 최적 제어(Stochastic Optimal Control, SOC) 문제로 정식화합니다. 핵심 혁신은 시뮬레이션이 필요 없는(simulation-free) 학습을 가능하게 하기 위해 Adjoint Matching (AM), 구체적으로는 **Reciprocal Adjoint Matching (RAM)**을 적용한 것입니다.

1. 이론적 기초

이 방법은 정책 개선 단계를, 종단 시간 $\tau=1$에서의 확산 과정의 주변 분포(marginal distribution)가 최적의 MaxEnt 정책 $\pi^*(a|s) \propto \exp(Q^\pi(s,a)/\alpha)$와 일치하도록 하는 드리프트 함수 $u$를 찾는 것으로 정의합니다.

• 슈뢰딩거 브릿지(Schrödinger Bridge): 이 문제는 제어 비용(kinetic cost)을 최소화하면서 경계 조건을 만족하는 슈뢰딩거 브릿지 문제로 구성됩니다.
• 상호 적응형 어드조인트 매칭(Reciprocal Adjoint Matching): 확산 경로를 통해 역전파하는 대신(이는 중간 상태를 저장해야 함), AMDP는 제어된 프로세스를 제어되지 않은 참조 프로세스(uncontrolled reference process)로 상호 투영(reciprocal projection)합니다. 이를 통해 조건부 상호 프로세스 $\Pi^0_{\tau|1}$를 사용하여 종단 상태 $X_1$으로부터 중간 상태 $X_\tau$를 직접 샘플링함으로써 손실을 최소화할 수 있습니다.
• 시뮬레이션 프리 학습(Simulation-Free Training): 리플레이 버퍼에 저장된 종단 행동 $X_1$으로부터 $X_\tau$를 샘플링함으로써, 알고리즘은 그래디언트 업데이트 중에 순방향 SDE를 시뮬레이션할 필요가 없습니다. 이는 명시적인 우도 추정(likelihood estimation)과 비용이 많이 드는 확산 체인 역전파를 피할 수 있게 해줍니다.

2. 주요 알고 알고리즘 구성 요소

• 행동 공간 적응 (오차 함수 스쿼싱, Error Function Squashing):
  유계된 행동 공간(예: $A = [-1, 1]^d$)을 처리하기 위해, 저자들은 미분 가능한 변환 $f: \mathbb{R}^d \to [-1, 1]^d$를 도입합니다. 표준적인 요소별 $\tanh$ 스쿼싱과 달리, AMDP는 스케일링된 오차 함수(error function, erf) 스쿼싱을 사용합니다:
  $$f_{\text{erf}}(x) = (\text{erf}(k x_i)_{i=1}^d)$$
  스케일링 인자 $k$는 참조 프로세스의 가우시안 주변 밀도와 변환의 자코비안 행렬식(Jacobian determinant)이 정확히 일치하도록 선택됩니다. 이 특정한 선택은 참조 밀도와 자코비안 행렬식의 로그 비율을 상수로 만들어, 회귀 목적 함수를 크게 단순화하고 수치적 안정성을 향상시킵니다.

• 정책 반복(Policy Iteration):

  • 평가(Evaluation): 수정된 "소프트" 벨만 백업(soft Bellman backup)이 사용됩니다. 확산 모델의 주변 엔트로피는 다루기 어렵기 때문에, 저자들은 제어 비용으로부터 유도된 다루기 쉬운 엔트로피 하한을 제안하며, 이는 수렴 특성을 보존합니다.
  • 개선(Improvement): 정책은 고정점 반복(fixed-point iteration)인 AMDP 손실을 최소화함으로써 업데이트됩니다. 이 손실의 유일한 고정점은 최적 제어 $u^*$에 해당합니다.

• 신뢰 영역 업데이트(Trust-Region Updates):
  타겟 분포가 빠르게 변하는 온라인 환경에서 안정성을 보장하기 위해, 이전 반복값 $u_{\text{old}}$로부터 새로운 드리프트 $u$의 편차를 페널티로 부여하는 신뢰 영역 제약이 추가됩니다. 이는 분포 붕괴나 Q-함수 발산을 방지합니다. 이는 완화된 라그랑주 형식론(relaxed Lagrangian formalism)을 통해 해결됩니다.

주요 기여

1. 확장 가능한 학습 알고리즘: AMDP는 확산 과정을 역전파하지 않고 확산 정책을 최적화하는 매우 효율적인 알고리즘을 도입하여 메모리 요구 사항을 획기적으로 줄였습니다.
2. 시뮬레이션 프리 최적화: 상호 적절 어드조인트 매칭을 활용함으로써, AM-DP는 리플레이 버퍼에 저장된 종단 행동만을 사용하여 최적화할 수 있으며, SDE를 재시뮬레이션하지 않고도 샘플당 여러 번의 그래디언트 업데이트를 가능하게 합니다.
3. 수치적으로 안정적인 행동 스쿼싱: 오차 함수 스쿼싱 제안은 $\tanh$에 대한 안정적인 대안을 제공하며, 손실 함수의 특정 항들을 상쇄시켜 목적 함수를 단순화합니다.
4. 안정적인 정책 업데이트: 신뢰 영역 손실의 통합은 어드조인트 매칭 프레임워크의 이론적 보증을 변경하지 않으면서 온라인 환경에서 안정적인 정책 반복을 보장합니다.

실험 결과

저자들은 온-폴리(on-policy) 설정에서는 MuJoCo Playground, ManiSkill, HumanoidBench의 63개 환경에서, 오프-폴리(off-policy) 설정에서는 7개의 고차원 개 및 휴머노이드 환경(DMC suite)에서 AMDP를 평가했습니다.

• 온-폴리 성능:

  • AMDP는 강력한 가우시안 베이스라인(REPPO, PPO, SPO) 및 최첨단 확산 방법(DIME, DPPO, FPO)과 대등하거나 이를 능가합니다.
  • 복잡하고 고차원적인 작업(HumanoidBench)에서 AMDP는 Best-of-N(BoN) 샘플링 없이도 모든 베이스라인을 크게 앞섭니다.
  • AMDP는 MuJoCo DMC 및 휴머노이드 로코모션 작업에서 REPPO보다 샘설 효율성 측면에서 더 빠르게 수렴합니다.
  • BoN 샘플링(16개 샘플 중 최적의 것 선택)과 결합했을 때, AMDP는 모든 벤치마크에서 우수한 성능을 달성합니다.

• 오프-폴리 성능:

  • DMC 개 및 휴머노이드 작업에서, AMDP는 역 KL을 최적화하는 DIME와 대등한 성능을 보이며 약간 더 빠르게 수렴하여, 오프-폴리 환경에서 어드조인트 매칭 목적 함수의 경쟁력을 입증했습니다.

• 계산 효율성:

  • 실제 실행 시간(Wall-clock Time): AMDP는 확산 단계 수와 관계없이 REPPO와 같은 매우 효율적인 가우시안 방법 대비 약 10%의 작은 상수 오버헤드만을 발생시킵니 다.
  • 확장성: 역 KL 역전파를 사용하는 방법들은 16단계 확산 시 10배, 128단계 확산 시 최대 72~80배 더 느립니다.
  • 모델 스케일링: 저자들은 AMDP를 사용하여 대규모 확산 정책(106M 파라미터, 표준 모델보다 약 300배 큼)을 성공적으로 학습시켰으며, 이는 메모리 제약으로 인해 역전파 기반 방법들이 실패할 수 있는 복잡한 모델로의 확장 능력을 보여줍니다.

의의 및 주장

본 논문은 AMDP가 생성적 정책(generative policies)을 실제 RL 애플리케이션에 배치하기 위한 이론적으로 근거 있고, 계산적으로 효율적이며, 견고한 경로를 제공한다고 주장합니다.

• 메모리 병목 현상 극克服: 확산 체인을 통한 역전파를 제거함으로써, AMDP는 메모리와 계산 자원이 제한된 온라인 RL 시나리오에서 확산 정책 학습을 실용적으로 만듭니다.
• 이론적 보장: 휴리스틱한 Q-스코어 방법과 달리, AMDP는 고정점 보증을 가진 확률적 최적 제어에서 유도되어 정책이 최적의 MaxEnt 분포로 수렴함을 보장합니다.
• 실용적 생존력: 시뮬레이션 프리 학습, erf-스쿼싱을 통한 수치적 안정성, 그리고 신뢰 영역 업데이트의 결합은 기존의 가우시안 정책을 대체할 수 있는 강력한 대안으로서 확산 정책의 주요 장벽(확장성, 안정성, 효율성)을 해결합니다.

저자들은 이 접근 방식이 복잡한 연속 제어 작업을 위해 전통적인 가우시안 정책의 보다 실용적인 대안으로서 확산 정책을 만드는 데 있어 계산 오버헤드를 크게 줄인다고 결론짓습니다.