Recovering Hidden Reward in Diffusion-Based Policies

"Recovering Hidden Reward in Diffusion-Based Policies" 논문 (ENERGYFLOW 도입) 에 대한 상세한 기술적 요약입니다.

1. 문제 제기

확산 기반 정책 (Diffusion-based policies) 은 복잡한 다중 모드 행동 분포를 모델링할 수 있는 능력 덕분에 전문가 시연으로부터 로봇 조작 기술을 학습하는 데 있어 최첨단 패러다임이 되었습니다. 그러나 이러한 정책들은 일반적으로 **행동 복제 (Behavior Cloning, BC)**에 의존하는데, 이는 근본적인 의도나 보상 구조를 명시적으로 모델링하지 않고 행동을 모방하는 방식으로 학습합니다.

• 표준 확산의 한계: 명시적인 에너지 표현이 부족하여 훈련 분포에서 벗어난 테스트 시간 상황 (Out-of-Distribution 또는 OOD) 에서 실패하기 쉽습니다. 불안정한 적대적 학습이나 계산 비용이 많이 드는 역강화학습 (Inverse Reinforcement Learning, IRL) 에 의존하지 않고서는 하류 강화학습 (Reinforcement Learning, RL) 을 위한 보상 신호를 쉽게 제공할 수 없습니다.
• 격차: 생성적 행동 모델링과 보상 복구를 통합하여 고품질 모방과 정책 정제를 위한 유효한 보상 신호 추출을 모두 가능하게 하되, 일반화에 대한 이론적 보장을 유지하는 프레임워크가 필요합니다.

2. 방법론: ENERGYFLOW

저자들은 스칼라 에너지 함수 $E_\theta(s, a)$를 매개변수화하는 ENERGYFLOW를 제안합니다. 여기서 이 함수의 기울기는 행동 생성을 위한 탈노이즈 필드로 작용합니다.

핵심 아키텍처

• 스칼라 에너지 매개변수화: 표준 확산 정책이 벡터 값 노이즈 예측 $\epsilon_\theta(s, a)$를 직접 회귀하는 것과 달리, ENERGYFLOW 는 스칼라 에너지 함수 $E_\theta(s, a, t)$를 학습합니다.
• 보수적 벡터 필드: 점수 함수 (탈노이즈 필드) 는 자동 미분을 통해 유도됩니다:
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  구성상 이 벡터 필드는 보수적 (회전이 0 인, $\nabla \times S = 0$) 이며, 즉 스칼라 포텐셜의 기울기임을 보장합니다. 이는 적분 가능성을 강제하여 행동 선호도가 전달적이며 물리적으로 실현 가능하도록 합니다.
• 학습: 모델은 **탈노이즈 점수 매칭 (Denoising Score Matching, DSM)**을 사용하여 학습됩니다. 목적 함수는 학습된 기울기와 실제 노이즈 방향 사이의 차이를 최소화합니다:
  $$\mathcal{L} = \mathbb{E}_{t, a_0, \epsilon} \left[ \sigma^2(t) \left\| -\nabla_{a_t} E_\theta(a_t, s, t) + \frac{\epsilon}{\sigma(t)} \right\|^2 \right]$$
• 추론: 행동은 학습된 에너지 함수의 기울기를 사용하여 확률 흐름 ODE 를 적분함으로써 생성됩니다.

보상 추출

이 프레임워크는 학습된 에너지 지형에서 직접 보상 신호를 추출합니다. 상태 의존적 적분 상수 $c(s)$ (이는 서로 다른 상태 간 원시 에너지 값을 모호하게 만듦) 를 처리하기 위해 저자들은 **중심화된 형성 (Centered Shaping)**을 제안합니다:
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
이는 상태 의존적 기준선 (몬테카를로 샘플링을 통해 추정됨) 을 빼서 보상을 상태 내에서의 상대적 행동 선호도만 반영하도록 정규화합니다.

3. 주요 이론적 기여

이 논문은 확산 점수 매칭과 보상 복구 사이의 연결을 확립하는 엄격한 이론적 증명을 제공합니다:

1. 점수 - 보상 동등성: 최대 엔트로피 최적성 가정 하에, 확산 모델이 학습한 점수 함수는 전문가의 소프트 Q-함수의 기울기에 비례함을 ($\nabla_a Q^* \propto S^*$) 증명합니다. 따라서 점수 매칭은 적대적 min-max 최적화 없이 보상 기울기를 암묵적으로 복구합니다.
2. 보수적 제약을 통한 복잡성 감소: 저자들은 가설 공간을 스칼라의 기울기인 보수적 필드로 제한함으로써 비제약 벡터 필드에 비해 Rademacher 복잡성을 엄격히 감소시킨다고 증명합니다. 구체적으로, 보수적 필드의 복잡성 상한은 표현의 매끄러움 ($L$) 에 비례하는 반면, 비제약 필드는 출력 차원 ($\sqrt{d}$) 에 비례합니다.
3. OOD 일반화 상한: 감소된 복잡성을 활용하여 논문은 Out-of-Distribution (OOD) 시나리오에 대한 더 엄격한 일반화 상한을 유도합니다. 보수적 제약은 기하학적 정규화제 역할을 하여 모델이 임의의 벡터 방향을 학습하는 것을 방지하고 에너지 지형의 모양을 외삽하도록 강제함으로써 분포 변화 하에서 더 나은 견고성을 이끌어냅니다.
4. 식별 가능성과 오차 전파: 논문은 절대 Q-값은 (알려지지 않은 상태 의존적 상수 때문에) 고유하게 식별될 수 없지만, 상태 내 행동 순위는 정확함을 확립합니다. 또한, 점수 추정 오차가 행동 선호도로 전파되는 것을 제한하여 오차가 행동 간 거리에 비례하여 선형적으로 증가함을 (리프시츠 연속성) 보여줍니다.

4. 실험 결과

저자들은 RoboMimic 및 Meta-World 벤치마크와 실제 로봇 작업에서 ENERGYFLOW 를 평가했습니다.

• 모방 성능 (RQ1): ENERGYFLOW 는 최첨단 (SOTA) 성능을 달성했습니다.
  • RoboMimic: 평균 성공률 93.8% (Diffusion Policy 의 91.2% 대비).
  • Meta-World: 평균 성공률 92.5% (Diffusion Policy 의 90.7% 대비).
  • 다른 에너지 기반 방법 (Implicit BC 등) 과 IRL 베이스라인을 크게 능가했습니다.
• 실제 배포 (RQ2): "Bottle" 및 "Drawer" 작업을 위해 AGIBOT G1 로봇에 배포되었습니다. ENERGYFLOW 는 베이스라인에 비해 부드러운 궤적과 더 적은 망설임으로 100% 성공률을 달성했습니다.
• 하류 RL 을 위한 보상 품질 (RQ3): Soft Actor-Critic (SAC) 에이전트를 훈련시키기 위한 보상 신호로 사용될 때:
  • "Centered Energy" 보상은 SAC 가 오라클에 가까운 성능에 도달하도록 하여 희소 보상 및 원시 에너지 보상보다 크게 우월했습니다.
  • 이는 초기 훈련의 정체기를 방지하는 밀집된 단계별 지침을 제공했습니다.
• OOD 일반화 (RQ4): 초기 위치 교란이 증가하는 조건에서 ENERGYFLOW 는 비제약 확산 및 흐름 정책이 크게 저하되는 동안 견고한 성공률을 유지했습니다. 이는 보수적 제약의 이론적 이점을 확인시켜 줍니다.
• 효율성 (RQ6): ENERGYFLOW 는 명시적 에너지 함수의 이점을 제공하면서도 기존 에너지 기반 모델 (EBM) 의 느린 샘플링 시간을 피하면서 비에너지 기반 Flow Policies 와 유사한 추론 속도 (~10ms 지연) 를 달성했습니다.

5. 중요성과 영향

• 모방과 RL 의 통합: ENERGYFLOW 는 행동 복제와 강화학습 사이의 격차를 해소합니다. 정책을 모방을 통해 훈련하면서도 추가 정제를 위한 유효한 보상 신호를 동시에 노출시켜 수동 보상 설계나 불안정한 적대적 IRL 의 필요성을 제거합니다.
• 확산에 대한 이론적 보장: 이 연구는 확산 정책에서 보수적 필드 (적분 가능성) 를 강제하는 것이 왜 유익한지에 대한 최초의 이론적 근거를 제공하며, 이를 감소된 가설 복잡성과 향상된 OOD 일반화와 연결합니다.
• 실용적 유용성: 이 방법은 계산적으로 효율적이며 고차원 행동 공간으로 확장 가능하고 실제 로봇 환경에서 효과적이어서 제한된 시연으로부터 복잡한 조작 기술을 학습하기 위한 견고한 해결책을 제공합니다.

요약하자면, ENERGYFLOW는 유효한 보상 추출에 필요한 구조적 제약 (보수적 필드) 이 확산 기반 정책의 일반화와 견고성을 향상시키는 강력한 귀납적 편향으로 동시에 작용함을 보여줍니다.