PaperlayerView exact revision on arXiv

Source-linked AI summary

EBT-Policy: Energy Unlocks Emergent Physical Reasoning Capabilities

Travis Davies, Yiqi Huang, Alexi Gladstone, Yunxin Liu, Xiang Chen, Heng Ji, Huxian Liu, Luhui Hu

arXiv:2510.27545v1cs.ROcs.AI

TL;DR

Diffusion policies can be computationally costly and susceptible to exposure bias, while earlier energy-based policies have faced scalability and stability challenges. EBT-Policy uses Energy-Based Transformers to generate robot action trajectories by minimizing energy. Across simulated and real-world benchmarks, it outperforms Diffusion Policy with fewer training epochs and inference iterations, and exhibits retry behavior after failure.

  • Problem

    Diffusion policies can require many denoising steps and incur high computational cost, while chain-structured generation can propagate errors; earlier EBM policies have faced scalability and training-stability challenges.

    Source: §1 · p. 2

  • Method

    EBT-Policy models visuomotor action distributions with an energy landscape and generates trajectories through iterative energy minimization using Energy-Based Transformers.

    Source: §3 · p. 3, §3.1 · p. 3

  • Results

    Across simulated and real-world benchmarks, EBT-Policy consistently outperforms Diffusion Policy while using fewer training epochs and as few as 2 inference iterations, up to 50× fewer than DP.

    Source: §7 · p. 9

  • Takeaways & Limitations

    EBT-Policy exhibits retry and action adjustment after failure, with scalar energy serving as an uncertainty signal for distinguishing successful from failed contacts.

    Source: §7 · p. 9

    Diffusion Policy continues to outperform EBT-Policy in training stability and sampling under highly multimodal action distributions, which EBT-Policy’s refinements only partially mitigate.

    Source: §6 · p. 9

Abstract

from arXiv · show

Implicit policies parameterized by generative models, such as Diffusion Policy, have become the standard for policy learning and Vision-Language-Action (VLA) models in robotics. However, these approaches often suffer from high computational cost, exposure bias, and unstable inference dynamics, which lead to divergence under distribution shifts. Energy-Based Models (EBMs) address these issues by learning energy landscapes end-to-end and modeling equilibrium dynamics, offering improved robustness and reduced exposure bias. Yet, policies parameterized by EBMs have historically struggled to scale effectively. Recent work on Energy-Based Transformers (EBTs) demonstrates the scalability of EBMs to high-dimensional spaces, but their potential for solving core challenges in physically embodied models remains underexplored. We introduce a new energy-based architecture, EBT-Policy, that solves core issues in robotic and real-world settings. Across simulated and real-world tasks, EBT-Policy consistently outperforms diffusion-based policies, while requiring less training and inference computation. Remarkably, on some tasks it converges within just two inference steps, a 50x reduction compared to Diffusion Policy's 100. Moreover, EBT-Policy exhibits emergent capabilities not seen in prior models, such as zero-shot recovery from failed action sequences using only behavior cloning and without explicit retry training. By leveraging its scalar energy for uncertainty-aware inference and dynamic compute allocation, EBT-Policy offers a promising path toward robust, generalizable robot behavior under distribution shifts.

1. Introduction

Diffusion policies face scheduler dependence, costly inference, and exposure bias, while earlier EBM policies have struggled with scalability and training stability. EBT-Policy applies Energy-Based Transformers to action generation, and the paper reports faster training and inference, strong benchmark performance, and retry behavior without explicit retry training.

Motivation

  • Diffusion and flow-based policies depend on externally defined noise schedulers, which can contribute to slow convergence, high training cost, many inference steps, and exposure bias.Errors early in their chain-structured generation can propagate downstream, complicating training and interpretation under distribution shifts.Source: §1 · p. 2
  • EBMs can avoid specific noise schedulers and use energy-guided MCMC to re-evaluate intermediate proposals, potentially reducing cascading mistakes and schedule sensitivity.This iterative correction contrasts with committing to a fixed denoising trajectory.Source: §1 · p. 2

Contribution

  • EBT-Policy combines regularized EBM losses, Transformers, energy-landscape regularization, and sampling refinements to address the scalability and stability challenges of implicit EBM policies.The refinements include energy-scaled MCMC step sizes, pre-sample normalization, Nesterov-accelerated gradients, and scaled Langevin Dynamics.Source: §1 · p. 2

Reported results

  • Across simulated and real-world tasks, EBT-Policy trains and infers faster than Diffusion Policy, and it outperforms DP by as much as 24% on robomimic Squre and Tool Hang.For the same task and success rate, EBT-Policy can use 2 inference steps versus DP's 100.Source: §1 · p. 2
  • EBT-Policy learns retry behavior without explicit retry data or supervision, alongside uncertainty-aware action adjustment.The paper presents this as an emergent capability of behavior cloning.Source: §1 · p. 2

2. Related Work

Related work contrasts diffusion-based behavior cloning with explicit energy-based policies and situates EBT-Policy’s retry behavior among prior robotics approaches. It also illustrates energy-based uncertainty through per-frame energies and a failure followed by a successful retry.

2.1. Implicit Behavior Cloning with Diffusion

  • Diffusion-based behavior cloning became a dominant approach, with high training stability and capacity for multimodal, high-dimensional action distributions.These methods amortize score prediction rather than differentiating an explicit scalar energy function.Source: §2.1 · p. 2

2.2. Implicit Behavior Cloning with Energy-Based Models

  • Earlier energy-based policies faced training instability from the need for exponentially many negative samples as data dimensionality increased, while newer EBM methods offer improved scalability.EBMs learn explicit energy landscapes whose equilibrium dynamics define the policy.Source: §2.2 · p. 2

2.3. Reasoning and Emergent Behavior

  • Energy-based models have also been applied to reasoning across discrete and continuous modalities, while robotics VLA approaches can draw on pretrained vision-language knowledge.The cited figure illustrates uncertainty-guided continuation or retry in a manipulation sequence.Source: §2.3 · p. 3, Figure 2 · p. 3
  • Prior robotics retry methods use planning, goal-conditioning, or hierarchical execution; EBT-Policy instead demonstrates retry from instantaneous action prediction without retry training data.Figure 2 associates lower per-frame energy with higher certainty and marks a failure at Step 7 followed by successful correction at Step 11.Source: §2.3 · p. 3, Figure 2 · p. 3

3. EBT-Policy Formulation

EBT-Policy represents visuomotor action distributions with an energy landscape and generates trajectories through iterative energy minimization. It conditions action sequences on recent observations and optionally language, refining candidates with energy-scaled updates and adaptive inference.

3. EBT-Policy Formulation

  • EBT-Policy represents multimodal action distributions as an energy landscape over actions, framing learning and inference as continuous optimization and sampling.The formulation is intended to support implicit, unnormalized action likelihoods while maintaining stability and scalability in high-dimensional visuomotor tasks.Source: §3 · p. 3

3.1. Energy-Based Transformers

  • EBMs map inputs to scalar energies, and the unnormalized formulation omits the partition function while treating lower energy as greater compatibility between inputs and candidate predictions.Generation then corresponds to minimizing energy, often by gradient descent.Source: §3.1 · p. 3
  • EBTs iteratively update candidate predictions by gradient descent on predicted energy, with an energy-scaled step size and Langevin Dynamics noise.The update procedure is approximately equivalent to MCMC sampling via Langevin Dynamics.Source: §3.1 · p. 3
  • Inference steps are chosen adaptively according to current energy, while the step size is scaled as a function of energy to support stable convergence.This lets the policy vary computation during iterative energy minimization.Source: §3.1 · p. 4

3.2. EBTs for Robot Learning

  • The policy generates joint- or Cartesian-space action sequences conditioned on recent multimodal observations and optionally natural-language instructions.The observation window contains RGB frames and proprioceptive state histories.Source: §3.2 · p. 4
  • The policy output is a sequence of continuous actions over a prediction horizon, planned to achieve the task described by the language instruction.The sequence covers the next n steps.Source: §3.2 · p. 4
  • Rather than learning an explicit policy, EBT-Policy assigns low energy to trajectories consistent with sensory observations and language, then refines sampled actions toward low-energy regions.Inference thereby samples from the manifold of task-consistent action sequences.Source: §3.2 · p. 4

3.3. Diffusion vs. Energy-Based Policies

EBTs learn an explicit energy landscape rather than a diffusion model’s score field, then refine actions through energy descent. The paper attributes their embodied reasoning capabilities to equilibrium dynamics, uncertainty-aware computation, and direct operation on continuous sensory inputs.

3.3. Diffusion vs. Energy-Based Policies

  • Diffusion policies learn a score field for reversing a noise process, whereas EBTs learn an explicit energy landscape and descend it toward low-energy, high-likelihood actions.The paper describes this as explicit density modeling with equilibrium dynamics.Source: §3.3 · p. 4

3.4. Why EBTs Work for Robot Policies

  • Equilibrium dynamics bring out-of-distribution points closer to the data manifold, while energy-based uncertainty modeling supports adaptive sampling and a built-in plausibility verifier.Sampling continues until energy converges or the gradient falls below a threshold; low energy indicates plausibility and high energy rejection.Source: §3.4 · p. 4

3.5. Towards Embodied Reasoning in Robotics

  • EBTs are presented as a unified framework for embodied System 2 reasoning, integrating perception, reasoning, and control in one energy-minimization objective.This formulation is contrasted with approaches relying on discrete tokenization, fixed compute budgets, or externally imposed reasoning modules.Source: §3.5 · pp. 4–5
  • EBTs operate directly on continuous sensory representations, including proprioceptive and camera data, rather than reasoning over discrete space such as written text.Source: §3.5 · p. 5
  • Inference steps are allocated dynamically according to task complexity and internal energy gradients, allowing computational effort to vary with uncertainty in observations and goals.Source: §3.5 · p. 5
  • Training is described as self-supervised cognitive optimization that minimizes energy dissipation for consistent action trajectories.Source: §3.5 · p. 5

3.6. EBT-Policy Ingredients

EBT-Policy addresses multimodal action data and unstable iterative refinement with training mechanisms for exploration and optimization stability. These include randomized MCMC sampling, noise and step-size controls, and gradient stabilization.

3.6. EBT-Policy Ingredients

  • Multimodal demonstrations can draw energy models into local minima, while iterative MCMC refinement can create exploding gradient chains; the paper introduces design tweaks to stabilize learning.Source: §3.6 · p. 5

3.6.1. Learning Multimodal Action Distributions

  • Randomized MCMC step counts, annealed Langevin noise, randomized step sizes, and Nesterov acceleration are used to explore multiple action modes and traverse complex energy surfaces.Higher-energy regions receive stronger noise for broad exploration, while lower-energy regions use smaller noise for precise convergence; step-size randomization also mitigates exploding gradients.Source: §3.6.1 · p. 5

3.6.2. Model Stability

  • Energy-scaled step sizes modulate update magnitudes in proportion to predicted energy to prevent overshooting or undershooting minima.Source: §3.6.2 · p. 6
  • Action trajectories are normalized before sampling, and gradient clipping caps the global gradient norm at 1.0 to limit uncontrolled growth and unstable updates.The paper reports gradient clipping as its most critical component for stable, reliable training.Source: §3.6.2 · p. 6

3.7. EBT-Policy Training

The training procedure starts from a noisy target trajectory and conditions the model on sensory inputs and, when available, language instructions. It refines trajectories through randomized MCMC sampling and accumulates losses against demonstrations.

  • Each training iteration initializes the target trajectory as pure noise and encodes multimodal sensory observations, with language instructions included when available.Source: §3.7 · p. 6
  • Training performs randomized MCMC sampling with trajectory normalization and Langevin noise injected according to an annealed schedule.Source: §3.7 · p. 6
  • The model predicts energy, randomizes and scales the step size in proportion to that energy, then accumulates loss by comparing denoised trajectories with demonstrations.Source: §3.7 · p. 6

3.8. Dynamic Inference in EBT-Policy

EBT-Policy uses adaptive energy descent at inference rather than a fixed number of sampling steps. It stops when the gradient norm falls below a threshold or the maximum step limit is reached.

  • Inference dynamically selects the number of MCMC updates, stopping when the maximum step limit is reached or the energy-gradient norm falls below threshold τ.Source: §3.8 · p. 6
  • Low-energy, high-certainty regions require fewer inference steps, while uncertain states receive more, allowing computation to balance precision and runtime without manual tuning.Source: §3.8 · p. 6

4. Experiments

The experiments evaluate EBT-Policy variants in simulation and physical robotic environments, using task suites and comparisons designed to assess performance and real-world applicability.

4. Experiments

  • EBT-Policy-S supports controlled simulation studies, while EBT-Policy-R is designed for real-world, language-conditioned, multitask policy learning.The variants serve distinct experimental goals across the evaluation.Source: §4 · p. 6, Table 2 · p. 7

4.2. Experiment Setup

  • Diffusion-based baselines use the same demonstration data, and the real-world comparison holds the parameter budget near 100M for fairness.Billion-parameter VLA and robot foundation models are omitted because internet-scale pretraining obscures a fair comparison of policy objectives.Source: §4 · p. 7, §4.2 · p. 7

4.1. Task Design

  • Simulation uses four robomimic manipulation tasks, while physical evaluation uses dual-arm tasks including towel folding and pan placement.The simulation tasks are Lift, Can, Square, and Tool Hang; physical tasks also include a two-arm pick-and-place task.Source: §4.1 · p. 7
  • Physical tasks assess multitask learning, sensor-noise robustness, goal understanding, spatial generalization, and placement precision.The tasks include collaborative towel folding, pan placement at varying positions and orientations, and a two-arm pick-and-place sequence.Source: §4.1 · p. 7

4.2. Experiment Setup

  • Simulation comparisons control architecture, hyperparameters, and augmentations, while testing diffusion variants with 10 and 100 inference steps.Simulation uses robomimic datasets and 50 episodes per rollout; real-world data comes from a teleoperated dual-arm, four-camera tabletop setup.Source: §4.2 · p. 7

5. Results

EBT-Policy outperforms Diffusion Policy across the simulated tasks and reaches strong performance with substantially less training and inference computation. It also shows recovery after failed actions in an out-of-distribution state without explicit retry training.

Simulation results

  • In simulation, EBT-Policy outperforms Diffusion Policy on all four tasks; Diffusion Policy was found to perform well only at 100 inference steps, often tenfold the iterations required by EBT-Policy.Source: §5 · p. 8

Training results

  • 100% success within 30 epochs for EBT-Policy, while the highest-performing Diffusion Policy variant stays below 60% after 100 epochs.The comparison reports the diffusion variant using up to 100 iterations.Source: §5 · p. 8

Computational efficiency

  • Approximately 50 times fewer iterations and 55% fewer epochs let EBT-Policy achieve results comparable to Diffusion Policy.The authors attribute faster convergence to focusing on plausible action sequences rather than generating them from scratch.Source: §5 · p. 8

Inference-step comparison

  • At 10 iterations, Diffusion Policy reaches 0% success, while EBT-Policy maintains high performance at 2 iterations.Reducing the iteration parameter significantly degrades Diffusion Policy performance.Source: §5 · p. 8

Emergent retry behavior

  • After an initial failed hook attempt, EBT-Policy retries successfully from an out-of-distribution state without having seen such data or explicit recovery training.The authors hypothesize that this emergent retry behavior arises from the energy function's equilibrium dynamics.Source: Figure 6 · p. 8

6. Discussion

In Tool Hang simulation, EBT-Policy recovers from out-of-distribution hook rotations that cause Diffusion Policy to fail. Its recovery tracks iterative energy minimization, while training stability and multimodal sampling remain challenges.

6. Discussion

  • EBT-Policy repositions the tool and completes Tool Hang after an out-of-distribution hook rotation, while Diffusion Policy loses stability and fails to recover.The authors report this retry capability across repeated runs and say it has not previously been observed in vanilla behavior-cloning policies.Source: §6 · p. 8
  • After a failed attempt, predicted energy first rises in high-energy states and then falls as EBT-Policy approaches a stable, task-completing configuration.The authors connect this trajectory to uncertainty modeling, energy-based optimization, and distinguishing successful from failed contact events.Source: §6 · p. 9
  • In simulation, maximum success rates are somewhat below those in the original Diffusion Policy paper, possibly because of hyperparameter differences.The authors matched released real-world parameters as closely as possible but note minor differences may have affected simulated outcomes.Source: §6 · p. 9
  • EBT-Policy's main challenges are training stability and sampling under highly multimodal action distributions; Diffusion Policy still performs better on both.The introduced hyperparameter strategies and architectural refinements only partially mitigate these issues.Source: §6 · p. 9

7. Conclusion

EBT-Policy replaces iterative denoising with energy minimization over action trajectories. Across simulated and real-world benchmarks, it outperforms Diffusion Policy with fewer training epochs and inference iterations, and exhibits retry behavior after failure.

7. Conclusion

  • Across simulated and real-world benchmarks, EBT-Policy outperforms Diffusion Policy while using fewer training epochs and as few as two inference iterations.The paper reports up to 50× fewer inference iterations than Diffusion Policy.Source: §7 · p. 9
  • EBT-Policy exhibits retry and adjustment after failure, using scalar energy as an uncertainty signal to distinguish successful from failed contacts and adapt its actions.The authors describe this as a foundation for energy-based world models in manipulation, with one energy function measuring consistency across observations, language, actions, and dynamics.Source: §7 · p. 9
  • The authors identify scaling EBT-Policy architectures, improving optimization schemes, and using larger, more diverse datasets as future directions.These directions are intended to further enhance performance and generalization.Source: §7 · p. 9
Loading 2510.27545v1…