Recovering Hidden Reward in Diffusion-Based Policies

This paper introduces EnergyFlow, a framework that unifies generative action modeling with inverse reinforcement learning by parameterizing a scalar energy function whose gradient is the denoising field. We establish that under maximum-entropy optimality, the score function learned via denoising score matching recovers the gradient of the expert's soft Q-function, enabling reward extraction without adversarial training. Formally, we prove that constraining the learned field to be conservative reduces hypothesis complexity and tightens out-of-distribution generalization bounds. We further characterize the identifiability of recovered rewards and bound how score estimation errors propagate to action preferences. Empirically, EnergyFlow achieves state-of-the-art imitation performance on various manipulation tasks while providing an effective reward signal for downstream reinforcement learning that outperforms both adversarial IRL methods and likelihood-based alternatives. These results show that the structural constraints required for valid reward extraction simultaneously serve as beneficial inductive biases for policy generalization. The code is available at https://github.com/sotaagi/EnergyFlow.

📋 한 줄 요약

**[모방학습 / 역강화학습]** 확산 정책의 디노이징 필드를 스칼라 에너지 함수의 그라디언트로 재해석해 적대적 학습 없이 보상까지 회복하는 EnergyFlow를 제안한다.

🎯 핵심 기여도

디노이징 score matching으로 학습된 score 함수가 최대 엔트로피 가정 하에서 전문가의 soft Q-함수의 그라디언트를 회복함을 이론적으로 증명
학습된 필드를 보존장(conservative)으로 제약하면 가설 복잡도가 줄고 OOD 일반화 경계가 좁아짐을 보장
회복된 보상의 식별 가능성과 score 추정 오차가 액션 선호도에 전파되는 양을 정량적으로 특성화
다양한 매니퓰레이션 태스크에서 SOTA 모방 성능을 달성하면서 다운스트림 RL 학습용으로도 효과적인 보상 신호 제공

💡 핵심 아이디어

확산 정책의 디노이징 필드는 보통 단순한 행동 생성 도구로만 쓰였지만, 실제로는 어떤 잠재적 에너지 함수의 그라디언트로 해석될 수 있다. 이 에너지 관점에서 보면 모방 학습과 역강화학습이 자연스럽게 통합되어 적대적 IRL 없이도 보상을 추출하고, 동시에 보존성 제약이 일반화에 도움이 되는 귀납 편향으로 작동한다.

🔬 기술적 접근법

**모델/방법론**: EnergyFlow — 디노이징 필드를 스칼라 에너지 함수의 그라디언트로 파라미터화하는 통합 생성·IRL 프레임워크
**핵심 기법**: 디노이징 score matching으로 score를 학습해 행동을 생성하면서, 동시에 그 score가 전문가의 soft Q-함수 그라디언트와 동치임을 이용해 적대적 학습 없이 보상 함수를 추출. 필드를 보존장으로 강제하는 구조적 제약이 가설 클래스를 좁혀 OOD 일반화 한계를 이론적으로 개선

📊 주요 결과

다양한 매니퓰레이션 벤치마크에서 모방 학습 SOTA 달성
회복된 보상으로 다운스트림 RL을 학습했을 때 적대적 IRL과 likelihood 기반 대안을 모두 능가
보존성 제약이 정책 일반화에도 도움이 되는 귀납 편향임을 실증

💭 의의 및 한계

**의의**: 확산 기반 정책에 정량적 의미(에너지·보상)를 부여해 모방학습과 IRL의 경계를 허물고, 다운스트림 RL과 자연스럽게 연결되는 새로운 패러다임을 제시한다. **한계**: 최대 엔트로피·보존성 가정이 깨지는 환경에서는 보상 회복의 정확도가 떨어질 수 있고, 고차원·다중 모달 행동에서의 수치 안정성에 대한 추가 검증이 필요하다.

🚀 실용적 활용

로봇 매니퓰레이션의 모방학습과 후속 RL 미세조정을 잇는 통합 파이프라인
명시적 보상 함수가 없는 데모 데이터에서 비적대적 IRL로 보상을 자동 회복
안정적이고 해석 가능한 에너지 기반 정책을 요구하는 안전·검증 중요 도메인