기술 요약: 오프라인 강화 학습을 위한 디퓨전 정책의 노이지 스페이스 정책 경사 (Noisy-Space Policy Gradient)

1. 문제 정의

오프라인 강화 학습(Offline Reinforcement Learning, RL)은 추가적인 환경 상호작용 없이 고정된 데이터셋으로부터 효과적인 정책을 학습하는 것을 목표로 한다. 이 설정에서 주요한 과제는 "분포 변화(distributional shift)" 문제이다. 즉, 데이터셋의 서포트(support)를 벗어나 외삽(extrapolate)하는 정책은 편향된 가치 추정치를 갖게 되어 성능이 저하된다. 오프라인 데이터에 존재하는 복잡하고 다봉형(multimodal)인 행동 분포를 처리하기 위해, **디퓨전 정책(diffusion policies)**과 같은 표현력이 풍부한 생성 모델이 등장했다. 디퓨전 정책은 노이즈가 섞인 잠재 변수(latent variables)를 최종적인 깨끗한 행동(clean action)으로 매핑하는 반복적인 디노이징 과정을 통해 행동을 생성한다.

그러나 디퓨전 정책을 가치 기반 RL과 통합하는 과정에는 근본적인 구조적 불일치가 발생한다:

• 최적화 공간(Optimization Space): 정책 최적화는 노이즈가 섞인 잠재 공간(디노이징 과정 중의 변수 시퀀스 $x^{(k)}$)에서 일어난다.
• 가치 정의(Value Definition): 강화 학습 목적 함수(가치 함수, 어드밴티지, 그래디언트)는 실행된, 깨끗한 행동($a \in \mathcal{A}$)에 대해 정의된다.

이 간극을 메우려는 기존 방식들은 다음과 같은 상당한 단점을 가진다:

1. 행동 공간 샘플링(Action-Space Sampling): 디퓨전 모델을 단순히 샘플러로 취급하고 행동에 따라 가중치를 다시 부여하는 방식은 잠재 공간 가이던스(latent-space guidance)를 우회하게 된다.
2. 엔드 투 엔드 역전파(End-to-End Backpropagation): 전체 디노이징 궤적을 통해 가치 그래디언트를 역전파하여 정책을 최적화하는 것은 긴 확률적 계산 그래프를 생성하여 계산 오버헤드와 불안정성을 초래한다.
3. 직접적인 잠재 변수 평가(Direct Latent Evaluation): 중간 단계의 노이즈 섞인 변수에 가치 가이던스를 직접 주입하기 위해 Q-함수를 평가하는 것은, 해당 변수들이 보상을 결정하는 실제 행동이 아니라는 점을 간과한다.

본 논문은 노이즈 섞인 잠해 공간 최적화와 행동 공간 리턴(return) 사이를 원칙적이고 의미론적으로 일관된 방식으로 연결하는 가치 정식화(value formulation)를 정의하는 데 있어 발생하는 공백을 식별한다.

2. 방법론

저자들은 **노이지 스페이스 행동-가치 함수(Noisy-Space Action-Value Function)**와 그로부터 유도된 **노이지 스페이스 정책 경사(Noisy-Space Policy Gradient, NSPG)**를 중심으로 한 프레임워크를 제안한다.

2.1 노이지 스페이스 행동-가치 함수 ($Q_{noisy}$)

핵심적인 이론적 기여는 노이즈 자체에 내재적 가치를 부여하는 것이 아니라, 노이즈가 유도하는 깨끗한 행동의 분포에 가치를 부여하는 노이즈 섞인 잠재 변수를 위한 가치 함수를 정의하는 것이다.

• 정의: 상태 $s$, 디퓨전 단계 $k$, 그리고 노이즈 섞인 잠재 변수 $x^{(k)}$에 대하여, 노이지 스페이스 행동-가치 함수는 $(s, x^{(k)}, k)$에 조건화된 디퓨전 정책에 의해 유도되는 기대 행동 공간 리턴으로 정의된다:
  $$Q_{noisy}^\pi(s, x^{(k)}, k) \triangleq \mathbb{E}_{a \sim \pi(\cdot|s, x^{(k)}, k)} [Q^\pi(s, a)]$$
• 의미론적 정당성: 저자들은 $Q_{noisy}$가 정책이 $(s, x^{(k)}, k)$에서 시작되어 환경에서 실행될 때 얻어지는 기대 할인 리턴과 정확히 일치함을 증명한다. 이는 가치 함수가 깨끗한 행동만을 관찰하는 환경 인터페이스와 일관됨을 보장한다.
• 대리 의미론(Surrogate Semantics): 정책의 역시간 역학(reverse-time dynamics)을 변경하지 않으면서 $x^{(k)}$에 대한 조건화를 잘 정의하기 위해, 저자들은 **대리 결합 분포(surrogate joint distribution)**를 도입한다. 이는 깨끗한 행동 $a$와 노이즈 섞인 잠재 변수 $x^{(k)}$ 사이의 관계를 순방향(forward) 디퓨전 과정(깨끗한 행동을 오염시키는 과정)을 통해 모델링한다. 이를 통해 정책의 생성 메커니즘을 변경하지 않고도 그래디언트를 유도할 수 있다.

2.2 노이지 스페이스 정책 경사 (NSPG)

대리 의미론을 사용하여, 저자들은 $x^{(k)}$에 대한 $Q_{noisy}$의 그래디언트에 대한 폐쇄형(closed-form) 표현식을 유도한다.

• 그래디언트 정식화:
  $$\nabla_{x^{(k)}} Q_{noisy}^\pi(s, x^{(k)}, k) = \mathbb{E}_{a \sim \pi} \left[ A_{noisy}^\pi(s, a, x^{(k)}, k) \nabla_{x^{(k)}} \log q_k(x^{(k)} | a) \right]$$
  여기서:
  • $A_{noisy}^\pi = Q^\pi(s, a) - Q_{noisy}^\pi(s, x^{(k)}, k)$ 는 노이지 스페이스 어드밴티지이다.
  • $\nabla_{x^{(k)}} \log q_k(x^{(k)} | a)$ 는 순방향 디퓨전 커널의 스코어 함수(score function)이다.
• 주요 차이점: 노이즈 섞인 잠재 변수를 하나의 디코딩된 행동 $\hat{a}$로 매핑하고 $Q(s, \hat{a})$를 미분하는 방식과 달리, NSPG는 호환 가능한 깨끗한 행동들의 기대 가치를 미분한다. 이를 통해 그래디언트는 가장 확률이 높은 행동이 반드시 최고 가치를 갖지 않더라도, 높은 가치를 가진 호환 가능한 행동들과 관련된 노이즈 공간의 방향을 강화하고 낮은 가치의 방향은 억제할 수 있다.
• 역전파 없음: 결정적으로, 이 그래디언트는 깨끗한 행동 공간의 가치 추정치($Q^\pi$)와 순방향 디퓨전 커널만을 사용하여 계산된다. 이는 전체 디노이징 궤적을 통한 역전파를 피함으로써 안정성과 효율성을 크게 향상시킨다.

2.3 노이지 스페이스 액터-크리틱 (NSAC)

저자들은 이 프레임워크를 실질적인 오프라인 RL 알고리즘인 NSAC로 구현한다.

1. 크리틱(Critic): 오버에스티메이션(overestimation)과 분포 변화를 완화하기 위해, 오프라인 데이터와 하한 신뢰 구간(Lower-Confidence Bound, LCB) 타겟을 사용하는 보수적인 행동 공간 Q-함수($Q_\phi$)를 학습한다. 크리틱은 노이즈 섞인 잠재 변수에 대해 결코 평가되지 않는다.
2. 액터(Actor): 표준 디노이징 손실(행동 복제)과 NSPG로부터 유도된 가치 가이던스 항을 결합한 회귀 목적 함수를 사용하여 디퓨전 정책을 업데이트한다.
   • 액터 손실은 $\lambda \cdot \epsilon_\theta^\top \nabla_{x^{(k)}} Q_{noisy}$ 항을 포함한다.
   • 이를 통해 정책은 데이터 분포에 묶여 있으면서도 잠재 공간에서 개선될 수 있다.

3. 주요 기여

1. 불일치의 식별: 본 논문은 디퓨전 기반 RL에서 노이즈 섞인 잠재 공간 최적화와 행동 공간 가치 함수 사이의 의미론적 간극을 공식적으로 식별하였다.
2. 이론적 정식화: $Q_{noisy}$를 도입하여, 노이즈 섞인 잠재 변수에 대해 오직 그것이 생성하는 깨끗한 행동의 분포를 통해서만 가치를 부여함으로써 MDP와 의미론적 일관성을 보장한다.
3. NSPG 유도: 깨끗한 행동 공간의 가치 추정치만을 사용하여 노이즈 섞인 잠재 변수를 최적화하는 정책 경사를 유도하여, 디노이징 과정에 대한 역전파의 필요성을 제거하였다.
4. NSAC 알고리즘: NSPG를 구현하여 디퓨전 정책을 오프라인 환경에서 효과적으로 학습할 수 있음을 보여주는 안정적인 액터-크리틱 알고리즘을 개발하였다.

4. 실험 결과

저자들은 상태 기반 D4RL 벤치마크(locomotion, AntMaze, Adroit)와 비전 기반 OGBench 태스크에서 NSAC를 평가한다.

• 성능: NSAC는 다양한 환경에서 최첨단(SOTA) 또는 경쟁력 있는 성능을 달성한다.
  • Locomotion: HalfCheetah, Hopper, Walker2d에서 기존 방법들(예: CQL, IQL, Diffuser)과 대등하거나 더 뛰어난 성능을 보인다. 특히, HalfCheetah medium-replay에서 약 10%의 성능 향상을 달성했다.
  • AntMaze: 기존 방법들, 특히 DAC보다 상대적 성능이 50% 이상 높은 대규모 미로 환경에서 상당한 개선을 보여준다.
  • Adroit: 평가된 방법들 중 Pen-Cloned-v1에서 가장 강력한 성능을 달상한다.
  • Vision (OGBench): 시각적 태스크에서 플로우 기반 정책 및 DAC를 능가하는 가장 높은 총점(271.5)을 기록한다.
• 절제 연구(Ablation Studies):
  • 가이던스 강도 ($\lambda$): 이 방법은 다양한 가이던스 강도 범위에서 견고하며, 중간 정도의 값이 최적의 성능을 낸다. 지나치게 공격적인 가이던스는 정책을 데이터의 서포트 밖으로 밀어내어 성능을 저하시킨다.
  • 몬테카를로 근사: 노이즈 공간의 기댓값은 적은 수의 샘플(예: 2-4개)로 효율적으로 추정할 수 있으며, DAC와 비교했을 때 계산 오버헤드가 매우 적다.
  • 비관주의(Pessimism): 알고리즘은 광범위한 비관주의 계수 범위에서도 견고함을 유지한다.

5. 의의 및 주장

본 논문은 NSPG가 디퓨전 정책을 훈련하기 위한 원칙적이고 효과적인 기초를 제공한다고 주장한다. 노이즈 섞인 잠재 공간 최적화와 행동 공간 가치 사이의 불일치를 해결함으로써, 이 방법은 엔드 투 엔드 역전파의 불안정성이나 노이즈 섞인 변수에 크리틱을 평가할 때 발생하는 의미론적 비일관성을 피한다.

저자들은 자사의 접근 방식이 다음을 특징으로 한다고 강조한다:

• 트랙터빌리티(tractability)를 희생하지 않으면서 다단계 디퓨전 정책의 표현력을 보존한다.
• 가치 추정과 정책 최적화가 밀접하게 결합된 도전적이고 긴 호흡의 다봉형 설정에서, 안정성과 성능을 완만하지만 유의미하게 향상시킨다.
• 이 프레임워크가 온라인 및 오프라인-투-온라인 학습으로 자연스럽게 확장될 수 있음을 시사하지만, 이는 향후 연구 과제로 남겨둔다.

본 연구는 모든 오프라인 RL 문제를 해결한다고 주장하는 것이 아니라, 특히 디퓨전 정책과 가치 기반 목적 함수를 통합하는 문제를 다루며, 휴리스틱한 샘플링이나 불안정한 엔드 투 엔드 훈련에 대한 이론적으로 근거 있는 대안을 제시한다.