Recovering Hidden Reward in Diffusion-Based Policies

1. Problem Statement

Diffusion-based policies have become a state-of-the-art paradigm for learning robotic manipulation skills from expert demonstrations due to their ability to model complex, multi-modal action distributions. However, these policies typically rely on Behavior Cloning (BC), which learns to imitate actions without explicitly modeling the underlying intent or reward structure.

• Limitations of Standard Diffusion: They lack an explicit energy representation, making them prone to failure when test-time situations deviate from the training distribution (Out-of-Distribution or OOD). They cannot easily provide a reward signal for downstream Reinforcement Learning (RL) without resorting to unstable adversarial training or computationally expensive Inverse Reinforcement Learning (IRL).
• The Gap: There is a need for a framework that unifies generative action modeling with reward recovery, enabling both high-quality imitation and the extraction of a valid reward signal for policy refinement, all while ensuring theoretical guarantees on generalization.

2. Methodology: ENERGYFLOW

The authors propose ENERGYFLOW, a framework that parameterizes a scalar energy function $E_\theta(s, a)$ whose gradient serves as the denoising field for action generation.

Core Architecture

• Scalar Energy Parameterization: Unlike standard diffusion policies that directly regress a vector-valued noise prediction $\epsilon_\theta(s, a)$, ENERGYFLOW learns a scalar energy function $E_\theta(s, a, t)$.
• Conservative Vector Field: The score function (denoising field) is derived via automatic differentiation:
  $$S_\theta(a, s, t) = -\nabla_a E_\theta(a, s, t)$$
  By construction, this ensures the vector field is conservative (curl-free, $\nabla \times S = 0$), meaning it is the gradient of a scalar potential. This enforces integrability, ensuring that action preferences are transitive and physically realizable.
• Training: The model is trained using Denoising Score Matching (DSM). The objective minimizes the difference between the learned gradient and the true noise direction:
  $$\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]$$
• Inference: Actions are generated by integrating the probability-flow ODE using the gradient of the learned energy function.

Reward Extraction

The framework extracts a reward signal directly from the learned energy landscape. To handle the state-dependent integration constant $c(s)$ (which makes raw energy values ambiguous across different states), the authors propose Centered Shaping:
$$\tilde{r}(a, s) = -\left( E_\theta(a, s) - \mathbb{E}_{a' \sim \mathcal{N}}[E_\theta(a', s)] \right)$$
This subtracts a state-dependent baseline (estimated via Monte Carlo sampling), normalizing the reward to reflect only relative action preferences within a state.

3. Key Theoretical Contributions

The paper provides rigorous theoretical proofs establishing the connection between diffusion score matching and reward recovery:

1. Score-Reward Equivalence: Under the Maximum Entropy Optimality assumption, the paper proves that the score function learned by diffusion models is proportional to the gradient of the expert's soft Q-function ($\nabla_a Q^* \propto S^*$). Thus, score matching implicitly recovers the reward gradient without adversarial min-max optimization.
2. Complexity Reduction via Conservative Constraints: The authors prove that constraining the hypothesis space to conservative fields (gradients of scalars) strictly reduces the Rademacher complexity compared to unconstrained vector fields. Specifically, the complexity bound for conservative fields scales with the smoothness of the representation ($L$), whereas unconstrained fields scale with the output dimension ($\sqrt{d}$).
3. OOD Generalization Bounds: Leveraging the reduced complexity, the paper derives tighter generalization bounds for Out-of-Distribution (OOD) scenarios. The conservative constraint acts as a geometric regularizer, preventing the model from learning arbitrary vector directions and forcing it to extrapolate the shape of the energy landscape, leading to better robustness under distribution shifts.
4. Identifiability and Error Propagation: The paper establishes that while the absolute Q-value cannot be uniquely identified (due to the unknown state-dependent constant), the within-state action ranking is exact. Furthermore, it bounds the propagation of score estimation errors to action preferences, showing that errors scale linearly with the distance between actions (Lipschitz continuity).

4. Experimental Results

The authors evaluated ENERGYFLOW on RoboMimic and Meta-World benchmarks, as well as real-world robot tasks.

• Imitation Performance (RQ1): ENERGYFLOW achieved State-of-the-Art (SOTA) performance.
  • RoboMimic: 93.8% average success rate (vs. 91.2% for Diffusion Policy).
  • Meta-World: 92.5% average success rate (vs. 90.7% for Diffusion Policy).
  • It significantly outperformed other energy-based methods (like Implicit BC) and IRL baselines.
• Real-World Deployment (RQ2): Deployed on an AGIBOT G1 robot for "Bottle" and "Drawer" tasks. ENERGYFLOW achieved 100% success rate with smooth trajectories and fewer hesitations compared to baselines.
• Reward Quality for Downstream RL (RQ3): When used as a reward signal to train a Soft Actor-Critic (SAC) agent:
  • The "Centered Energy" reward enabled SAC to reach near-oracle performance, significantly outperforming sparse rewards and raw energy rewards.
  • It provided dense, step-by-step guidance that prevented early training plateaus.
• OOD Generalization (RQ4): Under increasing initial position perturbations, ENERGYFLOW maintained robust success rates where unconstrained diffusion and flow policies degraded significantly. This confirms the theoretical benefit of the conservative constraint.
• Efficiency (RQ6): ENERGYFLOW achieved inference speeds comparable to non-energy-based Flow Policies (~10ms latency) while providing the benefits of an explicit energy function, avoiding the slow sampling times of traditional Energy-Based Models (EBMs).

5. Significance and Impact

• Unification of Imitation and RL: ENERGYFLOW bridges the gap between Behavior Cloning and Reinforcement Learning. It allows a policy to be trained via imitation while simultaneously exposing a valid reward signal for further refinement, eliminating the need for manual reward design or unstable adversarial IRL.
• Theoretical Guarantees for Diffusion: The work provides the first theoretical justification for why enforcing a conservative field (integrability) in diffusion policies is beneficial, linking it to reduced hypothesis complexity and improved OOD generalization.
• Practical Utility: The method is computationally efficient, scalable to high-dimensional action spaces, and effective in real-world robotic settings, offering a robust solution for learning complex manipulation skills from limited demonstrations.

In summary, ENERGYFLOW demonstrates that the structural constraints required for valid reward extraction (conservative fields) simultaneously serve as powerful inductive biases that enhance the generalization and robustness of diffusion-based policies.