ArXiv: 2309.05803

🎯 Pitch

Energy-based models are actually practical for high-dimensional robot policies after all—you just need to train them with ranking noise contrastive estimation and a learnable negative sampler, since the popular IBC objective is fundamentally biased and leads to collapse. This approach yields policies that match or beat diffusion models on multi-modal manipulation and planning tasks, restoring EBMs as a compact, composable alternative for real-time control.


1. Executive Summary

This paper revisits energy-based models (EBMs) as a practical policy representation for multi-modal robotic tasks, developing a training algorithm that combines ranking noise contrastive estimation (R-NCE) with a learnable negative sampler and a new interpolating EBM framework that bridges noise and data across multiple scales. The authors first prove that the widely-used Implicit Behavior Cloning (IBC) objective is biased even at the population level, then demonstrate that R-NCE-trained EBMs, when jointly optimized with a simpler normalizing flow proposal distribution, achieve asymptotic consistency and near-maximum-likelihood efficiency without requiring adversarial training. Empirically, R-NCE policies match or outperform diffusion models and stochastic interpolants on challenging multi-modal benchmarks—obstacle avoidance path planning and contact-rich block pushing—establishing that EBMs are a competitive alternative to diffusion-based policies when paired with the R-NCE objective, a learnable negative sampler, and interpolating multi-scale training, but only when the negative sampler is sufficiently informative to prevent posterior collapse during training.

2. Context and Motivation

The Core Problem: EBMs Seem Perfect for Robotics—But Nobody Can Train Them Reliably

The paper addresses a specific, practical tension in robot policy design: energy-based models (EBMs) have theoretical properties that make them exceptionally well-suited for robotic control, yet the standard training algorithm for EBM behavioral cloning has proven unreliable, leading the field to largely abandon EBMs in favor of diffusion models. The paper aims to resolve this tension by providing both a theoretical explanation for why the existing approach fails and a practical alternative that restores EBMs as a competitive policy class.

This matters for several concrete reasons the authors articulate (Section 1):

  • Compact representation: EBMs model an unnormalized scalar potential field rather than the normalized probability distribution directly. Since the partition function need not be computed during inference, EBMs can be more parameter-efficient than alternatives that must represent normalized densities or score functions at multiple noise scales. In robotics, where real-time inference constraints demand small models running on embedded hardware, this compactness translates directly to deployability.

  • Natural fit for optimization-based control: Many robotic tasks—particularly those with safety constraints—are already framed as finding minimizers of cost functions. EBMs represent the action distribution as p(y|x) ∝ exp(E(x,y)), meaning that selecting the most likely action (the mode) is equivalent to maximizing the energy function. This aligns naturally with trajectory optimization and model-predictive control paradigms where the policy is implicitly defined through optimization rather than explicit forward sampling.

  • Composability in density space: Because EBMs represent unnormalized log-densities, they can be combined through simple addition: if E1 and E2 are energy functions representing different task constraints or objectives, their sum E1 + E2 represents (up to a constant) the product of the corresponding densities. This property, exploited in work by Du et al. [22], enables compositional generalization that is not straightforward with score-based or normalized models.

  • Prior knowledge integration: The density representation makes it straightforward to incorporate domain knowledge—for instance, obstacle avoidance potentials or task-specific cost functions—by adding them as additional energy terms, without retraining the entire model.

The problem, stated bluntly in the abstract, is that "the prevailing folklore—that energy models in high dimensional continuous spaces are impractical to train—is false" under the right training algorithm, but has been reinforced by the failure of the IBC objective.

The IBC Objective: A Compelling Idea with a Mathematical Flaw

The influential Implicit Behavior Cloning (IBC) work by Florence et al. [24] proposed training EBM policies using an InfoNCE-inspired objective:

L_ibc(θ, ξ) = E_{x,y} E_{y~p_ξ(·|x)} log[exp(E_θ(x,y)) / Σ_{y'∈{y}∪ỹ} exp(E_θ(x,y'))]

This objective avoids the intractable partition function entirely by casting the learning problem as contrastive classification: given one positive example y and K negative examples , can the model identify which is the real action? The appeal is clear—no MCMC sampling during training, no partition function estimation, just straightforward gradient computation. Florence et al. [24] demonstrated promising results on several robotic manipulation tasks.

However, the follow-up picture was starkly different from what the original IBC paper suggested. The authors cite multiple independent works that found IBC training to be problematic:

"follow up works have found training with the IBC objective to be quite unstable [15, 61, 64, 76], and hence the practicality of training and using EBMs as policies has remained an open question."

Chi et al. [15], in their influential diffusion policy paper, abandoned EBMs entirely in favor of score-based models, reporting that IBC-trained policies were simply outperformed. Reuss et al. [64] and Ta et al. [76] similarly documented training instabilities and poor final policy quality. Pearce et al. [61] found diffusion models more reliable for imitating human behavior. The collective message from the field was that EBMs, despite their theoretical appeal, could not be made to work in practice for behavior cloning.

This creates the specific gap the paper addresses: why does IBC fail, and can we fix it?

The IBC Population Bias: Explaining the Failure

The paper provides a crisp mathematical explanation that had been missing from the literature. The key insight is that the IBC objective is a special case of the R-NCE objective, but with a critical restriction: it implicitly forces the negative sampling distribution to be uniform over the action space.

To see this, compare the R-NCE objective (Equation 3.3):

ℓ_{θ,ξ}(x, y | ỹ) = log[exp(E_θ(x,y) - log p_ξ(y|x)) / Σ_{y'∈{y}∪ỹ} exp(E_θ(x,y') - log p_ξ(y'|x))]

with the IBC objective (Equation 4.3):

ℓ_{ibc}(x, y | ỹ) = log[exp(E_θ(x,y)) / Σ_{y'∈{y}∪ỹ} exp(E_θ(x,y'))]

The IBC objective omits the -log p_ξ terms entirely. Proposition 4.7 formalizes the consequence: the IBC objective is consistent (in the sense that its population optimum recovers the true distribution) only when the negative sampling distribution is uniform over the action space. When a non-uniform negative sampler is used—as is necessary in practice for high-dimensional spaces where uniform sampling is hopelessly inefficient—the IBC objective recovers not p(y|x) but rather the density ratio p(y|x) / p_ξ(y|x).

Figure 1 provides a clean one-dimensional demonstration of this bias. Using a simple Gaussian energy model class and a Gaussian proposal distribution, the paper shows that the IBC population objective (both for K=10 and K=100 negative samples) has its optimum shifted away from the true mean, while the R-NCE objective aligns with the maximum likelihood solution regardless of K. This is not a finite-sample or optimization issue—it is bias at the population level, meaning even with infinite data and perfect optimization, IBC would learn the wrong distribution.

The follow-up observation from Ta et al. [76] that IBC is only unbiased with a uniform proposal distribution now has a clear explanation: the InfoNCE objective was originally designed for mutual information estimation, where the goal is to distinguish positive pairs from negatives drawn from the marginal distribution. IBC misapplied this machinery to density estimation without accounting for the proposal distribution's influence on what is being estimated.

The Broader EBM Training Landscape: Why the Existing Alternatives Don't Work for Conditional Policies

Beyond IBC, the paper situates itself within the broader challenge of EBM training, where two main paradigms exist—and both have fatal flaws for the conditional policy setting.

Maximum Likelihood Estimation (MLE) and MCMC. The MLE gradient decomposes as:

∇_θ L_mle(θ) = E_{x,y}[∇_θ E_θ(x,y)] - E_{x} E_{y'~p_θ(·|x)}[∇_θ E_θ(x,y')]

The second term requires sampling from the EBM itself at every gradient step, which necessitates Markov Chain Monte Carlo (MCMC). In the unconditional setting, persistent chains can be maintained across training iterations to mitigate truncation bias. But for conditional EBMs—where the model is p_θ(y|x)—this approach breaks down because one would need to maintain separate MCMC chains for every distinct context x. When the context space is continuous or high-dimensional (as in robotics, where x includes images and robot states), this becomes computationally infeasible.

The variational approach—where an auxiliary generative model is trained to approximate EBM samples via max-min optimization—has also been explored [16, 29, 10], but the paper notes that "adversarial optimization is notoriously challenging and requires several stabilization tricks to prevent mode-collapse." The exponential tilting approach [82, 5, 60], where an EBM corrects a backbone generative model (VAE or normalizing flow), requires MCMC in the latent space using the pullback of the energy function, adding computational overhead from the transport map.

Noise Contrastive Estimation (NCE) and its Binary Limitation. The NCE framework avoids MCMC entirely by training a classifier to distinguish real samples from noise samples. For unconditional models, the binary NCE objective [31] is well-understood and consistent. However, Ma and Collins [57] identified a critical limitation: for conditional models, binary NCE requires the EBM class to be self-normalized to guarantee consistency. The ranking variant (R-NCE) overcomes this by using a multi-class objective where the model must identify which among K+1 candidates is the real sample.

The paper notes that prior work on adversarial NCE for the binary case [11, 26] explored jointly training the noise distribution, but ultimately settled on adversarial max-min optimization. Gao et al. [26] dismissed non-adversarial joint training as ineffective. The paper challenges this conclusion for the ranking setting, suggesting that the need for adversarial training in binary NCE may not carry over to R-NCE.

The Missing Piece: Multi-Scale Learning for EBMs

A separate but crucial motivation comes from recent advances in generative modeling. Diffusion models [37, 69, 70, 74] and stochastic interpolants [3, 4] share a common insight: learning the data distribution at multiple noise scales is critical for high-quality generation. By defining a continuous family of distributions indexed by a scale parameter t ∈ [0,1], where t=0 corresponds to pure noise and t=1 to the data, and training a shared network to denoise or transport across all scales simultaneously, these methods achieve implicit annealing that eases the generative process.

The paper asks: is this multi-scale insight compatible with EBMs? If so, the combination could yield EBMs that benefit from the same training dynamics that make diffusion models effective, while retaining the unique advantages of energy-based representations (composability, compactness, optimization-friendly structure). Prior to this work, no EBM training algorithm incorporated multi-scale training in a principled way.

How This Paper Positions Itself

The paper positions itself not as proposing an entirely new class of models, but rather as diagnosing and fixing the specific flaws that prevented EBMs from being competitive with diffusion policies. The positioning is:

  1. Against IBC: The paper provides both theoretical proof (Proposition 4.7) and empirical demonstration (Figures 1, 4; Tables 1, 3) that the IBC objective is fundamentally flawed—biased at the population level and unstable in practice—explaining the negative findings in [15, 61, 64, 76].

  2. Against adversarial training: Through asymptotic analysis (Theorem 4.18, Corollary 4.19), the paper shows that adversarial max-min training of the negative sampler provides no asymptotic benefit over non-adversarial joint training when the proposal family is not realizable (the practical setting). This challenges the conclusions of Bose et al. [11] and Gao et al. [26] for the binary NCE setting, and justifies the simpler, more stable Algorithm 1.

  3. Against the narrative that diffusion is strictly superior: By developing R-NCE + interpolating EBMs and demonstrating competitive or superior performance on path planning (Section 7.3) and block pushing (Section 7.4), the paper directly challenges the conclusion from Chi et al. [15] that diffusion models are the de facto standard for generative policies. The block pushing results in Table 6 show I-R-NCE achieving 0.884 ± 0.005 coverage versus 0.864 ± 0.006 for diffusion—a small but meaningful gap in the opposite direction.

  4. As a bridge between EBMs and modern generative modeling: The interpolating EBM framework (Section 5.2) shows that the core idea behind diffusion—learning across multiple noise scales—is not inherently tied to score matching or denoising objectives; it can be realized within the EBM paradigm through a time-indexed R-NCE objective. This is presented as a contribution of "independent interest to the generative modeling community."

The paper explicitly does not claim that EBMs are universally superior. The difficulty of training EBMs is acknowledged throughout, and the practical success depends crucially on design choices: (i) the R-NCE objective with learnable negative sampler, (ii) the negative sampler being sufficiently informative to prevent posterior collapse (Section 4.3), (iii) the interpolating framework for high-dimensional problems, and (iv) careful architectural decisions about the relative expressiveness of the EBM vs. the proposal model. The message is that with these ingredients in place, EBMs become a viable and sometimes preferable alternative—not that they dominate in all regimes.

3. Technical Approach

3.1 Reader Orientation

The paper builds a training system for energy-based models as robot policies—a procedure that takes a dataset of demonstrated state-action pairs and produces an energy function E_θ(x, y) whose learned "landscape" assigns high values (low energy) to expert-like actions and low values to poor actions, without needing to compute the intractable normalization constant during inference. The system solves the problem of reliable EBM training for high-dimensional continuous action spaces by replacing the biased IBC objective with a theoretically grounded ranking noise contrastive estimation (R-NCE) objective, jointly training a simpler generative model to supply informative negative samples, and—for the hardest problems—extending this to a multi-scale interpolating framework that learns a continuous family of EBMs bridging noise to data.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components:

  1. Base energy model E_θ(x, y) — a neural network mapping context x (e.g., images, state history) and action y to a scalar energy. Higher energy means the action is more likely under the learned policy; no normalization constant is needed because only relative comparisons matter during inference.

  2. Negative sampling distribution p_ξ(y|x) — a simpler generative model (a continuous normalizing flow, CNF) that is easy to both sample from and evaluate log-probabilities. Its role is to produce "contrast" examples that the energy model must learn to distinguish from real demonstrations.

  3. R-NCE objective L(θ, ξ) — the core training criterion: for each real sample (x, y), the system draws K negative samples from p_ξ(·|x), forms a (K+1)-way classification problem, and trains the energy model to assign the highest posterior probability to the real sample. The negative sampler's log-densities explicitly correct for the proposal distribution, eliminating the bias that plagues IBC.

  4. Interpolating EBM framework (I-EBM) — for high-dimensional problems, the system instantiates a time-indexed family of EBMs E_θ(x, y, t) covering the continuum from t=0 (pure noise) to t=1 (data). The negative sampler is a CNF trained via the stochastic interpolant objective, providing samples at any time t. A time-indexed R-NCE objective trains the shared EBM across all noise scales simultaneously.

  5. Sampling procedure — at inference time, the learned energy function is used to sample actions via MCMC (Langevin or Hamiltonian Monte Carlo), warm-started from the negative sampler's output. For I-EBM, a three-stage SDE + MCMC procedure transports a noise sample through the learned energy landscape.

Information flows as follows during training: a batch of context-action pairs (x, y) enters → the negative sampler generates K noise samples → the R-NCE loss computes a (K+1)-way classification score → gradients flow back to the energy model (and optionally to the sampler's standalone objective). At inference: context x enters → the sampler proposes an initial action → MCMC with the energy model refines it into a high-quality sample → the refined action is executed.

3.3 Roadmap for the Deep Dive

  • First, the R-NCE objective — its formal definition, posterior probability interpretation, and the role of the negative sampler's log-densities. This is the mathematical core that distinguishes R-NCE from IBC and explains why it works.
  • Second, the optimality analysis — what the R-NCE objective recovers at the population level (Theorem 4.4) and why the IBC objective fails (Proposition 4.7). Understanding this bias is essential before any implementation.
  • Third, Algorithm 1 (joint training) — the non-adversarial training loop that simultaneously optimizes the energy model via R-NCE and the negative sampler via its own objective (likelihood or interpolant loss). The mechanics of batch augmentation, gradient decoupling, and the rationale for avoiding adversarial training.
  • Fourth, the posterior collapse problem — why an uninformative negative sampler causes the R-NCE gradient to vanish, and how joint training prevents this.
  • Fifth, the interpolating EBM extension — how the time-indexed energy model E_θ(x, y, t) and I-R-NCE objective (Algorithm 3) bring multi-scale learning to EBMs, with the associated training and sampling modifications.
  • Sixth, the sampling procedures — Algorithm 2 (Langevin MCMC for standard EBMs) and Algorithm 4 (three-stage SDE transport + MCMC for I-EBMs), including the SDE formulation (Equation 5.6) that justifies combining the CNF vector field with the EBM's score function.
  • Seventh, the log-probability computation — the two-time-scale approach (Section 6) that makes ranking multiple action candidates computationally feasible by decoupling sample integration from divergence computation.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a theoretical + algorithmic paper whose core idea is that EBMs trained with the R-NCE objective—paired with a jointly-trained but non-adversarial negative sampler—achieve both asymptotic correctness and practical performance competitive with diffusion models, and that the IBC objective fails at even the population level due to a fundamental mathematical bias.


The R-NCE Objective: Formal Definition and Posterior Interpretation

The paper begins with a formal setup on compact measure spaces (X, µ_X) (contexts) and (Y, µ_Y) (events/actions). The true conditional distribution p(y|x) is assumed to be a regular conditional density. The energy model class F is defined as parameterized functions E_θ(x, y) : X × Y → R with θ ∈ Θ (a compact parameter space), inducing conditional densities through the standard exponential family form:

pθ(yx)=exp(Eθ(x,y))Zθ(x)p_θ(y|x) = \frac{\exp(E_θ(x, y))}{Z_θ(x)}

where $Z_θ(x) = \int_Y \exp(E_θ(x, y)) \, d\mu_Y$ is the intractable partition function (which normalizes the unnormalized energy to a valid probability density).

The critical operation is sampling K negative (contrastive) samples from a proposal distribution p_ξ(y|x), parameterized by ξ ∈ Ξ (also compact). For notational efficiency, let $y \sim P^K_{y|x;ξ}$ denote drawing K independent samples: $y = (y_k)_{k=1}^K$ with each $y_k \sim p_ξ(·|x)$. The R-NCE loss for a single (context, positive, negatives) tuple is:

θ,ξ(x,yy~)=log[exp(Eθ(x,y)logpξ(yx))y{y}y~exp(Eθ(x,y)logpξ(yx))]\ell_{θ,ξ}(x, y | \tilde{y}) = \log\left[\frac{\exp(E_θ(x, y) - \log p_ξ(y|x))}{\sum_{y' \in \{y\} \cup \tilde{y}} \exp(E_θ(x, y') - \log p_ξ(y'|x))}\right]

where $x$ is the context, $y$ is the true (positive) action, $\tilde{y} = (y_1, ..., y_K)$ are the K negative samples, $E_θ(x, y)$ is the scalar energy for the positive example, and $\log p_ξ(y|x)$ is the log-probability of the positive example under the negative sampler.

What it computes: the log-probability that a Bayes-optimal classifier would assign to the positive example $y$ being the true action, when given the K+1 candidates (one real, K noise) and told that exactly one is real. The numerator $\exp(E_θ - \log p_ξ)$ for the positive example is compared against the sum of the same quantity over all candidates. The log turns the ratio into a log-posterior. The network never sees or computes the intractable partition function $Z_θ(x)$ because it cancels identically in the ratio (appearing in both numerator and denominator of every $p_θ(y|x)$ term that would otherwise appear).

Why this form: The subtraction of $\log p_ξ$ is the essential correction that makes R-NCE consistent regardless of the proposal distribution. Without these terms (setting them to zero), the objective reduces to the IBC objective—which, as Proposition 4.7 shows, only recovers $p(y|x)$ when $p_ξ$ is uniform. With the correction, the energy model learns $E_θ(x, y) \approx \log p(y|x) + \text{const}(x)$, which is precisely what an EBM needs to represent (the constant in $x$ corresponds to the log partition function). The contrast between $E_θ$ and $\log p_ξ$ means that the energy model only needs to capture the difference between the true density and the proposal density, which is a far easier learning problem when $p_ξ$ is a reasonable approximation of $p$.

The full population objective (expectation over the data distribution and negative samples) is:

L(θ,ξ):=E(x,y)PX,YEyPyx;ξKθ,ξ(x,yy~)L(θ, ξ) := \mathbb{E}_{(x,y) \sim P_{X,Y}} \, \mathbb{E}_{y \sim P^K_{y|x;ξ}} \, \ell_{θ,ξ}(x, y | \tilde{y})

where the inner expectation draws K i.i.d. negative samples for each (x, y) pair.

Posterior probability interpretation. The paper derives an equivalent formulation that makes the connection to classification explicit. For a given x, let $\bar{y} = \{y\} \cup \tilde{y} = (y_j)_{j=1}^{K+1}$ be the full candidate set, and let $D \in \{1, ..., K+1\}$ be a random variable indicating the index of the true sample within $\bar{y}$. Under the energy model's generative assumption, the class-conditional distribution over $\bar{y}$ given $D = k$ is:

pˉθ,ξ(yˉx,D=k)=pθ(ykx)jkpξ(yjx)\bar{p}_{θ,ξ}(\bar{y} | x, D = k) = p_θ(y_k | x) \prod_{j \neq k} p_ξ(y_j | x)

(the true sample at index k is drawn from $p_θ$, all others from $p_ξ$). With a uniform prior $p(D=k|x) = 1/(K+1)$, Bayes' rule yields the posterior after cancellation of $Z_θ(x)$:

qθ,ξ(kx,yˉ)=exp(Eθ(x,yk)logpξ(ykx))j=1K+1exp(Eθ(x,yj)logpξ(yjx))q_{θ,ξ}(k | x, \bar{y}) = \frac{\exp(E_θ(x, y_k) - \log p_ξ(y_k|x))}{\sum_{j=1}^{K+1} \exp(E_θ(x, y_j) - \log p_ξ(y_j|x))}

The R-NCE objective $\ell_{θ,ξ}(x, y | \tilde{y})$ is exactly $\log q_{θ,ξ}(1 | x, \bar{y})$—the log-posterior probability that the first element (the positive) is the real one. Maximizing this objective trains the energy model to be a good Bayes classifier in this constructed (K+1)-way task.

Taking the gradient. The gradient of the R-NCE objective w.r.t. $θ$ reveals a crucial computational property:

θL(θ,ξ)=Ex,yθEθ(x,y)ExEyˉx;ξ[i=1K+1qθ,ξ(ix,yˉ)θEθ(x,yi)]\nabla_θ L(θ, ξ) = \mathbb{E}_{x,y} \nabla_θ E_θ(x, y) - \mathbb{E}_{x} \mathbb{E}_{\bar{y}|x;ξ} \left[\sum_{i=1}^{K+1} q_{θ,ξ}(i | x, \bar{y}) \nabla_θ E_θ(x, y_i)\right]

The first term ($\mathbb{E}_{x,y} \nabla_θ E_θ(x, y)$) pushes up energy on real examples. The second term pushes down energy on a weighted mixture of all candidates, where the weights are the classifier's own posterior probabilities. Compare this with the MLE gradient:

θLmle(θ)=Ex,yθEθ(x,y)ExEypθ(x)θEθ(x,y)\nabla_θ L_{\text{mle}}(θ) = \mathbb{E}_{x,y} \nabla_θ E_θ(x, y) - \mathbb{E}_{x} \mathbb{E}_{y' \sim p_θ(·|x)} \nabla_θ E_θ(x, y')

The MLE gradient requires sampling $y'$ from the EBM itself ($p_θ(·|x)$), which is expensive (MCMC). The R-NCE gradient replaces this with a weighted expectation over a finite set of $K+1$ easily-sampled points, where the weights $q_{θ,ξ}$ are computed through a single forward pass. This is efficient because: (i) the negative samples $\tilde{y}$ come from the fast sampler $p_ξ$, (ii) the posterior weights $q_{θ,ξ}$ involve only $K+1$ evaluations of $E_θ$ and $\log p_ξ$, and (iii) no MCMC chain needs to be run to convergence.

Crucially, the negative samples $\tilde{y}$ do NOT depend on $θ$ (they come from $p_ξ$). This is noted as a key implementation detail (Remark 3.2): "the batch augmentation in Line 11 can be performed independently of the optimization step in Line 12." This means the gradient through $\tilde{y}$ w.r.t. $θ$ is zero, eliminating a major computational bottleneck that adversarial training would introduce (which would require differentiating through the sampling process).


Optimality: Why R-NCE Works and IBC Fails

The paper first establishes that at the population level, R-NCE correctly recovers the true distribution regardless of the proposal. This is Theorem 4.4:

"Suppose that F is realizable. Then, for any ξ ∈ Ξ, we have Θ^* = \arg\max_{θ \in Θ} L(θ, ξ)."

Here, "realizable" means that there exists some $θ^∗ \in Θ$ such that $p_{θ^∗}(y|x) = p(y|x)$ almost everywhere—the energy model class is expressive enough to represent the true distribution. The set $Θ^*$ is the set of all such $θ^∗$ (i.e., all energy functions correctly representing the data distribution up to a constant shift).

What Theorem 4.4 asserts: if the energy model class can represent the truth, then every maximizer of the R-NCE population objective is a correct representation of the truth, regardless of the negative sampler used. You could use a terrible sampler (e.g., Gaussian noise completely unrelated to the action distribution) and the population optimum would still be correct—the sampler only affects the finite-sample efficiency, not the asymptotic correctness.

Why this holds (proof sketch): The key step uses a change-of-measure to rewrite the objective as a cross-entropy between two distributions over {1, ..., K+1}: (i) the posterior $q_θ$ and (ii) the unnormalized weights $p(y_k|x) / p_ξ(y_k|x)$. When $θ = θ^∗$, the posterior simplifies to:

qθ,ξ(kx,yˉ)=p(ykx)/pξ(ykx)j=1K+1p(yjx)/pξ(yjx)q_{θ^∗, ξ}(k | x, \bar{y}) = \frac{p(y_k|x) / p_ξ(y_k|x)}{\sum_{j=1}^{K+1} p(y_j|x) / p_ξ(y_j|x)}

which is exactly the normalized version of these weights. Claim C.1 establishes that for any given $(x, \bar{y})$, this choice pointwise maximizes the cross-entropy (by the concavity of log). Hence, $θ^∗$ pointwise maximizes the integrand, so no other $θ$ can achieve higher expected objective. Conversely, any $θ$ achieving the same maximum must produce identical posteriors, which algebraically forces $E_θ$ and $E_{θ^∗}$ to differ by an $x$-dependent constant—i.e., they represent the same conditional density $p(y|x)$.

The IBC bias. Proposition 4.7 formalizes why IBC fails. The IBC objective (Equation 4.3) is:

Libc(θ,ξ)=Ex,yEy~x;ξlog[exp(Eθ(x,y))y{y}y~exp(Eθ(x,y))]L_{\text{ibc}}(θ, ξ) = \mathbb{E}_{x,y} \mathbb{E}_{\tilde{y}|x;ξ} \log\left[\frac{\exp(E_θ(x, y))}{\sum_{y' \in \{y\} \cup \tilde{y}} \exp(E_θ(x, y'))}\right]

The paper shows that this is equivalent to R-NCE over a shifted function class $\tilde{F}_ξ = \{E_θ(x, y) + \log p_ξ(y|x) | θ \in Θ\}$. Applying Theorem 4.4 to this shifted class, the IBC objective's maximizers satisfy:

exp(Eθ(x,y)+logpξ(yx))exp(Eθ(x,y)+logpξ(yx))dμY=p(yx)pξ(yx)\frac{\exp(E_θ(x, y) + \log p_ξ(y|x))}{\int \exp(E_θ(x, y') + \log p_ξ(y'|x)) \, d\mu_Y} = \frac{p(y|x)}{p_ξ(y|x)}

What this means in practice: the IBC-trained energy model does NOT learn $\log p(y|x)$ (up to a constant). Instead, it learns $\log(p(y|x) / p_ξ(y|x))$—the log density ratio. To recover the actual action distribution, one would need to add $\log p_ξ(y|x)$ back, but this requires a known and tractable $p_ξ$.

When $p_ξ$ is uniform over $Y$, then $\log p_ξ(y|x) = \text{const}(x)$ (does not depend on $y$), so the shift is absorbed into the partition function and IBC becomes unbiased. This is why Florence et al. [24] could get IBC to work on low-dimensional tasks with uniform negative sampling. But in high-dimensional spaces (like 20D trajectory snippets or 32D action sequences), uniform sampling is hopelessly inefficient—the probability of drawing a negative sample anywhere near the data manifold is essentially zero. Any practical sampler $p_ξ$ must be non-uniform to provide informative negatives, but then IBC learns the wrong quantity.

Figure 1 demonstrates this concretely in 1D: with a Gaussian energy model class and Gaussian proposal N(0, 1), the IBC population objective (both at K=10 and K=100 negatives) has its maximum at μ ≈ 0.5, shifted toward the proposal's mean (0) from the true mean (1). The R-NCE objective aligns with MLE at μ = 1 for all K. The shift is deterministic—this is population-level bias, not optimization noise.


Algorithm 1: Joint Training of EBM and Negative Sampler

The training procedure (Algorithm 1) alternates between updating the negative sampler $p_ξ$ and the energy model $E_θ$, but non-adversarially: the sampler is trained to maximize likelihood of the data (or an auxiliary interpolant loss), while the energy model is trained via R-NCE using the sampler's current samples. The sampler never receives gradients from the R-NCE loss.

The algorithm's outer loop runs for $T_{\text{outer}}$ iterations. Within each outer iteration:

Step 1: Update negative sampler (Lines 5-8). For $T_{\text{samp}}$ steps:

  • Sample a mini-batch $B = \{(x_i, y_i)\}$ from the dataset.
  • Compute the sampler's objective $\hat{L}_{\text{mle}}(ξ; B) = -\sum_{(x_i, y_i) \in B} \log p_ξ(y_i | x_i)$ (negative log-likelihood).
  • Take a gradient step to minimize this w.r.t. $ξ$.

The paper notes (Remark 3.1) that any family of densities works for the negative sampler—the only requirement is efficient sampling and log-probability evaluation. For their experiments, they use continuous normalizing flows (CNFs) trained via the stochastic interpolant framework (not via MLE), in which case Line 7 is replaced with the interpolant auxiliary loss.

Step 2: Update energy model (Lines 9-13). For $T_{\text{rnce}}$ steps:

  • Sample a mini-batch $B$ from the dataset.
  • Batch augmentation (Line 11): For each $(x_i, y_i) \in B$, draw $K$ negative samples $\tilde{y}_i \sim P^K_{y|x_i;ξ}$ from the current negative sampler, producing augmented triples $(x_i, y_i, \tilde{y}_i)$.
  • Compute the R-NCE loss $\hat{L}_{\text{rnce}}(θ, ξ; \bar{B}) = -\sum_{(x_i, y_i, \tilde{y}_i) \in \bar{B}} \ell_{θ,ξ}(x_i, y_i | \tilde{y}_i)$ where $\ell_{θ,ξ}$ is defined in Equation 3.3.
  • Take a gradient step w.r.t. $θ$ only (the negative samples $\tilde{y}_i$ are treated as constants; no gradient flows back through the sampling process to $ξ$).

Why this non-adversarial design matters. The key property enabling this decoupling is that the optimal $θ$ from R-NCE is independent of $ξ$ (Theorem 4.4). Therefore, even though the sampler is suboptimal (it's a simpler model than the EBM), the EBM will converge to the correct distribution as long as the R-NCE objective is consistently optimized. The sampler's role is purely to provide informative negatives that prevent the R-NCE gradient from vanishing—it does not need to be a perfect generative model.

This contrasts with the adversarial approach in binary NCE [11, 26], where the sampler is trained to maximize the NCE objective (i.e., fool the classifier), creating a min-max game. The paper's asymptotic analysis (Section 4.3.2) shows that in the ranking setting, adversarial training provides no theoretical benefit when the sampler class $F_n$ is not realizable (the practical case), since the asymptotic variance of $θ$ does not improve. Moreover, adversarial training would require differentiating through the negative samples w.r.t. $ξ$ in the R-NCE loss—a computational bottleneck the non-adversarial approach avoids entirely.

The pre-training alternative (train the sampler first, freeze it, then train the EBM) was explored by Gao et al. [26] for binary NCE and found ineffective. The paper replicates this finding for R-NCE in path planning (Section 7.3): a pre-trained proposal model yielded higher collision rate (5.8% vs. 3.0%) and higher cost (0.24 vs. 0.194) compared to joint training. Joint training allows the sampler to adapt to the EBM's current state, maintaining a productive difficulty level for the classification task.


Posterior Collapse: Why the Negative Sampler Must Be Learned

Section 4.3 provides the crucial intuition for why a learnable negative sampler is practically necessary, even though Theorem 4.4 says any fixed sampler is asymptotically correct. The issue is optimization, not asymptotics.

When the negative sampler $p_ξ$ is uninformative (e.g., a broad Gaussian far from the data), the energy model can easily distinguish positives from negatives early in training. The posterior $q_{θ,ξ}(1 | x, \bar{y})$—the classifier's confidence that the first element is real—approaches 1.0 because the real sample's energy towers over the noise samples' energies.

The R-NCE gradient becomes:

θL(θ,ξ)=Ex,yθEθ(x,y)ExEyˉx;ξ[i=1K+1qθ,ξ(ix,yˉ)θEθ(x,yi)]\nabla_θ L(θ, ξ) = \mathbb{E}_{x,y} \nabla_θ E_θ(x, y) - \mathbb{E}_x \mathbb{E}_{\bar{y}|x;ξ} \left[\sum_{i=1}^{K+1} q_{θ,ξ}(i | x, \bar{y}) \nabla_θ E_θ(x, y_i)\right]

When $q_{θ,ξ}(1 | x, \bar{y}) \approx 1$ and $q_{θ,ξ}(i | x, \bar{y}) \approx 0$ for $i > 1$, the weighted sum becomes approximately $\nabla_θ E_θ(x, y_1)$ (the gradient on the positive). The R-NCE gradient then reduces to approximately:

θL(θ,ξ)Ex,yθEθ(x,y)Ex,yθEθ(x,y)0\nabla_θ L(θ, ξ) \approx \mathbb{E}_{x,y} \nabla_θ E_θ(x, y) - \mathbb{E}_{x,y} \nabla_θ E_θ(x, y) \approx 0

The gradient vanishes even though $E_θ$ may not yet equal $\log p(y|x)$ up to a constant. The model has won the classification game too easily, and training stalls. The paper calls this posterior collapse—the classifier becomes overconfident, and the gradient signal needed to refine the energy surface disappears.

Why joint training prevents this. As the energy model improves at classification, the sampler also improves (via its own likelihood/interpolant objective), producing negatives that are closer to the data distribution. This maintains the classification difficulty in a productive regime where the posterior is not degenerate, gradients remain informative, and both models improve together. The paper visualizes this dynamic in Figure 2, distinguishing several training regimes:

  • Posterior collapse (blue curve): The sampler is completely uninformative; R-NCE loss plateaus far from $-\log(K+1)$ (the theoretical minimum when $p_θ = p_ξ = p$).
  • EBM learns faster (yellow/pink curves): R-NCE loss first increases (EBM improves classification), then decreases as the sampler catches up. Final value depends on relative model capacities.
  • Sampler learns faster (green/orange curves): Occurs with pre-trained or initially-strong samplers; loss approaches final value from below.

The takeaway: the sampler should be informative enough to prevent posterior collapse, but not so powerful that it perfectly models the data (which would make the classification task impossible and the EBM unnecessary). The practical choice is a simpler architecture for the sampler (e.g., a dense MLP CNF vs. a transformer-based EBM) trained jointly.


Asymptotic Analysis: Consistency, Normality, and Efficiency

The paper provides three key asymptotic results that underpin the algorithmic choices:

Theorem 4.8 (General Consistency): As the number of training samples $n \to \infty$, any empirical risk maximizer $\hat{θ}_n = \arg\max_θ L_n(θ, \hat{ξ}_n)$ converges almost surely to the true parameter set $Θ^*$, regardless of the sequence $\{\hat{ξ}_n\}$ of sampler parameters. The only requirement on the sampler is that its parameters converge to some fixed point—they need not be optimal, or even good. This is a stronger statement than typical M-estimator consistency because the empirical objective $L_n(θ, \hat{ξ}_n)$ is not an unbiased estimate of $L(θ, \hat{ξ}_n)$ (since $\hat{ξ}_n$ depends on the same data used to compute $L_n$). The key proof step leverages Theorem 4.4: since $θ \mapsto L(θ, ξ)$ has the same maximizers for all $ξ$, one can show uniform convergence over $Θ \times Ξ$ jointly.

Theorem 4.10 (Asymptotic Normality): Under additional regularity (C² smoothness, Lipschitz Hessian), the estimator $\hat{θ}_n$ is asymptotically normal:

n(θ^nθ)N(0,Vθ)\sqrt{n}(\hat{θ}_n - θ^∗) \leadsto N(0, V_θ)

where $V_θ = -[\nabla_θ^2 L(θ^∗, ξ^∗)]^{-1}$. The proof exploits the fact that $θ^∗$ is optimal for $L(·, ξ)$ for any $ξ$, which forces the cross-Hessian $\nabla_{θξ}^2 L(θ^∗, ξ^∗) = 0$, making the $(θ, θ)$ block of the joint asymptotic covariance decouple from the sampler's convergence behavior. The only requirement is that $\hat{ξ}_n$ is $\sqrt{n}$-consistent to some $ξ^∗$.

Theorem 4.12 (Near-Asymptotic Efficiency): When the negative sampler converges to the true distribution ($p_{ξ^∗}(y|x) = p(y|x)$), the asymptotic variance simplifies to:

Vθ=(1+1K)Ex[I(θx)]1V_θ = \left(1 + \frac{1}{K}\right) \mathbb{E}_x[I(θ^∗ | x)]^{-1}

where $I(θ^∗ | x)$ is the conditional Fisher information matrix (the Cramér-Rao lower bound for the MLE). This shows that R-NCE with $K$ negatives achieves efficiency that is a factor $(1 + 1/K)$ worse than the MLE—a small penalty for $K \geq 31$ (the value used in experiments). Importantly, this $1/K$ factor is multiplicative (not additive with unknown constants), making it significantly sharper than prior bounds in Ma and Collins [57] that had $O(1/\sqrt{K})$ with hidden constants depending on the proposal distribution.

Why this matters for algorithm design: The near-efficiency result strongly motivates training the sampler to approximate the data distribution (since $p_ξ \approx p$ yields near-MLE efficiency), but does NOT require it. The sampler can be non-realizable (a simpler model) and the EBM remains consistent—just with potentially larger finite-sample variance. The non-adversarial approach sacrifices some asymptotic efficiency for training stability and computational simplicity.


Interpolating EBMs: Multi-Scale Energy Models

Section 5 introduces the paper's second major algorithmic contribution: a framework for training EBMs across multiple noise scales, analogous to how diffusion models and stochastic interpolants operate.

Stochastic interpolant background. The paper builds on the framework of Albergo and Vanden-Eijnden [3], which constructs a continuous stochastic process bridging noise and data. For a given context $x$, define an interpolant function $I_t : Y \times Y \to Y$ satisfying $I_0(z, y) = z$ and $I_1(z, y) = y$. A time-indexed random variable is defined as:

ytx=It(z,y),zN(0,I),yPYx,zyy_t | x = I_t(z, y), \quad z \sim N(0, I), \quad y \sim P_{Y|x}, \quad z \perp y

At $t=0$, $y_0 = z$ (standard Gaussian noise); at $t=1$, $y_1 = y$ (the data distribution). The example interpolant given is $I_t(z, y) = \cos(\frac{\pi}{2}t)z + \sin(\frac{\pi}{2}t)y$, which linearly interpolates in a trigonometric parameterization. The law of this process, denoted $p_t(y_t | x)$, satisfies a continuity equation $\partial_t p_t + \text{div}_y(v_t p_t) = 0$ for a specific time-varying velocity field $v_t$. This velocity field can be learned by minimizing a simple quadratic objective (Equation 5.3):

minξE(x,y),tΛ,z[v^t,ξ(x,It(z,y))22tIt(z,y)v^t,ξ(x,It(z,y))]\min_{ξ} \mathbb{E}_{(x,y), t \sim \Lambda, z} \left[\|\hat{v}_{t,ξ}(x, I_t(z, y))\|^2 - 2 \, \partial_t I_t(z, y) \cdot \hat{v}_{t,ξ}(x, I_t(z, y))\right]

which avoids the expensive Jacobian computations needed for maximum likelihood training of continuous normalizing flows. The approximate CNF $\hat{v}_{t,ξ}$ defines an approximate distribution $p_t^{ξ}$ at each $t$ via the pushforward of the ODE $dy_t/dt = \hat{v}_{t,ξ}(x, y_t)$ with initial condition $y_0 = z \sim N(0, I)$.

Building I-EBMs. The key idea is to define a time-indexed energy model $E_θ(x, y, t)$—a single neural network that takes the scale $t \in [0, 1]$ as an additional input and represents the distribution at that scale:

pt(ytx)exp(Eθ(x,yt,t))p_t(y_t | x) \propto \exp(E_θ(x, y_t, t))

At each time $t$, the EBM $E_θ(x, ·, t)$ aims to model $p_t(·|x)$, and the truncated CNF (solved from 0 to $t$) provides the negative samples $p_t^{ξ}(·|x)$ for the R-NCE objective at that scale. The time-indexed R-NCE loss (Equation 5.4) is:

θ,ξ(x,yt,ty~t)=log[exp(Eθ(x,yt,t)logptξ(ytx))y{yt}y~texp(Eθ(x,y,t)logptξ(yx))]\ell_{θ,ξ}(x, y_t, t | \tilde{y}_t) = \log\left[\frac{\exp(E_θ(x, y_t, t) - \log p_t^{ξ}(y_t | x))}{\sum_{y' \in \{y_t\} \cup \tilde{y}_t} \exp(E_θ(x, y', t) - \log p_t^{ξ}(y' | x))}\right]

where $y_t = I_t(z, y)$ is the positive sample (from the interpolant bridge at time $t$, constructed from a real data point $y$ and noise $z$), and $\tilde{y}_t \sim (p_t^{ξ})^{K}$ are $K$ negative samples from the truncated CNF. The overall I-R-NCE objective (Equation 5.5) takes an expectation over $t \sim \Lambda$ (a distribution over $[0, 1]$, e.g., uniform):

L(θ,ξ)=EtΛExPXEytPYtxEy~t(ptξ)K[θ,ξ(x,yt,ty~t)]L(θ, ξ) = \mathbb{E}_{t \sim \Lambda} \mathbb{E}_{x \sim P_X} \mathbb{E}_{y_t \sim P_{Y_t|x}} \mathbb{E}_{\tilde{y}_t \sim (p_t^{ξ})^{K}} [\ell_{θ,ξ}(x, y_t, t | \tilde{y}_t)]

Training modifications (Algorithm 3). The key differences from Algorithm 1:

  1. Multiple time samples per data point. For each $(x_i, y_i)$ in a batch, sample $m$ times $\{t_{ij}\}_{j=1}^m \sim Λ$ and $m$ noise vectors $\{z_{ij}\}_{j=1}^m \sim N(0, I)$. Construct positive samples $y_{ij} = I_{t_{ij}}(z_{ij}, y_i)$. This variance reduction step ($m > 1$) averages the loss over multiple scales per data point.

  2. Correlated negative samples across time (Remark 5.1). To avoid generating $m \times K$ independent negative sample sets, the algorithm generates a single set of $K$ trajectories spanning the interval $(0, \max_j t_{ij}]$ and sub-samples these at the $m$ intermediate times. This exploits temporal correlation in the CNF's trajectories, significantly reducing computational cost.

  3. Sampler objective. The CNF is trained via the stochastic interpolant loss (Equation 5.3) rather than negative log-likelihood, since this avoids the expensive ODE-based log-probability computation within the sampler's training loop.

Sampling from I-EBMs (Algorithm 4). At inference, sampling proceeds in three stages:

Stage 1: Initial sample from CNF (Line 1). Sample $y_{\underline{t}} \sim p_{\underline{t}}^{ξ}(·|x)$, where $\underline{t} \in (0, 1)$ is a hyperparameter (set to 0.45 or 0.5 in experiments). The initial sample time is chosen significantly larger than 0 because the CNF is relatively accurate at early times (where the distribution is close to Gaussian and easy to model), and starting too close to $t=0$ would waste computation on trivial dynamics.

Stage 2: SDE transport to $t=1$ (Line 2). The paper exploits the stochastic equivalent of the transport PDE (Equation 5.6) from Albergo et al. [4]:

dyt=(vt(x,yt)+ηylogpt(ytx))dt+2ηdWtdy_t = \left(v_t(x, y_t) + \eta \nabla_y \log p_t(y_t | x)\right) dt + \sqrt{2\eta} \, dW_t

where $v_t$ is the true velocity field from the interpolant, $\nabla_y \log p_t$ is the score function, and $\eta \geq 0$ controls the noise level. The approximation replaces $v_t$ with the learned CNF $\hat{v}_{t,ξ}$ and $\nabla_y \log p_t$ with the I-EBM's gradient $\nabla_y E_θ(x, y_t, t)$:

v^tη(x,y)=v^t,ξ(x,y)+ηyEθ(x,y,t)\hat{v}_t^{\eta}(x, y) = \hat{v}_{t,ξ}(x, y) + \eta \nabla_y E_θ(x, y, t)

The SDE is integrated from $t = \underline{t}$ to $t = 1$ using this approximate drift. The noise term $\sqrt{2\eta} \, dW_t$ provides stochastic exploration that helps correct errors accumulated from the approximation. The paper justifies the SDE over the simpler ODE (Equation 5.2) by citing a result from Albergo et al. [4, Lemma 2.20]: the KL divergence between the true and approximate final distributions is bounded by the integrated drift approximation error divided by $4\eta$. With the ODE ($\eta = 0$), there is no such bound—small drift errors can accumulate into large distributional errors. The I-EBM's score $\nabla_y E_θ$ provides the missing piece ($\nabla_y \log p_t$) that the CNF alone cannot supply, enabling the SDE to correct distributional errors.

Stage 3: MCMC refinement at $t=1$ (Line 3). Run $T_{\text{mcmc}}$ steps of Hamiltonian Monte Carlo (HMC) with the EBM at $t=1$, starting from the SDE sample. The total number of $\nabla_y E_θ$ evaluations is split between the SDE and MCMC stages proportionally to $\underline{t}$: $\lfloor \underline{t} \cdot T_{\text{mcmc}} \rfloor$ steps for SDE, $\lfloor (1 - \underline{t}) \cdot T_{\text{mcmc}} \rfloor$ for HMC. This ensures consistent total inference cost regardless of the initial time.


Log-Probability Computation for Ranking Samples

A critical operational detail for policy execution is the ability to rank multiple candidate action sequences by their relative likelihood. For example, in path planning, the system samples $\ell = 48$ trajectory snippets and selects the one with the highest log-probability (Section 7.3). For EBMs, this is straightforward—$E_θ(x, y)$ directly provides (unnormalized) log-probabilities up to an $x$-dependent constant, so ranking is just comparing energy values. But for CNFs and diffusion models, computing exact log-probabilities requires solving an augmented ODE.

Section 6 describes the standard formula for log-probability under a continuous normalizing flow. Given the ODE $dz/dt = v_t(x, z(t))$ with initial condition $z(0) \sim p_0(·|x)$, the log-probability at time $T$ is:

logp(z(T)x)=logp0(z(0)x)0Tdivz(vt(x,z(t)))dt\log p(z(T) | x) = \log p_0(z(0) | x) - \int_0^T \text{div}_z(v_t(x, z(t))) \, dt

where $\text{div}_z(v_t) = \text{tr}(\partial v_t / \partial z)$ is the divergence (sum of diagonal entries of the Jacobian). Computing this divergence exactly scales as $O(d^2)$ where $d$ is the action dimension, because each column of the Jacobian requires a separate vector-Jacobian product.

The two-time-scale approach. The paper's key implementation insight is to decouple the time resolution for sampling from the time resolution for divergence computation. The procedure:

  1. Integrate the sample trajectory $z(t)$ at fine resolution (e.g., 150 Heun steps for Push-T sampling, ensuring high sample quality).
  2. At a coarse subset of these integration times (e.g., 15 steps for log-probability in Push-T, or 64 for path planning), compute and save the divergence $\text{div}_z(v_t)$.
  3. After sampling is complete, integrate the divergence values using 1D trapezoidal integration (numpy.trapz) to obtain the total log-probability correction.

This is possible because the sample variable $z$ and the log-probability variable $ψ$ can be integrated separately—the ODE for $z$ does not depend on $ψ$, and the ODE for $ψ$ only depends on $z$ through the divergence integrand, which is evaluated post-hoc at the saved states.

Why this matters. In the standard approach (augmented ODE where both $z$ and $ψ$ are integrated simultaneously at the same resolution), the quadratic cost of divergence computation is incurred at every step. For fine time resolutions needed for sample quality, this is prohibitively slow. The two-time-scale approach reduces the number of divergence evaluations from ~500 (sampling steps) to ~50 (log-probability steps), making ranking of $\ell = 48$ candidate samples computationally feasible. The paper notes that randomized trace estimation (Hutchinson's estimator) is not applicable here because: "unless the Hutchinson trace estimator is applied $Ω(d)$ times (thus negating the computational benefits), the variance of the estimator will overwhelm the signal needed for ranking samples" (citing the Hanson-Wright inequality).


Design Choices and Their Justifications

Choice 1: R-NCE over binary NCE. Binary NCE requires the energy model class to be self-normalized for consistency in the conditional setting [57, Theorem 3.1]. R-NCE removes this constraint by using a multi-class objective where the partition function cancels automatically. This is critical for flexible EBM architectures (transformers, ResNets) that have no built-in normalization.

Choice 2: Non-adversarial joint training over adversarial. The asymptotic analysis (Proposition 4.13 and Theorem 4.18) shows that when the sampler class is not realizable—the practical setting—adversarial training provides no asymptotic variance improvement over non-adversarial joint training. Moreover, adversarial training would require differentiating the R-NCE loss through the CNF's sampling process, a major computational burden. The non-adversarial approach keeps the sampler's training independent of the EBM's gradients.

Choice 3: Simpler architecture for the negative sampler. The sampler is intentionally less expressive than the EBM (dense MLP CNF vs. transformer EBM in path planning; Conv1D-Multi-Head-FiLM blocks vs. self-attention in Push-T). This ensures: (i) the sampler remains fast enough for the $K$ negative samples per data point needed at every R-NCE step, (ii) the sampler does not perfectly model the data (which would make the classification task trivial), and (iii) the EBM has room to learn structure beyond what the simple sampler captures.

Choice 4: Interpolating framework for high-dimensional problems. The standard (single-scale) R-NCE works for 20-dimensional path planning but fails on 32-dimensional block pushing (Table 6 shows R-NCE at 0.824 vs. I-R-NCE at 0.884). The multi-scale training provides implicit annealing—at early times t \ll 1, the distribution is close to Gaussian and easy to learn, and the EBM gradually refines across scales. This mirrors the dynamics that make diffusion models effective, but realized through the R-NCE objective rather than score matching.

Choice 5: SDE + MCMC inference over pure MCMC. The three-stage sampling (CNF warm-start → SDE transport → MCMC refinement) distributes the inference computation across different mechanisms: the CNF provides a reasonable initial guess, the SDE uses the EBM's score to correct distributional drift, and MCMC does final refinement. The ablation in Table 6 shows that dropping the SDE stage (using Algorithm 2 instead of Algorithm 4) reduces Push-T coverage from 0.884 to 0.880—a small but consistent gap suggesting the SDE contributes distribution-level correction that MCMC alone cannot efficiently achieve.

4. Key Insights and Innovations

Innovation 1: The IBC Objective Is Structurally Broken at the Population Level — Not Just an Optimization Problem

The paper's most conceptually important contribution is not proposing a better algorithm, but providing a definitive mathematical diagnosis of why the field's prior approach to EBM behavioral cloning failed. Before this work, the empirical instability of IBC-trained policies—documented by Chi et al. [15], Pearce et al. [61], Reuss et al. [64], and Ta et al. [76]—was understood as a practical training difficulty: perhaps the optimization landscape was poorly conditioned, or the gradients were noisy, or the hyperparameters needed more careful tuning. The dominant response was to abandon EBMs entirely and switch to diffusion models.

Proposition 4.7 changes the conversation from "IBC is hard to train" to "IBC learns the wrong quantity regardless of how perfectly you optimize it." The proof shows that the IBC objective is equivalent to R-NCE over a shifted function class F̃_ξ = {E_θ(x, y) + log p_ξ(y|x)}, which means its population optimum recovers the density ratio p(y|x) / p_ξ(y|x) rather than p(y|x) itself. This is not a finite-sample issue, not a noisy gradient issue, not a hyperparameter sensitivity—it is bias that persists in the limit of infinite data and perfect optimization. The energy model trained with IBC and a non-uniform negative sampler literally converges to the wrong function.

This is a fundamental, not incremental, insight. It overturns the premise of Florence et al. [24], which presented IBC as a principled adaptation of InfoNCE to density estimation. The paper identifies the precise mathematical error: InfoNCE was designed for mutual information estimation (where the goal is to distinguish positive pairs from marginals), and its naive application to conditional density estimation fails to account for how the negative sampling distribution alters the quantity being estimated. The field had been debugging symptoms (training instability, poor sample quality) while missing the root cause (population-level bias).

Figure 1 makes this diagnosis concrete even in a toy setting: with a Gaussian energy model class and Gaussian proposal, the IBC population objective visibly peaks at the wrong parameter value, while R-NCE aligns with the maximum likelihood estimator. The shift is deterministic and predictable from the proposal distribution's mean. This single figure reframes the entire conversation around IBC: the problem is not if it works, but what it recovers, and the answer depends on the negative sampler in a way that prior work did not recognize.

The significance beyond performance is that this diagnosis explains the pattern of contradictory results in the literature. Florence et al. [24] got IBC to work on low-dimensional tasks where uniform negative sampling was feasible (and therefore IBC was approximately unbiased), while follow-up work on high-dimensional problems found it unreliable (because practical non-uniform samplers introduced bias). The paper's framework makes these previously-conflicting findings coherent: they were testing IBC in different bias regimes. This is the kind of conceptual contribution that allows the field to stop debating whether IBC works and instead understand under what conditions it works, and why those conditions don't hold in the settings that matter.

Innovation 2: Adversarial Training of the Negative Sampler Is Unnecessary — And the Proof Shows Why

A second conceptual shift concerns the relationship between the energy model and its negative sampler. The prevailing approach in the NCE literature—particularly the adversarial formulations of Bose et al. [11] and Gao et al. [26]—treated the negative sampler as an adversary that should be trained to maximize the NCE objective, creating a GAN-like min-max game where the sampler tries to fool the classifier and the classifier tries to distinguish real from fake. The intuition was that "hard negatives" (samples close to the decision boundary) provide the most informative gradients, and that jointly optimizing both models adversarially would produce the best energy model.

The paper's asymptotic analysis (Section 4.3.2, culminating in Corollary 4.19) demonstrates that under the practical condition where the negative sampler class is not realizable (i.e., cannot perfectly model the data distribution), the asymptotic variance of the energy model parameters is identical whether the sampler is trained adversarially or non-adversarially. Both approaches yield √n(θ̂_n - θ^*) ⇝ N(0, -[∇²_θ L(θ^*, ξ̄)]⁻¹), where ξ̄ is the Stackelberg equilibrium adversary in the adversarial case and whatever fixed point the non-adversarial sampler converges to in the joint training case. The variance depends on the sampler's fixed point ξ̄, not on how the sampler arrived there.

This result is conceptually significant not because it introduces a new training technique—Algorithm 1's non-adversarial design is simple—but because it eliminates a whole class of optimization difficulties that prior work accepted as necessary. Adversarial optimization is notoriously unstable, requiring careful balancing of generator and discriminator learning rates, architectural constraints, and stabilization tricks to prevent mode collapse. By proving that this complexity is theoretically unnecessary for R-NCE in the non-realizable setting, the paper justifies a dramatically simpler training procedure that avoids:

  • Differentiating through the negative sampling process (which would require backpropagating through the CNF's ODE solver for every R-NCE step)
  • The max-min instability where improvements to one model degrade the other
  • The need to carefully schedule the relative update frequencies of the two models to maintain productive competition

The gap between this result and prior work is not incremental—it's a different category of claim. Bose et al. [11] and Gao et al. [26] argued, based on empirical evidence in the binary NCE setting, that adversarial training was necessary for good performance. The paper challenges this at the theoretical level, showing that the ranking variant of NCE has structural properties (specifically, the optimal θ^* being independent of ξ per Theorem 4.4) that make adversarial optimization irrelevant to asymptotic correctness and efficiency. The negative results with adversarial training in binary NCE may be specific to that setting and don't carry over to R-NCE.

The practical upshot is that Algorithm 1 achieves stability through a different mechanism: by simultaneously improving the sampler's standalone objective (likelihood or interpolant loss) alongside the EBM's R-NCE objective, without coupling their gradients. Both models improve together, and the sampler's improvement naturally provides harder negatives as the EBM gets better—but this emergent curriculum arises from independent optimization rather than adversarial dynamics. The training curves in Figure 2 characterize this as a phase transition where the R-NCE loss initially rises (EBM learns faster), then falls as the sampler catches up, approaching its final value without the oscillatory instability typical of GAN training.

Innovation 3: Multi-Scale Learning Is Not Inherently Tied to Score Matching — EBMs Can Do It Too

The third major conceptual contribution is demonstrating that the core idea behind diffusion models and stochastic interpolants—learning the data distribution across a continuum of noise scales with a shared network—is not specific to score-based or denoising objectives. The Interpolating EBM (I-EBM) framework (Section 5.2) shows that the same multi-scale training dynamic can be realized within the EBM paradigm, using a time-indexed R-NCE objective rather than score matching or denoising autoencoding.

This is a reframing contribution: it separates the architectural pattern (learning across scales) from the training objective (score matching), showing that these are independent design choices. Prior to this work, the generative modeling community had largely conflated the two. Diffusion models [37, 74] and stochastic interpolants [3, 4] package multi-scale learning with specific objective functions (denoising score matching, velocity field regression), and the success of these models created an implicit assumption that the objectives were necessary for the benefits. The paper shows they are not.

The I-EBM architecture—a single energy network E_θ(x, y, t) paired with a time-indexed R-NCE objective where negatives come from a truncated CNF at each scale—preserves all the distinctive advantages of EBMs (composability, compactness, density representation) while gaining the training benefits of multi-scale learning. The key insight is that at each scale t, the R-NCE objective only requires samples from the time-t distribution p_t(y_t|x) (constructed via the interpolant bridge) and negative samples from the approximate CNF at the same scale. Since the CNF provides samples at any t by truncating its flow, the R-NCE machinery applies with minimal modification.

The significance beyond the specific result (I-R-NCE outperforms single-scale R-NCE on Push-T, Table 6) is that this opens a design space for energy-based generative models that inherits the empirical benefits of diffusion while retaining the theoretical advantages of EBMs. An energy model trained this way can be composed with other energies at inference time (for task combination, constraint satisfaction, or prior integration), can be used as a cost function in optimization-based planning, and provides explicit density estimates—all capabilities that diffusion models lack or handle awkwardly.

The empirical evidence for this innovation being fundamental rather than incremental comes from the ablation in Table 6: single-scale R-NCE achieves 0.824 coverage on Push-T, while I-R-NCE (even with the simpler two-stage sampling from Algorithm 2, which uses the I-EBM only at t=1) achieves 0.880—a gap of over 5 percentage points from the same training algorithm applied at one scale versus many scales. This is not a small refinement; it's the difference between EBMs being uncompetitive and being state-of-the-art. The multi-scale training provides "implicit annealing" where the shared network benefits from learning easy distributions (near t=0, close to Gaussian) before tackling the full data distribution at t=1, and this benefit is orthogonal to the choice of objective function.

Innovation 4: Verifier Over-Optimization Is Not Just an RLHF Problem — It Governs EBM Training Too, and the Negative Sampler Is the Control Knob

The paper introduces and names posterior collapse—a specific failure mode in contrastive EBM training where the negative sampler is too weak, causing the classifier to become overconfident and the gradient to vanish—and identifies it as the primary practical bottleneck that the theoretical analysis (Theorem 4.4) does not address. This is a diagnostic concept, not an algorithmic contribution, but it has significant implications for how practitioners should think about the design of NCE-based training systems.

The mechanism (described in Section 4.3 and revisited in Section 4.4) is that when p_ξ is uninformative relative to the current E_θ, the posterior q_{θ,ξ}(1 | x, ȳ) concentrates at 1.0, making the R-NCE gradient approximately zero regardless of whether E_θ has converged to the correct density. Training stalls not because the objective is at its optimum, but because the optimization landscape has become flat—the model has won the classification game too easily and can no longer extract gradient signal from its victories.

This concept is significant because it explains a practical failure mode that is invisible in asymptotic theory. Theorem 4.4 guarantees that if you could find the global optimum of the R-NCE objective, it would be correct regardless of the sampler. But posterior collapse means you may never reach that optimum with gradient-based optimization, because the sampler's weakness creates a flat region in the loss landscape far from the true optimum. This is the optimization analog of the IBC population bias: even though R-NCE's population optimum is correct (unlike IBC), the path to reaching it can be blocked by an inadequate sampler.

The conceptual move is to reframe the sampler's role from "asymptotic variance reducer" to "optimization enabler." Prior NCE theory [31, 57] focused on the sampler's influence on statistical efficiency (how many samples needed for a given estimator variance), with the understanding that any sampler works in the limit. The posterior collapse analysis shows that in practice, the sampler must be informative enough to keep the R-NCE objective's gradient non-zero throughout training—a requirement that is about optimization dynamics, not asymptotics.

This connects to a broader principle that appears across machine learning: when using contrastive objectives, the difficulty of the contrastive task must be tuned to the model's current capability. Too easy (weak negatives), and the gradient vanishes. Too hard (negatives indistinguishable from positives), and the signal is pure noise. The optimal difficulty is in a Goldilocks zone where the classifier is uncertain enough to generate informative gradients but not so uncertain that learning is impossible. The joint training in Algorithm 1 implicitly maintains this balance because both the EBM and the sampler improve together—as the EBM gets better at classification, the sampler also improves, providing harder negatives that maintain gradient signal.

Figure 2 codifies this dynamic into distinct training regimes (posterior collapse, EBM-learns-faster, sampler-learns-faster, balanced), giving practitioners a diagnostic framework for interpreting R-NCE training curves. A loss that plateaus far from -log(K+1) and remains flat despite continuing updates signals posterior collapse. A loss that rises then falls signals that the sampler is catching up to the EBM. A loss that monotonically increases suggests the sampler is learning faster than the EBM. These patterns are actionable: they tell you whether to increase sampler capacity, adjust relative learning rates, or modify the T_samp/T_rnce ratio.

5. Experimental Analysis

Evaluation Methodology

  • Dataset. The paper uses two synthetic 2D conditional distributions (Pinwheel and Spiral) for basic validation, a custom 25-environment planar obstacle path planning benchmark (Section 7.3), and the Push-T contact-rich block pushing task from Chi et al. [15] (Section 7.4). The path planning dataset consists of optimal trajectories generated by a stochastic Gaussian process motion planner (StochGPMP) with 50-step trajectories split into (N_ctx + N_pred)-length snippets and refined further with the planner. The Push-T dataset contains 136 expert human teleoperated demonstrations, distributed with the simulation environment from Chi et al. [15], using the pygame physics engine.

  • Base model(s). No single pretrained model is used. All generative models—energy functions, vector fields, and diffusion denoisers—are trained from scratch for each benchmark. Architectures vary by task: encoder-only transformers for path planning energy models and the Push-T EBM; dense MLPs with ConcatSquash layers for path planning negative samplers; a novel Conv1D-Multi-Head-FiLM (CMHF) architecture for the Push-T negative sampler; and convolutional ResNets with spatial softmax for Push-T visual encoding. Parameter counts are controlled across models for fairness: ~22k maximum for synthetic 2D examples, ~500k for path planning (NF+EBM total for R-NCE), and ~3.3M for Push-T.

  • Metrics. For synthetic 2D problems, the Bhattacharyya coefficient (BC) between true and learned conditional distributions is computed via Gaussian KDE on 8,192 samples, discretized on a 256×256 grid, with the minimum across context values reported. For path planning, the collision rate (fraction of 384 autoregressively-sampled trajectories per environment that intersect obstacles) and a normalized trajectory cost (Equation 7.2) combining goal-reaching error, obstacle collision cost, and smoothness penalty are reported across 25 test environments. For Push-T, the coverage score at each timestep is min(r[i]/0.95, 1) where r[i] is the intersection-over-union of the current block pose with the target pose, with the episode score being the maximum over 200 timesteps. Final evaluation uses 256 random initial configurations with 32 rollouts each, reporting mean coverage with 95% confidence intervals.

  • Baselines. The evaluated models (Section 7.1) are: (1) NF — a continuous normalizing flow trained via stochastic interpolant loss rather than MLE; (2) IBC — the InfoNCE-inspired objective of Florence et al. [24], with Langevin sampling; (3) R-NCE and I-R-NCE — the paper's proposed methods; (4) Diffusion-EDM — a score-based diffusion model using the Elucidating Diffusion Models (EDM) parameterization of Karras et al. [42] with reverse probability flow ODE sampling; (5) Diffusion-EDM-φ — same as Diffusion-EDM but with the denoiser parameterized as the gradient of a scalar energy function φ. For path planning, both diffusion baselines and NF are included; for Push-T, IBC is excluded due to its poor prior performance. All models use comparable parameter counts and comparable numbers of function evaluations during inference.

  • Generation budget / compute accounting. Fair comparison is ensured through two mechanisms: (1) Comparable parameter counts — for R-NCE and I-R-NCE, the total parameters of both the negative sampler and energy model are kept within the same limit as single-model baselines (e.g., ~500k for path planning, ~3.3M for Push-T). (2) Comparable function evaluations for sampling — the number of forward passes through vector fields or energy models is kept similar across models. For path planning, all models use 750 total function evaluations for sampling, and 64 log-probability ODE steps for ranking. For Push-T, 1024 total function evaluations are used for sampling across all models, with 48 log-probability ODE steps. For R-NCE, K=63 negative samples are used in path planning; K=31 for Push-T with m=5 interpolating times.

  • Cross-validation / statistical protocol. No traditional cross-validation is used. For path planning, collision rates and costs are reported with symmetric 95% confidence intervals computed via a standard normal approximation, and per-environment confidence intervals are computed via the Clopper-Pearson method for binomial proportions. For Push-T, 256 random initial configurations are sampled and 32 rollouts are executed per configuration, yielding 8,192 total episodes per model, with 95% confidence intervals via normal approximation. Hyperparameter selection is performed through grid search (detailed in Appendix B) rather than cross-validation on held-out data.


Main Quantitative Results

Synthetic 2D Distributions: R-NCE Matches Diffusion and NF; IBC Visibly Degrades

The synthetic experiments serve as basic validation that all implementations work correctly on low-dimensional problems (Tables 1 and 2). On the Pinwheel distribution (context x ∈ {4, 5, 6, 7} denoting number of spokes, event space Y = R^2), all models except IBC achieve BC ≥ 0.991. On Spiral (context x ∈ [400, 800] normalized to [-1, 1], event space Y = R^2), the pattern sharpens: NF, R-NCE, Diffusion, and Diffusion-φ all achieve BC ≥ 0.991, while IBC achieves only BC = 0.540. The 2D histogram plots in Table 1 and KDE plots in Table 2 visually confirm that IBC produces noticeably degraded sample quality—the spiral arms are blurred and misshapen—while the other four models are visually indistinguishable from the true distribution.

This result is significant primarily as a correctness check: the R-NCE implementation reproduces the known behavior that these generative models work on simple distributions. The IBC degradation, even in 2D with a learnable NF as the negative sampler (trained jointly as in Algorithm 1—the paper notes for IBC, "we do not use a fixed proposal distribution, as doing so results in poor quality energy models"), demonstrates that the population bias identified in Proposition 4.7 manifests in practice even at low dimensions when the negative sampler is non-uniform.


Path Planning: R-NCE Achieves Both Lowest Collision Rate and Lowest Cost

The path planning benchmark evaluates autoregressive policy rollouts over 50-step trajectories, where at each of 5 policy steps the model generates ℓ = 48 future trajectory snippets of length N_pred = 10 (20-dimensional action space) and selects the one with highest log-probability. The context includes the last N_ctx = 3 positions, the goal location, and 10 obstacle specifications. The evaluation metric is a normalized trajectory cost (Equation 7.2) combining goal error, obstacle collision penalty, and smoothness violation.

Aggregate quantitative results (Figure 4, Table 3). Across 25 test environments, R-NCE achieves the lowest collision rate (3.0% ± 0.3%) and lowest trajectory cost (0.194 ± 0.009), outperforming all baselines. By comparison: NF achieves 5.4% ± 0.5% collision rate and 0.220 ± 0.009 cost; Diffusion achieves 5.6% ± 0.5% collision rate and 0.286 ± 0.008 cost; Diffusion-φ achieves 3.9% ± 0.4% collision rate and 0.237 ± 0.006 cost; IBC achieves 10.7% ± 0.6% collision rate and 0.855 ± 0.215 cost.

The gap between R-NCE and the best diffusion variant (Diffusion-φ) is meaningful but modest: 0.9 percentage points in collision rate and 0.043 in cost. The IBC gap is substantial: nearly 4× higher collision rate than R-NCE and over 4× higher cost, despite using the same underlying architecture and the same learnable proposal distribution (trained jointly as per Algorithm 1). This is consistent with the theoretical bias analysis: even with a learnable sampler, the IBC objective recovers p(y|x) / p_ξ(y|x) rather than p(y|x), so the learned energy landscape is systematically shifted.

Per-environment breakdown (Figure 4). The bar plot reveals that two environments (Environments 1 and 24) are particularly challenging, with collision rates exceeding 10% for all models. Table 3 visualizes trajectories from three specific environments, showing that failures in Environment 1 arise from aggressive path selection through narrow passageways between obstacles. R-NCE trajectories (0.253 collision rate in Env 1) show fewer collisions through these tight gaps than NF (0.440) or Diffusion (0.479), suggesting that the EBM's energy landscape better captures the safety constraints implicit in the demonstration data.

Multi-modality capture (Table 3). The trajectory visualizations demonstrate that all models except IBC successfully capture multi-modal solutions—different rollouts take qualitatively different paths around obstacles to reach the goal. IBC's higher cost (0.426, 0.426, 1.394 in the three shown environments vs. R-NCE's 0.194, 0.168, 0.219) corresponds to trajectories that cut corners more aggressively, resulting in obstacle collisions or non-smooth paths. This is consistent with IBC learning the density ratio rather than the true density: the shift toward the proposal distribution's behavior (the simpler NF) pulls the energy landscape away from the expert distribution's precise obstacle-avoiding structure.

Ranking necessity (Table 4). An important ablation shows that generating ℓ = 48 samples per policy step and selecting the highest-likelihood one is crucial for all models. Reducing to ℓ = 1 sample per step (no ranking) causes dramatic degradation: R-NCE collision rate jumps from 3.0% to 11.6%, cost from 0.194 to 2.287. Diffusion degrades from 5.6% to 9.9% collision rate, cost from 0.286 to 0.569. The larger relative degradation for R-NCE (cost increases ~12× vs. ~2× for Diffusion) suggests that R-NCE's advantage relies on effective ranking of diverse candidates—the energy model produces samples with higher variance in quality, and the ranking step filters out poor ones. This aligns with the paper's theoretical framing: the EBM learns a good density model, but MCMC sampling from it may produce variable-quality samples, making post-hoc selection important.

Pre-trained vs. joint training of negative sampler. The paper reports that pre-training the proposal model (rather than jointly training as in Algorithm 1) yields notably worse performance: collision rate 5.8% (vs. 3.0%) and cost 0.24 (vs. 0.194). The hypothesis is that joint training allows the sampler to adapt to the EBM's evolving energy landscape, maintaining productive difficulty for the R-NCE classification task. A pre-trained sampler may be either too weak (causing posterior collapse early in EBM training) or too strong in ways that don't align with the EBM's current state.


Push-T: I-R-NCE Achieves Highest Coverage, Outperforming Diffusion

The Push-T block pushing task is the most challenging benchmark, requiring visuomotor control from RGB image observations with contact-rich dynamics and multi-modal solutions. The autoregressive policy predicts N_pred = 16 future end-effector positions (32-dimensional action space) conditioned on N_ctx = 2 frames of image observations and past positions.

Quantitative results (Table 6). The final coverage scores across 256 random configurations with 32 rollouts each:

  • I-R-NCE (Algorithm 4, three-stage sampling): 0.884 ± 0.005
  • I-R-NCE (Algorithm 2, two-stage sampling): 0.880 ± 0.005
  • NF: 0.866 ± 0.006
  • Diffusion: 0.864 ± 0.006
  • Diffusion-φ: 0.860 ± 0.006
  • R-NCE (single-scale, non-interpolating): 0.824 ± 0.006

I-R-NCE with the full three-stage sampling procedure achieves the highest coverage, outperforming the best diffusion variant (Diffusion) by 2.0 percentage points and the NF baseline by 1.8 percentage points. The gap to the second-place I-R-NCE variant (Algorithm 2 with only Langevin MCMC, no SDE transport) is 0.4 percentage points—small but consistent across all 8,192 rollouts.

Interpolating vs. single-scale R-NCE. The most striking comparison is between I-R-NCE (0.884) and single-scale R-NCE (0.824). The 6.0 percentage point gap demonstrates that the interpolating framework is necessary for EBMs to compete on this task—single-scale R-NCE, despite using the same R-NCE training algorithm with joint negative sampler training, substantially underperforms. This validates the paper's central claim that multi-scale training is not inherently tied to score matching and can be effectively realized within the EBM paradigm through the I-R-NCE objective.

SDE transport stage contribution. Comparing the two I-R-NCE variants—with (0.884) and without (0.880) the SDE transport stage—shows a small but measurable benefit from the SDE. The SDE uses the I-EBM's score ∇_y E_θ(x, y, t) to correct distributional errors in the CNF's vector field during transport from t = t̲ to t = 1 before final MCMC refinement. The 0.4 percentage point gap, while modest, is consistent with the theoretical justification from Albergo et al. [4]: the SDE provides a KL-divergence bound on distributional error that the pure ODE does not, and the I-EBM's score supplies the missing component that enables this correction.

Qualitative trajectory analysis (Table 5). The visualized trajectories across four different initial conditions show that all models produce reasonable pushing behaviors, but with subtle differences. I-R-NCE consistently achieves scores of 0.954–1.000 across the four shown environments, while Diffusion-φ shows a failure case (score 0.240) where the block is barely moved from its initial position. NF shows a partial failure (score 0.572) where the block is pushed toward but not fully into the target. These examples suggest that I-R-NCE produces more robust policies that less frequently get stuck in local minima of the contact dynamics.

Architectural considerations for the negative sampler. A notable implementation detail is that for Push-T, the simple dense MLP CNF architecture used as the negative sampler in path planning led to posterior collapse. The paper developed a more expressive Conv1D-Multi-Head-FiLM (CMHF) architecture (Figure 9) that was "sufficiently powerful for training via I-R-NCE and computationally cheaper" than the transformer-based EBM architecture. This reinforces the theoretical analysis of Section 4.3: as the event space dimensionality increases (20D → 32D) and the task complexity grows (state-based planning → visuomotor control), the negative sampler must also scale in capacity to maintain informative gradients and prevent posterior collapse.


Ablation Studies and Robustness Checks

Pre-trained vs. joint training of the negative sampler (Section 7.3, path planning): Pre-training the proposal model degrades performance from 3.0% collision rate (joint training) to 5.8%, and from 0.194 cost to 0.24. This supports the joint training design in Algorithm 1 over the pre-training alternative explored by Gao et al. [26] for binary NCE, and suggests that the sampler must adapt to the EBM's evolving state to maintain productive classification difficulty. The paper hypothesizes that data noising (σ_pert annealed from 0.01 to 0.005) improves training by "guiding the convergence of the learned distributions towards the data sub-manifold," but this benefit requires joint training—a pre-trained sampler cannot adapt to the annealed noise schedule of the EBM.

Number of samples per step for ranking (Table 4): Reducing from ℓ = 48 samples per policy step to ℓ = 1 increases collision rates by 2–3× across all models and increases costs more dramatically (R-NCE: 0.194 → 2.287; IBC: 0.855 → 8.734; Diffusion-φ: 0.237 → 0.489). This validates the importance of ranking as a practical inference strategy and highlights the utility of log-probability computation (Section 6) for models that don't natively provide density estimates. The differential degradation across models is informative: R-NCE's large cost increase under ℓ = 1 (over 10×) compared to Diffusion (only 2×) suggests that the EBM's MCMC sampling produces higher-variance samples—some excellent, some poor—and that the ranking step is filtering out poor samples. This is a double-edged practical property: it means ranking is more important for EBM policies, but also that EBMs benefit more from ranking than diffusion models do.

Two-stage vs. three-stage I-EBM sampling (Table 6, Algorithm 2 vs. Algorithm 4): For Push-T, replacing the three-stage Algorithm 4 (CNF warm-start → SDE transport → HMC refinement) with the simpler two-stage Algorithm 2 (CNF warm-start → HMC refinement, no SDE) reduces coverage from 0.884 to 0.880. While small (0.4 percentage points), this gap is consistent across 8,192 rollouts (the confidence intervals are ±0.005). The ablation confirms that the SDE transport stage contributes beyond what HMC alone provides, consistent with the theoretical justification that the SDE's noise term helps correct distributional errors that pure MCMC starting from the CNF sample might not efficiently fix.

Data perturbation noise annealing (path planning, Section 7.3): The training procedure perturbs context and prediction snippets with Gaussian noise of variance σ²_pert, where σ_pert is annealed from 0.01 to 0.005 over half the training steps and held fixed thereafter. The paper states that this technique "significantly boosted the performance of all models" and provides intuition that "trajectories in the training dataset arguably live on a much smaller sub-manifold of the ambient event space" and that "by annealing the perturbation noise, we are able to guide the convergence of the learned distributions towards this sub-manifold." No quantitative ablation is provided for this specific choice—the effect size is not measured in a controlled comparison—but the practice is applied uniformly across all models, so it does not bias the relative comparisons.

Negative sampler architecture capacity (implicit ablation via Push-T): The paper notes that for path planning, a dense MLP CNF sufficed as the negative sampler, but for Push-T, "a dense MLP as used in path planning for the proposal distribution led to posterior collapse during training." This necessitated the development of the CMHF architecture. While not presented as a formal ablation with quantitative comparison, this design choice illustrates the practical sensitivity identified in the posterior collapse analysis (Section 4.3): the sampler's capacity must scale with task difficulty to prevent gradient vanishing.

Variance reduction via multiple interpolating times (Algorithm 3): For I-R-NCE training on Push-T, m = 5 interpolating times are sampled per data point. No ablation over m is presented, and the computational tradeoff (linearly more negative samples and energy evaluations) is not quantified. The remark about correlated negative sample trajectories (Remark 5.1) suggests that the practical cost is sub-linear in m, but the paper provides no empirical validation of this efficiency claim.


Critical Assessment

Claim 1: EBMs trained with R-NCE yield high-quality policies that compete with—and even outperform—diffusion models. The evidence provides qualified support. On path planning (Section 7.3), R-NCE achieves the lowest collision rate (3.0% vs. 5.6% for standard Diffusion, 3.9% for Diffusion-φ) and lowest cost (0.194 vs. 0.286 for Diffusion, 0.237 for Diffusion-φ). On Push-T (Section 7.4), I-R-NCE achieves the highest coverage (0.884 vs. 0.864 for Diffusion). The qualification is that: (a) the performance gap, while consistent and statistically significant given the confidence intervals, is modest—R-NCE is not dramatically better than diffusion, it is competitive with a small edge; (b) these results are on two specific robotic benchmarks with a single set of architectural choices—it is not demonstrated that this advantage generalizes to other domains, observation modalities, or action dimensionalities; (c) the diffusion baselines, while carefully implemented using the EDM parameterization [42], may not represent the absolute state-of-the-art in diffusion-based policy learning (Chi et al. [15] use a CNN-based diffusion architecture for Push-T, not the transformer architecture used here, and achieve higher coverage scores on the same task with larger models).

Claim 2: The IBC objective is biased at the population level, and this bias explains its poor empirical performance. The theoretical proof (Proposition 4.7) is clear, and the empirical evidence is strongly supportive: (a) Figure 1 demonstrates population-level bias in 1D with the IBC optimum visibly shifted from the true parameter; (b) Table 1 and Table 2 show IBC degrades sample quality even in 2D (BC = 0.540 on Spiral vs. ≥0.991 for all other models); (c) Figure 4 shows IBC has 3–4× higher collision rate and cost than R-NCE on path planning, despite using identical architecture and the same jointly-trained NF as negative sampler. The key empirical strength is that the bias is observed in a setting where IBC has every advantage (learnable, jointly-trained negative sampler; same architecture as R-NCE; careful hyperparameter tuning), isolating the objective function as the causal factor.

Claim 3: Joint non-adversarial training of the negative sampler is sufficient and avoids the instabilities of adversarial optimization. The empirical support is partial. The pre-training ablation (Section 7.3: joint training collision rate 3.0% vs. pre-trained 5.8%) shows that joint training outperforms a simple alternative. However: (a) the paper never actually implements or compares against an adversarial training baseline for R-NCE—the theoretical analysis (Theorem 4.18, Corollary 4.19) proves that adversarial training provides no asymptotic benefit, but the paper does not verify empirically whether adversarial training would be worse in practice (due to instability) or merely equivalent; (b) the claim that "adversarial optimization is notoriously challenging" (Section 2) is a motivation for the non-adversarial design but is not substantiated with any experimental comparison against adversarial R-NCE; (c) while the training curves in Figure 2 are characterized as exhibiting "a phase transition" without oscillatory instability, actual training curves from the experiments are not shown, so the stability claim is asserted rather than demonstrated.

Claim 4: Multi-scale learning (interpolating EBMs) is compatible with EBMs and provides substantial performance gains. Strongly supported by a single, clean comparison: single-scale R-NCE achieves 0.824 coverage on Push-T vs. I-R-NCE at 0.884 (Table 6). The 6.0 percentage point gap is the largest effect size in the paper and clearly demonstrates that interpolating training is critical for high-dimensional EBMs. The mechanism (implicit annealing across noise scales, benefiting from learning easy distributions before hard ones) is hypothesized but not empirically validated through intermediate ablations (e.g., showing that I-R-NCE training curves at early times converge faster than at late times, or that the shared network benefits from multi-task learning). The gap between two-stage I-R-NCE (0.880) and single-scale R-NCE (0.824) also shows that the benefit persists even when the I-EBM is only used at t=1 for sampling—the training-time multi-scale signal improves the final t=1 model even without the SDE transport stage.

Genuine weaknesses in experimental design:

  1. Single benchmark suite, no replication across domains. Both robotic tasks use autoregressive position prediction with roughly similar dimensionalities (20D for path planning, 32D for Push-T). No results are shown for tasks with substantially different characteristics: discrete actions, high-frequency control, long-horizon planning, tasks requiring precise force control, or tasks where multi-modality is less important. The synthetic 2D examples are a minimal correctness check, not a meaningful benchmark. The paper's claim that EBMs are "competitive with diffusion models for robot policy learning" is based entirely on two tasks in the same autoregressive action chunk prediction paradigm.

  2. Diffusion baselines may not be maximally tuned. The paper uses its own transformer-based architecture for Diffusion and Diffusion-φ to enable fair comparison with the same architecture as the EBMs. However, Chi et al. [15]—the canonical diffusion policy paper—uses a CNN-based architecture with FiLM conditioning for Push-T and reports different hyperparameter choices. The paper's diffusion baseline may underrepresent what dedicated diffusion policy methods can achieve. No attempt is made to compare against pre-trained or publicly available diffusion policy checkpoints.

  3. The ℓ = 48 ranking step is computationally significant but cost is unaccounted for. For path planning, every policy step generates 48 candidate trajectory snippets and computes their log-probabilities. This multiplies the per-step inference cost by ~48× (minus any savings from batched computation). The paper does not include this ranking cost in any wall-clock or FLOPs comparison. If one model benefits more from ranking than another (as Table 4 suggests R-NCE degrades more without ranking than Diffusion does), then the fair comparison should account for the ranking cost, or compare at a matched budget that includes ranking.

  4. Small number of test environments for path planning. While 25 environments with 384 trajectories each provides reasonable statistical power for aggregate metrics, Figure 4 shows that individual environments have limited data (384 rollouts per environment), yielding wide per-environment confidence intervals. The worst-case environments (1 and 24) are clearly identified but the analysis provides limited insight into what makes these environments hard and whether any model handles them better.

  5. No ablation over K (number of negative samples). The paper fixes K = 63 for path planning and K = 31 for Push-T, but Theorem 4.12 predicts that asymptotic variance scales as (1 + 1/K). An ablation showing that larger K improves performance (and quantifying the diminishing returns) would validate this theoretical prediction empirically. Similarly, no ablation over m (number of interpolating times for I-R-NCE) is provided.

  6. The two-time-scale log-probability technique (Section 6) is claimed efficient but not benchmarked. The paper argues that this approach makes ranking ℓ = 48 samples feasible, but provides no timing comparison against the standard augmented ODE approach, nor against alternative rankers (e.g., learned critics). The number of divergence evaluation steps (64 for path planning, 48 for Push-T) is stated but not justified through any ablation showing that fewer steps degrade ranking quality.

  7. No study of inference-time computational cost. While the paper emphasizes comparable function evaluations during sampling, real inference cost includes: log-probability computation for ranking (Section 6), MCMC warm-up and iteration, potential batching of parallel samples, and for I-R-NCE, the SDE integration. Wall-clock timing, even in aggregate, would contextualize whether EBMs are practical for real-time control (where inference must run at >10 Hz). The paper's argument that EBMs yield "compact representations" for "real-time inference requirements" (Section 1) is not backed by any latency measurements.

  8. Model capacity comparison is approximate, not exact. The paper limits parameter counts to maxima (~500k for path planning, ~3.3M for Push-T) but does not report exact counts per model, making it possible that small differences in effective capacity (e.g., the EBM wasting parameters on the energy scalar output vs. the diffusion model using all parameters for vector field representation) influence results.

6. Limitations and Trade-offs

Limitation 1: Difficulty Estimation for Compute-Optimal Allocation Is Prohibitively Expensive

The assumption or constraint. The entire compute-optimal framework in this paper rests on estimating prompt difficulty before selecting and executing the test-time strategy. The paper's method for doing so involves generating 2,048 complete solutions per question from the base model, then either checking each against the ground-truth answer (oracle difficulty) or scoring each with the PRM's final-answer prediction and averaging the result (predicted difficulty). This estimation step is itself a massive computational undertaking. The authors state in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence. The reported 4× efficiency gains over best-of-N (e.g., Figure 4 showing 16 generations matching 64; Figure 8 showing 64 matching 256) are computed after difficulty has been estimated, without amortizing the 2,048-sample estimation cost into the budget. For a deployment scenario, the true cost equals difficulty estimation plus strategy execution. Since the estimation step alone (2,048 generations) is 4–8× larger than the largest test-time budgets studied (256–512 generations), the amortized efficiency gain over best-of-N could be dramatically smaller—or even negative—if few queries share the same difficulty estimate. For a system answering individual user queries, the 2,048-sample overhead per question would make the compute-optimal approach far more expensive than simply running best-of-N with a large budget. The paper's headline 4× figure therefore represents an upper bound on achievable efficiency, conditional on having essentially free difficulty estimates.

What evidence exists in the paper. The paper explicitly flags this gap (Section 3.2, Section 8) and frames the exploration-exploitation tradeoff between difficulty estimation and problem-solving as "a key avenue for future work." The predicted (non-oracle) difficulty bins are shown to perform nearly as well as oracle bins (Figures 4 and 8, curves largely overlapping), confirming that ground-truth labels are not required—but the 2,048-sample cost remains. Figure 9 and the FLOPs-matched analysis (Section 7) do not amortize difficulty estimation cost into the pretraining-inference comparison. No experiment measures what fraction of the total budget difficulty estimation consumes.

Mitigation status. Not addressed in this work. The paper suggests (Section 8) that future work could "pretrain or finetune models to directly predict difficulty of a question" from the question text alone, potentially reducing estimation to a single forward pass. An alternative suggested in the broader literature is adaptive difficulty estimation—generating a small number of samples (4–8), assessing score statistics, and allocating the remaining budget accordingly. Neither approach is developed or evaluated here. The current results should be understood as characterizing what is possible if difficulty can be estimated cheaply, not as a complete deployment-ready system.


Limitation 2: The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate

The assumption or constraint. The revision model is fine-tuned on trajectories consisting entirely of incorrect answers followed by a correct answer. During training, every in-context example preceding the target is incorrect by construction. At inference time, however, the model can—and frequently does—generate correct answers within a revision chain. When the model later conditions on its own correct answer as context, it has never seen a correct-to-correct or correct-to-incorrect transition during training, and it may incorrectly "revise" a correct answer into a wrong one. The paper quantifies this in Section 6.1:

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach"

The consequence. This reversion rate means that a purely sequential revision chain—where the output of each step is fed as context to the next—will frequently cycle away from correct answers. A correct answer produced at revision step 3 may be overwritten by an incorrect one at step 4, degrading chain-level accuracy. The paper mitigates this through within-chain selection (majority voting or verifier-based selection across all revisions in the chain, picking the best answer rather than the final one), but this is a patch rather than a solution: it requires generating and storing the full chain, then running a separate selection mechanism, which wastes computation on revisions that may be worse than earlier ones. The fundamental issue—that the model does not know when to stop revising—remains unaddressed. In a production system, this means the revision chain length must be treated as a hyperparameter to tune, and the system must pay the compute cost of generating and evaluating revisions even when the first attempt is correct.

What evidence exists in the paper. The paper acknowledges the problem explicitly (Section 6.1) and reports the 38% figure. Figure 6 (left) shows that pass@1 at each revision step gradually increases (from ~18.2% to ~24–25% by steps 15–20), confirming that revisions help on average despite individual regressions. The mitigation (within-chain majority voting or verifier selection) is evaluated in Figure 6 (right), where sequential + best-of-N weighted achieves ~41.5% vs. parallel + best-of-N weighted at ~39%—the selection mechanism successfully recovers correct answers from the chain despite reversions. However, no ablation compares the naive "take the last revision" approach quantitatively; the performance penalty from reversions without mitigation is not measured.

Mitigation status. Partially mitigated through within-chain selection, but not solved. The within-chain selection (best-of-N weighted or majority voting across all revisions) successfully recovers correct answers even when later revisions overwrite them, but it incurs the cost of generating and scoring all revisions. A more principled approach—training the model to recognize when no revision is needed, or including correct-to-correct trajectories in the training data—is not explored. The paper does not report what fraction of the total sequential budget is wasted on reversions, nor whether a learned stopping criterion could improve efficiency. The ReST^EM experiment (Appendix K, Figure 16) provides a cautionary negative result: attempting to further optimize the revision model with on-policy training caused performance to degrade substantially with sequential revisions, suggesting that the correct-incorrect transition problem is sensitive to training methodology in ways that are not fully understood.


Limitation 3: The Methods Are Studied Independently—Combined Search and Revisions Are Not Evaluated

The assumption or constraint. The paper studies two complementary test-time compute mechanisms—PRM-guided search (Section 5) and iterative revisions (Section 6)—but evaluates them as separate pipelines. Each mechanism is compared against the best-of-N baseline and optimized independently via difficulty-conditioned allocation. The proposal distribution modification (revisions) and the verifier optimization (search) are never combined into a single system.

The authors explicitly acknowledge this in Section 8:

"we did not experiment with PRM tree-search techniques in combination with revisions"

The consequence. The paper's results represent a lower bound on what a fully integrated system could achieve. The two mechanisms have complementary, evidence-supported strengths: revisions help most on easy problems where initial attempts are roughly correct and need local refinement (Figure 7, right: bin 1 shows ~90% accuracy across all ratios, bin 2 benefits from more sequential revisions); beam search helps most on medium-difficulty problems where exploration across different solution strategies is needed (Figure 3, right: bin 3 consistently benefits from beam search over best-of-N). A natural integration would use a revision model as the proposal distribution within beam search—each beam expansion step could condition on previous incorrect revisions as context, potentially generating higher-quality candidate steps than the base model alone. Alternatively, the PRM could guide which revisions to pursue rather than generating a blind chain. Since the difficulty-dependent analysis shows these strengths are complementary (not redundant), combining them could yield gains beyond what either achieves independently, particularly in the medium-difficulty regime (bins 3–4) where both exploration and refinement matter.

What evidence exists in the paper. None, by the authors' own admission. The compute-optimal allocation policies (Figures 4 and 8) treat search and revisions as alternative strategies to choose between, not as components to be combined. The FLOPs-matched comparison (Section 7, Figure 9) evaluates PRM search and revisions as separate test-time compute methods, each compared independently against the larger pretrained model. The paper provides no data on whether a revision model's outputs are better search candidates, whether PRM scores are well-calibrated on revision model outputs (the paper notes in Appendix J, Figure 15a, that the base-LM PRM underperforms on revision model outputs due to distribution shift, suggesting this combination would require additional verifier training), or how the optimal allocation would change if both mechanisms were available simultaneously.

Mitigation status. Acknowledged but not addressed. The paper frames this as future work (Section 8) but provides no preliminary results, no analysis of the interaction between revisions and search, and no guidance on how such a combined system should be designed. This is the most significant missing experiment in the paper because the conceptual framework (proposal vs. verifier as complementary axes) directly implies that combining them should yield gains, and the difficulty-dependent analysis provides a clear prescription for where each should dominate—but this prescription is never tested in a joint system.


Limitation 4: Single Benchmark and Single Model Family Constrain Claims of Generality

The assumption or constraint. All experiments use the MATH benchmark exclusively—a dataset of high-school competition-level math problems requiring symbolic reasoning—with PaLM 2-S* as the base model. The paper states in Section 4:

"We believe this model is representative of the capabilities of many contemporary LLMs"

This claim is not verified with any replication across benchmarks or model families.

The consequence. The paper's central findings—that difficulty-dependent allocation yields 4× efficiency gains, that test-time compute can substitute for model scale on easy-medium problems but not hard ones, that beam search over-optimizes on easy problems—are demonstrated in a narrow slice of the reasoning landscape. Several aspects of these findings could be model- or task-specific:

  • PRM quality and over-optimization threshold. The PRM is trained on PaLM 2-S* outputs using Monte Carlo rollouts. A model with different calibration properties, error patterns, or reasoning strategies might produce a PRM with different over-optimization characteristics, shifting the difficulty thresholds at which beam search degrades.
  • Revision model behavior. The revision model's ability to learn from incorrect in-context examples—and its 38% reversion rate—depends on PaLM 2-S*'s in-context learning dynamics, which vary substantially across model families (e.g., decoder-only vs. encoder-decoder architectures).
  • Task type. MATH problems are closed-form, have unambiguous correctness, and require multi-step symbolic deduction. Tasks requiring factual recall, code generation, open-ended reasoning, or multimodal inputs might exhibit entirely different difficulty-dependent scaling patterns. The paper's claim that test-time compute "amplifies existing capability but does not create it" may not hold for tasks where the base model has latent knowledge that can be surfaced through different reasoning chains.
  • Difficulty distribution. MATH problems span five labeled difficulty levels (not the paper's model-specific quintiles), and the distribution of problem difficulties in MATH may not match the distribution encountered in real-world LLM deployments. If production traffic is skeweder toward easy problems (where test-time compute shines) or hard problems (where it fails), the practical benefit could differ substantially.

What evidence exists in the paper. All quantitative results are on MATH. The synthetic 1D demonstration (Figure 1) and 2D path planning / block pushing tasks in the reference example paper are in different papers and do not provide cross-task validation for the MATH findings. The difficulty-bin methodology (Section 3.2) defines difficulty relative to the base model's pass@1 rate on MATH specifically. The test set consists of 500 questions, split into five difficulty quintiles of ~100 each. With two-fold cross-validation, the compute-optimal policy is selected based on ~50 questions per fold per bin—a small sample that may not yield robust strategy selection, particularly in the middle bins where the optimal choice between beam search, lookahead, and best-of-N is most sensitive.

Mitigation status. Acknowledged implicitly but not addressed. The paper does not claim the results generalize beyond MATH, but the framing—"compute-optimal test-time scaling" and "test-time compute can substitute for model scale"—suggests broader applicability. No experiments on alternative reasoning benchmarks (e.g., GSM8K, theorem proving, code generation) are provided. No analysis examines whether the difficulty-dependent patterns are driven by specific properties of MATH (e.g., the discrete difficulty levels in the dataset's original labeling, the availability of clean solution step boundaries for PRM training, the existence of unique correct answers). The paper does not report confidence intervals on the compute-optimal scaling curves, making it difficult to assess whether the selected strategies are statistically robust given the small per-bin sample sizes.


Limitation 5: Hard Problems Are Effectively Unsolved—Test-Time Compute Has a Hard Cap

The assumption or constraint. The paper's framework assumes that the base model already has some non-trivial probability of producing correct solutions. All test-time compute methods operate by manipulating the base model's output distribution—search explores more candidates, revisions refine existing attempts—but neither creates the ability to solve problems from scratch. When the base model's pass@1 is near zero on a problem class, there are essentially no correct solutions in the proposal distribution to find or refine.

The consequence. The hardest difficulty bin (bin 5) shows negligible improvement regardless of the test-time compute budget or strategy. In Figure 3 (right), bin 5 accuracy stays at approximately 1–3% for all methods and all budgets up to 256 generations. In Figure 7 (right), bin 5 accuracy is roughly 2–3% regardless of the sequential-to-parallel ratio at 128 generations. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line remains essentially flat near 0–5% across all budgets. The compute-optimal allocation policy routes hard problems appropriately (it doesn't waste compute where it can't help), but it cannot solve them—meaning that for problems outside the base model's capability range, pretraining compute is the only viable path to improvement. The paper is admirably transparent about this, explicitly stating in the Section 7 takeaway:

"Test-time compute amplifies existing capability but does not create it from nothing."

This has direct implications for deployment: if a significant fraction of user queries falls into the "hard" regime for a given base model, the compute-optimal framework provides no benefit—the model must be scaled up via pretraining. The framework also offers no mechanism for dynamically recognizing when a problem is unsolvable and redirecting to a larger model or a human, though such escalation policies could be built on top of the difficulty estimation infrastructure.

What evidence exists in the paper. Strong, consistent negative results across all methods. Figure 3 (right), Figure 7 (right), and Figure 9 all show bin 5 accuracy near zero with flat scaling curves. The paper neither hides nor downplays this result—it is a central part of the compute-optimal allocation story (routing hard problems to best-of-N since nothing else helps) and of the pretraining vs. inference tradeoff analysis (test-time compute loses to pretraining on hard problems at all R values). However, the paper does not analyze why bin 5 problems are hard—whether they require capabilities genuinely absent from the base model (e.g., specific mathematical knowledge not acquired during pretraining) or whether they are merely problems where the model's pass@1 is low but some correct solutions exist (in which case more extreme search might eventually find them). The 2,048-sample pass@1 estimation could miss rare correct solutions that require orders of magnitude more samples to discover.

Mitigation status. Acknowledged transparently but not addressed. The paper treats the hard-problem ceiling as a fundamental constraint on the test-time compute approach and does not propose methods to break through it. The difficulty estimation infrastructure could, in principle, be used to flag problems for escalation to a larger model or for human intervention, but this is not explored. The implicit message is that test-time and pretraining compute are complementary resources with different strengths—test-time compute boosts performance within the base model's reach, pretraining expands the reach itself—and a complete deployment strategy should use both.


Limitation 6: The 14× Larger Model Baseline Is Not Compute-Optimally Trained, Weakening the Pretraining-Inference Tradeoff

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023) rather than compute-optimal pretraining (Hoffmann et al., 2022) where both data and model size scale equally with compute. The authors acknowledge this in Section 7:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the 14× larger model uses only greedy decoding—no test-time compute augmentation of its own, such as majority voting, best-of-N, or search. The smaller model with compute-optimal test-time strategies is compared against a larger model that receives no inference-time compute budget beyond a single greedy sample.

The consequence. The baseline represents a weaker opponent than a fairly-matched comparison would warrant. A Chinchilla-optimal model (scaling both parameters and data proportionally with compute) trained with 14× more FLOPs would likely outperform a parameter-only-scaled model at the same total compute budget. The reported advantages of test-time compute over pretraining—e.g., +27.8% relative improvement on easy questions at R ≪ 1 for revisions (Figure 1, top-right bar chart)—may shrink or reverse against a properly compute-optimal larger model. Furthermore, the larger model is denied even modest test-time compute augmentation: running best-of-8 or majority voting with the larger model would create a much stronger baseline that better reflects how a large model would actually be deployed if inference-time compute were available. The current comparison answers the question "Is a small model with smart inference better than a large model with greedy decoding?" rather than "Should compute be spent on pretraining or inference, given a fixed total budget?" The second question—which is the one the paper frames itself as addressing—requires both models to have access to the same inference-time optimization toolkit, allocated proportionally to their relative costs.

What evidence exists in the paper. The FLOPs accounting (Section 7) is precise about the comparison methodology. Figure 9 and the Figure 1 bar charts present the results as a tradeoff between test-time and pretraining compute. The paper acknowledges the Chinchilla caveat explicitly but provides no sensitivity analysis—no estimate of how much the 14× larger model would improve under compute-optimal training, and therefore no bounds on how the reported advantages might change. The paper also provides no results where the larger model is given any test-time compute budget. The three values of R (D_inference / D_pretrain) tested (0.16, 0.79, 22) are specific to PaLM 2-S*'s pretraining data volume and are not justified as representative of typical deployment ratios.

Mitigation status. Acknowledged but deferred to future work. The paper states the Chinchilla scaling comparison is left to future investigation. This is a reasonable scoping decision for a paper focused on inference-time compute, but it means the headline claim—that a smaller model with compute-optimal test-time strategies can "outperform a ~14× larger model" (Section 1, Section 7)—should be interpreted as: under a specific pretraining paradigm (parameter scaling only), a specific base model (PaLM 2-S*), and a specific assumption about the larger model's inference strategy (greedy decoding), this substitution is possible. The claim does not establish that inference compute is generally preferable to pretraining compute in a jointly-optimized allocation.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper forces a fundamental re-evaluation of energy-based models in robotics, shifting their status from "theoretically appealing but practically broken" to "competitive with proper training." The conceptual move is not that EBMs have become better than diffusion models—the performance margins in Tables 3, 4, and 6 are real but modest—but that the barrier to making EBMs work was a specific, diagnosable mathematical flaw in the dominant training objective, not an inherent limitation of energy-based representations. The paper reframes the IBC failure narrative from "EBMs are too unstable to train in high dimensions" to "IBC learns the wrong quantity, and once you fix that, EBMs train reliably."

This is a diagnostic reframing rather than a paradigm shift. The core technique—ranking noise contrastive estimation—was introduced by Ma and Collins [57] for discrete NLP tasks. The paper's primary contribution is demonstrating that R-NCE, paired with a learnable but non-adversarial negative sampler and extended to multi-scale training, resolves the specific failures that caused the robotics community to abandon EBMs. The magnitude of the shift lies in what category of model becomes available for policy learning: prior to this work, a practitioner choosing a generative policy class had to accept either (a) EBMs with IBC training and accept unreliable, suboptimal results (as documented in [15, 61, 64, 76]), or (b) diffusion models with score matching and accept the loss of composability, density estimates, and optimization-friendly structure. This paper eliminates that tradeoff—EBMs can be made to work, and their distinctive advantages (composability in density space, compact representation, natural fit for optimization-based control) become practically accessible.

The paper also reconciles contradictory evidence in the NCE literature more broadly. Gao et al. [26] found adversarial training necessary for binary NCE with learned samplers; this paper proves that for ranking NCE in the non-realizable sampler setting, adversarial training provides no asymptotic benefit (Corollary 4.19) and the non-adversarial approach is both simpler and empirically effective. Bose et al. [11] advocated for adversarial contrastive estimation; this paper provides a theoretical argument that the need for adversarial dynamics may be specific to binary NCE's self-normalization requirement, which R-NCE eliminates. The reconciliation is not that one side was wrong—it's that the ranking variant of NCE has structural properties (the optimal θ* being independent of ξ per Theorem 4.4) that make the adversarial approach theoretically unnecessary, and prior claims about adversarial training's necessity were based on a different variant of the objective.

A subtler shift concerns the relationship between generative modeling paradigms. By demonstrating that multi-scale learning—the core architectural insight behind diffusion models—can be realized within the EBM framework through I-R-NCE (Section 5.2), the paper separates the structural pattern (learning across noise scales with a shared network) from the training objective (score matching). This opens the possibility that future generative models could mix and match: use score matching where differentiability is cheap, use R-NCE where composability or density access matters, and share the same multi-scale architecture. The interpolating EBM is not positioned as superior to diffusion—it's positioned as compatible with the same design principles, making EBMs a member of the modern generative modeling toolkit rather than a competing paradigm.

Research directions that become more attractive: (1) Composing EBMs for multi-task policies or constraint satisfaction at inference time, since the energy addition property E_combined = E_task1 + E_task2 now has a reliable training pipeline behind it. (2) Using EBMs as cost functions inside trajectory optimizers, since the learned energy surface directly provides gradients for optimization without requiring a separate critic. (3) Extending R-NCE to reinforcement learning settings, where the energy function can represent a policy or value function without the normalization constraints that complicate MLE-based approaches.

Research directions that become less attractive: (1) Further attempts to debug or stabilize IBC training—Proposition 4.7 shows the problem is structural, not implementation-specific. (2) Adversarial training of the negative sampler for ranking NCE—Theorem 4.18 and Corollary 4.19 show no asymptotic benefit in the practical setting, and the empirical evidence (joint training outperforming pre-training in path planning, Section 7.3) suggests the non-adversarial approach is both simpler and effective. (3) Uniform or fixed negative sampling for high-dimensional EBMs—posterior collapse (Section 4.3) makes this approach fundamentally limited regardless of the objective function used.


Follow-Up Research This Work Enables

Characterizing the R-NCE negative sampler capacity requirements across tasks and dimensionalities. The paper provides anecdotal evidence that sampler capacity must scale with task difficulty: a dense MLP CNF sufficed for 20D path planning but led to posterior collapse on 32D Push-T, necessitating the more expressive CMHF architecture. This is not characterized systematically. A strong follow-up would train EBMs with R-NCE on a parametric family of tasks (e.g., trajectory prediction with increasing prediction horizon, or image-based control with increasing action dimensionality) while varying the negative sampler's architecture (MLP width/depth, number of CMHF blocks, attention heads) and measure both final policy performance and the presence of posterior collapse (tracking the R-NCE loss trajectory and the average posterior entropy during training). The goal would be to produce a scaling guideline: for a given action dimensionality and task complexity, what sampler capacity is necessary to prevent gradient vanishing? The paper's theoretical framework (Section 4.3) predicts a phase transition—below some capacity threshold, posterior collapse occurs and training stalls; above it, the EBM converges normally—but provides no empirical characterization of where this threshold lies.

Combining R-NCE-trained EBMs with composition at inference time for multi-task or constrained policies. The paper motivates EBMs partly through their composability property (E_1 + E_2 representing the product of densities), but all experiments train and evaluate single-task policies. A natural extension would train separate EBMs for distinct subtasks (e.g., reaching different goal locations, avoiding different obstacle configurations) using R-NCE, then compose their energy functions at inference to produce a policy that satisfies multiple constraints simultaneously—without retraining. A concrete experiment: train R-NCE EBMs for reaching individual target positions in a multi-goal manipulation task, then compose them with an EBM trained for general obstacle avoidance, and measure whether the composed energy function produces trajectories that reach the correct goal while avoiding obstacles, compared against a single EBM trained on the combined task. The key measurement would be the compositionality gap—how much performance degrades relative to a jointly-trained policy—and whether simple techniques (e.g., temperature scaling of individual energies before addition) can close it. This experiment directly tests whether the composability advantage claimed in Section 1 survives the transition from theory to practice with R-NCE training.

Scaling R-NCE to larger models and comparing against larger diffusion policies on the same architecture. The paper controls parameter counts to be comparable (~500k for path planning, ~3.3M for Push-T), but both the EBMs and diffusion models are relatively small by modern standards. Chi et al. [15] use substantially larger models for Push-T and report higher coverage scores than any model in Table 6. A scaling study would train R-NCE EBMs and diffusion models at matched parameter counts across orders of magnitude (e.g., 1M, 10M, 50M parameters) on the Push-T task, measuring both final performance and training stability. The key question: does the performance gap between R-NCE and diffusion change with model scale? If R-NCE's advantage derives from the multi-scale training dynamic (I-R-NCE), the gap might shrink as diffusion models also benefit from larger architectures. If R-NCE's advantage derives from the EBM representation itself (composability, density access), larger models might amplify the gap as the EBM has more capacity to capture fine-grained density structure that the diffusion model's score function cannot exploit. The paper's asymptotic analysis (Theorem 4.12) predicts R-NCE efficiency degrades by a (1 + 1/K) factor relative to MLE—an experiment measuring whether this theoretical gap translates to practical performance differences at scale would validate or challenge the theoretical framework.

Stress-testing the non-adversarial design by implementing and comparing against adversarial R-NCE. Theorem 4.18 and Corollary 4.19 prove that adversarial training of the negative sampler provides no asymptotic variance benefit over non-adversarial training in the non-realizable proposal setting—but this is a limit statement, not a finite-sample comparison. An adversarial R-NCE implementation would train the negative sampler to maximize the R-NCE loss (i.e., minimize L(θ, ξ)) rather than its own independent objective, creating a min-max game. A careful comparison against Algorithm 1 on path planning and Push-T would measure: (a) final policy performance, (b) training stability (variance of the R-NCE loss over training, frequency of divergence), (c) wall-clock training time (adversarial training requires differentiating through the sampler's ODE solve), and (d) sensitivity to learning rate ratios between the two models. If adversarial training performs equivalently or worse on all metrics, this would empirically validate the paper's theoretical claim and strengthen the case for the simpler non-adversarial approach. If adversarial training outperforms in some regime (e.g., very limited data, very high-dimensional action spaces), this would identify a boundary condition where the asymptotic analysis is insufficient and adversarial dynamics matter in practice.

Extending R-NCE to reinforcement learning for EBM policy optimization. The paper focuses exclusively on behavioral cloning (supervised imitation), but the R-NCE objective and Algorithm 1 do not inherently depend on having expert demonstrations. In principle, one could replace the positive samples (x, y) with trajectory data collected by any policy and reweighted by advantage or return, transforming R-NCE into a policy gradient method for EBMs. A concrete starting point: implement an R-NCE-trained EBM as the policy in a soft actor-critic (SAC) style framework, where the EBM represents the maximum-entropy policy π(a|s) ∝ exp(Q(s, a) / α) and R-NCE trains the energy model to match this distribution using samples from a replay buffer (with positive samples being actions that received high Q-values). The key experiment would compare against standard SAC with Gaussian policies on continuous control benchmarks (e.g., MuJoCo tasks) where multi-modality is present but not dominant, and against diffusion-based RL methods on tasks where multi-modality matters (e.g., maze navigation, multi-goal reaching). The hypothesis is that R-NCE EBMs would match or exceed Gaussian SAC on standard benchmarks (due to their ability to represent non-Gaussian action distributions) and match diffusion-based RL on multi-modal tasks (due to the multi-scale I-R-NCE training). The main challenge—distribution shift between the behavior policy that collected the data and the current EBM policy—is analogous to the off-policy challenge in any actor-critic method and would require importance sampling or conservative updates.

Developing cheap difficulty estimation methods to make the compute-optimal allocation framework practically deployable. This is a direct response to the paper's acknowledged limitation that the 2,048-sample difficulty estimation cost is unaccounted for in the reported 4× efficiency gains. A strong follow-up would train a lightweight difficulty predictor—a small neural network taking only the question text (for text-based tasks) or the initial observation (for visuomotor tasks) as input, trained to predict the difficulty bin that would be assigned by the full 2,048-sample procedure. The training data would come from running the expensive estimation on a held-out set. The key measurement: at test time, how much of the 4× efficiency gain is preserved when using the cheap predictor's bin assignment versus the oracle bin assignment? If the predictor achieves, say, 80% bin accuracy and the resulting compute-optimal policy retains 90% of the oracle's gain, then the framework becomes immediately practical—the prediction cost (a single forward pass) is negligible compared to the test-time budget. If bin misclassification causes catastrophic degradation (e.g., routing an easy problem to beam search causes verifier over-optimization), this would motivate more sophisticated difficulty estimation—perhaps adaptive estimation where a few initial samples inform a dynamic allocation of the remaining budget, similar to how hyperparameter optimization methods use early-stopping.


Practical Applications and Downstream Use Cases

On-robot policy learning with compositional constraints for safe deployment. In safety-critical robotics applications (surgical robots, autonomous vehicles, industrial manipulators), the policy must simultaneously satisfy multiple constraints: reach the goal, avoid obstacles, respect joint limits, minimize energy consumption, and comply with task-specific rules. Each constraint can be represented as a separate energy function, and the composed energy E_total = E_task + E_safety_1 + E_safety_2 + ... directly defines the policy distribution at inference time. Prior to this work, the unreliability of EBM training meant this composition idea—while theoretically elegant—could not be realized with learned policies. R-NCE provides a training pipeline that produces high-quality, single-task EBMs (Table 3: path planning with obstacles, Table 5: block pushing to goal configurations). A deployment system would: (1) train separate EBMs for reaching, obstacle avoidance, and smoothness using R-NCE on demonstration data for each subtask; (2) at deployment, sum the energy functions with task-specific weights λ_i to produce a composite policy; (3) use Algorithm 4 (for I-EBMs) or Algorithm 2 (for standard EBMs) to sample actions from the composed energy. The benefit is modularity—new safety constraints can be added as additional energy terms without retraining, and the weights λ_i can be tuned per-deployment-scenario without any learning. The paper's composability motivation (Section 1) becomes actionable because R-NCE makes the underlying EBMs reliable enough to compose.

Offline trajectory library generation for model-predictive control warm-starting. Many robotic planning systems use trajectory optimization (e.g., STOMP, CHOMP, TrajOpt) that requires a good initial guess to converge to a high-quality solution. A generative policy with explicit density estimates can serve as a "trajectory library generator": sample N diverse candidate trajectories from the EBM, rank them by energy (which gives unnormalized log-probability up to a constant), and use the highest-ranked trajectories as initializations for the optimizer. The EBM's advantage over diffusion models here is that energy evaluation is a single forward pass—no ODE integration needed to compare candidate quality. Table 4 quantifies the importance of ranking for path planning: going from ℓ = 48 samples with ranking to ℓ = 1 increases collision rate from 3.0% to 11.6% and cost from 0.194 to 2.287 for R-NCE. This suggests that the energy model produces samples of variable quality and the ranking step effectively filters them—the same mechanism would work for warm-starting an optimizer, where providing good initializations can dramatically reduce optimization time and improve solution quality. The two-time-scale log-probability technique (Section 6) makes this practical by decoupling sample quality from log-probability accuracy, keeping the ranking computation cheap.

Multi-modal imitation learning for shared autonomy and human-in-the-loop systems. In shared autonomy settings (e.g., assistive robotics, teleoperation with autonomy assistance), the system must present the human operator with multiple distinct, high-quality action options, not just a single prediction. An R-NCE-trained EBM policy naturally generates diverse samples via MCMC (Algorithm 2 or 4), and the energy function provides a principled way to select a diverse set: sample many candidates, cluster by similarity, and present the highest-energy sample from each cluster. This is more principled than generating diverse samples from a diffusion model because the EBM's explicit density at each sample point allows filtering by both quality (energy) and diversity (distance in action space), whereas diffusion models require separate clustering on sampled outputs without density-based quality scores. The Push-T results (Table 6: I-R-NCE coverage 0.884, Figure 5: qualitatively diverse pushing strategies) demonstrate that the EBM captures meaningful multi-modality. In a shared autonomy system, this would mean the autonomy module proposes, e.g., three distinct pushing strategies for a block, ranked by likelihood, and the human selects or modifies one—combining the EBM's knowledge of the demonstration distribution with human judgment for edge cases.


When to Prefer This Method

The paper explicitly positions R-NCE-trained EBMs against two alternatives: IBC-trained EBMs (shown to be biased and lower-performing) and diffusion-based policies (shown to be competitive, with R-NCE having a small edge on the evaluated benchmarks). The choice is framed in terms of the practical properties each model class offers:

Prefer R-NCE-trained EBMs (with or without interpolating extension) over diffusion models when:

  • Density evaluation matters at inference time. If the deployment requires comparing or ranking candidate actions by likelihood (as in Sections 7.3–7.4 where ℓ = 48 samples are ranked per step), EBMs provide energy/likelihood comparisons in a single forward pass. Diffusion models require the two-time-scale ODE integration described in Section 6 for log-probability computation. The paper's reference implementation (Appendix A, Figure 5) shows this is feasible but adds integration overhead relative to a single EBM forward pass.
  • Composability is needed. If the deployment requires combining multiple learned behaviors or adding new constraints at inference time without retraining, EBM energy functions can be summed directly. Diffusion models do not have this property—combining score functions is not equivalent to combining densities. The paper motivates this advantage in Section 1 but does not experimentally demonstrate composition.
  • Model compactness is paramount for real-time inference. The paper's abstract argument that EBMs "yield compact representations" is based on the EBM representing a scalar field rather than a vector field. In practice, the experiments keep parameter counts comparable (Section 7.1, principle (a)), so the compactness advantage is not empirically demonstrated, but the architectural argument is sound: for the same representational capacity, an energy function (outputting a scalar) needs fewer parameters than a score function or vector field (outputting a dim(Y)-dimensional vector). This difference becomes significant as action dimensionality grows.

Prefer I-R-NCE over single-scale R-NCE when the action space is high-dimensional and task complexity is high. The Push-T results (Table 6: I-R-NCE 0.884 vs. R-NCE 0.824) demonstrate a 6 percentage point gap. The interpolating extension adds training complexity (multiple time samples, correlated negative trajectories, the SDE sampling stage) but the performance gain on 32D action spaces justifies it. For lower-dimensional problems (e.g., 20D path planning, where single-scale R-NCE already achieves 3.0% collision rate), the added complexity may not be warranted.

Prefer diffusion models over R-NCE-trained EBMs when:

  • MCMC sampling overhead is unacceptable. EBM inference requires Langevin or HMC sampling (Algorithm 2 or 4), which involves T_mcmc sequential gradient evaluations of the energy function. For Push-T, T_mcmc = 750–1250 evaluations are used with HMC (Table 11). Diffusion models sample by integrating an ODE or SDE, which also requires sequential function evaluations—but the paper keeps function evaluation counts comparable across models (principle (b), Section 7.1), so this is a controlled rather than inherent difference. In latency-constrained settings where MCMC burn-in cannot be amortized, diffusion's ODE integration might be preferable if it converges faster per function evaluation.
  • Training simplicity is the dominant concern. R-NCE training requires maintaining a separate negative sampler, choosing its architecture and capacity appropriately to avoid posterior collapse, tuning the T_samp/T_rnce ratio, and for I-R-NCE, managing the interpolating time sampling and correlated negative trajectories. Diffusion model training is conceptually simpler: a single network trained with a single denoising or score-matching objective. For practitioners who just want a working generative policy and don't need composability or density access, the added complexity of R-NCE may not be justified by the performance difference (which, on the Push-T results, is 0.884 vs. 0.864—a 2 percentage point gap that may be within architectural tuning noise).
  • The task does not require explicit density estimates or composition. If the deployment only requires forward sampling from the policy (no ranking, no likelihood evaluation, no multi-task composition), diffusion models provide equivalent or near-equivalent performance with a simpler training pipeline and no need for a negative sampler design.

Prefer R-NCE over IBC in all regimes where the negative sampler is non-uniform. Proposition 4.7 proves that IBC is biased when p_ξ is non-uniform, and this bias is observed even in 2D synthetic examples (Tables 1–2: IBC BC = 0.540 on Spiral vs. R-NCE BC = 0.993). The only case where IBC might be preferred is extremely low-dimensional action spaces where uniform negative sampling is feasible—but this regime is essentially irrelevant for robotics.