PaperlayerView exact revision on arXiv

Source-linked AI summary

Recovering Hidden Reward in Diffusion-Based Policies

Yanbiao Ji, Qiuchang Li, Yuting Hu, Shaokai Wu, Wenyuan Xie, Guodong Zhang, Qicheng He, Deyi Ji, Yue Ding, Hongtao Lu

arXiv:2605.00623v2cs.RO

TL;DR

Diffusion policies often imitate actions without recovering the preferences behind them, limiting reward-guided selection beyond demonstrations. EnergyFlow learns a scalar energy whose gradient drives action generation and extracts reward-like action preferences, matching or exceeding state-of-the-art diffusion policies while improving out-of-distribution generalization and supporting downstream reinforcement learning.

  • Problem

    Diffusion policies typically match demonstrated actions without modeling underlying task preferences, making reliable action selection under distribution shift difficult and reward design for reinforcement learning challenging.

  • Method

    EnergyFlow parameterizes a scalar observation–action energy, trains it with denoising score matching, and derives a conservative generative field and reward signal from its action-gradient.

  • Results

    EnergyFlow matches or exceeds state-of-the-art diffusion policies, provides effective downstream reinforcement-learning shaping, and improves out-of-distribution generalization over unconstrained flow policies.

  • Takeaways & Limitations

    EnergyFlow shows that conservative structure can jointly support reward extraction from demonstrations and policy generalization without adversarial optimization.

  • Takeaways & Limitations

    From demonstrations, EnergyFlow recovers soft advantage or soft Q-function up to an unknown state-dependent integration constant, not global state value.

Abstract

from arXiv · show

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.

1. Introduction

Diffusion policies capture diverse expert behaviors through iterative denoising but typically imitate actions without modeling task preferences, limiting robustness outside demonstrations. EnergyFlow addresses this by learning a scalar energy whose action-gradient generates actions and provides a reward signal, with integrability supporting generalization and preference guarantees.

  • Motivation: Diffusion policies generate actions by iteratively denoising corrupted samples conditioned on state and capture complex, multi-modal expert behaviors.These properties make them well-suited to learning diverse manipulation skills from demonstrations.
  • Motivation: Behavior cloning imitates demonstrations without explicitly modeling why actions are desirable, which can limit robustness and extrapolation when test situations differ from demonstrations.Matching action likelihood alone may not provide a sufficient basis for behavior outside the demonstration distribution.
  • Method: EnergyFlow parameterizes a scalar energy over observation–action pairs, using its action-gradient to generate actions and its energy landscape to provide a Boltzmann-aligned reward signal.The energy is trained through denoising score matching and supports probability-flow ODE derivations connecting training and sampling.
  • Theory: The integrability constraint acts as implicit regularization by reducing hypothesis complexity and tightening generalization bounds.EnergyFlow enforces integrability by construction through a scalar energy and its action-gradient.
  • Theory: EnergyFlow bounds how score-matching errors propagate to recovered action preferences when the learned energy is used as a reward signal.This connects estimation accuracy to the reliability of extracted preferences.
  • Empirical results: Experiments show that the learned energy provides an effective downstream-RL shaping signal and that integrability improves out-of-distribution generalization over unconstrained flow policies.The reported gains are attributed to energy-based reward extraction and the integrability constraint, respectively.

2. Preliminaries

This section introduces denoising score matching, its multi-scale extension in score-based generative models, and the conditional score formulation used by diffusion-based policies. It also describes how noise-conditioned scores are trained and integrated to generate actions.

  • Denoising Score Matching: Denoising score matching estimates a data distribution’s score by perturbing samples with Gaussian noise and learning to denoise them.The objective is equivalent to explicit score matching up to a constant.
  • Denoising Score Matching: The denoising objective reduces to predicting the scaled noise direction, −ε/σ, because the perturbation-kernel score is −(x̃ − x0)/σ2.This follows from ∇x̃ log qσ(x̃|x0) = −(x̃ − x0)/σ2.
  • Score-Based Generative Models: Score-based generative models extend denoising score matching across noise scales using a monotonically increasing schedule σ(t) and a noise-conditional score network.The network approximates ∇x log pt(x), while λ(t) = σ2(t) provides uniform contribution across noise levels.
  • Diffusion-Based Policies: Diffusion-based policies model πθ(a|s) with a conditional score network trained by noise prediction and generate actions by integrating the probability-flow ODE conditioned on s.Inference starts from aT ∼ N(0, σ2(T)I), with at = a0 + σ(t)ε during training.

3. Theoretical Analysis

The theoretical analysis shows that, under maximum-entropy optimality, diffusion scores encode gradients of the expert’s soft Q-function, enabling reward recovery through scalar energy modeling. Conservative parameterization ensures valid action preferences, reduces complexity, improves OOD control, and clarifies identifiability and score-error effects.

  • Score–Reward Equivalence: Under maximum-entropy optimality, the expert policy’s action score is proportional to the gradient of its soft Q-function.Differentiating the Boltzmann policy removes the state-only partition function, directly linking score estimation to reward gradients.
  • Score–Reward Equivalence: An energy function whose negative action gradient matches the expert score recovers the soft Q-function up to a state-dependent constant.Equivalently, the recovered quantity represents the soft advantage up to state-dependent terms.
  • Conservative Structure: Conservative score fields are necessary for well-defined rewards because nonconservative fields can induce cyclic, intransitive action preferences.Parameterizing a scalar energy and defining the score as its negative action gradient enforces conservativeness.
  • Generalization: Conservative constraints strictly reduce empirical Rademacher complexity relative to arbitrary linear vector fields, with the gap becoming pronounced in high-dimensional action spaces.The result treats conservative fields as gradients of potentials over the same feature representation.
  • Generalization: For OOD generalization, unconstrained-field error bounds grow with action dimensionality, whereas conservative-field bounds remain controlled by smoothness.The global scalar-potential constraint couples prediction errors across dimensions and constrains extrapolation in unseen regions.
  • Identifiability and Error: Recovered rewards identify exact within-state action rankings, but state-dependent offsets can alter sequential policies unless they have the potential-based reward-shaping form.For same-state comparisons, the offset cancels; centered shaping removes it explicitly.

4. Methodology

EnergyFlow parameterizes a scalar energy and obtains its action score by automatic differentiation, enforcing a conservative field. It trains this energy with denoising score matching and centers extracted rewards to preserve relative action preferences while reducing state-dependent offsets.

  • Action Generation: Action generation samples noisy actions and iteratively updates them using the energy gradient, returning the final action and its endpoint energy.The procedure uses K discretized steps and endpoint γ.
  • Architecture: EnergyFlow parameterizes a scalar energy function Eϕ over actions, states, and noise levels, then obtains the score through automatic differentiation rather than directly regressing a vector field.This architectural choice satisfies the conservative-field requirement.
  • Architecture: The resulting action score is curl-free, ensuring learned preferences remain transitive and physically realizable.The construction enforces ∇a × Sϕ ≡ 0.
  • Training Paradigm: Denoising score matching uses a variance-exploding noise schedule with σmin = 0.01, σmax = 10.0, and T = 1.0.As t → 0, minimizing the objective is equivalent to recovering the maximum-entropy reward gradient.
  • Reward Extraction: Centered shaping subtracts the expected energy under a reference distribution, normalizing state-dependent offsets so rewards reflect only within-state action preferences.The ODE endpoint is γ = 10−3, and the baseline uses M = 16 samples from N(0, I).
  • Reward Extraction: For fixed reference samples, baseline computation is deterministic, yielding a low-variance reward signal for policy-gradient updates.This contrasts with methods requiring stochastic trace estimation for CNF log-likelihoods.

5. Experiments

Experiments show that ENERGYFLOW preserves strong benchmark performance, transfers successfully to real-world manipulation, and provides a reward signal for reinforcement learning. Its conservative structure improves robustness under perturbations, while reward extraction remains stable across a broad range of γ values and inference remains efficient.

  • Simulation Benchmarks: ENERGYFLOW achieves the highest average success rates on RoboMimic (93.8%) and Meta-World (92.5%), outperforming Diffusion Policy and Flow Policy.On RoboMimic ToolHang, it reaches 84.2% versus 77.2% for Diffusion Policy.
  • Real-World Transfer: ENERGYFLOW obtains a 100% success rate on both real-world Bottle and Drawer tasks across three initial-position changes and 20 rollouts each.The evaluation uses a 7-DoF arm, parallel-jaw gripper, fixed RGB camera, and 20 expert demonstration trajectories.
  • Reward Extraction: The learned energy enables SAC training without ground-truth rewards, evaluated for 200k environment steps on RoboMimic Square and Transport.Centered shaping is compared with sparse task rewards, raw energy rewards, and oracle dense rewards; sparse rewards delay learning, while raw energy rewards can plateau early.
  • Robustness: ENERGYFLOW maintains superior stability as initial-position perturbations increase, outperforming Diffusion and Flow Policy at medium and large perturbation levels.The results support the curl-free constraint as a geometric regularizer that improves extrapolation under spatial variability.
  • Sensitivity and Efficiency: Reward-extraction performance remains robust across γ ∈[10^-4, 10^-2], with degradation only at γ ≥0.1 as the noised distribution diverges from the data distribution.ENERGYFLOW also achieves superior success rates with latency comparable to Flow Policy, avoiding the expensive Langevin sampling required by Implicit BC.

6. Related Work

Related work spans generative behavior cloning, inverse reinforcement learning, and energy-based imitation. EnergyFlow is positioned against methods that either model actions generatively, recover rewards through adversarial or optimization-heavy procedures, or use energies without establishing reward identifiability.

  • Behavior cloning: Generative behavior cloning captures expert action distributions directly, with diffusion policies achieving strong results in contact-rich manipulation and later extending to 3D visual manipulation.The cited progression includes expressive generative models, Diffusion Policy, and subsequent 3D visual manipulation extensions.
  • Inverse reinforcement learning: IRL recovers latent rewards rather than directly imitating actions, while classical maximum-entropy and Bayesian methods faced computational intractability from repeated policy optimization.Adversarial methods such as GAIL and AIRL improve scalability through occupancy-measure matching but inherit training-related limitations.
  • Energy-based imitation: Implicit Behavioral Cloning and related methods learn state-action energies for action prediction, improving stability through contrastive training, refined negative sampling, transformer scaling, and iterative inference.EBT-Policy is described as achieving strong robustness with fewer inference steps than diffusion policies.
  • Energy-based imitation: Policy-oriented energy methods use energies as decision scores rather than identifiable rewards, and their iterative optimization dynamics need not define a well-defined potential.On the reward-learning side, EBIL estimates expert energy via score matching and uses it as a reward for downstream maximum-entropy reinforcement learning.

7. Conclusion · A. Proofs

ENERGYFLOW bridges diffusion-based imitation learning and inverse reinforcement learning through energy-based parameterization. Its theory links optimal-policy scores to soft-Q gradients for adversary-free reward recovery and shows that conservative vector fields reduce hypothesis complexity.

  • 7. Conclusion: ENERGYFLOW bridges diffusion-based imitation learning and inverse reinforcement learning through energy-based parameterization.
  • 7. Conclusion: The framework uses an energy-based parameterization to connect diffusion policy modeling with inverse reinforcement learning.
  • 7. Conclusion: The score function of an optimal policy encodes the gradient of its soft Q-function.
  • 7. Conclusion: This score–soft-Q relationship enables reward recovery via score matching without adversarial optimization.
  • 7. Conclusion: Constraining the learned vector field to be conservative provably reduces hypothesis complexity.
  • 7. Conclusion: The theoretical analysis identifies three key results, including score-based reward recovery and the benefits of conservative field constraints.

A.1. Proof of Theorem 3.6 · A.2. Proof of Lemma 3.8

Theorem 3.6 proves that conservative vector fields have tighter complexity scaling than arbitrary linear fields, especially in high-dimensional action spaces with smooth representations. Lemma 3.8 extends this advantage to OOD risk bounds, where unconstrained complexity diverges as d →∞ while conservative complexity remains controlled by L.

  • A.1. Proof of Theorem 3.6: Theorem 3.6 compares arbitrary linear vector fields Func with conservative gradient fields Fcons under bounded feature norm B and Jacobian norm L.The conservative class consists of fields f(x) = Jϕ(x)⊤w generated by potentials E(x) = w⊤ϕ(x).
  • A.1. Proof of Theorem 3.6: For high-dimensional action spaces with smooth representations, ˆRS(Fcons) ≪ ˆRS(Func).The condition given is L ≪B d.
  • A.1. Proof of Theorem 3.6: The unconstrained complexity scales with the square root of output dimension d, whereas conservative complexity scales with representation smoothness L.The comparison follows from the bounds for arbitrary linear fields and conservative gradient fields.
  • A.1. Proof of Theorem 3.6: Because neural representations learn smooth manifolds whose tangent-space volume grows more slowly than ambient dimension, conservative constraints provide structurally superior generalization guarantees.The tangent-space volume is captured by L.
  • A.2. Proof of Lemma 3.8: Lemma 3.8 bounds target-domain risk for a conservative estimator with probability at least 1 −δ when the ground-truth field h∗ belongs to Fcons.The analysis assumes uniformly bounded hypotheses and uses squared L2 risk.
  • A.2. Proof of Lemma 3.8: The conservative source complexity bound is ˆRS(Fcons) ≤ΛL √n, compared with the unconstrained bound ˆRS(Func) ≤ΛB d √n.These bounds are substituted into the domain-adaptation and empirical-to-population generalization inequalities.
  • A.2. Proof of Lemma 3.8: As d →∞, the unconstrained-field OOD bound diverges, while the conservative bound remains controlled by smoothness parameter L.The conservative estimator also benefits from negligible ideal joint error because h∗ belongs to Fcons.

A.3. Proof of Proposition 3.9 … C.5.1. AUTOREGRESSIVE AND GENERATIVE POLICIES

EnergyFlow exactly ranks actions within each state, while cross-state energy comparisons remain ambiguous; score errors induce bounded errors in relative preferences. The implementation uses a scalar, twice-differentiable energy head with shared baseline evaluation and specified generative-policy baselines.

  • A.3. Proof of Proposition 3.9: Within any fixed state, minimizing learned energy exactly recovers the expert-optimal action ranking.The state-dependent integration constant cancels, yielding arg min_a Eϕ(a, s) = arg max_a Q∗(s, a).
  • A.3. Proof of Proposition 3.9: Cross-state energy comparisons are ambiguous because c(s)−c(s′) is unknown from demonstrations.Demonstrations reveal within-state preferences, not relative values across states.
  • A.4. Proof of Theorem 3.11: If score error is uniformly bounded by ϵ, relative preference error is at most ϵ · ∥a −a′∥2.The bound applies to two actions evaluated at the same state.
  • B. Baselines: ENERGYFLOW is evaluated against autoregressive, diffusion, flow, energy-based, and nonadversarial IRL baselines.The suite includes LSTM-GMM, Diffusion Policy, Flow Policy, Implicit BC, EBT-Policy, EBIL, NEAR, and IQ-Learn.
  • C. Additional Implementation Details: The implementation uses PyTorch, AdamW training, shared MLP state encoders, temporal position embeddings, and grid-searched learning rates and batch sizes.The grid uses learning rates {10−3, 10−4, 10−5} and batch sizes {128, 256}.
  • C.2. Network Architecture: EnergyFlow adapts Diffusion Policy’s 1D Conditional U-Net, replacing its vector noise head with a scalar energy head conditioned on state and time.The scalar output is produced by GlobalAveragePooling1D followed by a 3-layer MLP, while FiLM injects conditioning into convolutional blocks.
  • C.3. Differentiable Training Infrastructure: Mish activations make the network C2, enabling the second derivatives required for score matching and double backpropagation.Automatic differentiation constructs the gradient computation graph, and spectral normalization stabilizes energy magnitudes.
  • C.5.1. AUTOREGRESSIVE AND GENERATIVE POLICIES: Baseline implementations include a 2-layer, 256-unit LSTM-GMM with K = 5 components, an identical-backbone Diffusion Policy, and a 4-layer RealNVP Flow Policy.Diffusion Policy uses T = 100 diffusion steps and 20 DDIM denoising steps; Flow Policy uses ReLU MLP coupling networks.

C.5.2. ENERGY-BASED METHODS … C.7. OOD Perturbation Implementation

The implementation compares energy-based and inverse-reinforcement-learning baselines, trains recovered rewards with SAC, stabilizes learning through replay and entropy tuning, and evaluates OOD robustness under progressively stronger perturbations.

  • C.5.2. ENERGY-BASED METHODS: IBC uses a 3-layer, 512-unit ReLU MLP with discontinuous energy and derivative-free inference rather than differentiable optimization.Training uses an InfoNCE-style loss with uniform negative samples, while inference uses the original autoregressive derivative-free optimizer.
  • C.5.2. ENERGY-BASED METHODS: EBT-Policy parameterizes tokenized state-action sequences with a 4-layer, 4-head Transformer Encoder of embedding dimension 128, trained by NCE.Inference uses Langevin Dynamics for K = 100 steps with step size 0.01.
  • C.5.3. INVERSE REINFORCEMENT LEARNING (IRL): IRL baselines recover rewards and train policies with Soft Actor-Critic as the underlying reinforcement-learning optimizer.The specific implementations include IQ-Learn, EBIL, and NEAR, with their respective Q-function, adversarial energy, and noise-conditioned architectures.
  • C.6.1. SOFT ACTOR-CRITIC ALGORITHM: SAC learns a stochastic policy with clipped double-Q critics, Gaussian reparameterized actions, tanh squashing, and linear rescaling to valid action ranges.Automatic entropy tuning uses target entropy Htgt = −dim(a), and target critics are updated by Polyak averaging.
  • C.6.1. SOFT ACTOR-CRITIC ALGORITHM: SAC minimizes the soft Bellman error using successor actions sampled from the policy and target critics with discounting and entropy regularization.The target incorporates the success indicator, discount factor γ, entropy temperature α, and target critics Q¯θi.
  • C.6.2. EXPERIENCE REPLAY BUFFER: Experience replay stabilizes SAC by uniformly sampling transitions from a fixed-capacity FIFO buffer collected under the current policy.Environment termination signals are stored as dt, while time-limit truncations are treated as timeouts.
  • C.7. OOD Perturbation Implementation: OOD evaluation adopts S, M, and L perturbation levels corresponding to D0, D1, and D3 datasets with progressively broader object and target-position changes.S and M perturb initial positions of contacted objects, whereas L additionally perturbs target positions or objects.

D. Experiment Tasks … E.1.2. META-WORLD TASKS

The experiments span RoboMimic and Meta-World simulation tasks covering object manipulation, articulated-object interaction, grasping, assembly, and tool use. EnergyFlow successfully executes the demonstrated tasks on both benchmarks.

  • D.1. Simulation Tasks: RoboMimic includes Can, Lift, Square, and Transport tasks involving soda-can relocation, cube lifting, nut–peg fitting, and two-arm hammer transfer.Transport requires one arm to retrieve and pass a hammer while the other transfers it to a bin.
  • D.1. Simulation Tasks: Meta-World evaluates ButtonPress and DrawerOpen through switch actuation and constrained-joint manipulation.These tasks assess interaction with articulated objects.
  • D.1. Simulation Tasks: Meta-World’s BinPicking task tests robust grasping and retrieval from a confined volume.
  • D.1. Simulation Tasks: Meta-World’s Assembly and Hammer tasks require tight-tolerance nut–peg alignment and tool use to drive a nail into a target surface.Hammer requires grasping the hammer accurately before driving the nail.
  • E.1.1. ROBOMIMIC TASKS: Figure 6 illustrates successful execution of each RoboMimic benchmark task with the EnergyFlow policy.
  • E.1.2. META-WORLD TASKS: Figure 7 illustrates successful execution of each Meta-World benchmark task with the EnergyFlow policy.

E.2. Real Robot Experiment

The real-robot experiments use ten teleoperated demonstrations per task and a vision-based policy operating at 10Hz. The setup combines augmented training crops, center-crop inference, and an ImageNet-pretrained ResNet-18 encoder.

  • Experimental setup: Each task uses 10 teleoperated demonstrations, with random-crop augmentation during training and a static center crop during inference.The inference crop matches the training crop size.
  • Experimental setup: The policy receives 226 × 226 RGB images at 10Hz and outputs 8-dimensional actions comprising seven joint velocities and a gripper command.This defines the real-robot observation and action interfaces.
  • Experimental setup: The visual backbone is an ImageNet-pretrained ResNet-18 encoder, consistent with prior work.The cited prior work includes Chi et al. (2023) and Zhao et al. (2023).
  • Task demonstrations: Figure 6 visualizes rollout sequences for different RoboMimic tasks.Each row presents a task demonstration sequence.
  • Task demonstrations: Figure 7 visualizes rollout sequences for different Meta-World tasks.Each row presents a task demonstration sequence.
Loading 2605.00623v2…