기술 요약: 확률적 MeanFlow 정책 (SMFP)

1. 문제 제기

고차원 연속 제어 분야의 온라인 오프-정책 강화 학습 (RL) 은 정책 표현력과 최적화 실용성 사이의 근본적인 트레이드오프에 직면해 있습니다.

• 정책 표현: 표준 가우시안 정책은 빠른 추론과 실용적인 엔트로피 추정을 제공하지만, 복잡하고 다중 모달 (multimodal) 인 행동 분포를 모델링하지 못합니다. 반면, 생성 정책 (예: 확산 또는 흐름 기반) 은 풍부한 다중 모달 표현력을 제공하지만, 비실용적인 가능도, 어려운 엔트로피 추정, 그리고 반복적 샘플링으로 인한 높은 추론 지연 시간을 겪습니다.
• 최적화 방법론: 광범위한 탐색을 장려하지만 불안정성을 초래할 수 있는 엔트로피 정규화 탐색(예: Soft Actor-Critic, SAC) 과 업데이트를 이전 정책과 가깝게 유지함으로써 안정적이고 보수적인 정책 개선을 보장하는 미러 디센트 (MD) 업데이트 사이에는 긴장 관계가 존재합니다.
• 핵심 갈등: SAC 스타일의 엔트로피 최대화와 MD 제약을 결합하면 다중 모달 타겟 분포(가치 유도 볼츠만 분포와 이전 정책의 복합체) 가 생성됩니다. 표준 단모달 가우시안 정책은 이러한 타겟을 표현할 수 없어 표현 불일치가 발생합니다. 표현력 있는 생성 정책이 이론적으로 이 간극을 메울 수 있지만, 일반적으로 명시적 엔트로피 제어를 위한 실용적인 메커니즘이 부족하여 엔트로피 주도 탐색과 미러 디센트 안정성의 결합은 실제로 불가능합니다.

2. 방법론: 확률적 MeanFlow 정책 (SMFP)

저자들은 MeanFlow 기반의 노이즈 - 행동 매핑과 가우시안 재매개화를 결합하여 위의 갈등을 해결하도록 설계된 **확률적 MeanFlow 정책 (SMFP)**이라는 원스텝 생성 정책 프레임워크를 제안합니다.

2.1 확률적 재매개화

SMFP 는 결정론적 MeanFlow 항등식을 확률적 매핑으로 재구성합니다. 순수한 결정론적 속도장 대신, 정책 $g_\theta$는 평균 행동과 조건부 표준 편차를 모두 예측합니다:
$$g_\theta(a_t, b, t) = \underbrace{a_t - u_\theta(a_t, b, t)}_{\text{평균 } \mu} + \underbrace{\sigma_\theta(a_t, b, t) \odot \epsilon}_{\text{확률성 } \sigma}$$
여기서 $a_t$는 노이즈와 행동 사이의 보간값이며, $u_\theta$는 학습된 평균 속도이고, $\epsilon$은 표준 가우시안 노이즈입니다. 이 구조를 통해 정책은 **단일 순전파 (one-step inference)**로 행동을 출력하면서도 조건부 가우시안 분포를 유지할 수 있습니다.

2.2 실용적 엔트로피 대리 변수

주요 혁신은 실용적 엔트로피 대리 변수의 유도입니다. 생성 정책의 주변 엔트로피는 실용적이지 않으므로, SMFP 는 잠재 노이즈 $e$에 대한 조건부 엔트로피를 기반으로 한 하한을 최대화합니다.

• 목적 함수는 예측된 가우시안의 기대 로그 표준 편차를 사용합니다: $\mathbb{E}[\sum \log \sigma^{(i)}_\theta]$.
• 분산 붕괴를 방지하면서도 무제한의 엔트로피 증가를 유도하지 않기 위해 힌지 스타일 엔트로피 정규화가 적용됩니다. 이는 정규화된 로그 표준 편차가 임계값 $\kappa_\sigma$ 미만으로 떨어질 때만 활성화되어 탐색을 위한 최소 노이즈 크기를 보장합니다.

2.3 오프-정책 미러 디센트 최적화

SMFP 는 이 엔트로피 제어를 오프-정책 미러 디센트와 통합합니다. 최적화 목적 함수는 다음을 결합합니다:

1. 가치 최대화: 표준 Q-학습 손실.
2. 엔트로피 정규화: 위에서 설명한 실용적 조건부 엔트로피 대리 변수.
3. 미러 디센트 제약: 이전 정책 $\pi_{\theta_{old}}$에 업데이트를 고정하는 KL-발산 항.

미러 디센트의 타겟 분포는 비모수적이고 정규화되지 않았으므로, SMFP 는 가중 회귀 방식을 사용합니다. 이전 정책에서 행동을 샘플링하고, 이들을 이점 기반 항 (근사적으로 $\max(0, Q(s,a) - V(s))$) 으로 가중치하여 생성 정책을 최적 타겟에 투영합니다. 이는 타겟의 명시적 밀도 평가를 필요로 하지 않습니다.

2.4 가치 유도 행동 선택

분산을 더 줄이기 위해 SMFP 는 Best-of-K 샘플링을 활용합니다. 주어진 상태에 대해 $K$개의 행동 후보를 병렬로 생성 (JAX vmap 사용) 하고, 가장 높은 Q-값을 가진 행동을 선택합니다. 이는 무시할 수 있는 추론 오버헤드로 경량 거부 필터 역할을 합니다.

3. 주요 기여

1. SMFP 아키텍처: MeanFlow 표현력과 가우시안 재매개화를 결합한 원스텝 확률적 생성 정책의 도입으로, 단일 스텝 추론을 가능하게 하면서 명시적 엔트로피 제어를 지원합니다.
2. 실용적 엔트로피 제어: 생성 정책과 함께 엔트로피 정규화 미러 디센트 RL 을 실제로 효과적으로 만드는 조건부 엔트로피 하한 유도.
3. 통합 목적 함수: 엔트로피 주도 탐색, 가치 가중 미러 디센트, 원스텝 생성 추론을 공동으로 최적화하는 통합 손실 함수로, 다중 모달 타겟과 단모달 정책 간의 표현 불일치를 해결합니다.
4. 실증적 검증: 가우시안 전용 설정에서의 결과와 달리, 표현력 있는 생성 정책이 지원될 때 SAC 스타일 탐색과 미러 디센트의 결합이 유익함을 입증했습니다.

4. 실험 결과

저자들은 30 개의 무작위 시드를 사용하여 7 개의 MuJoCo 연속 제어 벤치마크(Hopper, Walker2D, Ant, HalfCheetah, Humanoid, HumanoidStandup, Swimmer) 에서 SMFP 를 평가했습니다.

• 성능: SMFP 는 모든 벤치마크에서 최고의 집계 성능을 달성하여 고전적 가우시안 베이스라인 (PPO, TD3, SAC) 과 최첨단 생성 베이스라인 (DIME, QVPO 와 같은 확산 기반; SAC Flow, FPMD 와 같은 흐름 기반) 을 모두 능가했습니다.
• 효율성: 16~20 개의 샘플링 스텝이 필요한 반복적 확산 정책과 달리, SMFP 는 **단일 스텝 추론 (NFE=1)**을 유지합니다. NVIDIA RTX 5090 에서 SMFP 는 무시할 수 있는 추론 지연 시간 (0.42ms) 으로 강력한 성능을 달성했는데, 이는 표준 가우시안 정책과 유사하며 반복적 생성 방법보다 훨씬 빠릅니다.
• 절대 실험 (Ablation Studies):
  • 엔트로피 정규화기를 제거하면 성능이 저하되어 확률적 탐색의 필요성이 확인되었습니다.
  • 미러 디센트 항을 제거하면 특히 Humanoid 와 같은 고차원 작업에서 안정성이 감소했습니다.
  • 이 방법은 엔트로피 온도 ($\alpha$) 와 미러 디센트 계수 ($\lambda$) 의 하이퍼파라미터 변화에 대해 강건함을 보였습니다.

5. 중요성과 주장

이 논문은 SMFP 가 표현력 있는 생성 모델링과 안정적인 엔트로피 정규화 최적화 사이의 간극을 성공적으로 연결함으로써 온라인 오프-정책 RL에서 중요한 진전을 이루었다고 주장합니다.

• 실용적 효과성: 저자들은 SMFP 가 "엔트로피 정규화 미러 디센트 RL 을 실제로 효과적으로 만든 첫 사례"라고 명시합니다. 다중 모달 타겟을 표현할 수 없다는 한계로 인해 가우시안 매개화 하에서 이러한 목적을 결합하려는 이전 시도는 실패했습니다. SMFP 는 그 생성적 특성으로 이를 극복합니다.
• 균형 잡힌 트레이드오프: 이 연구는 생성 정책의 다중 모달 표현력과 미러 디센트의 안정성을 실시간 제어에 필요한 추론 속도를 희생하지 않고 유지할 수 있음을 보여줍니다.
• 한계: 저자들은 겸손하게도 실용적 대리 변수는 MeanFlow 의 재매개화된 구조에 특정적으로 의존한다고 지적합니다. 이를 다른 생성 클래스로 확장하려면 새로운 엔트로피 경계가 필요할 수 있습니다. 또한, 현재 구현은 $\alpha$와 $\lambda$에 대해 고정된 하이퍼파라미터를 사용하므로, 적응형 가중치 방식은 향후 과제로 남습니다.

요약하자면, SMFP 는 표현력 있는 생성 정책의 상호 보완적 장점과 효율적인 가우시안 기반 엔트로피 제어를 보존하는 원칙적인 프레임워크를 제공하여, 복잡한 연속 도메인에서 안정적이고 고성능인 온라인 RL 을 가능하게 합니다.