Hypergradient-based Bilevel Reinforcement Learning with Improved Sample Complexity

기술 요약: 개선된 샘플 복잡도를 가진 하이퍼그레이디언트 기반 바이레벨 강화학습

문제 정식화

본 논문은 메타 학습, 계층적 작업 분해, 그리고 인간 피드로백 기반 강화학습(RL-HF)과 같은 문제를 공식화하는 데 사용되는 프레임워크인 **바이레벨 강화학습(Bilevel Reinforcement Learning, Bilevel RL)**을 다룹니다. 이 설정에서 외곽 레벨(outer-level) 목적 함수 $\phi(x, \theta^*(x))$는 파라미터 $x$(예: 보상 함수 파라미터)에 대해 최적화되며, 내부 레벨(inner-level) 목적 함수 $J(\theta, x)$는 정책 파라미터 $\theta$에 대해 최적화됩니다. 내부 레벨 문제는 일반적으로 비볼록(non-convex)이며, 외곽 레벨 목적 함수는 내부 문제의 최적해 $\theta^*(x)$에 의존합니다.

기존 방식들은 다음과 같은 상당한 한계에 직면해 있습니다:

1. 페널티 기반 방법(예: Gaur et al., 2025)은 종종 근사된 그래디언트에 의존하며 외곽 레벨 목적 함수에 대한 폴야크-로자시츠(Polyak-Łojasiewicz, PL) 조건을 요구하므로, $\tilde{O}(\epsilon^{-3})$의 하위 최적 샘플 복잡도를 가집니다.
2. 헤시안 기반 하이퍼그레이디언트 방법(예: Chakraborty et al., 2024)은 PL 가정을 피하지만, 헤시안 행렬을 계산하고 역행렬을 구해야 하는 필요성 때문에 확장성 문제를 겪으며, 이는 거대한 파라미터 공간에서 계산적으로 불가능합니다.
3. 유한 공간 방법(예: Yang et al., 2025)은 볼츠만 정책 최적성을 활용하여 헤시안 프리(Hessian-free) 그래디언트를 도출하지만, 연속 공간에서의 파라미터화된 비실현적(non-realizable) 정책 클래스로는 직접 확장되지 않습니다.

핵심 과제는 외곽 레벨 목적 함수에 대해 PL 가정을 요구하지 않거나 실현 가능한 정책 클래스(정책 클래스가 모든 볼츠만 정책을 표현할 수 있는 경우)를 요구하지 않으면서도, 최첨단 샘플 복잡도를 달성할 수 있는 확장 가능한 헤시안 프리 바이레벨 RL 알고리즘을 개발하는 것입니다.

방법론

저자들은 근사 하이퍼그레이디언트 최적화(Approximate Hypergradient Optimization, AHO) 알고리즘을 제안합니다. 이 방법론은 세 가지 핵심적인 이론적 및 알고리즘적 구성 요소에 의존합니다:

1. 볼츠만 정책 최적성 활용

내부 레벨 목적 함수는 엔트로피 정규화된 할인 보상 함수로 정의됩니다. 본 논문은 이 특정 목적 함수에 대한 볼츠만 정책($\pi_B$)의 알려진 최적성을 활용합니다. 정책 클래스 $\Theta$는 일반적으로 비실현적(즉, 최적의 볼츠만 정책을 완벽하게 표현할 수 없음)이지만, 저자들은 볼츠만 정책의 최적성을 사용하여 그래디언트를 도출할 수 있는 조건을 확립합니다.

2. 그래디언트 시프트 가치 함수 (Gradient-Shifted Value Functions)

하이퍼그레이디언트 식에서 흔히 나타나는 헤시안 항(Implicit Function Theorem을 통해 유도됨)을 피하기 위해, 저자들은 그래디언트 시프트 가치 함수($U^\pi$ 및 $W^\pi$)를 도입합니다.

• $U^\pi(s, \theta, x)$와 $W^\pi(s, a, \theta, x)$는 볼츠만 정책 하에서의 보상 그래디언트 $\nabla r(s, a, x)$에 대한 장기 할인 합계로 정의됩니다.
• 이 함수들은 볼츠만 가치 함수($\nabla V^{\pi,B}$)와 Q-가치 함수($\nabla Q^\pi$)의 대리물(surrogates) 역할을 합니다.
• 근사 오차: 정책 클래스의 비실현성으로 인해, 이러한 대리 함수와 실제 그래디언트 사이에는 근사 오차가 존재합니다. 이 오차는 정책 클래스가 실현 가능하다면 0이 되는 $\epsilon_{kl}$(KL 발산 관련) 및 $\epsilon_{fd}$(그래디언트 차이 관련) 항들에 의해 경계가 정해집니다.

3. 근사 하이퍼그레이디언트 정리

저자들은 외곽 레벨 목적 함수에 대한 근사 하이퍼그레이디언트를 도출합니다. 실제 그래디언트를 그래디언트 시프트 함수로 대체함으로써, 헤시안 역행렬 항을 제거하는 폐쇄형(closed-form) 식을 얻습니다. 결과적인 하이퍼그리디언트는 다음으로 구성됩니다:

• 선호도 손실(preference loss)에 대한 $x$의 직접적인 그래디언트.
• 온도 계수 $\tau$에 의해 스케일링된, 그래디언트 시프트 Q-가치 함수와 가치 함수의 차이를 포함하는 항.

4. 알고리즘 (AHO)

AHO 알고리즘(Algorithm 1)은 중첩 루프(nested loop)로 작동합니다:

• 내부 루프: 추정된 Q-가치 함수를 사용하여 정책 파라미터 $\theta$를 업데이트하는 표준 정책 그래디언트 방법을 사용합니다.
• 외곽 루프: 근사 하이퍼그레이디언트를 사용하여 보상 파라미터 $x$를 업데이트합니다. 이는 수집된 궤적(trajectories)을 기반으로 한 오프 폴리시(off-policy) Q-러닝 유형의 업데이트(Algorithm 3 및 4)를 사용하여 그래디언트 시프트 함수($U^\pi_{off}$ 및 $W^\pi_{off}$)를 추정하는 과정을 포함합니다.

주요 기여

1. 비실현적 클래스를 포함한 볼츠만 최적성: 본 논문은 파라미터화된 비실현적 정책 클래스에서도 볼츠만 정책의 최적성을 적용할 수 있는 조건(Assumption 3.1)을 제공하며, 그로 인한 근사 오차가 수렴에 미치는 영향을 논의합니다.
2. 확장 가능한 헤시안 프리 알고리즘: 제안된 AHO 알고리즘은 완전한 헤시안 프리 방식으로, 헤시안 기반 방법들이 실패하는 고차원 파라미터 공간에서도 확장 가능합니다.
3. 개선된 샘플 복잡도: 알고리즘은 $O(\epsilon^{-1})$의 반복 복잡도와 $\tilde{O}(\epsilon^{-2})$의 샘플 복잡도를 달성합니다. 이는 페널티 기반 방법(Gaur et al., 2025)의 기존 최첨단 샘플 복잡도인 $\tilde{O}(\epsilon^{-3})$를 개선한 것입니다.
4. 완화된 가정: 기존의 최첨단 연구(Gaur et al., 2025)와 달리, AHO는 외곽 레벨 목적 함수에 대한 PL 가정을 제거했습니다. 또한 내부 문제의 유일한 최소해(unique minimizer) 요구 사항을 완화하여, 국소적으로 고립되고 비퇴화된(non-degenerate) 최소해만을 요구합니다.

결과 및 분석

이론적 보장

완만한 정칙성 조건(Lipschitz 연속성, 유계 보상/로그 확률, 그리고 내부 레벨에 대한 PL 조건) 하에서, 저자들은 평균 제곱 외곽 레벨 그래디언트가 근사 오차 및 샘플링 파라미터에 의존하는 바운드로 수렴함을 증명합니다:
$$\frac{1}{T} \sum_{t=0}^{T-1} \|\nabla \Phi(x_t)\|^2 \leq O(\epsilon) + O(\epsilon_{approx}) + O(\epsilon_{fd}) + O(\epsilon_{kl})$$
$T = \Theta(\epsilon^{-1})$, $B = \Theta(\epsilon^{-1})$, $K = \Theta(\log(\epsilon^{-1}))$, $H = \Theta(\log(\epsilon^{-1}))$로 설정함으로써, 샘플 복잡도는 $\tilde{O}(\epsilon^{-2})$로 확립됩니다.

실험적 평가

저자들은 선호도 기반 바이레벨 RL 문제로 정식화된 DM Control 스위트의 연속 제어 작업(walker-walk 및 cheetah-run)에서 AHO를 평가합니다.

• 베이스라인: AHO를 Gaur et al. (2025)의 페널티 기반 바이레벨 RL 알고리즘과 비교합니다.
• 비교 제약 사항: 저자들은 PARL(Chakraborty et al., 2024)의 헤시안 기반 접근 방식이 사용된 네트워크 크기에서 계산적으로 불가능하기 때문에(약 $10^{12}$개의 항목을 가진 헤시안 역행렬 필요) PARL과는 비교하지 않는다고 명시합니다.
• 결과: AHO는 두 작업 모두에서 Gaur et al. (2025) 베이스라인보다 높은 총 평가 보상(total evaluation reward)을 달성하여, 샘플 복잡도의 이론적 개선을 경험적으로 뒷받침합니다.

의의

본 논문은 AHO가 확장성과 샘플 효율성 사이의 트레이드오프를 해결함으로써 바이레벨 RL 분야에서 중요한 진전을 이루었다고 주장합니다.

• 확장성: 헤시안을 제거함으로써, 이 방법은 선호도 기반 RL에서 사용되는 현대의 대규모 신경망 정책에도 적용 가능해졌습니다.
• 효율성: 엔트로피 정규화된 목적 함수의 구조와 볼츠만 최적성을 활용함으로써, 페널티 방법의 고유한 문제인 그래디언트 근사와 샘플링 오차 간의 트레이드오프를 피하고 더 나은 샘플 복잡도($\tilde{O}(\epsilon^{-2})$ vs. $\tilde{O}(\epsilon^{-3})$)를 달もの합니다.
• 이론적 견고성: 외곽 레벨에 대한 PL 조건을 제거함으로써, RL-HF 및 메타 학습에서 흔히 나타나는 다양한 비볼록 외곽 목적 함수에 대한 알고리즘의 적용 범위를 넓혔습니다.