기술 요약: 점수 기반 단일 단계 MeanFlow 정책 최적화 (SOM)

1. 문제 제기

연속 행동 공간에서의 강화 학습 (RL) 은 단일 모드 가우시안 정책의 표현력 한계를 극복하기 위해 확산 (diffusion) 및 흐름 매칭 (flow matching) 과 같은 생성 모델을 increasingly 채택하고 있습니다. 이러한 모델은 복잡하고 다중 모드인 행동 분포를 포착하여 더 풍부한 탐색과 다양한 최적 행동을 가능하게 합니다. 그러나 중요한 병목 현상이 남아 있습니다. 표준 확산 및 흐름 매칭 정책은 단일 행동을 생성하기 위해 **반복적인 다단계 탈노이즈 (iterative, multi-step denoising)**를 필요로 합니다. 이는 추론 시 상당한 계산 오버헤드를 초래하며, 특히 정책 추론이 데이터 수집과 지속적으로 교차하는 온라인 RL에서 치명적입니다.

MeanFlow는 평균 속도 장 (average velocity field) 을 학습하여 샘플링을 단일 네트워크 평가로 줄이는 해결책으로 제안되었지만, 온라인 RL 에 직접 적용하는 것은 근본적인 요구 사항으로 인해 방해받습니다. 즉, 목표 속도 장을 구성하려면 일반적으로 목표 분포 (최적 정책) 에서의 샘플이 필요합니다. 온라인 RL 에서는 에이전트가 정책과 동시에 목표 분포를 학습해야 하므로, 이러한 실제 샘플 (ground-truth samples) 을 사용할 수 없습니다. 이를 우회하려는 기존 접근법들 (예: Best-of-N 샘플링) 은 행동 차원이 증가함에 따라 확장성 문제를 겪는 경우가 많습니다.

2. 방법론

저자들은 명시적인 목표 분포 샘플 없이 완전한 온라인 RL 에서 MeanFlow 학습을 가능하게 하는 **점수 기반 단일 단계 MeanFlow 정책 최적화 (SOM)**라는 액터-크리틱 알고리즘을 제안합니다. 핵심 혁신은 Q-함수를 통해 점수 추정 (score estimation) 으로 목표 속도 장을 직접 구성하는 데 있습니다.

2.1 시간 의존적 에너지 함수

이 방법은 음수 Q-함수를 목표 분포의 에너지로 간주합니다. 엔트로피 정규화 RL 프레임워크 하에서 최적 정책은 볼츠만 분포를 따릅니다:
$$\pi^*(a|s) \propto \exp(\alpha Q(s, a))$$
여기서 $\alpha$는 온도 매개변수입니다. 저자들은 확산 시간 단계 $t \in [0, 1]$ 전반에 걸쳐 이 볼츠만 구조를 유지하는 시간 의존적 에너지 함수 $Q_t(s, a_t)$를 정의합니다. 임의의 시간 단계 $t$에서 분포의 점수 (score) 는 다음과 같이 유도됩니다:
$$\nabla_{a_t} \log p_t(a_t|s) = \nabla_{a_t} Q_t(s, a_t)$$
중요하게도, $Q_t$를 폐쇄형으로 계산하는 것은 불가능하지만, 이 방법은 근사 가능한 그 점수 (gradient) 만을 필요로 합니다.

2.2 iDEM 적응을 통한 점수 추정

목표 샘플 없이 점수 $\nabla_{a_t} Q_t(s, a_t)$를 추정하기 위해 SOM 은 반복 탈노이즈 에너지 매칭 (Iterated Denoising Energy Matching, iDEM) 추정기를 적응합니다. 실제 목표 분포에서 샘플링하는 대신, 이 방법은 현재 노이즈 상태 $a_t$를 중심으로 한 가우시안 제안 (proposal) 에서 후보 행동을 샘플링합니다. 이러한 후보들은 Q-값에 기반하여 재가중치되어 사후 기대값을 근사합니다. 결과적으로 도출된 점수 추정기는 다음과 같습니다:
$$\nabla_{a_t} Q_t(s, a_t) \approx \nabla_{a_t} \log \sum_{i=1}^K \exp(\alpha Q(s, a^{(i)}_{0|t}))$$
여기서 $a^{(i)}_{0|t}$는 $a_t$에 조건부인 가우시안 분포에서 추출된 샘플들입니다. 이는 노이즈를 통한 탐색을 유지하면서 행동 업데이트를 고가치 영역으로 효과적으로 편향시킵니다.

2.3 속도 장 구성

추정된 점수는 목표 속도 장을 유도하기 위해 **확률 흐름 ODE (Probability Flow ODE, PF-ODE)**에 통합됩니다. 분산 보존 (Variance Preserving, VP) SDE 의 경우, 속도 장은 다음과 같이 정의됩니다:
$$v_t(a_t|s) = -\frac{1}{2}\beta(t) \left( a_t + \nabla_{a_t} Q_t(s, a_t) \right)$$
이 속도 장은 노이즈를 현재 Q-함수에 의해 유도된 고가치 모드 (modes) 로 직접 수송합니다. 그런 다음 MeanFlow 정책 $u_\theta$는 이 구성된 목표 속도를 사용하여 MeanFlow 항등식을 만족하도록 훈련되며, 다음 손실을 최소화합니다:
$$\mathcal{L}(\theta) = \mathbb{E} \left[ \| u_\theta(a_t, r, t, s) - \text{sg}(u^{\text{new}}_{\text{tgt}}) \|^2_2 \right]$$
여기서 $u^{\text{new}}_{\text{tgt}}$는 PF-ODE 속도에서 유도됩니다.

2.4 구현 세부 사항

• 정규화: 점수 기울기는 단위 노름으로 정규화되고, 다양한 보상 규모에 걸쳐 훈련을 안정화하기 위해 조정 가능한 계수 $w$로 재조정됩니다.
• 크리틱 훈련: 크리틱은 트윈 Q-네트워크를 사용하여 분산형 벨만 목적함수 (distributional Bellman objective) 로 업데이트됩니다.
• Best-of-N (BoN): 주요 기여는 점수 기반 속도이지만, 저자들은 크리틱 업데이트 동안 BoN 샘플링을 통합하여 여러 샘플 중 가장 높은 Q-값을 가진 후보를 선택함으로써 목표 값 추정을 가속화합니다.

3. 주요 기여

1. 병목 현상 식별: 이 논문은 목표 분포 샘플의 부재를 온라인 RL 에 MeanFlow 를 적용하는 주요 장벽으로 식별합니다.
2. 점수 기반 속도 구성: 명시적인 목표 샘플의 필요성을 우회하기 위해 Q-함수 기울기에서 직접 목표 속도 장을 구성하는 새로운 방법을 제안합니다.
3. SOM 알고리즘: MeanFlow 를 온라인 RL 에 통합하여 단일 단계 행동 샘플링으로 표현력 있는 정책 학습을 달성하는 액터-크리틱 알고리즘 SOM 을 도입합니다.
4. 실증적 검증: SOM 이 반복적 확산 및 흐름 매칭 기준 방법들에 비해 훈련 및 추론 시간을 크게 단축하면서 보행 작업에서 최첨단 성능을 달성함을 입증합니다.

4. 실험 결과

저자들은 MuJoCo v4 연속 제어 작업 (Ant, HalfCheetah, Hopper, Humanoid, Walker2d, Swimmer, Pusher, Reacher, InvertedDoublePendulum 포함) 9 개에서 SOM 을 평가했습니다.

• 성능: SOM 은 5 개 보행 벤치마크 중 4 개에서 가장 높은 평균 반환 (mean return) 을 달성했으며, Hopper-v4 에서는 두 번째로 높은 성능을 보였습니다. 개선은 Humanoid-v4(높은 행동 차원성) 에서 가장 두드러졌는데, SOM 이 두 번째로 좋은 방법보다 명확한 차이로 우세하여 기울기 기반 접근법이 차원에 따라 유리하게 확장됨을 시사합니다.
• 효율성: 표 2 에所示된 바와 같이, SOM 은 HalfCheetah-v4 에서 0.218 ms의 단일 단계 추론 시간을 달성하는 반면, 반복적 기준 방법들 (10100 단계의 탈노이즈 필요) 은 0.61.5 ms 가 소요됩니다. 훈련 시간도 크게 단축되었습니다.
• 정책 분석: 시각화 (t-SNE 및 Q-지형도) 는 SOM 이 행동 샘플을 고 Q-모드 주변에 밀집시키는 반면, 경쟁하는 단일 단계 방법들 (MFP 등) 은 단일 모드로 붕괴되거나 확산적으로 분포하는 경향이 있음을 나타냅니다.
• 강건성: 노이즈가 있는 Q-함수를 사용한 실험은 SOM 의 가우시안 평활화 점수 추정기가 거친 크리틱 오류를 효과적으로 평균화하면서도 고가치 모드 근처의 국소 구조를 보존하여 강건하게 유지됨을 보여주었습니다.

5. 중요성 및 주장

이 논문은 SOM 이 생성 정책의 표현력과 온라인 RL 의 계산 제약 사이의 긴장을 해소한다고 주장합니다. 목표 속도를 Q-함수에서 직접 유도함으로써 SOM 은 목표 샘플이 없는 완전한 온라인 환경에서 MeanFlow 정책을 훈련할 수 있게 합니다.

저자들은 SOM 이 단일 생성 단계로 최첨단 성능을 달성하여 추론 지연 시간이 중요한 실시간 제어 응용 분야에 실용적인 해결책을 제공한다고 강조합니다. 현재 이 방법은 연속 제어 작업 (MuJoCo) 으로 제한되어 있으며 학습된 크리틱의 품질에 의존하지만, 이는 효율적이고 표현력 있는 정책 최적화를 향한 중요한 단계라고 지적합니다. 이 연구는 향후 방향이 더 높은 차원 또는 부분 관측 가능 환경 (예: 시각 제어) 으로 SOM 을 확장하고 더 고급 불확실성 인지 크리틱을 통합하는 것을 포함할 것이라고 시사합니다.