Wasserstein Proximal Policy Gradient

이 논문은 엔트로피 정규화 강화학습을 위한 새로운 알고리즘인 'Wasserstein 근사 정책 경사 (WPPG)'를 제안하며, 최적 수송과 가우시안 합성곱을 번갈아 적용하여 명시적 확률 정책 없이도 전역 선형 수렴을 보장하고 연속 제어 벤치마크에서 경쟁력 있는 성능을 달성함을 보여줍니다.

핵심

1. 배경: 로봇이 배울 때 겪는 두 가지 문제

상상해 보세요. 로봇이 새로운 춤을 배우려 합니다.
기존의 학습법 (KL 발산 기반) 은 로봇에게 **"네가 지금 추는 춤과 다음에 추는 춤이 얼마나 다른지"**를 계산하게 합니다. 이때 중요한 전제는, 로봇이 **"내 춤 동작의 확률 분포를 수학적으로 완벽하게 설명할 수 있어야 한다"**는 것입니다.

하지만 현실은 어떻습니까?

• 문제 1 (복잡한 춤): 로봇이 너무 복잡하고 자유로운 춤 (암시적 정책, Implicit Policy) 을 추고 싶다면, 그 춤을 수학 공식으로 설명하는 건 불가능에 가깝습니다.
• 문제 2 (탐험의 필요성): 로봇이 너무 똑똑해져서 같은 동작만 반복하면 (수익은 좋지만), 새로운 것을 시도하지 못해 더 좋은 춤을 찾을 수 없습니다. 그래서 '엔트로피 (무작위성)'를 넣어주어 조금씩 엉뚱한 동작을 시도하게 해야 합니다.

2. WPPG 의 핵심 아이디어: "물방울의 흐름"과 "소금물"

이 논문은 이 문제를 해결하기 위해 **'워터스틴 거리 (Wasserstein Distance)'**라는 개념을 도입합니다.

비유 1: 물방울의 이동 (워터스틴 거리)

기존 방식은 "두 확률 분포가 얼마나 다른가?"를 단순히 숫자 차이로 봅니다. 하지만 WPPG 는 **"이 물방울 (현재 행동) 을 저 물방울 (목표 행동) 로 옮기려면 얼마나 많은 에너지를 써야 하는가?"**를 생각합니다.

• 기존: "내 춤과 네 춤은 50% 다릅니다." (단순 비교)
• WPPG: "네가 지금 오른손을 들었는데, 목표는 왼손을 드는 거야. 오른손을 왼쪽으로 부드럽게 이동시키는 데 드는 '비용'을 계산하자." (기하학적 이동)

이 방식은 행동 공간 (Action Space) 의 기하학적 구조를 자연스럽게 반영합니다. 예를 들어, '왼쪽'과 '오른쪽'은 거리가 멀지만, '왼쪽'과 '왼쪽 약간'은 가깝습니다. WPPG 는 이런 **친밀감 (Proximity)**을 잘 이해합니다.

비유 2: 소금물과 열기 (엔트로피 정규화)

학습 과정에서 로봇이 너무 경직되지 않게 하려면 '엔트로피 (무작위성)'를 추가해야 합니다. 보통은 수학적으로 복잡한 계산을 통해 이 무작위성을 추가합니다.
하지만 WPPG 는 아주 간단한 방법을 씁니다.

• 방법: 로봇이 다음 행동을 결정할 때, 약간의 '소금물 (가우시안 잡음)'을 섞어줍니다.
• 효과: 이 소금물이 섞이는 과정은 마치 뜨거운 물이 식으면서 퍼지는 것처럼, 자연스럽게 무작위성을 만들어냅니다. 이 과정을 수학적으로는 '열 방정식 (Heat Equation)'이라고 하는데, WPPG 는 이를 **가우시안 합성곱 (Gaussian Convolution)**이라는 쉬운 연산으로 처리합니다.

3. 혁신적인 점: "설명할 필요 없음"

가장 큰 장점은 수학적 설명 (확률 밀도 함수) 이 필요 없다는 점입니다.

• 기존 방식: "네가 이 행동을 할 확률이 얼마야? 그 로그를 계산해 줘." (수학 공식이 있어야 함)
• WPPG: "네가 이 행동을 할 때, 그 행동이 얼마나 좋은지 (Q 값) 알려주고, 그 방향으로 조금씩 움직여. 그리고 약간의 잡음을 섞어."

마치 요리사를 생각하세요.

• 기존 요리사: "이 요리의 레시피 (공식) 를 정확히 적어줘야 내가 다음 요리를 만들 수 있어."
• WPPG 요리사: "맛있는 요리를 만드는 '방향'을 알려주고, 재료를 섞는 '손맛 (잡음)'만 추가해. 레시피는 몰라도 돼!"

이 덕분에 **매우 복잡하고 표현력 있는 신경망 (암시적 정책)**을 사용하여 로봇이 훨씬 더 다양하고 창의적인 행동을 배울 수 있습니다.

4. 결과: 더 빠르고 안정적인 학습

논문은 이 방법이 수학적으로 **선형 수렴 (Linear Convergence)**한다는 것을 증명했습니다. 쉽게 말해, "이 방법을 쓰면 실수가 줄어들고 최적의 해답에 도달하는 속도가 일정하게 보장된다"는 뜻입니다.

실제 실험 (MuJoCo 시뮬레이션) 에서도:

• 기존 방법 (PPO, SAC): 잘하지만, 복잡한 환경 (Humanoid 등) 에서는 학습이 불안정하거나 느립니다.
• WPPG (특히 암시적 정책 버전): 복잡한 환경에서도 더 높은 점수를 기록하며, 더 안정적으로 학습했습니다.

5. 요약: 한 줄로 정리하면?

"WPPG 는 로봇이 복잡한 춤을 배울 때, 수학 공식으로 설명할 필요 없이 '행동의 흐름'과 '적당한 무작위성 (소금물)'만 섞어주면, 더 빠르고 창의적으로 최고의 춤을 추게 해주는 새로운 학습법입니다."

이 기술은 인공지능이 더 유연하고 강인하게 세상을 이해하고 행동하는 데 중요한 발걸음이 될 것입니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 배경: 강화학습 (RL) 에서 정책 경사 (Policy Gradient, PG) 방법은 복잡한 의사결정 문제를 해결하는 핵심 도구입니다. 기존 방법들은 주로 매개변수 공간의 유클리드 기하학 (Euclidean geometry) 이나 KL 발산 (Kullback-Leibler divergence) 을 기반으로 한 자연 정책 경사 (Natural Policy Gradient) 나 TRPO/PPO 와 같은 신뢰 영역 (Trust Region) 방법을 사용합니다.
• 한계점:
  • KL 기반 방법의 제약: KL 발산은 행동 공간의 기하학적 구조를 무시하고 행동을 독립적인 범주로 취급합니다. 또한, 명시적인 확률 밀도 함수 (log-density) 와 그 기울기 (score function) 를 계산할 수 있는 정책 (Explicit Policy) 에만 적용 가능합니다.
  • 암시적 정책 (Implicit Policy) 의 부재: 복잡한 행동 분포를 표현할 수 있는 암시적 정책 (예: 푸시포워드 맵을 통해 정의된 정책) 은 밀도 함수를 직접 계산하기 어렵기 때문에 기존 KL 기반 방법이나 SAC 와 같은 엔트로피 정규화 방법의 적용이 제한적입니다.
  • 수렴성 이론의 부재: 와세르스타인 (Wasserstein) 거리를 활용한 정책 최적화 연구는 존재하지만, 연속 행동 공간에서 매개변수화된 정책 (입자 근사 외) 에 대한 전역 수렴성 (Global Convergence) 보장은 아직 미해결 과제로 남아 있었습니다.

2. 제안 방법: Wasserstein Proximal Policy Gradient (WPPG)

이 논문은 와세르스타인 기하학 (Wasserstein geometry) 을 기반으로 한 새로운 정책 업데이트 프레임워크인 WPPG를 제안합니다.

핵심 아이디어

• 와세르스타인 근사 업데이트: 정책 업데이트를 와세르스타인 공간에서의 근사 (Proximal) 문제로 정의합니다. 이는 행동 공간의 기하학적 구조를 보존하면서 행동 가치 함수 (Action-Value Function) 를 최대화하는 방향으로 정책을 이동시킵니다.
• 연산자 분할 (Operator-Splitting) 기법: 복잡한 와세르스타인 근사 업데이트를 두 단계로 분해하여 효율적으로 해결합니다.
  1. 와세르스타인 수송 단계 (Transport Step): 행동 가치 함수 ($Q$) 의 기울기를 따라 행동을 이동시켜 기대 보상을 증가시킵니다.
  2. 열 흐름 단계 (Heat Flow Step): 엔트로피 정규화를 처리하기 위해 가우시안 노이즈를 주입합니다. 이는 열 방정식 (Heat Equation) 의 해가 가우시안 합성곱 (Convolution) 임을 이용합니다.
• 암시적 정책 지원: 이 접근법은 정책의 로그 밀도 (log-density) 나 그 기울기를 계산할 필요가 없습니다. 대신, 행동 가치 함수의 행동에 대한 기울기 ($\nabla_a Q$) 만을 사용하여 생성 모델 (Generator) 의 매개변수를 업데이트합니다. 이는 **암시적 정책 (Implicit Policies)**을 직접 최적화할 수 있게 합니다.

알고리즘 흐름

1. ** critic 업데이트:** Double-Q 학습을 사용하여 행동 가치 함수 $Q$를 추정합니다.
2. actor 업데이트 (WPPG):
   • 현재 정책에서 샘플링된 행동에 대해 $Q$의 기울기를 계산합니다.
   • 이 기울기를 따라 행동을 이동시킨 후, 엔트로피 항을 처리하기 위해 가우시안 노이즈를 추가합니다.
   • 생성기 (Generator) 의 매개변수를 업데이트하여 새로운 행동 분포가 이 목표 방향으로 이동하도록 합니다.
3. 엔트로피 추정: 암시적 정책의 경우, 가우시안 합성곱을 통해 유도된 정책 분포에 대한 엔트로피를 플러그인 (Plug-in) 혼합 모델 (Mixture Model) 을 사용하여 추정합니다.

3. 주요 기여 (Key Contributions)

1. 새로운 최적화 프레임워크: 와세르스타인 기하학을 기반으로 한 WPPG 알고리즘을 제안하여, 명시적/암시적 정책 모두에 적용 가능한 정책 최적화 방법을 개발했습니다.
2. 밀도 함수 불필요: 기존 방법들과 달리 정책의 로그 밀도나 스코어 함수 (Score function) 를 요구하지 않아, 복잡한 구조를 가진 암시적 정책 (Pushforward maps) 을 직접 학습할 수 있습니다.
3. 전역 선형 수렴성 증명 (Global Linear Convergence):
   • 엔트로피 정규화 하에서 WPPG 의 전역 선형 수렴 속도를 수학적으로 증명했습니다.
   • 정확한 가치 함수 평가와 Actor-Critic 기반의 근사 평가 (Approximation error 포함) 모두에 대해 수렴성을 보장합니다.
   • 이 분석은 와세르스타인 거리 ($W_2$) 와 운송 - 정보 부등식 (Transportation-Information inequality, $T_2$) 을 활용하여 KL 기반 분석 (Mirror Descent 등) 과는 다른 새로운 이론적 통찰을 제공합니다.
4. 실증적 성능: MuJoCo 연속 제어 벤치마크 (Hopper, Humanoid 등) 에서 기존 SOTA 방법들 (PPO, SAC, WPO) 과 비교하여 우수한 성능을 입증했습니다. 특히 암시적 정책을 사용하는 WPPG-I 은 모든 작업에서 가장 높은 보상을 기록했습니다.

4. 실험 결과 (Results)

• 벤치마크: MuJoCo 의 6 가지 연속 제어 작업 (Hopper, Walker2d, HalfCheetah, Reacher, Swimmer, Humanoid) 에서 평가 수행.
• 비교 대상: PPO (KL 기반), SAC (엔트로피 정규화), WPO (Wasserstein 기반).
• 성능:
  • WPPG (명시적 정책): SAC 와 유사한 성능을 보이며, KL 기반 기하학이 와세르스타인 기하학과 경쟁력 있음을 입증했습니다.
  • WPPG-I (암시적 정책): 모든 베이스라인을 압도하는 성능을 기록했습니다. 특히 Humanoid 와 같은 고차원/어려운 작업에서 안정적으로 수렴하며 높은 보상을 달성했습니다. 이는 복잡한 행동 분포를 표현하는 암시적 정책의 잠재력을 보여줍니다.
  • WPO: Humanoid 나 Swimmer 와 같은 어려운 환경에서 수렴 불안정성을 보였으며, Reacher 에서는 학습에 실패하기도 했습니다.
• 분석 (Ablation Study):
  • 엔트로피 계수 ($\tau$): 적절한 노이즈 주입 ($\tau$) 은 탐색을 촉진하지만, 과도한 노이즈는 학습을 저해합니다.
  • 잠재 변수 차원: 암시적 정책의 잠재 변수 차원이 너무 작으면 탐색이 부족하고, 너무 크면 학습이 느려집니다. 상태 차원의 약 1/3 정도가 최적의 균형을 이룹니다.
  • Double-Q: Double-Q 기법을 적용하면 학습 안정성과 성능이 크게 향상됩니다.

5. 의의 및 결론 (Significance)

• 이론적 의의: 와세르스타인 기하학을 사용하여 연속 행동 공간에서 매개변수화된 정책에 대한 전역 수렴성 보장을 최초로 제시한 연구 중 하나입니다. 이는 기존 KL 기반 이론의 한계를 넘어 새로운 수렴 분석 도구를 제공합니다.
• 실용적 의의:
  • 암시적 정책의 활용: 밀도 함수 계산이 불가능한 복잡한 생성 모델 (Implicit Policies) 을 강화학습에 효과적으로 통합할 수 있는 길을 열었습니다. 이는 고차원 행동 공간이나 멀티모달 (Multimodal) 행동 분포가 필요한 로봇 제어 및 복잡한 의사결정 문제에 매우 유용합니다.
  • 간단한 구현: 로그 밀도 계산 없이 행동 기울기만으로 구현 가능하여, 기존 RL 프레임워크에 쉽게 통합될 수 있습니다.

요약하자면, 이 논문은 와세르스타인 기하학을 활용하여 암시적 정책을 직접 최적화할 수 있는 WPPG를 제안하고, 이에 대한 이론적 수렴성과 실제 성능을 입증함으로써 강화학습 정책 최적화의 새로운 지평을 열었습니다.