Revisiting Energy Based Models as Policies: Ranking Noise Contrastive Estimation and Interpolating Energy Models
Authors: Sumeet Singh
Stephen Tu
Vikas Sindhwani
Venue: Accepted by TMLR ()

Abstract: A crucial design decision for any robot learning pipeline is the choice of policy representation: what type of model should be used to generate the next set of robot actions? Owing to the inherent multi-modal nature of many robotic tasks, combined with the recent successes in generative modeling, researchers have turned to state-of-the-art probabilistic models such as diffusion models for policy representation. In this work, we revisit the choice of energy-based models (EBM) as a policy class. 

We show that the prevailing folklore---that energy models in high dimensional continuous spaces are impractical to train---is false. We develop a practical training objective and algorithm for energy models which combines several key ingredients: (i) ranking noise contrastive estimation (R-NCE), (ii) learnable negative samplers, and (iii) non-adversarial joint training. We prove that our proposed objective function is asymptotically consistent and quantify its limiting variance. On the other hand, we show that the Implicit Behavior Cloning (IBC) objective is actually biased even at the population level, providing a mathematical explanation for the poor performance of IBC trained energy policies in several independent follow-up works. We further extend our algorithm to learn a continuous stochastic process that bridges noise and data, modeling this process with a family of EBMs indexed by scale variable. In doing so, we demonstrate that the core idea behind recent progress in generative modeling is actually compatible with EBMs. Altogether, our proposed training algorithms enable us to train energy-based models as policies which compete with---and even outperform---diffusion models and other state-of-the-art approaches in several challenging multi-modal benchmarks: obstacle avoidance path planning and contact-rich block pushing.

Published in Transactions on Machine Learning Research (09/2024)
Revisiting Energy Based Models as Policies: Ranking Noise Contrastive Estimation and Interpolating Energy Models
Sumcet Singh
saumect@@google.com
Google DeepMind
Stephen Tu
clephen.lu@usee edu
Ming Hsich Department of Electrical ond Computer Engineering University of Southenn California
Vikas Sindhwani
sindhwani@google.com
Google DeepMind
Reviewed on OpenReview: https: opw.reviee........mmmeeeee net/ formm? id= JmKAY67100
Abstract
crucial design decision for any robot learning pipeline is the choice of policy representa- tion: what type of model should be used to generate the next set of robot actions? Owing to the inherent multi-modal nature of many robotic tasks, combined with the recent successos in generative modeling, researchers have turned to state-of-theeart probabilistic models such diffusion models for policy representation. In this work, revisit the choice of energy- based models (EBM) as policy class. We show that the prevailing folklore that energy models in high dimensional continuous spaces are impractical to train-is false. We develop a practical training objective and algorithm for energy models which combines several key ingredients: (i) ranking noise contrastive estimation (R-NCE), (fi) learnable negative sam- plers, and (iii) non-adversarial joint training. We prove that our proposed objective function asymptotically consistent and quantify its limiting variance. On the other hand, we show that the Implicit. Behavior Cloning (IBC) objective is actually biased even at the population level, providing mathematical explanation for the poor performance of IBC trained energy policies in several independent follow-up works. We further extend our algorithm to learn a continuous stochastic process that bridges noise and data, modeling this process with a family of EBMs indexed by scale variable. In doing so, we demonstrate that the core idea. behind recent progress in generative modeling actually compatible with EBMs. Alto- gether, our proposed training algorithms enable us to train energy- models as policies which compete
even outperform ddiffuoo models and other state-of-the-art ap- proaches in several challenging multi-modal benchmarks: obstacle avoidance path planning and contact-rich block pushing.
1 Introduction
Many robotic lasks e.g. grasping, manipulation, and trajectory planning- are inherently multi-modal: at any point during task exccution, there may be multiple actions which yield task-optimal behavior. Thus, a fundamental question in policy design for robotics how to capture such multi-modal behaviors, especially when learning from optimal demonstrations. A natural starting point is to treat policy learning distri- bution learning problem: instead of representing policy as deterministic map To(r), utilize conditional
* Work performed while author affiliated with Google DeepMind.

---

Published in Transactions on Machine Learning Research (09/2024)
generative model of the form po | z)) to model the entire distribution of actions conditioned on the current robot state.
This approach of modeling the entire action distribution, while very powerful, leaves open a key design question: which type of generative model should one use? Due to the recent advances in diffusion models across a wide variety of domains (Dhariwal & Nichol, 2021; Blattmann et al., 2023; Kong et al., 2021; Saharia et al., 2023; Song et al., 2022; Yang et al., 2022), it is natural to directly apply diffusion model for the policy representation (Chi et al., 2023; Janner et al., 2022; Reuss et al., 2023). In this work, however, we revisit the use of eneryy-base based models (EBMs) policy representations (Florence et al., 2022). Energy- based models have a mumber of appcaling propertios in the context of robotics. First, by modeling a scalar potential field without any normalization constraints, EBMs yield compact representations; model size is an important consideration in robotics due to real-time inference requirements. Second, action selection for many robotics tasks (particularly with strict safety and performance considerations) can naturally be expressed as finding a minimizing solution of a particular loss surface (Adams et al., 2022); this implicit structure is suceinetly caplured through the sampling procedure of EBMs. Furthermore, reasoning in the density space is more amenable to capturing prior task information in the model (Urain et al., 2022b), and allows for straightforward composition which is not possible with score-based models (Du et al., 2023).
Unfortunately, while the advantages of EBMs are clear, FBMs for continuous, high-dimensional data can be difficult to train due to the intractable computation of the partition function. Indeed, designing efficient algorithms for training EBMs in general is still an active area of research (see e.g., Du & Mordatch (2019); Dai et al. (2019); Song & Kingma (2021); Arbel et al. (2021)). Within the context of imitation learning, recent work on implicit behavior cloning (IBC) (Florence et al., 2022) proposed the use of an InfoNCE (van den Oord et al., 2018) inspired objective, which we refer to as the IBC objective. While IBC is considered state of the art for EBM behavioral cloning, follow up works have found training with the IBC objective to be quite unstable (Ta et al., 2022; Reuss et al., 2023; Chi et al., 2023; Pearce et al., 2023), and hence the practicality of training and using EBMs as policies has remained an open question.
In this work we resolve this open question by designing and analyzing new algorithms based on ranking noise constrastive estimation (R-NCE) (Ma& & Collins, 2018), focusing on the behavioral cloning setting. Our main theoretical and algorithmic contributions are summarized as follows:
The population level IBC objective is biased: We show that even in the limit of infinite data, the population level solutions of the IBC objective are in general not correct. This provides a mathematical explanation as to why policies learned using the IBC objective often exhibit poor performance (Ta et al., 2022).
Ranking noise contrastive estimation with learned sampler is consistent: We utilize the ranking noise contrastive estimation (R-NCE) objective of Ma & Collins (2018) to address the shortcomings of the IBC objective. We further show that jointly learning the negative sampling distribution is compatible with R-NCE, preserving the asymptotic normality properties of R-NCE. This joint training turns out to be quite necessary in practice, as without it the optimization land- scape of noise contrastive estimation objectives can be quite ill-conditioned (Liu et al., 2022; Lee el al., 2023).
EBMs are compatible with multiple noise resolutions: A key observation behind recent gen- erative models such as diffusion (Ho et al., 2020; Song et al., 2021b) and stochastic interpolants (Al- bergo et al., 2023; Tipman et al., 2023; Liu et al., 2023) is that learning the data distribution at multiple noise scales is critical. We show that this concept is actually compatible with EBMs and contrastive training, by introducing a framework that models a continuum of EBMs which we term interpolating EBMs, trained jointly using R-NCE. We believe this contribution to be of independent interest to the generative modeling community.
EBMs trained with R-NCE yield high quality policies: We show empirically on several multi- modal benchmarks, including an obstacle avoidance path planning benchmark and a contact-rich block pushing task, that EBMs trained with R-NCE are competitive with and can even outper- form IBC and diffusion-based policies. To the best of our knowledge, this is the first result in the
2

---

Published in Transactions on Machine Learning Research (09/2024)
literature which shows that EBMs provide a viable alternative to other state-of-the-arr generative models for robot policy learning, and can be stably trained even with modern transformer-basel backbones.
2 Related work
Generative Models for Controls and RL. Recent advances in generative modeling have inspired many applications in reinforcement learning (RI. (Heess et al., 2013; Haarnoja et al., 2017; 2018; Levine, 2018; Liu et al., 2021; Ho & Ermon, 2016), trajectory planning (Du et al., 2020; Urain et al., 2022b; Ajay et al., 2022; Janner et al., 2022), robotic policy design (Florence et al., 2022; Pearce et al., 2023; Chi et al., 2023; Reuss et al., 2023), and robotic grasp and motion generation (Urain et al., 2022a). In this paper, we focus specifically on generative policies trained with behavior cloning, although many of our technical contributions apply more broadly learning energy-based models. Most related to our work is the influential implicit behavior cloning (IBC) work of Florence et al. (2022), which advocates for using a noise contrastive estimation objective (which we refer to as the IBC objective) based on InfoNCE (van den Oord et al., 2018) in order to train energy-based policies. As discussed previously, many works (Ta et al., 2022; Reuss et al., 2023; Chi et al., 2023; Pearce et al., 2023) have independently found that the IBC objective is numerically unstable and does not consistently yield high quality policies. One of our main contributions a theoretical explanation of the limitations of the IBC objective, and an alternative training procedure based on ranking noise cont rastive estimation. Another closely related work is Chi et al. (2023), which shows that based diffusion models (Sohl-Dickstein et al., 2015; Ho et al., 2020; Song et al., 2020; 2021b) yield state-of-the the art generative policies for many multi-modal robotics tasks. In light of this work and the poor empirical performance of the IBC objective, it is natural to conclude that score-based diffusion methods are now the de facto standard for learning generative policies to solve complex robotics tasks. One of our contributions is to show that this conclusion is false; energy-based policies trained via our proposed R-NCE algorithms are actually competitive with diffusion- policies in terms of performance.
We focus the rest of the related work discussion on the training of EBMs, as this is where most of our technical contributions apply. Given the considerable scope of this topic and the extensive literature (see e.g.; Song & Kingma (2021); Lecun et al. (2006) for thorough literature reviews), we concentrate our discussion on the two most common frameworks: maximum likelihood estimation (MLE) and noise contrastive estimation (NCE).
Training EBMs via Maximum Likelihood. A key challenge in training EBMs via. MLE is in computing the gradient of the MLE objective, which involves an expectation of the gradient of the energy function w.r.t. samples drawn from the EBM itself (cf. Section 3). This necossitates the use of expensive Markov Chain Monte Carlo (MCMC) techniques to estimate the gradient, resulting in potentially significant truncation bias that can cause training instability, A common heuristic involves using persistent chains (Tieleman, 2008; Du & Mordatch, 2019; Du et al., 2021) over the course of training to better improve the quality of the MCMC samples. However, such techniques cannot be straightforrardly applied to the conditional distribution setting, as one would require storing context-condittonel chains, which becomes infeasible when the context space is contimuous.
Another common technique for scalable MLE optimization is to approximate the EBM model's samples using an additional learned generative model that casier to sample from. Re-writing the partition function in the MLE objective using a change-of- lower bound and thus, max-min optimization problem where the sampler is optimized to tighten the bound while the EBM is optimized to maximize likelihood (Dai et al., 2019; Grathwohl et al., 2021). For example, Dai ct al. (2019) propose to learn base distribution followed by using a finite number of Hamiltonian Monte Carlo (HMC) leapfrog innegration steps to sample from the EBM. Grathwohl et al. (2021) use a Gaussian sampler with a learnable mean function along with variational approximations for the inner optimization. More examples of this max-min optimization approach can be found in e-g., Song & Kingma (2021); Bond- Taylor et al. (2022). Unfortunately, adversarial optimization is notoriously challenging and requires several stabilization tricks to prevent mode-collapse (Kumar et al., 2019).
argument and applying Jensen's inequality yields a variational
3

---

Published in Transactions on Machine Learning Research (09/2024)
Alternatively, one may instead leverage an EBM to correct backbone latent variable-based generative model, commonly referred to as exponential tilling. For instance, this has been accomplished by using VA (Xiao et al., 2021) and normalizing flows (Arbel et al., 2021; Nijkamp et al., 2022) as the backbone. Leveraging the inverse of this model, one may perform MCMC sampling within the latent space, using the pullback of the exponentially tilted distribution. Ideally, this pullback distribution is more unimodal, thereby allowing faster mixing for MCMC. Pushing forward the samples from the latent space via the backbone model yields the necessary samples for defining the MLE objective. Thus, the EBM acts as an implicit "generative correction" of the backbone model, as opposed to standalone generative model. The challenge however is that the overall energy function that needs to be differentiated for MCMC sampling now involves the transport map as well, which, for non-trivial backbone models (e.g., continuous normalizing flows (Chen et al., 2018)) yields a non-negligible computational overhead.
We next outline NCE as form of "discriminative correction", drawing a natural analogy with Generative Adversarial Networks (GANs).
Training EBMs via Noisc Contrastive Estimation. Instead of directly trying to maximize the like- lihood of the observed samples, NCE leverages contrastive learning (Le-Khac et al., 2020) and thus, is composed of two primary parts: (i) the contrastive sample generator, and (ii) the classification-bnsed critic, where the latter serves as the optimization objective. The general formulation was introduced by Gutmann & Hyvärinen (2012) for unconditional generative models, whereby for each data sample, one samples K "contrast. (also referred to as "negative") examples from the noise distribution, and formulates a binary classification problem based upon the posterior probability of distinguishing the true sample from the syn- thesized contrast samples.
For conditional distributions, Ma & Collins (2018) found that the binary classification approach is severely limited: consistency requires the EBM model class to be self-normalized. Instcad, they advocate for the ranking NCE (R-NCE) variant (Jozelowicz el al., 2016). R-NCE posits a multi-class classification objective as the critic, and overcomes the consistency issues with the binary classification objective in the conditional sctting (Ma & Collins, 2018), at least for discrete probability spaces. In this work, focus on the R-NCE formulation, and extend the analysis of Ma & Collins (2018) to include jointly optimized noise distribution models over arbitrary probability spaces (i.e., beyond the discrete setting). In particular we study various properties such as: (i) asymptotic convergence, (ii) the pitfalls of weak contrastive distributions, (iii) jointly training the contrastive generative model, either via an independent objective or adversarially, (iv) extension to a multiple noise scale framework by defining a suitable time-indexed family of models and castrastive losses, and (v) sampling from the combined contrastive generative model and EBM.
Despite leveraging similar dassification--aseed objective to GANs, NCE does not need adversarial training in order to guarantee convergence. In particular, the contrastive model in NCE is not the primary genera- tive model being learned; it is merely used to provide the counterexamples needed to score the EBM-basod classifier. Furthermore, our asymptotic analysis shows that a fixed noise distribution suffices for consistency and asymptotic normality. In practice, however, the design of the noise distribution is the most critical factor the success of NCE methods (Gutmann & Hirayama, 2011). An overly simple fixed noise distri- bution will lead to trivial classification problem, which is problematic from an optimization perspective due to exponentially flat loss landscapes (Liu et al., 2022; Lee et al., 2023). Thus, a key part. of our al- gorithmic contributions is in how the negative sampler is learned. We propose simultaneously training a simpler normalizing flow model as the sampler, jointly with the EBM. This is in contrast with adversarial approaches (Gao el al., 2020; Bose et al., 2018) which require max-min training." Without a need for ad- versarial optimization, our joint training algorithm is numerically stable and yields a pair of complementary generative models, whereby the more expressive model (EBM) bootstraps off the simpler model (normalizing flow) for both training and inference-ting time sampling.
"aoo et al. (2020) actually consider similar approach in the binary NCE setting, but ultimately dismiss it in favor of adversarial training. We delve further by analyzing the statistical properties of adversarial training and question its necessity in regards to the R-NCE objective.
4

---

Published in Transactions on Machine Learning Research (09/2024)
3
Ranking Noise Contrastive Estimation
3.1 Notation
Let the context space (X, (x) and event space (Y, JY) be compact measure spaces (for simplicity, we assume that the event space is not a function the context). Equip the product space x Y with a probability measure Px,y = Px x Pylx, and assume that this probability measure is absolutely continuous w.r.t. the base product measure px X py. Suppose furthermore that the conditional distribution Pylx x is regular. Let p(æ, w), p(x), and ply æ) denote the joint densily, the marginal density, and the conditional density, respectively (all densities are with respect to the their respective base measures).
Let ④ be a compact subset of Euclidean space, and consider the function class of conditional energy models:
(Eo(1,)) v) X x Y R e E).
(3.1)
This family of energy functions induccs conditional densitios in the following way.?
po x) exp(Eo(xx,)) ())
Zo(x) / (xp(E duy.
Zz(
Next, let also be a compact subset of Euclidean space, and consider the parametric class of conditional densitics for the contrastive model:
: (2)
(3.2)
While our theory will be written for general parametric density classes, for our proposed algorithm be practical we require that both computing and sampling from pę(y x) is efficient. Some examples of models which satisly these requirements include normalizing flows (Grathwohl el al., 2019) and stochastic interpolants (Albergo & Vanden-Eijuden, 2023; Albergo et al., 2023).
We globally fix positive integer K E N4. For a given I, let PK ylnng denote the product conditional sampling y/æ;ş denotes a random vector y (yx) Yk with y P=( x) for k 1,..., K. Now, for a given (ar, y) ~ Px,,, y PK yluge and parameters 0 E θ, & E we define:
distribution over YK where y pk
exp(Eo(x,)) log x)) Co yy) log Evellyyy exp(Ex(x,)) y¹) lo x))
(3.3)
With this notation, the Ranking Noise Contrastive Estimation (R-NCE) population mazimization objective I(0. {) is:
L(),E) := B(e Eyupg
Vlaie f,,e(x, ly).
(3.4)
For notational brevity, we will write this as:
L(O,E) := E lo, y yy)
where the length K of the vector y implied from context.
Intuition. Let us build some intuition for the ranking objective (3.3). For any (x, v) Px.Y and y Pk let : Uy = ( yK+1, Let the random variable D E (l,....,k, K 1} denote the index of the true sample y within y. We will construct Bayes classifier for D using the EBM. In particular, leveraging the EBM as the generative model for the true distribution Prix, can define the "class-conditional" distribution over y as:
y
K+1 II 3 1ijkk
py x. D. D= k)) Po(y x) PelUs x))
While EBM would traditionally be represented po exp( Es(a, v)) (sco c.g., Locun al. (2006)), utilize the non-negated version notational convenience and consistency with referenced works such Gutmann Hyvärinen (2012); Ma Collins (2018), upon which base our theoretical development.

---

Published in Transactions on Machine Learning Research (09/2024)
Then, the posterior p(D = 1o x. x) - g0.5(k: | x. may be derived via Bayes rule as:
po ly x. D = k)
g0.5 x.
k+ اド=1 po.g (g x. D
Now, remains to specify a prior distribution p(D = æ). A simple, yet natural, choice reflecting the desire that no index is to be preferred over any other indices is to choose the uniform prior distribution p(D D k I 1/(K + 1) for all & (1,...,, K + 13. With this prior, the posterior probability simplifies as:
pa (x) I1 PE (3)) x))
( x) - Ek+1 (o( ()) I14 Pe(m x)
(3.5)
Po(u x)/pp (M1 x))
(3.6)
j=1 po(U5 1x)/pg (Uy xx)
Substituting po(v z) exp(E6 vl)l.,)))/aoo (z)), and noting that the partition function Zo cancels:
exp(Ee(<; m&) log bogge(us x))
q0 x Eyen exp(E6 (z. y) Pe x))
(3.7)
Thus, the objective in (3.3) denotes the posterior log- bog-rrbbability (conditioned on both the data (2, and the contrastive samples y) that the sample y is drawn from the energy model & and the remaining samples K samples in y are all drawn iid from pę(- æ). Based on this interpretation, we adopt the terminology that y the posilive example, y contains the negative examples, and pe(- x) is the negalive proposal distribution.
Given this posterior log-probability interpretation, it will prove useful to re-state the R-NCE objective follows. For k E {1..... K + 1}, let Pyleee denote the true "class-coniitional distribution over Yk+1, That is, for given y (w)) omitted, this refers to the setting of k = 1: Pylm: = Py æ:g.1- By symmetry then, for any & E {1....,K K +1,, the population objective in (3.4) can be cquivalently written as:
Pylme,, we have yk p(- x)) and 8 ~ P x) for j # k. When the index k is
L(D,E) Ex-Px Byn log log x. y).
(3.8)
For notational brevity, we will write this as:
L(0,E) == EyEgii log do.
x)
This notation will prove useful for subsequent analysis.
Comparison to Maximum Likelihood Estimation. Let us take moment to compare the R-NCE objective to the standard maximum likelihood estimation (MLE) objective:
Lmte (0) = EEr,y v log Pe xx)
(3.9)
Letting y {y}U Uy, and taking the gradient of both Lmtt (0) and L(O,E) (assume for now the validity of exchanging the order of expectation and derivative),
(eLmle(() Exr,y VeE Ey (-x())))))))))))))))))))>)))))>) VoEo(z,
(MLE)
K+1
VoL(O,E) = Ex,,y VnEo Ey un:
Σ co (
(R-NCE)
Comparing the two expressions, the main difference is in the gradient correction term on the right. For MLE, this gradient computation requires sampling from the energy model po | x), which is prohibitively expensive during training. On the other hand, the R-NCE gradient computationally efficient to compute, provided the negative sampler is efficient to both sample from and compute log-probabilities
6

---

Published in Transactions on Machine Learning Research (09/2024)
Algorithm 1: Ranking noise-contrastive estimation with learnable negative sampler
Data: Dataset D = {(0,, WO number of outer steps Touter, number of sampler steps per outer step Tummp number of R-NCE steps per outer step maee, number of nogative samples K. Result: Energy model parameters 0, negative sampler parameters c. Initialize 0 and &.
Set (fmle(; B) = -- - L(e JEB log leg xz). (Oa (Ti, Vi ys).
Set (0,
for Touter iterations do
for Teamp iterations do
B - batch of D.
< t minimize using (E: // Or use auxiliary loss, see Remark 3,
end
for Tine iterations do
10
B batch of D. // Augment batch with negative samples.
B M.) (Ia, vi) E B), with y, pk
11
yla;
12
< minimize using 60
13
end
14
return (8
15 end
Algorithm 2: Sample from po(y (r) for Euclidean Y
Data:
Context X, energy model parameters 8, negative sampler parameters &, Langevin steps Tucmue; Langevin step size n7.
Result: Sample y po(y x:) Set P:(y z).
// Run Langevin sampling, could alternatively run HMC sampling.
for / == ......
:.....Tmemm Tmmme 1 do
Set M1+1 y/2 + EV 2ipe, with uy ~ No 1).
end
return yT
3.2 Algorithm
Our main proposed R-NCE algorithm with learnable negative sampler is presented in Algorithm 1. We remark that one could alternatively pre-train the negative sampler distribution Pt (i.e., move Lines 5-8 after Line 3 and set Tsamp large enough so that the negative sampler converges), as opposed to jointly training the negative sampler concurrently with the energy model. Using a pre-trained negative sampler distribution was previously explored in Gao et al. (2020), where it was empirically shown to be an ineffective idea. for binary NCE. We will also revisit this choice within our experiments.
Remark 3.1. As mentioned earlier, any family of densities Fn suffices in Algorithm 1. For specific families such as stochastic interpolants (Albergo & Vanden- Landen-Fijnden, 2023) where auziliary loss is used instead of negative log-likelihood, Algorithm 1, Line 7 is replaced with the corvesponding auriliary loss.
Remark 3.2. Note that in Algorithm 1, the proposal model is optimized independently of the EBM; the latler is oplimized purely vio the R-NCE objective. Crucially, the negative samples y do not depend upon the EBM's parameters 0. Hence, no re-parameterization tricks are required since the batch augnentation Line 11 does not introduce any gradient dependencies on for the optimization step in Line 12.
Algorithm outlines how sample from the combined proposal Pş and EBM pe models. In particular, leverage Ps to generate initial sample in Line and warm-start finite-step MCMC chain with the

---

Published in Transactions on Machine Learning Research (09/2024)
EBM in Lines 2-4. We outline the simplest possible MCMC implementation (Langevin sampling), however one may leverage any additional variations and tricks, including Hamiltonian Monte Carlo (HMC) and Metropolis- Hastings adjustments.
4 Analysis
We now present our analysis of ranking noise contrastive estimation. First, we will require the following regularity assumptions: Assumption 4.1 (Continuity). Assume the following assumptions are satisfied:
1. For a.e. vI, the map 0 + Ex(2, v) is continuous, and for all 0, the map (I,z) + Eo(x, y) is measurable.
2. For a.e. (I, the map & + Pe x) is continuous, and for all &, the map (0,z) + pe(y x) is measurable.
We additionally require the following boundedness assumption: Assumption 4.2 (Uniform Integrability). Assume that:
30. Eo e = s.t. Eny lase sup #Ee. ÇEE
loo.c0 (x, 1y)l co.
Furthermore, define the random variable (o (2, y), measurable on X x Y:
(,.e = Bylng l0.e yy)
Then, by the continuity of (0, + (0, y) for a.e. (x, and the Lebesgue Dominated Convergence
theorem, it follows that: (i) for a.e. (x, v), the map (0,E) + (0 y) is continuous, and (ii) the map ().g) + L(B,E) {) is continuous. For all results in this section, the proofs are provided in Appendix C.
4.1 Optimality
The first step in analyzing the R-NCE objective is to characterize the set of optimal solutions. To do so, we introduce the following key definition:
Definition 4.3 (Rcalizability). We say that F realizable if there exists a 8, Q such that for almost every (I. y):
Exp(Eo,(x,V)) ply x)).
po, (y x))
(4.1)
Zo. (x)
We let O denote the of 0, such that the above condition (4.1) holds. If O2 is a singleton, we say that F is identifiable.
Realizability stipulates that the energy model & is sufficiently expressive to capture the true underlying distribution. We now establish the following optimality result: Theorem 4.4 (Optimality). Suppose that realizable. Then, for any & E E, we have that:
O. = argmax L(0,E)-
(4.2)
8EB
We note that the proof of Theorem 4.4 mostly follows that of Ma & Collins (2018, Theorem 4.1), but includes the weasure-theoretcc arguments necessary to extend it to our more general setting.
Remark 4.5. Notice that the optimality of 9, is independent of the negotive proposal distribution, thereby allowing us to use any generative model for negative samples y, provided easily samplable and the computation of log Pe(- x) is tractable. The question of which negative proposal distribution to use is addressed in Section 4.9.
8

---

Published in Transactions on Machine Learning Research (09/2024)
4.1.1 Comparison to the IBC Objective
In Florence et al. (2022), the following implicit behavior cloning (IBC) objective is proposed:
exp(Eo(z,))) y)) Eu exp(Eo(z,
(0, () = Ex. Eylng log
(4.3)
The IBC objective can be seen as a special case of the R-NCE objective, with the particular choice of negative sampling distribution Pe(y z) as the uniform distribution on Y for a.e. x E X. Note that the uniform distribution the only choice which renders the objective unbiasod (this observation was also made in Ta al. (2022)). Figure illustrates toy one dimensional setting with Gaussian proposal distribution, illustrating the bias inherent in the population objective when non-uniform proposal distribution is used. Shortly, we will characterize the optimizers of the IBC objective, quantifying the bias.
2
Objective
-4
-5
IBC (K= 10)
IBC (K 100)
R-NCE (K 10)
R-NCE (K = 100)
MLE
2
1
0
2
Figure The population objective landscape for the IBC objective (4.3) versus the R-NCE objective (3.4), for K E {10, 100). The true distribution p(y x)) is M(L1,) distribution (for simplicity, the context is not relevant), and the proposal distribution (N 1). The energy model function class is comprised of unit variance Gaussians, Le., & = ((2.)) y) + - [v--)(() 10²² E R}. Note that for ease of comparison, all objectives have been shifted so that the maximum values are zero. Furthermore, the maximum likelihood objective (3.9) has also been plotted for reference.
The fact that the IBC objective forces the use of uniform proposal distribution substantially limits its applicability in high dimension. Indeed, in Section 4.3 will that proposal distribution which is non-informative causes the gradient to vanish, and hence training to stall. This issue has been empirically observed in many follow up works (Ta et al., 2022; Reuss et al., 2023; Chi et al., 2023; Pearce et al., 2023).
Optimizers of the IBC Objective. The optimizers of the IBC objective can be studied using the optimality of the R-NCE objective (Theorem 4.4) in conjunction with change of variables. We first introduce the notion of g-realizability, which posits that the density ratio ply x)/Pe(y x) is representable by the function class F.
Definition 4.6 (f-realizability). Let & E E. We say that I is g-realizable if there erists 0, E θ such that for a.e. (I,y):
exp(Eo, (x,, u))
ply x)) pe (y x)
Sy exp(Eo, (x, y) u))e x) duy
We let Og denote the subset of θ such that the above condition holds.

---

Published in Transactions on Machine Learning Research (09/2024)
With this definition, can now characterize the optimizers of the IBC objective.
Proposition 4.7. Lel E and suppose that is g-renlizable (Definition 4.6). We have:
g = argmax (0,E).
e50
Proof. The proof follows immediately from Theorem 4.4 after simple transformation. In particular, we note that:
exp(Ee(x,)) y) lo x) - log P5 2))
Libe (0.E) Bx, Bylngg log
(y)yy exp(Ex(e, expe + log p x)) - log (v/ xx)
This is equivalent to the R-NCE objective over the shifted function class:
F (E v) log x)) G).
Note that realizability of Fg (Definition 4.3) is equivalent to g-realizability of (Definition 4.6), from which the claim now follows by Theorem 4.4.
Proposition 4.7 illustrates that optimizing the IBC objective results in an energy model which represents the density ratio p(y <>/pe(y z), explaining the bias in the IBC objective when the negative sampling distribution is non-uniform. The reason for this because the IBC objective is based on an incorrect application of InfoNCE (van den Oord al., 2018), which designed to maximize a lower bound on mutual information between contexts æ and events y, rather than extract an energy model (x(x, y) to model p(y æ) (Ta et al., 2022).
4.2 Asymptotic Convergence
In this section, study the asymptotics of the R-NCE joint optimization algorithm. We begin by defining the finite-sample version of the R-NCE population objective:
(., = Ex, (o.c(x.v) = Σ
(4.4)
We consider a sequence of arbitrary negative sampler parameters (Ex) C which are random variables depending on the corresponding prefix of data points ((11,11)). From these negative sampler parameters, we select energy model parameters via the optimization:
Ôn E argmax I.n (0, En)
(4.5)
BEB
We note that the definition of 0,, in (4.5) is a stylized version of Algorithm 1, where the training procedure (e.g., gradient-based optimization) is assumed to succeed. We now establish consistency of the sequence of estimators (0n) to the set O.. Let A(z,A) denote the set distance function between the set A and the point 2, i..,3
d(2, = inf || 2'11. Z'EAA
Theorem 4.8 (General Consistency). Suppose that. F is realizable, O2 is contained in the interior of O and Assumptions 4.1 and 4.2 hold. Let 0., argmaxxge (,, En) denote orbitrary empirical risk matimizer from n samples. Then, d(On, A.) m almost surely.
Throoghhut this work, -1] denotes the fg-norm.
10

---

Published in Transactions on Machine Learning Research (09/2024)
The proof of Theorem 4.8 requires care to deal with the fact that even for fixed 8, the empirical is not unbiased estimate of the population objective (.), since the negative sampler
objective (0, (.) parameters E. in general depend the data points states that for any &, the map 0 ~ [(),E) has the same set of maximizers θ. With this observation, we show that it suffices to establish uniform convergence jointly over è x E, i.e., supaee,ge= (0, g) - L(0, () - a.s. as n + 00.
The key, however, lies in Theorem 4.4, which
As with the optimality result in Theorem 4.4, the asymplotic consistency of the R-NCE estimator is indepen- dent of the sequence of noise distributions, parameterized by {En]. In Section 4.3, we study the algorithmic implications of the design of the noise distribution and its concurrent optimization. We now characterize the asymptotic normality of dn, under certain regularity assumptions:
A Assumption 4.9 (C=-Continuity). Assume the following conditions are salisfied:
1. Por a.e. (I, the maps 0 + Ed( and + pe(U x) are Cc-conlinuous.
2. For a.e. (x,z) and all (0,5) x
3 0 s.t. Bylas sup
(o y)lo < oc.
(4.6)
0| EEll
Then, by the Lebesgue Dominated Convergence theorem, the map (0,E) + (2,,(1,))) y) is C² (Pcootinnuous for a.e. (z. v). Additionally, ossume that for all (0,E) E E x:
3 s.t.
Exy
sup 110 -o
Vo (z.v)llop 00.
(4.7)
It follows then that the map (0,E) + L(O,E) is C2 conttiinuuuss
We are now ready to formalize the asymplotic distribution of On.
Theorem 4.10 (Asymptotic Normality). Suppose that F is realizable and Assumptions 4.1, 4.2, and 4.9 hold. Let {0m} be sequence of optimizers as defined via (4.5). Denote the joint parameter 7 == (0.E) and joint r.v. (x. v), and assume the following properties hold:
The negalive proposal distribulion parumelers } are Vn-consistent about a fixed point Ens i.e., Vn(En. {.) = =1p().
The set θ, = {0.} v is a singleton, and n, : E0,,E.) lies in the interior o E.E.
The Hessian of E, is Lipschitz- i.e.,
[V (=) (=)lloo & M(=)ll 7all, Eg [M()) oo, VT1, x 4. The (0, block of the Hessian L denoled VoL(O,, Ex), salisfies ViL(O.,E.) <00
Then, vn(( 0.) ~ NO0, Vo), where
Ve V3L(00.E.)-'' -E (36.
A few remarks are in order regarding this result. First, note that the fixed point for the negative distribution parameters &, need not be optimal in any sense, as long as it is approached by the sequence &. sufficiently fast. Second, the proof itself begins by using the standard machinery for proving asymptotic normality of M-estimators for the joint set of parameters == (0,E) (see e.g. Ferguson, 2017). In order to extract the marginal asymptotics for On, while employing minimal assumptions regarding the convergence of &n, again the key step is to leverage Theorem 4.4, which states that 0, is optimal for the population objective L(,5E) for any E E. This implies that the joint Hessian (O.,0,,) is block-diagonal, yielding the necessary simplification. Finally, we remark that the result in Ma & Collins (2018, Theorem 4.6) is a spocial case of Theorem 4.10 where (i) the negative distribution is held fixed, i.e., &n = & for all n, and (ii) the event spaces X and are finite.
4he matrix norm lp this work refers the induced + f2 operator norm, Lee, the largest singular value.
11

---

Published in Transactions on Machine Learning Research (09/2024)
4.3 Learning the Proposal Distribution
Given that the consistency and asymptotic normality for the R-NCE estimator are guaranteed under the very mild condition that the negative proposal distribution parameters converge to a fixed point (cf. Theo- rem 4.10), there is considerable algorithmic flexibility in generating the sequence of proposal models, indexed by }. We begin our discussion with first establishing the importance of having a learnable proposal dis- tribution. Using the posterior probability notation (3.7), we can write the R-NCE gradient as:
K+1 Σ
= Ex,y EE VoEo(x, Ey lovs
-
go,cli (i VoCo
Further examination gives the following key intuition about training: the proposal distribution pe(y) x) should not be significontly less informative than the energy model E(x, . If pe x)) is uninformative, then training will stall. This is because the posterior distribution qo,e(- v, will concentrate nearly all its mass on the positive example, as the energy of the true sample will be significantly higher than the energy of the negative samples. As a consequence, the R-NCE gradient approximately vanishes without necessarily having (zz y) x ply z), stalling the training procedure. We term this phenomenon posterior collapse. Indeed, this consistent with the findings Liu et al. (2022); Lee et al. (2023) the context of binary NCE, and motivates having a learnable (cf. Algorithm 1) negative sampling distribution Pę(y x) during training, avoid early gradient collapses that result in sub-optimal energy model. We next attempt to characterize the properties of an "optimal" negative sampler.
4.3.1 Nearly Asymptotically Efficient Negative Sampler
In principle, the asymptotic normality result Theorem 4.10 gives us a recipe to characterize an optimal proposal model: simply select the distribution P& which minimizes the trace of the asymptotic variance matrix Vo (which itself is function of pt). However, this is an intractable optimization problem in general (see e.g. Gutmann & Hyvärinen, 2012, Section 2.4). Nevertheless,
will show that considerable insight may be gained by setting the negative distribution (3 x)) equal to the conditional data distribution p(y x).
To begin, we first define the conditional Fisher information matrix, which generalizes the classical Fisher information matrix from asymptotic statistics.
Definition 4.11 (Fisher information matrix). The conditional Fisher information matric of a conditional density (y x) at point E X is defined as:
I(0) x) = Ex-po((=] Vo V logpely =.
Note that for our parameterization F have:
100 Φ) EE (V Ve (y [V
v)
Under suitable conditions similar those in Theorem 4.10, it well known that the inverse Fisher infor- mation matrix governs the asymptotic behavior of the maximum likelihood estimator:
Vn(@m - 0.) ~ NN, E, =))""), where @mle E argmax
1
log po xz).
PEO
Our next result. shows that by selecting pe ( x:) = p(y x)), the R-NCE estimator has similar asymptotic efficiency as the MLE.
Thcorem 4.12. Consider the setting of Theorem 4.10, and suppose that Pe. (y x)) = p(y x). Then, the asymptotic voriance matrix Vo is:
Ve = (1+ 1
E, [(0, xx)--
K
12

---

Published in Transactions on Machine Learning Research (09/2024)
Theorem 4.12
extends similar result for binary NCE in Gutmann & Hyvärinen (2012, Corollary 7) to the ranking NCE setting. As aside, we note that (g x) = ply x)) is actually not the optimal choice for minimizing asymptotic variance (Chehab et al., 2022). Nevertheless, even for reasonable values of K, the sub-optimality factor is fairly well controlled.
We now compare Theorem 4.12 to Ma & Collins (2018, Theorem 4.6). In Ma & Collins (2018, Theorem 4.6), it is shown that regardless of the (fixdd) proposal distribution Pt, one has that IIVe Eg (00, =))- < O(I/VK), where the (()) hides constants depending on the proposal distribution Pg. These hidden constants are not made explicit, and are likely quite pessimistic. On the other hand, in Theorem 4.12, we show that under the idealized setting when Pe æ) ply x), we have a sharper multiplicative bound w.r.t. the maximum likelihood estimator. Furthermore, since analysis and algorithm allows for learnable proposal distributions, Pe(y / x) p(v 2) is much more of reasonable assumption for theoretical analysis.
The near-asymptootcalll efficient result of Theorem 4.12 strongly motivates a learnable proposal model, but still leaves undetermined the algorithm by which such a result can be attained. We next investigate adversarial training for partial clues.
Is Adversarial Training the Answer? Rocently, several papers which study the binary version of the NCE objective have advocated for an adversarial "GAN-like" approach where & is optimized via concurrently minimizing the NCE objective (Gao al., 2020; Bose et al., 2018). This approach is justified from the viewpoint of "hard negative" sampling when, similar to our carlier discussion on postcrior collapse, the proposal distribution is easily distinguishable by a sub-optimal EBM. The following result establishes a close connection between the near-optimal asymptotic variance in Theorem 4.12 and adversarial optimization, under a realizability assumption for the set of proposal distributions.
Proposition 4.13. Suppose that I is realizable, and furthermore, the family of proposal distributions Fn is realizable. That is, there erists some non-emply E. C = s.t. for every & E E., Pe, (y x) = plu 2)) for a.e. (I. Then:
1
max min L(O,E) = L(O., log dEB EEE
(4.8)
K K+1
where (0.,E.) O, E..
Recall from Theorem 4.12 that if the fixed point for the proposal distribution Pc. (. () equals p(- | x), the asymptotic variance for On equals 1 + 1/K multiplicative factor of the Cramér-Rao optimal variance. The
result above corresponds precisely to such a scenario: given a realizable family of proposal distributions, the adversarial saddle-point for R-NCE is characterized by the equalities po. (- 2) = Pc. (-1 z) = p(- z))
Thus, when a sufficiently rich family of proposal distributions (i.e., F, is realizable) paired with a tractable convergent algorithm that ensures E.) "= 0, there is no additional benefit to using adversarial R-NCE: the asymptotic distribution of 0, is unaffected.
Notwithstanding, our preferred setup assumes, for both computational and conceptual reasons, that Fn is not realizable. This is motivated by the fact that the proposal distribution used both for sampling and likelihood evaluation within the R-NCE objective, and thus a significantly simpler model for computa- tional reasons. In the following, we investigate the asymptotic properties of adversarial R-NCE under such conditions.
4.3.2 Equilibria and Convergence for Adversarial R-NCE
Let us define the finite-sample maxmin objective as follows:
max min L,, == Br., (x,
(4.9)
OEO EE=
In order to study the statistical properties of this game, we first define two families of solutions: Nash and Stackelberg equilibria.
13

---

Published in Transactions on Machine Learning Research (09/2024)
Definition 4.14 (Maxmin Nash equilibria). Let f: xXE R. A pair (0..g.) 6 is Nash equilibrium W.T.l. i/:
((0.5.) & S(0.. < (( ((0.))
A Nash equilibrium is characterized as follows: taking the other "player" as fixed, each player has no incentive to change their "action". It straightforward verify from the proof of Proposition 4.13 that any pair (+) E O, E. is in fact Nash equilibrium w.r.t. the population objective L. However, for general nonconcave--nnnoomne global or local Nash equilibrium is not guaranteed (Jin et al., 2020). Thus, we need an alternative dcfinition of optimality, which comes from considering games as sequential.
objectives, e.g., the finite-sample game defined in (4.9), the existence of a
Definition 4.15 (Maxmin Stackelberg equilibria). Let f: xx A pair (0,5) E x is Stackelberg equilibrium w.r.l. if:
inf f(0,5") < 5(0 < ((0,5) v(0.)) E'E
That is, E minimizes ((0,)), :), whereas maximizes inferez f(0,g") In contrast to global Nash equílibrium, global Stackelberg equilibrium points are guaranteed exist provided continuity of f and compactness of θ x =\t (Jin ct al., 2020).
For what follows, we will study the knite-sample Stackelberg equilibrium points (On, Ex):
L, (0, & (nn, & (0n, AO EB EE
inf
(4.10)
Prior to proving a consistency statement about the convergence (0n, En), we first establish some basic facts regarding Stackelberg optimal pairs for the population game.
Proposition 4.16. For any Stackelberg optimal pair (0, E) to the population game (4.8), holds that E O.. Conversely, for any 0. O., there exists an E such that (0..E) E) is Stackelberg optimal for the population game (4.8).
We now establish convergence of the finite-sample Stackelberg optimal pairs (6n, Ex) to a Stackelberg optimal pair for the population objective.
Theorem 4.17 (Adversarial Consistency). Suppose that Assumption 4.1 and Assumption 4.2 hold. Let (On, E.) E θ E be a sequence of Stackelbery optimal pairs, as defined in (4.10), for the finite sample game (4.9). Then, the following results hold:
1. Suppose that O, is contained in the interior of O. We have that d(Om, "03 0 a.s.
Suppose furthermore that Assumption 4.9 holds. Let E := =E== E L(O,,E) infgrez L(O., {') for an arbilrary choice of 0, E O. (the set defonition is independent of this choice), and suppose E is contained the interior of E. We have that d(En, n 0 a.s.
We note that part (1) the consistency result perhaps not surprising in light of Theorem 4.8, which merely assumed that On maximizes the finite-sample objective Ln for any particular &.. In particular, no assumptions are placed upon the negative distribution parameter sequence {En]. Under a mild regularity assumption however, we are additionally able to show convergence the noise distribution parameters to a population Stackelberg adversary. In the setting where the family of proposal distributions Fn, is realizable, the set E is indeed the Nash optimal set E., defined in Proposition 4.13. Next, we characterize the asymptotic normality of the sequence of finite-sample
bacckelber--pttnnde pairs (n).
Theorem 4.18 (Adversarial Normality). Suppose that F is realizable and Assumptions 4.1, 4.2, and 4.9 hold. Let In = (0n.E.) O be sequence of Stackelbery optimal pairs, as defined in (4.10), for the fintt sample game (4.9), and define E in Theorem 4.17. Assume then that the following properties hold:
14

---

Published in Transactions on Machine Learning Research (09/2024)
1. The sets 0,,ã are singletons, i.e., θ {8.}, and E {{), and y (0.,E) lies in the interior of + x E.
2. The Hession of ey is Lipschits- i.e., IV.l (z)llo & M( - 7el1, E, [M[()] < co, Vn, E x
3. The (0,0) block of the Hessian (0., denoted (8,(0., satisfies ViI(O.,E) x 0, and the (E.E) block, denoted (2,,0,,)), satisfies (eL(0,,)))
Then, vn(fnn 7.) (O, V where
V = E2 [ (=)] Var. (V),,())) E. [V
(4.11)
A particularly important corollary of this result is as follows: Corollary 4.19 (Marginal Normality). Under the assumptions of Theorem, 4.18, the porameters {0x} satisfy
vn(6,- -.) NO, Vy²), where
V8 L.
The variance V is similar to Vo in Theorem 4.10, with the difference being the specification of particular fixed-point for the proposal distribution parameters: the Stackelberg adversary of 0.. In the special case when the family of proposal distributions is realizable, since the fixed-point involves the Nash adversary &. (cf. Proposition 4.13), these two variances coincide. However, when we lack realizability of the proposal distributions, we are left to compare the asymptotic variance Vo from Theorem 4.10 with the new asymptotic variance Vir Given the quite non-trivial relationship between the proposal parameters that a non-adversarial negative sampler converges versus the Stackelberg adversary, is not clear that a general statement can be made comparing these two variances. Therefore, the theoretical benefits of adversarial training are quite unclear.
Furthermore, actually finding Stackelberg solutions comes with its own set of challenges in practice. Com- putationally, differentiating the R-NCE objective with respect to the negative proposal parameters would involve differentiating both the negative samples y as well as the likelihoods log æ), a non-trivial com- putational bottleneck. contrast, differentiating the negative samples y with respect to the proposal model parameters & is not required in the non-adversarial setting (cf. Remark 3.2). For this work, we assume that the independent optimization of & via suitable objective for the generative model Pe is sufficient to optimize the EBM, without assuming realizability for Fn. We leave the explicit minimization of some function of Vo to future work.
Finally, we conclude by noting that the proofs of Theorem 4.17 (adversarial consistency) and Theorem 4.18 (adversarial normality) are based off of ideas from Biau et al. (2020), who study the asymptotics of GAN training. The main difference that we are able to exploit specific properties of the R-NCE objective derive simpler expressions for the resulting asymptotic variances.
4.4 Typical Training Profiles
The results of Section 4.2 and Section 4.3 focus on the asymptotic properties of R-NCE training, and do not consider the optimization aspects of the problem. Here, in Figure characterize typical training curves that one can expect when jointly training an EBM and negative proposal distribution with Algorithm 1.
4.5 Convergence of Divergences
To conclude this section, we translate asymptotic normality in the parameter space to convergence of both KL-divergence and relative Fisher information. For two measures and over the same measure space, the KL- KL-divergence and the relative Fisher information of μ w.r.t. v are defined as:
KT( || v)
E,p log (") ]
FI(n v) = E,p V log
(M)
15

---

Published in Transactions on Machine Learning Research (09/2024)
Posterior collapse
EBM Improves faster relative negative sampler
EBM capacity > negative sampler capacity
L10. (3..5
Balanced model capacities
-log(K + 1)
Iterations
Figure 2: An illustration of the different qualitative behaviors that R-NCE training curves can exhibit. The illustrates example of posterior collapse, where the negative sampler is completely un-informative, leading to a trivial classification problem. Both the and curves illustrate phase transition in the training: at first the EBM learns quicker relative to the negative sampler (and hence the classification accuracy increases), but then the negative sampler catches up in learning, resulting in drop in the R-NCE objective. The main difference between the and curves is the final value each curve converges towards. The curve tends toward - log(( + 1), which indicates that both the EBM and negative sampler have sufficient capacity to represent the true distribution p(y | z) (cf. Equation (3.6) in the case where pe, (: |x) = Pe. (+1 = p(. æ)). On the other hand, the curve tends towards final value larger than - log(K 1) but bounded away from of the negative sampler. The remaining EBM; this occurs when, for example, pre-trained negative sampler is used. Here, the training curve approaches its final value from below. As before, the difference between the final value both the and curves is due to negative sampler model capacity, with the curve denoting when the EBM has much larger capacity than the negative sampler, and the curve denoting when the model capacities are relatively balanced.
indicating that the EBM capacity exceeds that and curves represent when the negative sampler learns faster than the
Here, we have (a) overloaded the same notation for measures and densitics, and (b) in the definition of relative Fisher information we have assumed smoothness of the densities. The following propositions are straightforward applications of the second-order delta method.
Proposition 4.20. Suppose that 0, E O., and that {0,} is sequence of estimators satisfying vil( -0.) m
N( E.). Then:
E. (KL(p( ()) Pa_ (y x))] ~ 1 111 Z~ N'O,
E [[(, x)jp!/
Hence, the setting of Theorem 4.10, if pe. (y x)) p(u x) and if θ C TRP, then we have:
1 2/ 1+ K x2)
n E. [KL((l | x)) Pa, (y x))ه ~
(4.12)
where x²(p) is Chi-squared distribution with degrees of freedom. Now turning to the relative Fisher information, we have the following result.
16

---

Published in Transactions on Machine Learning Research (09/2024)
Proposition 4.21. Suppose that 0, o. and that {0,} is sequence of estimators satisfying v ((0, E.). Then:
n En Fl(pyy 2) ll Pa, (v x))] ~ |z||. ~~ NN,, X E. [[ | r)))
where J(0) [x) = R 100 v) 0o E v)].
The asymptotic variance from Proposition 4.21 is more difficult to interpret than the corresponding asymp- totic variance from Proposition 4.20. Unfortunately in general, there is no relationship between the two, as we will see shortly. In order to aid our interprelation, suppose the energy model follows the form of an exponential family, ie., Eg y) = (), for some fixed sufficient statistic vb. Then, a straightforward computation yields the following:
1(0 | x) (VO
x0) = Epotutn)00 (0,((,,))")
Hence, assuming as in (4.12) that Pe, (y |-) ply x), then the asymplotic variance is determined by the variance of the smoothness of the sufficient statistic (in the event variable y) relative to the covariance of the sufficient statistic itself. Let us now consider specific example. Suppose that po(y x) = N(u, o²) (so the context variable x is not used). Here, the sufficient statistic is (b( = v²). A straightforward computation yields that:
4n²
tr [L [T =E, x)))
04
Thus, depending on the specific values of the paramcters (H, o2, the asymptotic variance quantity above can be either much smaller much larger than 2.
5 Interpolating EBMs
Recent methods such as diffusion models (Sohl-Dickstein al., 2015; Ho al., 2020; Song et al., 2020; 2021b) and stochastic interpolants (Albergo & Vanden-Eijnden, 2023; Albergo et al., 2023) advocate for modeling family of distributions with single shared network, indexed by a single scale parameter lying in fixed interval. At one end of this interval, the distribution corresponds to standard Gaussian noise, while the other end models the true data distribution. Sampling from this generative model begins with sampling a "latent" vector 2 E Y from the base Gaussian distribution, and then propagating this sample through a sequence of pushforwards over the interval. This propagation can either refer a fixed discrete sequence of Gaussian models, or solving continuous-ting SDE or ODE. The strength of such a framework has been attributed to implicit annealing of the data distribution over the interval (Song & Ermon, 2019), casing the generative process by stepping through sequence of distributions instead of jumping directly from pure noise to the data.
5.1 Stochastic Interpolants
A useful instantiation of this framework that subsequently allows generalizing all other interval-based models as special cases, corresponds to the stochastic interpolants formulation (Albergo & Handen-Fijnden, 2023). In this setup, continuous stochastic process over the interval [0, 1] is implicitly constructed by defining an interpolant function I, Y x Y Y, E [o, 1]. where To(2, y) = 2 and [1 y) = y. For a given x, define the continuous-timm stochastic process Y; x as:
2 ~~,10,11, y ~~ Pyles yt x Ie(z.)),
z 1 y
By construction, the law of Y, at t 1 exactly Py An example interpolant function from Albergo & Handen-Eijdden (2023) given as le(z,)) = cos( mt)z + sin(f The first key result from Albergo & Vanden- -ijnden (2023) that law of the process Y x, denoted as Py, with associated timo-varying density pl (ve æ), satisfies the following continuity equation:
O.p' + divy(upp') 0.
(5.1)
17

---

Published in Transactions on Machine Learning Research (09/2024)
where V [0, x Y + Rdim(Y) is particular time-varying velocity field (described subsequently).
yielding the following continuous normalizing flow (CNF) generative model:
dy (z) (2 (z)) (z= =====~~00 ~~
(5.2)
dt
The distribution of the samples from this model parameterize any function approximator 8ggg approximate this vector field.
correspond the desired distribution Pylm- One
While CNFs generative modeling framework are certainly not new, their training has previously relied on leveraging the associated log-pro bog-probability ODE (Chen et al., 2018; Grathwohl et al., 2019) (cf. Section 6), corresponding to the change-of-variables formula, and using maximum likelihood estimation as the opti- mization objective. Unfortumately, solving this ODE involves computing the Jacobian of the parameterized vector field fr,g with respect to Wr gradients of the objective w..t.. therefore require 2nd gradients of 8ggg. As result of these computational bottlenecks, CNFs have been largely overtaken as framework for generative models. However, the second key result from Albergo & Vanden-Eijnden (2023) is that the truc vector field is given as the minimum of a simple quadratic objective:
(x(x,)) = argmin (e,)))
E, Ey (1) (((,,)))1== 20,1(22,1) 6,(2, [(z,,))] Vr,l.
This now enables the use of expressive models for oe,5, paired with the following population optimization objective:
min B(xx))xxx E [1 ((2,,))11 20,1(22,0) 0,,(2, Te(2, y))]
(5.3)
where A is fixed distribution over the interval [0, 1] (o.g., uniform).
5.2 Building Interpolating EBMs
A natural way of using the Interpolant-CNN model described previously is to use the pushforward distribution at t= ==1 1 the proposal generative model Pe(v) a.), with relatively simple parameterized vector field 6r.-- Paired with a parameterized energy function Eg(x, y), one can now proceed with optimizing both models as outlined in Algorithm 1.
However, by truncating the flow at any intermediate time E (0, the CNF gives an approximation (y æ) of the distribution pf (v, æ). Thus, can define time-indexed EBM Ea(2, y,t) to approximate this same distribution
p'(v x)) exp(Eo(2, yc, t)).
We refer to the above model as an Interpolating-E.MM. (I-EBM). Notice that a single shared network & is used to represent the distribution across all t [0, 1].
By pairing the EBM at time with the truncated CNF time t as the proposal model, we can straightfor- wardly define time-indexed R-NCE objective; Figure 3 for an illustration.
18

---

Published in Transactions on Machine Learning Research (09/2024)
Dt-0 2
Lif, a) ~ (
Interpolant Bridge
R-NCE
Approx CNF
dy fegge gjdt
31 p 166-11 1z)
I-EBM
U1 pbl- ()) or Egta,s
Figure 3: An illustration of the Interpolating-EIBM (L-EBM) framework applied motion planning problem. Here, the event space corresponds to the waypoint trajectory (green circles) of the robot in the obstacle environment. The interpolant function I, defines continuous time stochastic bridge process over [0, 1] between Gaussian noise at & = 0 and the true conditional distribution at == L; the time-varying density of this process is notated p- 2). The CNF with (a relatively simply parameterized) vector field ôr is an approximation of this continuous- time process, and is independently trained using (5.3); the truncation of the flow of the CNF at time f yields samples from the distribution ( z). The I-EBM defines a family of (more expressively parameterized) EBMs pa(: x), indexed by i, to better approximate p'. At given time i, we pair K contrastive samples from the CNF and a true sample from the interpolant bridge to define a time-indexed R-NCE objective for Eo.
In particular, let K E N4, the number of negative samples defining the R-NCE objective be fixed as before. For given x, let pK yilzgg denote the product conditional sampling distribution over YK where ye ~ y axf denotes a random vector yt = (y), ) YK with yt, ~ rf( x) for k I,...,,. K.
P
Now, for a given x ~ Px, L E (0, 1], posilive sample ~ Pysler negalive samples ~ PK Velaig and parameters
E O g E define:
exp(Ee(x, ye t) lo (2) Ey'Elu exp(Eo (z, - log [=))
(z Ve,t l = log
(5.4)
the R-NCE loss at time t. We thus define the overall population objective as:
L(0 = REAABEE Buu-Py, Px
wilng [Co. yr, yt)].
(5.5)
Notice that this is simply expectation of time-indexed population R-NCE objective. We refer to the objective above as the Interpolating-R-NOT (I-R-NCE) objective. Provided that the function class of time- indexed conditional energy models:
8: (Ee(z, v.t) + O).
is realizable (cf. Definition 4.3) for all t E 10, i.e., there exists 0, E O s.t. for all E (0, 1] and a.e. (z, v)):
exp(Eo, (x, Uc, t)) Jy exp(Ee, (z, y, t)) duy
1b.(1 (1 ())
('v xx)
then the population objective in (5.5) is maximized by 0, for each t.
5.2.1 Training
The adjustment to Algorithm for training straightforward for I-EBMs. In particular, for each (2,, v.) E B, sample m values of 6, ie., (1g (4 A, latents {2y 1m N/(0, 7), and negative-sample sets {y
19

---

Published in Transactions on Machine Learning Research (09/2024)
Yeg l The net R-NCE loss for (2t, V6) is then defined the average of the R-NCE loss over ((r, Mij. lij, )== where yij It (zij. The training pseudocode is presented in Algorithm 3, with the modifications to Algorithm 1 highlighted in teal. We note that choosing values of m > 1 serves as variance reduction, and we also apply a similar variance reduction trick to the stochastic interpolant loss in Line 8.
where Yiy P
Algorithm 3: Interpolatiig- R--NCB with learnable negative sampler
Data: Dataset D = (Z1, 3)) HE number of outer steps Tonter, number of sampler steps per outer step Tsmpp number of R-NCE steps per outer step Tince, mumber of negative samples K, interpolating time distribution A, number interpolating times m E N21- Result: Energy model paramcters 8, negative sampler parameters ç.
Initialize 0 and g.
Set (vr,e(z, 2, y, t) ||,,, 10 w) 20,1,(2,8) ve.e(z Le(2, y)). 1 Set Limterp(E (E: B) = Elz, (lx m Σド=1 lv (Ie, Zzfs yo tiz).
Set Lirnce (0,
1n D la (2 Yost bi vis).
for Touter ilerations do
for T Tammp iterations do
7
B + batch of D.
// Augment batch with latent variables and times.
B {(z,, Ve, Uzy. (1)))==)) (o,, 91) B}, with bi ~ A, By ~ 10,
8
minimize using Ve Vehinterp(E: (E;
10
end
for T mnce Ulerations do
B batch of D.
// Convert batch to time-indexed positive and negative samples.
13
B {(1,, ((1y, ty. tt)) (xi, M) E B), with try 24 ~ NoO, us Le,, (215. vi).
Yij ~ pk
Vuylems"
14
minimize using Vo
15
end
16
return (0,E)
17 end
Remark 5.1. naive implementation of variance reduction for the R-NCE loss Algoritho 3 would generate m independent sets of K negative samples, i.e., {y corresponding to the m intermediate times (tay) To make this more efficient, allow the negative sample sets to be correlated across time by generating single set of K trajectories, spanning the interval (0, max, tiyl, and sub-sampling these trajectories at m intermediate times.
5.2.2 Sampling
In follow-up work (Albergo al., 2023), the stochastic equivalent of the transport PDE (5.1) is derived using the Fokker-Planck equation, which yields the following, equivalent in law, SDE-based generative model:
(z) = ve(2, 21zz(2)) + V = (x.(4.))))))
x) dt + dw, (== = = (5.6)
where n >> 0 and W; is the standard Wiener process. Given the simplicity of the former ODE-based model in (5.2), a natural question is: why the SDE-based model necessary? The answer lies studying the divergence between the true conditional distribution p(y x) and the approximate distribution p (y x) induced by the learned vector field dg.- In particular, from Albergo et al. (2023, Lemma 2.19), for CNF
20

---

Published in Transactions on Machine Learning Research (09/2024)
parameterized by the approximate vector field de.e
KI(ply x)) ll Pē - x))
L (1V 1o
Vy logp'' x)) (0),((x,)) ue( ve(a, p' (v x)) duydl.
It follows that control over the vector field error, i.c., the objective minimized in (5.3). does not provide any direct control over the Fisher divergence, iee, error in the learned score function V, logr: x), thus preventing any control over the target KL divergence.
In contrast, let 87 (x. be an approximation of the drift term vf( (z. y) in (5.6). Define the following approx- imate generative model and associated Fokker-Planck transport PDE for the induced density f' =))
dz)]((2) = (, 3aja (z)) dt + d2n M(a(()) = ~~ 10, 0 + div,
(5.7)
, (0,p') = nAp'.
(5.8)
Then, Albergo et al. (2023, Lemma 2.20) states that:
4nj 1 / [10(0,,0) (7 v)) p'( x)) diy
KL(pl x) || (y x))
Thus, if one can control the approximation error of the SDE drift, then the SDE-based generative model should yield samples which are closer to the target distribution in KL. While error in the velocity field || - ll is controllable via the minimization of (5.3), the I-EBM framework gives us a handle the other component to the SDE drift: an approximation of the time-varying score function, 5 i.e., VyEa( yt, t) z v, log p'Ve z). Indeod, from Proposition 4.21, convergence to the global maximizer 8, yields controllable convergence (in distribution) on the relative Fisher information between p' (yt x) and pà. (ye æ). Algorithm outlines a throe-stage procedure for sampling with I-EBM.
Algorithm 4: I-EBM: Sample from po æ) for Euclidean Y
Data: Context x E X, energy model parameters 0, negalive sampler &, initial sample time & E (0, 1),
MCMC steps Tmome: MCMC step size r).
Result: Sample y ~ Po x))
Set ~ / x))
// Leverage the SDE in (5.7) to transport Ye from t=t to t ==
y' <- SDESolve(""(r,,)) = (,,((1))) + + V u,) [), // Finite-step MCMC with Ea(2, y. starting at y' (cf. Algorithm 2).
< MCMC(Tment 71,3 n1,5))
return y
Notice in Line we use the CNF parameterized by vector field ,,,g, with the flow truncated at time to form the initial sample. This is reasonable since for ( & 1, the proposal model pt becomes more accurate with the true process Y x tending towards isotropic Gaussian, dispelling the need for EBM. In Line 2, leverage the SDE generative model from (5.7) to transport the sample to t = 1, making use of both the proposal model vector field êτ,g and the I-EBM's score function VyEa to form the approximation of the drift term. Finally, in Line 3, we run T Tmeme steps of MCMC (either Langevin or HMC, cf. Algorithm 2) at fixed (or annealed) step-size of 77 with the EBM at t to generate the final sample.
We note that in Albergo al. (2023), the authors introduce stochustic extension the interpolant framework from Albergo & Sanden-Bijnden (2023) which clicits quadratic loss function for the time-varying score function, similar in spirit to the canouical denoising loss used for training diffusion models. Our framework however, agnostic the type of bridge created between the base distribution and true data. Notably, I-ESMs and the associated I-R-NCE objective are well-defined for any type of interpolating diffusion /diffuusioni bridge.
21

---

Published in Transactions on Machine Learning Research (09/2024)
5.2.3 Advantages
With the I-EBM framework in place, we now descríbe the key advantages behind the formulation compared with other seate-of-the-art generative models. As we will soon show experimentally (Section 7), it is im- portant for generative policies to be able to both generate and score samples efficiently. I-EBMs meet both requirements in ways that many other slate-of-the-arl models cannot, as we discuss below.
Comparison to stochastic interpolants: As a stochastic interpolant model is as its core a continuous normal- izing flow, computing sample likelihoods require solving the standard log-probability ODE (ef. Section 6), which is computationally bottlenecked by vector-Jacobian products. While we do employ various tricks to speed this procedure up (as discussed in Section 6), I-EBMs completely avoid this ODE computation for likelihood computation, and only require one forward pass through the network (1 y, 1.
Comparison to diffusion models: A similar computation issue arises when scoring samples with ODE-based diffusion models (Song et al., 2021b; Karras et al., 2022) as with interpolant models. This is because de- noiser models are typically parameterized as unconstrained vector-valued functions, and hence cannot be analytically integrated. Therefore, the most straight way to compute log-probability is treat a diffusion model as normalizing flow and use the log- bog-prbbbiylity ODE, inheriting the associated computa- tional bottlenecks. An alternative to the log-probability ODE is to use Girsanov's thcorem to derive a lower bound on the log-likelihood, which can then be estimated via Monte-Carlo sampling (Song el al., 2021a). A similar sampling- based approach to likelihood computation also applies to DDPM-style diffusion mod- els (Sohl-Dickstein et al., 2015; Ho et al., 2020), where Monte-Carlo sampling of the evidence lower bound can be used to estimate a lower bound on the log-likelihood. However, in addition to introducing an extra hyperparameter (the number of Monte-Carlo samplings used) which necessary to tune the best trade-off between accurancy and performance, comparing ELBO scores across samples is actually quite problematic as described next.
Comparison to conditional VAE models: Conditional variational autoencoders (CVAE) (Doersch, 2016) have been recently proposed by several authors (Zhao et al., 2023; 2021) as a model for generative policies. Unfortunately, the ELBO loss which CVAFs optimize is insufficient to support ranking samples. To see this, first recall that CVAEs utilize the following decomposition of the bog-pobabbility (Docrsch, 2016):
et al., 2020; Ivanovic et al.,
loge(y xx) - KL(((= | x, v) |l p(= x2,))) = Es (a Dogg(y x =)] - KL(g(= a2,v) |l p(= | x)).
(5.9)
The RHS of (5.9) lower bounds the log- log-probability log ply x) and is the ELBO loss used during training; it only references the encoder q(= | x, y), the decoder p(y ar, 2), and the prior p(> x)), and is therefore possible to optimize over. The LHS of (5.9) contains the desired logp(y [x), plus an extra error term o(x,)) v) == -KL(g(> x,,) || p(= x,,)), which cannot be computed due to its dependence on the true posterior p(= xr y). From this decomposition, we see that for a given shared context x and a set of samples V1...., Uk. the RHS of (5.9) cannot serve as reliable ranking score function, since the unknown error terms o(I v:) for i = 1, 1,...,k k are in general all different and not comparable. Note that this issue persists even in the limit of infinite samples used to compute the RHS of (5.9).
Efficient sampling from I-EBMs: We have thus far focused on the advantages of I-EBMs over other generative models when it comes to likelihood computation. We now discuss the sampling advantages of I-EBMs. Combining the vector field and energy function within the SDE sampling step (cf. Line 2 in Algorithm 4) significantly reduces the burden of MCMC-based sampling, which can be difficult to converge when starting far from the stationary distribution of the EBM. In this respect, I-EBMs bring the advantages of learning over multiple noise resolutions/ /annealed sampling to E.Mss.
6 Efficient Log-Probabilte Computation
In this section, we discuss efficient computation of bog-probability for contimuous normalizing flows. We first discuss the generic setup, and then describe where the setup is instantiated in our work.
22

---

Published in Transactions on Machine Learning Research (09/2024)
Setup. Let (z denote vector field, and consider the following continuous flow:
dæ
dt (1) = (,, (()), =(0) pol 1
(6.1)
For any T> O, a standard computation (Grathwohl al., 2019; Chen et al., 2018) shows that the random variable z(T) has log-probability given by the following formula:
log p(z(T) x)) = log Po (=(() x²) 1 //' (iv=(vr(x,=(t)) dt.
(6.2)
= Recall that for a vector field f RRd + Rd, we have div(ʃ) = -r(0))=== =;= af
Uz."
Applications. We make use of the computation described in (6.2) in several key places for our work. First, we use this formula to compute log Pe 2) inside the R-NCE objective (cf. (3.3)). Second, we use this formula to compute psy z) for a diffusion model where the denoiser is directly parameterized (instead of indirectly via the gradient of a scalar function); for diffusion models, this tog-probability calculation is necessary for ranking samples in next action selection (ef. Section 7.3). Here, we apply (6.2) by interpreting the reverse probability flow ODE as continuous normalizing flow of the form (6.1).
Efficient Computation. The main bottleneck of (6.2) is in computing the integrand, which involves computing the divergence of the vector field =)) w.r.t. the flow variable Assuming the flow variable E Rd, the computational complexity of a single divergence computation scales quadratically in d, i.., O(d)) The quadratic scaling arises due to the fact that for exact divergence computation, one needs to separately matcrialize cach column of the Jacobian, and hence O(d) computation is repeated times. The typical presentation of CNFs advocates solving an augmented ODE which simultaneously performs sampling and log-probability computation:
ue (1, z(t) dt (t) - (0)=-[-
d
(1), (
z(0) pol: x)) (0) log Po(2(0) z).
This augmented ODE suffers from one key drawback: it forces both the sample variable z and the log- probability variable & be integrated at the same resolution. Due to the quadratic complexity of computing the divergence, this integration becomes prohibitively slow for fine time resolutions. However, high sample quality depends on having relatively small discretization error.
We resolve this issue by using a two time-scale approach. In our implementation, we separately integrate the sample variable a and the log- bog-pobbaiiltyy variable sb. We first integrate z at a fine resolution. However, instead of computing the divergence of u at every integration timestep, we compule and save the divergence values at much coarser subset of timesteps. Then, once the sample variable is finished integrating, we finish off the computation by integrating the v variable with one-dimensional trapezoidal integration (e.g., using numpy. trapz). Splitting the computation in this way allows one to decouple sample quality from log- probability accuracy; empirically, we have found that using at most 64 tog-probability steps suffices even for the most challenging tasks we consider. A reference implementation is provided in Appendix A, Figure 5.
7 Experiments
In this section, we present experimental validation of R-NCE a competitive model for generative policies. We implement models in the jax (Bradbury et al., 2018) ecosystem, using flax (Heek et al., 2023) for training neural networks and diffrax (Kidger, 2021) for mumerical integration. All models are trained using the Adam (Kingma & Ba, 2015) optimizer from optax (Babuschkin et al., 2020).
applicable here, since need to be able compare and rank log-probabilitics across small batch of samples. Generically, unless the Hutchinson trace estimator is applied 22(d) times (thus negating the computational benefits), the variance of the estimator will overwhelm the signal needed for ranking samples; this consequence of the Hanson- Wright inequality (Meyer al., 2021).
Randomized methods such the standard Hutchinson trace estimator
23

---

Published in Transactions on Machine Learning Research (09/2024)
7.1 Models Under Evaluation and Overview of Results
Here, collect the types of generative models which evaluate in experiments. Note that the more complex experiments only evaluate subset of these models, as not all of the following models work well for high dimensional problems.
NF (normalizing flow): a CNF (Grathwohl et al., 2019), where for computational efficiency the vector field parameterizing the flow ODE learned via the stochastic interpolant framework (Albergo al., 2023) instead of maximum likelihood; see Section 5.1. The flow ODE is integrated with Heun. During training, we sample interpolant times from the push-forward mcasure of the uniform distribution [0, 1] transformed with the function 1 + (4/a, where is hyperparameter.
IBC (implicit behavior cloning): the InfoNCE (van den Oord et al., 2018) inspired IBC objective of Florence et al. (2022), defined in (4.3), for training EBMs. Langevin sampling (ef. Algorithm 2) is used to produce samples.
R-NCE and I-R-NCE (ranking noise contrastive estimation): the learning algorithm presented in Algorithm for training EBMs, and its interpolating variation, Algorithm 3, presented in Section 5.2 for training I-EBMs; both leveraging the stochastic interpolant CNF as the learnable negalive sam- pler. Sampling is performed using either Algorithm for EBM or Algorithm 4 for 1-EBM.
Diffusion-EDM:: standard score-based diffusion model. We base our implementation heavily off the specific parameterization described in Karras et al. (2022), including the use of the reverse probability flow ODE for sampling lieu of the reverse SDE, which we integrate via the Heun integrator. We also treat the reverse probability flow ODE as CNF for the purposes of log- probability computation (cf. Section 6).
Diffusion-EDM- Same as DDttusion-EDMM, except instead of directly parameterizing the denoiser function, represent the denoiser as the gradient of a parameterized energy model (Salimans & Ho, 2021; Du et al., 2023). This allows us to use identical architectures for models as IBC, R-NCE, and I-R-NCE. More details regarding recovering the relative likelihood model from the parameterization of Karras et al. (2022) are available in Appendix B.1. We note that this model has an extra hyperparameter, Orel, which indicates the noise level at which to utilize the energy model ø for relative likelihood scores.?
Note for both diffusion models, for brevity we often drop the EDEDM" label when it is clear from context that are referring to this particular parameterization of diffusion models.
Basic Principles for Model Comparisons. As this list represents wide variety of different methodolo- gics for generative modeling, some care is necessary in order to conduct fair comparison. Hcre, outline some basic principles which we utilize throughout our experiments to ensure the fairest comparisous:
(a) Comparable porameter counts: We keep the parameter counts between different models in similar ranges, regardless of whether or not a particular model parametorizes vector ficlds (NF, Diffusion) or energy models (IBC, R-NCE, I-R-NCE, Diffusion-é)). Furthermore, for R-NCE and I-R-NCE, we count the total number of parameters between both the negative sampler and the energy model. (b) Comparable functions evaluations for sampling: We keep the number of function evaluations made to either vector fields or energy models in similar ranges regardless of the model, so that inference times are comparable.
(c) Avoiding excessive hyperparumeter tuning: In order to limit the scope for comparison, we omit exploring many hyperparameter settings which apply equally to all models, but may give some
We find that, much akin how the diffusion reverse process for sampling is typically terminated at a small but non-zero time, this is also necessary for the numerical stability of the relative likelihood computation (cf. Appendix 1.1, Equation (B.1)). Furthermore, we find that the best stopping time for the latter typically an order of magnitude higher than the former. We leave further investigation into this for future work.
24

---

Published in Transactions on Machine Learning Research (09/2024)
marginal performance increases. A few examples of the types of optimizations omit include peri- odic activation functions (Sitzmann et al., 2020), pre vs. post normalization for self-attention (Xiong et al., 2020), and exponential moving average parameter updates (Yazııı et al., 2019; Song & Er- mon, 2020). While these types of optimizations have been reported in the litcrature to non-trivially improve performance, since they apply equally lo all models, we omit them in the interest of limiting the scope of comparison.
Overview of Results. We provide brief overview of the results to be presented. First, we consider synthetic two-dimensional conditional distributions (Section 7.2) as basic sanity check. In these two- dimensional examples, we see that all models perform comparably except for IBC, which has visually worse sample quality. The next two examples feature robotic tasks where we learn policies that predict multiple actions into the future. This multistep horizon prediction is motivated by arguments given by Chi et al. (2023, Section IV.C), namely temporal consistency, robustness to idle actions, and eliciting multimodal behaviors.
The first robotic task is high dimensional obstacle avoidance path planning problem (Section 7.3). Here, we see again that IBC yields policies with non-trivially higher collision rates and costs than the other models. More importantly, we start to see separation wet the other methods, with R-NCE ultimately yielding both the lowest collision rate and cost. The next and final task is our most challenging benchmark: a contact-rich block pushing task (Section 7.4). For this block pushing task, we utilize the full power of the interpolating EBM framework (cf. Section 5.2). Our experiments show that. I-R-NCE yields the policy with the highest final goal coverage. Thus, the key takeaway from our experimental evaluation that training EBMs with R-NCE does indeed yield high quality generative policies which are competitive with, and even outperform, diffusion and stochastic interpolant based policies.
7.2 Synthetic Two-Dimensiona. Examples
We first evaluate R-NCE and various baselines on several conditional two-dimensional problems, Pinwheel and Spiral.
The context space uniform distribution over X (4, 6, 7), denoting the number of spokes of a planar pinwheel. We apply 10 dimensional sinusoidal positional embedding embed X into R The event space Y = R², denoting the planar coordinates of the sample. The distribution is visualized in the upper left row of Table 1, with only x E (4,5,6) 6} visualized for clarity.
Pinwheel:
The context space is a uniform distribution over X [400, 800], denoting the length of a parametric curve representing a planar spiral. This range normalized to the interval [-1, 1] before being passed into the models. The event space Y R², again denoting the planar coordinates of the sample. The distribution is visualized in the upper right row of Table 1, with only {400, 500, 600} visualized for clarity.
Spiral:
The performance of the models listed in Section 7.1 is shown in Table 1 (samples) and Table 2 (KDE plots). The specific details of model architectures, training details, and hyperparameter values are given in Appendix B.2. For all models, we restrict the parameter count of models to not exceed ~22k; for R-NCE, this number is a limit on the sum of the NF+EBM parameters. We report the Bhattacharyya coefficient (BC) between the sampling distribution and the true distribution, computed as follows. First,, recall that the BC between true conditional distribution p(y x)) and a learned conditional distribution po(y 2)) is:
BC(p(y | po(y x)) = / p(y x)po(y x)duy
VI E X.
(7.1)
Note that the BC lies between [0, 1], with a value of one indicating that ply x) = poly x). We estimate both p(y a) and po(y æ² via the following procedure, which we apply for each context x that we evaluate separately. First, we compute a Gaussian KDE using 8,192 samples with scipy . stats. gaussian kde, invoked with the default parameters. Next, we discretize the square [-4,4) x [-4,4] 4] into 256 x 256 grid
25

---

Published in Transactions on Machine Learning Research (09/2024)
Table A two-dimensional histogram plot of the samples generated by each of the models listed Section 7.1. The samples are binned resolution of 64 bins coordinate. Each column represents different context value see Section 7.2 for the specific values for each problem. The Bhattacharyya coefficient (BC) is computed as described in Section 7.2.
Model
Pinwheel
Spiral
6
Actual
2
6
NF
BC-0.540
IBC
0,996
-C3933 0.333
R-NCE
BC-0.998
BC-0.993
Diffusion
BC 0.998
BC 0.994
Diffusion-o
points, and compute (7.1) via numerical integration using the KDE estimates of the density. Finally, we report the minimum of the BC estimate over all values that we consider.
The main findings Table and Table 2 are that the highest quality samples are generated by the R-NCE, Diffusion, and NF models, between which the quality is indistinguishable, followed by a noticeable drop off in quality for IBC. Note that for IBC we use a uniform noise distribution to generate negative samples. While ielatively simple, this example sufficient to validate our implementations, albeit still illustrating the limitations of IBC.
7.3 Path Planning
We next turn to a higher dimensional problem, one of optimal path planning around a set of obstacles. We study this task due to its inherent multi-modality. The environment is planar environment with 10 randomly sampled spherical obstacles, random goal location, and random starting point. The objective
26

---

Published in Transactions on Machine Learning Research (09/2024)
Table A Gaussian KDE estimate generated via 8,192 samples for each of the models listed in Section 7.1. The KDE is computed via scipy. stats. gaussian__aae kde with default values for the sw_method and weights parameters. the specific values each problem. The Bhattacharyya coefficient (BC) is computed described in Section 7.2.
Each column represents a different context value æ; see Section
Model
Pinwheel
Spiral
Actual
NF
2
IBC
0,996
0,333
R-NCE
BC90.998
8C
Diffusion
BC 0,998
BC -C 0.994
Diffusion-o
is learn an auloregressive navigation policy which, when conditioned on the scene (i.e., obstacles and goal) and the last Nax positions of the agent, models the distribution over the next Npred positions.
Dataset Generation. In order generate the demonstration optimal trajectories, Gaussian process motion planner (StochGPMP) of Urain et al. (2022b), in which we optimize samples generated from an endpoints conditioned Gaussian process (GP) prior with an obstacle avoidance cost. The GP prior captures dynamic feasibility (without considering the obstacles) and the following key properties: (i) closeness to the specified start-state with weight o, > 0, (ii) closeness to the specified goal-state with weight 0g > 0, and (iii) trajectory smoothness with weight og > O; see Urain et al. (2022b); Mukadam et al. (2018); Barfoot al. (2014) for details on the computation of the prior. Roughly speaking, given a trajectory ( start location 8, and goal location g, the primary components of the effective "cost
the stochastic
27

---

Published in Transactions on Machine Learning Research (09/2024)
function' associated with the GP prior scale as:
N +
| - s||2 [1blN]] g||2+
lolt + 1] y|i]12
In generating optimal trajectories from SLochGPMP, we set the trajectory length to N = 50. The obstacle cost Cobs is parameterized by the obstacle set O = {(ck r=) where cj. E R2² and rk E RR are the centers and radii of the spherical obstacles, respectively:
N-1 10
Cobs (tulil) 0)
Σ
Σ
2 22
dist (ult), z/i +
(vli). + 1)) & rr.).
1=1
The contribution to the StochGPMP planner from the obstacle cost scales as coba/o?. Following the gener- ation of global trajectories with the planner, we recursively split the trajectories into (Netx + length snippets, and refine these snippets further with the StochGPMP planner, using the endpoints of the snip- pets as the "start" and "goal" locations. The collection of these snippets form the training and evaluation datasets.
Policy and Rollouts.
The autoregressive policy is then defined as the following conditional distribution:
plyli .., s/i + Npred] vl ........... yli Nax + 11
For simplicity, we assume perfect state observation and perform this modcling in position space, i.c., y/i] E R denotes the planar coordinates of the Irajectory at time Furthermore, we assume perfect knowledge of the goal g and obstacles O. We leave relaxations of these assumptions to future work. We set Nexx = 3 and NNreed 10, yielding a 20-dimensional event space. In order generate a full length trajectory, we autoregressively sample and rollout the policy. Given prediction horizon of N Npred 10 and total episode length of N = 50 steps, we sample the policy 5 times over the course of an episode. In order to encourage more optimal trajectories from sampling, for each putext/pollccy-sep we sample 48 future prediction snippets, and take the one with the highest (relative) log-likelihodd is a straightforward forward pass through the energy model. On the other hand, for NF and Diffusion, this computation relics on the differential form of the evolution of log-probability described Section
For R-NCE and Diffusion-o, this computation
Evaluation. For evaluation, we compute the cost of the entire trajectory resulting from the autoregressive rollout the summation of three terms: (i) goal-reaching, (ii) smoothness, and (iii) obstacle collision cost. To appropriately scale these costs in a manner consistent with the demonstration data, the evaluation cost is modeled after the primary cost. terms in the GP cost plus the obstacle avoidance cost, used to generate the data:
N Σ
1 Cor Cote (1ul)))) 0)
Obj(y. 9. 0) = | g
yli 1111..
(7.2)
Here, og. oo, and o. are derived by normalizing the corresponding weights from the StochGPMP planner.
Architecture and Training. For both energy models and vector fields, we design architectures based on encoder-only transformers (Vaswani et al., 2017). Specifically, we lift all conditioning and event inputs to a common embedding space, apply a positional embedding to the lifted tokens, pass the tokens through multiple transformer encoder layers, and conclude with final dense accumulation MLP. However, as previously discussed in Section 4.3.1, for computational reasons we make the vector field parameterizing the proposal model for R-NCE dense MLP, which is significantly simpler compared to the EBM. Indeed, the proposal CNF on its own is insufficient to solve this task, but is powerful enough as a negative sampler. More specific architecture, training, and hyperparameter tuning details are found in Appendix B.3. Note that for IBC, we do not use a fixed proposal distribution, as doing so results in poor quality energy models. Instead, we use the same NF' proposal distribution as used for R-NCE, trained jointly as outlined in Algorithm 1.
For all models considered, we limit the parameter count to a maximum of ~500k (for IBC/R-NCE, this is a limit on the EBM+NF parameters). Finally, for all models, during training, we perturbed the context and
28

---

Published in Transactions on Machine Learning Research (09/2024)
prediction snippet, i.e, fylé Netx il,.... y/2 + N.pred]); using Gaussian noise with variance opert where pert is annealed starting from 0.01 down to 0.005 over half the training steps, and held fixed thereafter. We found that this technique significantly boosted the performance of all models, and provide some intuition in the discussion of the results.
NF (Col=0.054 0.005, Cost= 0.009)
Diftusion- & (Col=0.039 0.004, Cost=0.231 2 0.006)
R-NCE (Col=0.030 0.003, Cost= 00.00000000000004440000
IBC (Col=0.107 0.006. Cost=0.855 0.215)
Diffusion (Col=0.056 0.005, cost=0.286:000000 0.008
0.10
0.08
Colltsion Rate
0.06
0.04
0.00
10
Env
0.10
Collisoon Rate
0.06
0.04
0.02
0.00
12
13
17
18
19
20
21
22
23
24
Figure 4: A bar plot comparing the collision rates amongst different models, for each of the 25 test environments. For each euvironment, a (environment conditioned) 95% confidence interval for the collision rate is computed via the Clopper-Pearson method and shown via collision rate and the cost is shown for each model in the legend, computed via the standard normal approximation. Each color corresponds a different model, with the color legend shown in the top plot. Note that for several of the environments, the collision rate extends beyond the limits of the displayed plots. For all models, Environments and are particularly challenging; we illustrate the qualitative behavior of this failure
bar. Furthermore, symmetric 95% confidence interval for both the
Table 3.
Results.
We focus our evaluation on two metrics: (a) collision rate, and (b) the objective cost (7.2). Our main finding is that R-NCE yields the sampler with both the lowest collision rate and the lowest obstacle
29

---

Published in Transactions on Machine Learning Research (09/2024)
Table 3: Sampled trajectories for three separate path planning test environments (out of 25). The specific envi- ronments shown to highlight multi-modal solutions. Each environment contains 384 trajectories sampled autore- gressively, 10 obstacles, and one goal location indicated by the orange marker. 'The collision rate and cost values are shown for each environment, along with symmetric 95% confidence interval computed via standard normal approximation. The "Env 1" column highlights the failure cases for the best performing models (cf. Figure 4), where collisions arise due the models generating trajectories which attempt pass through narrow passageways.
Model
Env 1
Env 10
Env 18
col 0.440 0,050, cost 0287 0.080
col 0.010 0.000, cost 0216 0.050
cal 0.040 0.021, cost 0.244 0.033
NF
col 0.383 0. 049,, cost - 0. 0.426 0.054
col 0.161 + 0.037, cost 0.426 0.095
col 0.094 = 0.029, cosl 1..994 0.335
IBC
0,253 O 043,, 0.194
col 0..06622 0.012. cost 0.168 0.048
0o0=0.000= 0.010, cost 0,21 0.045
R-NCE
col 0.479 0.050, cost 0.265
0.062 + 0.024, cost 252 0.038
Col 0.023 0.015, cost 0.372 = 0.062
Diffusion
col 0.417 0.049, cost 0.022
col 0.005 0 005, cost 217 0.023
c 0.005 0.005, cost 0 0.025
Diffusion-ó
30

---

Published in Transactions on Machine Learning Research (09/2024)
cost. Figure 4 provides a quantitative evaluation, showing the collision rates for each model on all 25 test environments. Furthermore, in Table 3, we show the performance of the various models three test environments. Observe that in Figure 4, there are two environments (specifically Environments 1 and 24) for which the collision rate excceds 10% on every model. Example trajectories from Environment 1 are shown in Table 3 to rate this failure scenario (note that Environment 24 exhibits similar behavior). Here, the failures arise from the policy attempting to aggressively cut through very narrow slot between two obstacles.
For R-NCE, we additionally experimented with leveraging pre-trained proposal model; all other relevant training hyperparameters were held the same as joint training, other than holding pert fixed at 0.005. The resulting model's performance yielded an average collision rate of 5.8% and cost 0.24, notably worse that the joint training equivalent. We hypothesize that data noising improves the conditioning of the generative learning problem. Indeed, the trajectories the training dataset arguably live on a much smaller sub-manifold of the ambient event space. By anncaling the perturbation noise, we are able to guide the convergence of the learned distributions towards this sub-manifold. This is not possible, however, if the proposal model is pre-trained and the perturbation noise is annealed during EBM training.
In Table 4, we additionally ablated the number of negative samples K within the R-NCE objective. While higher K yields better overall closed-loop performance, likely due to better generative models, the drop-off in performance is relatively mild for this problem.
Table 4: Ablation over the number of negative samples K within the R-NCE objective. Increasing K yields higher quality policies w.r.t. closed-loop collision rate and cost. However, the drop-off in performance mild, suggesting some inherent robustness in the learning problem.
K
Col
Cost
7
0.061 + 0.005
0.20 + 0.009
31
0.052 + 0.004
0.212 + 0.01
63
0.03 + 0.003
0.194 + 0.009
Finally, in Table we show that generating multiple (E = 48) samples per step and selecting the one with highest log- log-probability has non-trivial effect on the quality of the resulting trajectories, thus demonstrating the importance of ranking samplings.
Table 5: A study to illustrate the necessity of generating multiple samples per timestep, and selecting the sample with the highest score. Here, we let the variable é denote the number of samples which are generated at every timestep. Symmetric 95% confidence interval for both the collision rate and the cost are computed via the standard normal approximation. Uniformly across all models, both the collision rate and trajectory cost rise significantly going from E = 48 samples per step down to e 1. Recall that for Diffusion, the log-probability computed by treating the model as a CNF (cf. Section 6).
Model
Col (E = 48)
Col (8 = 1)
Cost (e = 48)
Cost (E = 1)
NF
0.054 = 0.005
0.101 = 0.006
0.220 + 0.009
0.500 + 0.016
R-NCE
0.030 = 0.003
0.116 = 0.006
0.194 = 0.009
2.287 = 1.089
Diffusion
0.056 = 0.005
0.099 = 0.006
0.286 = 0.008
0.569 = 0.017
Diffusion-ó
0.039 = 0.004
0.087 = 0.006
0.237 = 0.006
0.489 = 0.015
IBC
0.107 = 0.006
0.203 = 0.008
0.855 = 0.215
8.734 = 1.677
7.4 Contact-Rich Block Pushing (Push-T)
Our final lask leatures contact-rich multi-modal planar manipulation which involves using a circular end effector to push a T-shaped block into a goal configuration (Chi et al., 2023; Florence et al., 2022). In this environment, the initial pose of the T-shaped block randomized. The agent reccives both an RGB image observation of the environment (which also contains rendering of the target pose) and its current end effector position, and is tasked with outputting a sequence of position coordinates for the end effector. The
31

---

Published in Transactions on Machine Learning Research (09/2024)
target position coordinates are then tracked via a PD controller. This task is the most challenging of the tasks we study, due to (a) contact-rich behavior and (b) reliance on visual feedback. Indeed, we find that successfully solving this task with R-NCE requires the I-R-NCE machinery introduced in Section 5.2.
Our specific simulation environment and training data comes from Chi et al. (2023), which uses the pygame physics engine for contact simulation, and provides a dataset of 136 expert human teleoperated demonstra- tions.
Policy and Rollouts. Similar to path planning (cf. Section 7.3), we use an autoregressive policy to perform rollouts:
Pluli - .l.......l sli + Npredd yld], [i]........................ ................................ , - Nexx + 1], oli - Notx 1J).
(7.3)
Here, v|i] denotes the end effector position at timestep i, and oli] denotes the corresponding RGB image observation. Following Chi et al. (2023), set Npred = 16 and Netx 2, yielding 32- 52-dimensional event space. Furthermore, during policy execution, even though we predict pred = 16 positions into the future, we only play the first 8 prodictions in open loop, followed by replanning using (7.3); Chi et al. (2023, Figure 6) showed that this ratio of prediction to action horizon yielded the optimal performance on this task. As with the path planning example, for each policy step, we sample 8 predictions and execute the sequence with the highest (relative) ligllikei
Evaluation. Each rollout is scored with the following protocol. At each timestep i, the score of the current configuration is determined by first computing the ratio r|i] of the area of the intersection between the current block pose and the target pose to the area of the block. The score s[i] is then set to s|i] = min(r(i)/0 /0.95, 1). The episode terminates when either (a) the score s[i] = 1, or (b) 200 timesteps have elapsed. The final score assigned to the episode is the maximum score over all the episode timesteps. Our final evaluation is done by sampling 256 random seeds (i.e., randomized initial configurations), and rolling out each policy 32 times within each environment (with different policy sampling randomness seeds).
Architecture and Training. We design two-stage architectures, which first encode the visual observations {oli].... oli Notx + +1} into latent representations via convolutional ResNet, with coordinates appended to the channel dimension of the input (Liu et al., 2018), and spatial softmax layers at the end (Levine et al., 2016). To further improve the parameter efficiency of R-NCE and I-R-NCE, both the EBM and NF models share this visual encoder. Updates to the encoder are then allowed to be propagated by either the EBM or the NF optimization step through an appropriate coop-gadiinn operator. For the EBM, the spatial features are flattened and combined with both the agent positions and time index (the time index of the generative model, not the trajectory timestep index i), and passed through encoder-only transformers similar to the architectures used in path planning (ef. Section 7.3). For the interpolant NF which forms the negative sampler for I-R-NCE, special care is needed to design an architecture which is simultaneously expressive and computationally efficient: a dense MLP as used in path planning for the proposal distribution led to postcrior collapse during training (cf. Section 4.3). We detail our design in Appendix B.4, in addition to various training and hyperparameter luning details. For all models, we limit the parameter count to 33.3 (for I-R-NCE, this limit applies to the sum of the NF+EBM parameters)." Finally, similar to the training procedure for path planning, we employed the same annealed data perturbation schodule to guide training, with noise applied only to the sequence of end effector positions.
Results. We show our results in Table 6 and Table 7. Note that we do not evaluate IBC here due to its previous poor performance both in path planning, and poor performance shown for this task in Chi et al. (2023). Table 6 contains, for each model, visualizations of policy trajectories in four different evaluation enviromments. In Table 7, we show the final evaluation score achieved by each of the models; we additionally include results for I-R-NCE paired with the simpler two-stage sampling method from Algorithm 2, and the simpler non-interpolating R-NCE (Algorithm 1). We note that the best performance is achieved by
Package website: https://gitaeei //github........ wwt.eo/////yeeeee pygane/pygans We find that models which are several orders of magnitude smaller than the models used Chi al. (2023) suffice.
32

---

Published in Transactions on Machine Learning Research (09/2024)
Table visualization of the policy trajectories each model four different initial conditions. The trajectory is color coded using time, with the beginning of the trajectory represented using the lightest shade and the end of trajectory with the darkest shade The desired goal configuration show in green color Above each environment, the score assigned to the rollout is shown (see the evaluation paragraph in Section 7.4 for the definition of the score).
Model
Examples
1.000
0.972
0.964
0.572
NF
r= 0.954
r= 1.000
r= 0.992
r= 0.988
I-R-NCE
1.000
0.958
r= 0.907
1.000
Diffusion
r= 0.967
r= 0.240
/= 1.000
r= 0.955
Diffusion-ó
I-R-NCE, paired with the three-stage sampling method from Algorithm 4. This is closely followed by 1- R-NCE with the two-stage sampling method from Algorithm 2: both I-R-NCE results notably outperform non-interpolating R-NCE. This suggests that training via I-R-NCE yields significantly superior models, even when the EBM only used at == 1 during sampling. reasonable hypothesis would be that training via I-RNCCE yields multi-scale variance reduction whereby the shared network & benefits from lcarning across all times within (0, 1], similar the benefits attributed time-varying score modeling. We anticipate that challenging
the gap between the two- and three-stage sampling algorithms should widen with
33

---

Published in Transactions on Machine Learning Research (09/2024)
distributions. Meanwhile, the non-interpolating variant of R-NCE does not inherit any multi-scale learning benefits. Finally, with path planning, I-R-NCE (with or without the SDE sampling stage) outperforms the NF, Diffusion, and Diffusion-ó models.
Table 7: The final evaluation score across 256 randomly sampled initial configurations, and 32 rollouts within each environment configuration. The reported score is scalar between [0, 1] which describes the maximum coverage of the target achieved by the policy; the precise definition is given in the evaluation paragraph of Section 7.4. A 95% symmetric confidence interval computed via a normal approximation is shown for each score. Here, we see that the model trained with T-R-NCE outperforms the other models. The difference between I-R-NCE, and I-R-NCE (Alg. 2) corresponds to the use of Algorithm for the former and Algorithm 2 for the latter during sampling.
NF
R-NCE
I-R-NCE (Alg. 2)
I-R-NCE
Diffusion
Diffusion-
0.866 + 0.006 0.824 + 0.006
0.880 + 0.005
0.884 + 0.005 0.864 + 0.006 0.860 + 0.006
8
Conclusion
We showed that in the context of policy representation, energy-based models are competitive alternative to other state-of-the-art generative models such as diffusion models and stochastic interpolants. This was done through several key ideas: (i) optimizing the ranking NCE objective in lieu of the InfoNCE objective proposed in the IBC work (Florence et al., 2022), (ii) employing a learnable negative sampler which is jointly, but non-adversariallly, trained with the energy model, and (iii) training a farily of energy models indexed by a scale variable to learn a stochastic process that forms a continuous bridge between the data distribution and latent noise.
This work naturally raises several broader points for further investigation. First, while we focused on theo- retical foundations in addition to a comprehensive evaluation of R-NCE trained EBMs compared with other generative models, in future work we plan to deploy these EBMs on real hardware platforms. Second, this work investigated the use of R-NCE for training EBMs in the supervised setting, ie., behavioral cloning. A promising future direction would be adapt our R-NCE algorithm and insights for training FBM policies via reinforcement learning (Levine, 2018). Finally, despite our motivations being targeted towards robotic ap- plications, many of our technical contributions actually apply generative modeling more broadly. Another direction for future research is to apply our techniques to her domains such as text-to-image generation, super-resolution and inverse problems in medical imaging, and see if the benefits transfer over to other domains.
References
Stephen Adams, Tyler Cody, and Peter A. Beling. A survey of inverse reinforcement learning. Artificial Intelligence Revieu, 55(6):4307-43466 2022.
Amurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is condi- tional generative modeling all you need for decision-making? ar preprint arXiv:2211.1563 2022.
Michael S. Albergo and Eric Vanden-FFjndenn. Building normalizing flows with stochastic interpolants. Tn International Conference on Learning Representations, 2023.
Michael S. Albergo, Nicholas M. Boffi, and Eric Vanden-Lijndeen. Stochastic interpolants: A unifying frame- work for flows and diffusions. arXiv preprint arXiv:2903.087 2023.
Michael Arbel, Liang Zhou, and Arthur Gretton. Generalized energy based models. In International Con- ference on Learning Representations, 2021.
Igor Babuschkin, Kate Baumli, Alison Bell, Surya Bhupatiraju, Jake Bruce, Peter Buchlovsky, David Bud- den, Trevor Cai, Aidan Clark, Ivo Danihelka, Antoine Dodieu, Claudio Fantacci, Jonathan Godwin, Chris Jones, Ross Hemsley, Tom Hennigan, Matteo Hessel, Shaobo Hou, Steven Kapturowski, Thomas Keck, Turii
34

---

Published in Transactions on Machine Learning Research (09/2024)
Kemaev, Michael King, Markus Kunesch, Lena Martens, Hamza Merzic, Vladimir Mikulik, Tamara Nor- man, George Papamakarios, John Quan, Roman Ring, Francisco Ruiz, Alvaro Sanchez, Rosalia Schneider, Eren Sezener, Stephen Spencer, Srivatsan Srinivasan, Wojeiech Stokowiec, Luyu Wang, Guangyao Zhou, and Fabio Viola. The DeepMind JAX Ecosystem, 2020. URL http: http:///itthnn /gthtt....uu.... htw.itttu..meemeeeeeeeeeeeeeeeee cwt/com///////eeeee com/deepmind.
Timothy D. Barfoot, Chi Hay Tong, and Simo Särkkä. Batch continuuuttimm exactly sparse gaussian process regression. In Proceedings of Robotics: Science and Systems, 2014.
trajectory estimation as
Gérard Biau, Benoît Cadre, Maxime Sangnier, and Ugo Taniclian. Some theorctical properties of GANS.
The Annals of Stabistics, 48(3):1539 1566, 2020.
Andreas Blattmann, Robin Rombach, Huan Ling, Tim Dockhorn, Seung Wook Kim, Sanja Fidler, and Karsten Kreis. Align your latents: High-resolution video synthesis with latent diffusion models. In Proceedings of the IEBE/CVF Conference on Computer Vision and Pattern Recogrition, pp. 22563 22575, 2023.
Sam Bond- Taylor, Adam Leach, Yang Long, and Chris G. Willcocks. Deep generative modelling: A compara- tive review of VAžs, GANs, normalizing flows, mnergy-based and autoregressive models. IEEE Transactions on Pattern Analysis and Machine Intelligence, 44 4A(11):7327-7347, 2022.
Avishek Joey Bose, Huan Ling, and Yanshuai Cao. Adversarial contrastive estimation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1021-1032, 2018.
James Bradbury, Roy Frostig, Peter Hawkins, Malthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne and Qiao Zhang. JAX: compos- able transformations of Python+NumPy programs, 2018. URL http://gi tttb//ttttuuu...////// com googgeo
Omar Chehab, Alexandre Gramfort, and Aapo Hyvärinen. The optimal noise noise-contrastive learn- ing is not what you think. In Proceedings of the Thirty-Eighth Conference on Uncertointy in Artificial Intelligence, volume 180, pp. 307-316. PMLR, 2022.
Ricky T. Q. Chen, Yulia Rubanova, Jesse Bettencourt, and David K. Duvenaud. Neural ordinary differential equations. In Advances in Neural Information Processing Systems, volume 31, 2018.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song, Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems, 2023.
Bo Dai, Zhen Liu, Hanjun Dai, Niao He, Arthur Gretton, Le Song, and Dale Schuurmans. Exponential family estimation via adversarial dynamics embedding. In Advances in Neural Information Processing Systems, volume 32, 2019.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthosis. In Advances in Newral Information Processing Systems, volume 34, 2021.
Carl Doersch. Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05900 05900, 2016.
Yilun Du and Igor Mordatch. Implicit generation and modeling with energy based models. Advances in Neural Information Processing Systems, volume 32, 2019.
Yilun Du, Torn Lin, and Igor Mordatch. Model-based planning with energy-based models. In Proceedings of the 3rd Conference on Robot Learning, volume 100, pp. 374-383. PMLR, 2020.
Yilun Du, Shuang Li, Joshua Tenenbaum, and Igor Mordatch. Improved contrastive divergence training of energy-basee models. In International Conference on Machine Learning, volume 139, pp. 2837-2848. PMLRR 2021.
35

---

Published in Transactions on Machine Learning Research (09/2024)
Yilun Du, Conor Durkan, Robin Strudel, Joshua B. Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. In International Conference on Machine Learning, volume 202, pp. 8489-8510. PMLR, 2023.
Thomas S. Ferguson. A course in large sample theory. Routledge, 2017.
Pete Florence, Corey Lynch, Andy Zeng, Oscar A. Ramirez, Ayzaan Wahid, Laura Downs, Adrian Wong, Johnny Lec, Igor Mordatch, and Jonathan Tompson. Implicit behavioral cloning. In Proceedings of the 5th Conference on Robot Leurning, volume 164, pp. 158 168. PMLR, 2022.
Gerald B. Folland. Real Analysis: Modern Techniques and Their Applications. Wiley, 1999.
Ruiqi Gao, Erik Nijkamp, Diederik Kingma, Zhen Xu, Andrew M. Dai, and Ying Nian Wu. Flow contrastive estimation of energy-basod models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pultern Recognition, pp. 7518 7528, 2020.
Sebastian Gomez-Gonzales Sergey Prokudin, Bernhard Schölkopf, and Jan Peters. Real time trajectory prediction using deep conditional generative models. IEFE Robotics and Automation Letters, 5(2):970- 976, 2020.
Will Grathwohl, Ricky T. Q. Chen, Jesse Bettencourt, and David Duvenaud. Scalable reversible generative models with free-form continuous dynamics. In International Conference on Learning Represenlations. 2019.
Will Grathwohl, Jacob Kelly, Milad Hashemi, Mohammad Norouzi, Kevin Swersky, and David Duvenaud. No MCMC for me: Amortized sampling for fast and stable training of energy-based models. In International Conference on Learning Representations, 2021.
Michael U. Gutmann and Jun-ichiro Hirayama. Bregman divergence as general framework to estimate unnormalized statistical models. In Proceedings of the Twenly-Seventh Conference on Uncertainly in Artificial Intelligence, pp. 283 290, 2011.
Michael U. Gutmann and Aapo Hyvärinen. Noiso-contrastive estimation of unnormalized statistical models, with applications to natural image statistics. Journal of Machine Learning Research, 13(11):307 361, 2012.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference Machine Learning, volume 70, pp. 1352 1361. PMLR, 2017.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Conference on Machine Learning, volume 80, pp. 1861-1870. PMLR, 2018
Jonathan Heek, Anselm Levskaya, Avital Oliver, Marvin Ritter, Bertrand Rondepierre, Andreas Steiner, and Mare van Zee. Flax: A neural network library and ecosystem for JAX, 2023. URL http://github com/google/ffax.
Nicolas Heess, David Silver, and Yee Whye Teh. Actor-critic reinforcement learning with energy-based policies. In Europeon Workshop on Reinforcement Learning, volume 24, pp. 45-56 48-58, PMLR, 2013.
Jonathan Ho and Stefano Ermon. Generative adversarial imitation learning. In Advances in Neural Infor- mation Processing Systems, volume 29, 2016.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. In Advances Neural Information. Processing Systems, volume 33, 2020.
36

---

Published in Transactions on Machine Learning Research (09/2024)
Boris Ivanovic, Karen Leung, Edward Schmerling, and Marco Pavone. Multimodal deep generative models for trajectory prediction: A conditional variational autoencoder approach. TEEF Robolics and Automation Letters, 6(2):295 302, 2021.
Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, volume 162, pp. 9902-9915. PMLR, 2022.
Chi Jin, Praneeth Netrapalli, and Michael Jordan. What is local optimality in nonconvex-noncoonaay minimax optimization? In International Conference on Machine Learning, volume 119, pp- 4880-4889. PMLR, 2020.
Rafal Jozefowicz, Oriol Vinyals, Mike Schuster, Noam Shazeer, and Yonghui Wu. Exploring the limits of language modeling. arXiv preprint arXiv:160..00 024 10, 2016.
Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine. Elucidating the design space of diffusion-based generalive models. In Advances in Newral Information Processing Systems, volume 35, 2022.
Patrick Kidger. On Neural Differential Equations. PhD thesis, University of Oxford, 2021.
Diederik P. Kingma and Jimmy Lei Ba. Adam: A method for stochastic optimization. In International
Conference Learning Represenlations, 2015.
Zhileng Kong, Wei Ping, Jiaji Huang, Kexin Zhao, and Bryan Catanzaro. Diffwave: versatile diffusion model for audio synthesis. In International Conference on Learning Representations, 2021.
Rithesh Kumar, Sherjil Ozair, Anirudh Goyal, Aaron Courville, and Yoshua Bengio. Maximum entropy generators for energy-based models. arXiv preprint arx:w:1001.000 1001.0800,, 2019.
Phuc H. Le-Khac, Graham Healy, and Alan F. Smeaton. Contrastive representation learning: A framework and review. IEEE Access, 8:193907 193934, 2020.
Yann Lecun, Sumit Chopra, Raia Hadsell, Mare Aurelio Ranzato, and Fu Jie Huang. tutorial on eneryy- based learning. MIT Press, 2006.
Holden Lee, Chirag Pabbaraju, Anish Prasad Sevekari, and Andrej Risteski. Pitfalls of gaussians as noise distribution in NCE. In International Conference on Learning Representations, 2023.
Sergey Levine. Reinforcement learning and control probabilistic inference: Tutorial and review, anXiv preprint 1 00909, 2018.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. Journal of Machine Learning Research, 17(39):1-41) 40, 2016.
Yaron Lipman, Ricky T. Q Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In International Conference Learning Representations, 2023.
Bingbin Liu, Elan Rosenfeld, Pradeep Kumar Ravikumar, and Andrej Risteski. Analyzing and improv- ing the optimization landscape of noise-contrastive estimation. In International Conference on Learning Representations, 2022.
Minghuan Liu, Tairan He, Minkai Xu, and Weinan Zhang. Energy-based imitation learning. In Proceedings of the 20th International Conference on Autonomous Agents and MultiAgent Systems, pp. 809 817, 2021.
Rosanne Liu, Joel Lehman, Piero Molino, Felipe Petroski Such, Eric Frank, Alex Sergeev, and Jason Yosinski. An intriguing failing of convolutional neural networks and the coordconv solution. In Advances in Neural Information Processing Systems, volume 31, 2018.
Xingchao Liu, Chengyue Gong, and Qiang Liu. Flow straight and fast: Learning to generate and transfer data with rectified flow. In International Conference on Learning Representations, 2023.
37

---

Published in Transactions on Machine Learning Research (09/2024)
Zhuang Ma and Michael Collins. Noise contrastive estimation and negative sampling for conditional models: Consistency and statistical efficiency. In Proceedings of the 2018 Conference on Empirical Methods in Natural Language Processing, pp. 3698 3707, 2018.
Raphael A. Meyer, Cameron Musco, Christopher Musco, and David Woodruff. Hutch++: Optimal stochastic trace estimation. In SIAM Symposium on Simplicity in Algorithms, pp. 142-155, 2021.
Mustafa Mukadam, Jing Dong, Xinyan Yan, Frank Dellaert, and Byron Boots. Continuous-tine gaussian process motion planning via probabilistic inference. The International Journal of Robotics Research, 37
(114:1319-1340, 2018.
Erik Nijkamp, Ruiqi Gao, Pavel Sountsov, Srinivas Vasudevan, Bo Pang, Song-Chun Zhu, and Ying Nian Wu. MCMC should mix: learning energy-based model with neural transport latent space MCMC. In International Conference on Learning Representations, 2022.
Tim Pearce, Tabish Rashid, Anssi Kanervisto, Dave Bignell, Mingfei Sun, Raluca Georgescu, Sergio Valcarcel Macua, Shan Zheng Tan, Ida Momennejad, Katja Hofmann, and Sam Devlin. Imitating human behaviour with diffusion models. In International Conference on Learning Representations, 2023.
Ethan Perez, Florian Strub, Harm de Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In AAAI Conference on Artificial Intelligence, 2018.
Prajit Ramachandran, Barret Zoph, and Quoc V. Le. Searching for activation functions. arXiv preprinl arXiv:1710.08 05940, 2017.
Moritz Reuss, Maximilian Li, Xiaogang Jia, and Rudolf Lioutikov. Goal-conditioned imitation learning using scoro-based diffusion policies. arXic preprint arXiu:234.002522 2023.
Walter Rudin. Function theory in the unit ball of C", Springer Science & Business Media, 2008.
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J. Fleet, and Mohammad Norouzi. Image super-resolution via iterative refinement. IEEE Transactions on Pattern Analysis and Machine Intelligence, (564):4713-4722) 2023.
Tim Salimans and Jonathan Ho. Should EBMs model the energy or the score? In International Conference on Learning Representations (aanennnsseeeerrr (Energy Based Models Workshop), 2021.
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Implicit neural representations with periodic activation functions. In Advances in Neural Information Processing Systems, volume 33, 2020.
Jascha Sohl-Dickstein Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thernodynamics. In International Conference on Machine Learning, volume 35, pp. 2256 2265. PMLR, 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In International Conference Learning Representations, 2020.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution. In Advances in Newral Information Processing Systems, volume 32, 2019.
Yang Song and Stefano Ermon. Improved techniques for training score- score-based generative models. In Advances in Neural Information Processing Systems, volume 33, 2020.
Yang Song and Diederik P. Kingma. How to train your energy-based models. arXiv:2101.0325. 2021.
arXiv preprint
Yang Song, Conor Durkan, Iain Murray, and Stefano Ermon. Maximum likelihood training of score-basod diffusion models. In Advances in Neural Information Processing Systems, volume 34, 2021a.
38

---

Published in Transactions on Machine Learning Research (09/2024)
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference Learning Representations, 2021b.
Yang Song, Liyue Shen, Lei Xing, and Stefano Ermon. Solving inverse problems in medical imaging with scoro-based generative models. In International Conference Learning Representations, 2022.
Duy-Nguyen Ta, Eric Cousineau, Huihua Zhao, and Siyuan Feng. Conditional energy- based models for implicit policies: The gap between theory and practice. arXie preprint arXiv:2207.0588 2022.
Tijmen Tieleman. Training restricted boltzman machines using approximations to the likelihood gradient. In International Conference on Machine Learning, 2008.
Julen Urain, Niklas Funk, Georgia Chalvatzaki, and Jan Peters. SEf3)-DiffusionFiedd:: Learning cost func- tions for joint grasp and motion oplimization through diffusion. ar preprint orXiv:2209. 03855, 2022a.
Julen Urain, An Le, Alexander Lambert, Georgia Chalvatzaki, Byron Boots, and Jan Peters. Learning implicit priors for motion optimization. In IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 7672-7679, 2022b.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arviv preprint
2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkorcit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Newral Information Processing Systems, volume 30, 2017.
Zhisheng Xiao, Karsten Kreis, Jan Kautz, and Arash Vahdat. VAEBM: A symbiosis between variational autoencoders and energy-based models. In International Conference on Learning Representations, 2021.
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In International Conference Machine Learning, volume 119, pp. 10524-10533. PMLR, 2020.
Ling Yang, Zhilong Zhang, Yang Song, Shenda Hong, Runsheng Xu, Yue Zhao, Wentao Zhang, Bin Cui, and Ming-Hsuan Yang. Diffusion models: comprehensive survey of methods and applications. arXiv preprint ar Xii::209.00077 2022.
Yasin YazIcı, Chuan-Sheng Foo, Stefan Winkler, Kim-Hui Yap, Georgios Piliouras, and Vijay Chandrasekhar. The umusual effectiveness of averaging in GAN training. In International Conference on Learning Repre- sentations, 2019.
Tony Z. Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fnc-grained bimanual manipulation with low-cost hardware. In Proceedings of Robolics: Science and Systems, 2023.
39

---

Published in Transactions on Machine Learning Research (09/2024)
Appendix
Table of Contents
A Reference Two Time-Scale Implementation of Joint Sample and Log- Nog-Probability Computation
41
B Experimental Details
41
B.1 Recovering the Energy Model Diffusion-
B.2 Synthetic Two- Pvo-Dimensional Examples
42
B.3 Path Planning
43
B.4 Push-T
44
Proofs
C.1 Optimality
47
C.2 Asymptotic Convergence
49
C.3 Learning the Proposal Distribution
54
40

---

Published in Transactions on Machine Learning Research (09/2024)
A Reference Two Time-Scale Implementation of Joint Sample and Nog-Probability Computation
Figure 5 provides a simple reference implementation of the two time-scale approach to jointly sampling and computing log-probability described in Section 6.
from typing import Callable
import jax
import import jax.scipy as jap import diffrax as dfx
as jnp
Array, KeyArray jax.Array jax. Time, Context, Event = float, Array, Array
VectorField = Callable [[Time, Event, Context], Event]
EVENT WVENT_DIM-
14
15
def latent
latent_sample and_log_prob KeyArray) tuple[Event; float]:
Isotropic Gaussian base distribution for simplicity.
z0 ja..ja..anaee return z0, jup. sum (jsp. stats ntram.norm.... mhorm.loge nww..ooopdmmccccccccooo.co.o. (z0))
(key, (EVENTT_D___
def sample_ane @ample_and__ _nd__oo_proo(((
91
vf: VectorField, key: KeyArray, ctx: Context, ode_ts: Array, 1p_ts: Array tuple [Event, float]:
22
def d_dt_psi (t, 2, x):
fn lambda
vf (t, 2, x)
., vjp_ fn = jax. jax.vip(fn, 2)
(dfdz,) = jax. vmap (vjp_fnl fn) (jnp. (EVENT_
-jnp.trace return -jnp. trace (dfdz)
98
srajasuubaavest dfx. SubSaveAt (ts=lp_ts, in-d_dt_psil
29
z0, psi0 latent sample and _og_prob (z,), d_ dldt__pat psis dfx. diffeqsolve( dfx. dfx.0DEEI (xEEDTerrm(v)). solver=dfx.Heun t0=ode t0=ode_ts[0]] [0], [soode_ts[[]]] dtO=None, y0=z0, args=ctx, stepsize wtt.:lleeeeddeeeeetotooooooo ta), saveat=dfx. SaveAt (subs= Stw..aaa..aeeetttddcaaaeaaa.a... (tataaaatttt:tttt)) (t1=True), traj. subsaveat]]) lp_ts)
85
return 2, psi0 jnp.
Figure 5: reference implementation the two time-scale approach described in Section for jointly computing samples and log-probabilities from continuous normalizing flow.
B Experimental Details
B.1 Recovering the Energy Model in Diffusion-ó
The score function Vlogp(y e) in the Diffusion-EDDM parameterization (Karras al., 2022) not represented directly by neural network, but rather indirectly through a sequence of affine transformations. In order recover (up to constant) the corresponding function ((x,y,)) ply 2,, (), need to
41

---

Published in Transactions on Machine Learning Research (09/2024)
integrate the score across the affine transformations. Specifically, letting f(z, y. denote the raw vector- valued network, Diffusion-EDD parameterizes the score as (Karras et al., 2022, Eq. 3 and Eq. 7):
v log p(y x,() vd(x, u))a(()²
d(x, v, = Cskip (4) + Cont (t) (Cim (1)3,, Cuoise (4))-
Now, if the network fx, y, 1) has special gradient structure f(x, y. t) = Vod(x, v, then a straightforward computation yields:
logp(y x2 1) r(x,,,1) + const
2/2 2/2(20 (Cakip (t) 1)]lvl1+ + Cout 6(I, en (t)y, Cnoise (t)).
1
r(a, y. =
(B.1)
(t)
where co,st(r,,t) t) a function that unknown, but only depends on x and t. Hence, for a fixed X and events {yi} C Y,, we can compute the relative likelihoods {r(z, &t,e), where is a value close zero (but not equal to zero as (B.1) becomes degenerate at 0). These relative likelihood values can be used to compare the likelihood of samples for a given context, and can be hence used to perform fast action selection.
B.2 Synthetic Two-Dimensional Examples
For each problem (i.e., pinwheel and spiral), the training data consists of 50,000 samples from the joint distribution Px...
Architectures and EBM Sampling. For all energy models (TBC, Diffusion-o, and R-NCE-EBM, as we do not use the interpolating variant of R-NCE for this benchmark), we model the energy function E(x, y) by first concatenating (z, y) and then passing through 8 dense layers with residual connections. Similarly for the NF vector field and Diffusion denoiser, which also a vector-valued map (1 y), we first embed the time argument through 2 fully connected MLP layers of width 10, and then pass the concatenated arguments (x,, y, MLP(()) into 8 dense layers of width 64 with residual connections. Finally for parameter efficiency reasons, the R-NCE-NF vector field first concatenates into 2 ConcatSquash 10 layers. Both here and throughout our experiments, we utilize the swish activation function (Ramachandran et al., 2017). EBM sampling (i.e., for IBC and R-NCE) is done using Algorithm 2 with Langevin MCMC. Table and Table 9 contains the hyperparameters we search over.
(z,y) and then passes the inputs (2,t)
Table 8: List of hyperparameters for the EBMs. LR denotes learning rate, WD denotes weight decay, Timme refers to the number of Langevin steps taken, n refers the Langevin step size, and K denotes the number of negative examples. Note that K is odd here since it was more computationally efficient to sample K ++ 1 negative samples and replace one with the positive sample. Note that all models considered have at most 111k NF EBM parameters in total.
Model
Width
WD
meme
71 (10-3,5
IBC
(64, 128)
[10-3 5. 10-4)
{0. 10-1)
(250, 500, 750, 1000)
10 10-1) (9,63, 63, 127, 255)
R-NCE-EBM
64
(10-3,5-10
(0. 10
(500, 1000)
[10 110- 10-4)
R-NCE-NF
(42 64)
10-
{0. 10-4
Training. We use a constant learning rate for all models, a batch size of 128 examples, and train the models for 20,000 total batch steps. For R-NCE, 20,000 R-NCE batch steps are taken, but the negative sampler varies between 20,000 or 50,000 interpolant loss batch steps depending on the specific hyperparameter sotting. In particular, we tune both the number of negative sampler updates (Tsamp in Algorithm Line 5) in addition to the R-NCE updates (Trnce in Algorithm 1, Line 9) in the main loop. Specifically, recalling that Touter denotes the total number of outer loop iterations (Algorithm 1, Line 4), we search over the grid:
4000 =5, Tmce==), (T (Touter=10000 Tramp- =3, T -2)}.
42

---

Published in Transactions on Machine Learning Research (09/2024)
Table List of hyperparameters for diffusion and NF models. The parameter a is described in Section 7.1, and the remaining parameters dasa, Fmax, Omin, and p refer to various Diffusion-EDM parameters from Karras et al. (2022. Table 1). All models considered have at most ~22 parameters. For sampling, all models use 768 Heun steps (1,024 total function evaluations).
Model
Width
LR
WD
Odata
Fmax
Omia
e
a
Diffusion
(10-3,5 10-4) (0, (0,10-4) (0.5, 13
{60, 80}
0.002
(7,9)
Diffusion-o
(64, 128)
NF
(2,3.4)
B.3 Path Planning
Architectures and EBM Sampling. As discussed in Section 7.3, train architectures based encoder-only transformers. Here, we give more detailed description. The imputs vector fields are interpolation/noisee /coise-scale time E R, obstacles 1000 goal location g E R2, previous positions T_ (yli - Netx + .........l.. v[t)) E RNaxx2 and future positions 7 (uli + iI.....,/ yii + NNread) E R Nace We fix an embedding dimension demb, and denote the combined past and future snippets as T = T+)- The architecture depicted Figure 6, where dout is equal to the dimensionality of the event space (2.prred for path planning) for vector fields and is oqual to one for energy models.
energy models and
XN
MLP(4dembe demb)
9
MLP Ademb denb)
MH-SA
Convld(1, demb))
Convld(1, demb)
MLP (demb dembe dont.)
Figure 6: The encoder-only transformer architecture we utilize for path planning models. Here, refers to the swish activation function (which is also the activation function used inside the MLP blocks). @ is the standard sinusoidal positional embedding, and the first argument to Convld denotes the filter size. MH-SA denotes a standard Multi-Head Self- Attention encoder block (Vaswani et al., 2017); use 4 self-attention heads in each block. The f tokens after the terminal block are flaitened and passed through the oulput MLP.
In addition, we also consider architectures which utilize convolutional "prior" which operates only on the trajectory T tokens. Namely, after F is embedded into Rdamb and positionally encoded, we pass the resulting trajectory tokens through the following two blocks:
Convld(3, demb) - GroupNorm(8) + a + Dense((dem)) + a Convid(3, demb) - GroupNorm(S) + o + Dense(dout) -
(First block)
(B.2a)
(Second block).
(B.2b)
The output of the second block then summed along the temporal axis, to produce the output of the prior. This is then added to the output of the final MLP, shown in Figure 6.
Furthermore, we remark that for computational and parameter efficiency, both IBC-NF and R-NCE-NF do not use the transformer encoder architecture described in Figure 6. Instead, all obstacle, goal, and trajectory inputs are concatenated together and fed into several ConcatSquash layers. While this produces a proposal distribution which is nowhere near optimal, turns out be sufficiently powerful negative sampler for this problem. Sampling during inference from the EBM models is done using Algorithm with Langevin sampling for the MCMC stage.
We use train batch size of 128. For both IBC and R-NCE, we use Touter 250,000 main loop steps; each loop consists of either Tsamp E (2, 4) negative sampler batch steps, and 2 IBC/R-NCE updates
Training.
For the IBC-EBM and EBM models, fixed constant value 1 as we do the I-EBM formulation here.
43

---

Published in Transactions on Machine Learning Research (09/2024)
List of hyperparameters for the IBC and R-NCE path planning models. Width denotes the width of the ConcatSquash layers, demb and N₂ are described in Figure 6, Tmeme refers to the number of Langevin steps taken, and y refers to the Langevin step size. The number of ConcatSquash layers for IBC-NF and R-NCE-NT set equal to the number of self-attention blocks No within the EBM; the constant for the proposal NEs was fixed at 2. Note that in our hyperparameter sweeps, we exclude combinations of Width and No which yield total (NF + EBM) parameter accounts exceeding ~500k.
Table 10:
Model
Width
demb
Nb
LR
WD
Tieme
1)
1BC-EBM
64
(3, (3,
5 10
(0,10-9 4)
(500, 750)
10-4
IBC-NF
{64, 256)
10-3
R-NCE-EBM
64
{3,4,5)
5 10-
{0, 10,10-
(500,
10-4
R-NCE-NF
{64, 128, 256}
10-3
steps. The IBC-NF TBC-NF/R-NCF-NT optimizers use a cosine decay learning rate without warmup, whereas the IBC-EBM/R-NCE-EBN optimizers use a warmup cosine decay rate with 5,000 warmup steps. For both IBC and R-NCE, we use fixed K = 63 negative samples. More hyperparanetters are listed in Table 10. On the other hand, for Diffusion, Diffusion-o, and NF, train our models for 500,000 batch steps and use a warmup cosine decay learning rate schedule, 12 with 5,000 warmup steps. The remaining hyperparameters are listed in Table 11.
List of hyperparameters for diffusion and NF path planning models. Note that in our hyperparameter sweeps, we exclude combinations of demb and No which yield parameter accounts exceeding ~500k. For sampling, all models use 375 Heun steps (750 total function evaluations), and 64 log-probability ODE steps (cf. Section 6).
Table 11:
Model
demb
No
LR
WD
datta
mmax
miia
Diffusion
Diffusion-
864,80}
{3, 4,5}
{10 1.008,.00005)
{0, 10 10-4)
{.5, 1]
{60, 80}
.002
{.01. 02, 02.03
(7. 9)
Model
demb
No
LR
WD
NF
{64, 80)
[3, 4,5)
(1005)
(0, 10 'I (2.3)
B.4 Push-T
Architectures and EBM Sampling. For all models, we use similar encoder-only transformer architec- ture as for path planning. The inputs to the energy models and vector fields consist of: interpolation/nott /noisee scale time 1 E TR, 13 past image observations o E IR Natx x x9xx9xx3 past agent positions T_ = (y]² - Nex + 1J......(i7) v(i)) E Naxx22 and future positions T = (yli + II...... wli + Npreaa) E X2
The image observations are first featurized into sequence of embeddings R R/ Xd, via a convolutional ResNet, with coordinates appended to the channel dimension of the images (Liu et al., 2018), and spatial softmax layers at the end (Levine et al., 2016). We use a 4-layer ResNet, with width dres, and 4 spatial- softmax output filters, giving d, = 8. The rest of the architecture resembles the self-attention archilecture from Appendix B.3 (see Figure 7 for completeness). As before, we choose an embedding dimension demb: the output dimension dou is equal to the dimensionality of the event space, 2N for vector fields, and is equal to one for energy models. As with path planning, we also add the output from a convolution prior, defined in (B.2), which operates only on the prediction takens T+ post positional embedding. The dimensions dres and demb were held fixed at 64 and 128 respectively, for all models.
The negative sampler NF model's architecture for I-R-NCE needed to be designed differently since the self-attention layers in Figure 7 were computationally prohibitive for negative sampling and log-probability computations, both of which are necessary for defining the I-R-NCE objective. As a result, we devised a novel MaayI--Mutt-Heea-FFILMM. (CMHF) block, inspired from Perez et al. (2018), which, while inferior to the general self-attention block, was sufficiently powerful for training via. T-R-NCE and computationally cheaper. The overall EBM+NF architecture and CMHF block are shown in Figures 8 and 9, respectively. Finally, drawing inspiration from the pre-conditioning function Cooise (1) presented in Karras et al. (2022), we
phtttps:/eptax: http:///tpxxxxrrtteeeeeeee rew.dheddcs.ee... www/en/aatest/pp........ ctapp..ttmlooptt... parpptt....rruuppp..mmmmmm www.up_araaaeeeeeeeepee cowiiee www.snoeeeccae.eeaeeeeeee lee_dd_echeeuullee schedule
Note that interpolation time e (0. 1] also an input to the ERM as we use the I-R-NCE formulation.
44

---

Published in Transactions on Machine Learning Research (09/2024)
Context Tokenization
M4ddeeemb demb)
X
MLP(4demth demb)
MH-SA
MLP(4demb dembi
T+
MLP(44emub demb)
MLP(2dembb demb: dout)
Figure 7: The encoder-only transformer architecture for Push-'T. The T tokens after the terminal block are flattened and passed through the output MLP. Each MH-SA block used 4 attention heads.
applied the non-linear transform + - log(1.1 t) on the interpolation time used within I-R-NCE-EBM and I-R-NCE-NF. The advantage of this transform to counteract the effect of the decreasing step-sizes from the transform t!/a defined in Section 7.1.
q
xNN
T+ MLP(d))
MLP(2dy, dy. dout)
MLP(d,)
Conv ID MH-FiLM
MLP(d,
Context Tokenization
MLP(4demmb dembi
xNN,
MLP(4demb dembi
MH-SA
(Ldeedd demb)
T4
MLP(4demis demb))
MLP(2damm dembi 1)
Figure 8: Combined NF (top) EBM (bottom) architecture I-R-NCE. To eliminate parameter redundancy, both models share the context tokenizer. However, only the sampler (cf. Algorithm 3, Line 9) the EBM (cf. Algo- rithm 3, Line 14) is allowed to update the context tokenizer; we swept over both choices. Post context tokenization, smaller feature embeddingg d, demb// 64. The design of the Couv L-D-Muli-Headd-PPP
the NF blocks block shown Figure 9. We use 4 heads within the EBMs self-attention blocks and 2 heads within the NFs CMHF blocks. The additional input q the CMHF blocks explained in Figure 9.
Sampling during inference from the EBM models done via Algorithm 4 with HMC sampling for the MCMC stage. We split the net amount evaluations of v (z, y.t) between the SDE and MCMC stages using the interpolation time lower-bound t; particular, It T meme] evaluations are dedicated to the MCMC stage, while [(1-t) Tmemel evaluations are given the SDE stage. This allows us control the net inference-time computational fidelity.
45

---

Published in Transactions on Machine Learning Research (09/2024)
MLP(dd)
MH-SA
Heads
MH
FLLM
q
MH-CA
Ym 80)
LayerNorm ConviD(3. dy)
7+
Figure 9: The ConYID-MH-FiLM block. The input X is the combined set of context tokens from o- and t. Each block has its own independent set of set query tokens q E IR treated as learnable parameters, that cross-attend (MH-CA) to the composite context tokens X produce FiLM weights {(x, B:)) where nr Br R The number of queries in each CMHF block ny was set to the number of heads (2). Each set of FiLM weights (7t, (3)) are broadcasted across the temporal dimension of the prediction tokens T+ and affinely transform the prediction tokens (see Perez et al. (2018)) to yield my copies of transformed prediction tokens. These are then concatenated (similar to a multi-head attention block) along the head-dimension, and projected back down to yield a set of output prediction tokens 7+. Separately, the composite context tokens X are propagated through a standard MH-SA block.
List of hyperparameters for I-R-NCE. The constant N₂ is described in Figure 8, Theme refers to the total number of VyEo evaluations, split between the SDE sampler and HMC leapfrog integrator steps, and n refers to both the SDE noise-scale (see (5.7) and HMC leapfrog integrator slep size; we used 50 leapfrog sleps per HMC momentuum sampling step. The value is the initial CNF sample time; see Algoritiu 4. The constant a for the proposal NFs was fixed at 2.5. Note that in our sweeps, we exclude combinations of No for the EBM and NF which yield total parameter accounts exceeding 3.3M.
Table 12:
Model
No
LR
WD
Timme
LR-NCE-EBM
[6,8)
2. 1010-
(10-b 5.
10
10 1010) (750, 1000, 1250)
[5 10 10 3]
(0.45,0.3
LR-NCEENF
(3.4)
10-3
10-6
Training. We use a train batch size of 128 for the diffusion, NF, and I-R-NCE-NE models, and 64 for ER-NCE-EBM. For Diffusion, Diffusion-o, and NF, we train our models for 200,000 batch steps and use warmup cosine decay learning rate schedule with 5,000 warmup steps. For I-R-NCE, we use Touter 100,000 main loop steps; each loop consists of Tamm= = 4 negative sampler batch steps, and Truce = 2 R-NCE batch steps. Both models used a cosine decay learning rate schedule with a warmup of 1,500 main loop steps. We fix the number of negative samples for I-R-NCE to K = 31 and the number of interpolating times m 5. For I-R-NCE, we also leveraged the computation trick from Section 6 by relying on 150 time-steps with a Heun integrator for sampling from the proposal model, and sub-sampling at only 15 time-steps for computing the log probability via trapezoidal integration. The remaining of hyperparameters for I-R-NCE are outlined in Table 12.
Finally, we conclude by discussing the remaining set of hyperparameters for NF, Diffusion, and Diffusion-
Φ. For all three models we hold the following hyperparameters fixed: (a) WD = 10-6 (b) the number of
sampling ODE steps equals 512 (so 1024 total mumber of function evaluations), and (c) the number log probability ODE steps equal to 48 (cf. Section 6). For Diffusion and Diffusion-o, we use the default scttings
of Jaata; miin mmax and p from Karras et al. (2022). We also fix dzel 0.02 for Diffusion--. Finally, for NF, we fix = 2. The remaining hyperparameter values sweep over listed in Table 13.
46

---

Published in Transactions on Machine Learning Research (09/2024)
Table List of hyperparameters for NF, Diffusion, and Diffusion-o. The constant N, is described in Figure The parameter pert governs the data noising, discussed in the main text. All models considered have parameter counts not exceeding 3..3M.
Model
No
LR
pert
NF Diffusion
(4,6,8)
[10-4, 2 1..
(0,0.01)
Diflusion-d
C Proofs
C.1
Optimality
Before we proceed, we state the following claim which will be used in the proof of Theorem 4.4. Claim C.1. Let (w))-1 be strictly positive vector in R2, and let Sa-1 denote the (d 1) -ddimensional simpler on Rd. The function F(q)= -' 74 log qi uniquely marimized over Sa-1 by q = 7/ Σi= %.
Proof. Follows from standard Lagrange multiplier argument combined with the strict concavity of log the positive reals.
We now re-state and prove Theorem 4.4.
Theorem 4.4 (Optimality). Suppose that f is realizable. Then, for any & E E, we have that:
θ, argmax L(0,E). 8ee
(4.2)
Proof. As noted in Section 4.1, the proof strategy follows Ma & Collins (2018, Theorem 4.1), but includes the mensure-theoretic arguments necessary in our sctting. First, we note our assumptions cnsure that L(0, E) is finite for every 0,5 E E. Second, since the index is does not affect the value in (3.8), we can take the average
K+1
L(,,E)
Σ
Ey log q x,
K
K+1
We now use the existence of regular conditional densities perform a change of measure argument follows:14
K+1
L(O,E) K +
E2
Bylnee log (0 x. x)
Fx
log c (k| Ir, F)p(uk x:) II Pe(us (x) dyn:
K+1 1
K+1
pluk xx)
K+1
E,
log g0 x)
II
Pe(us dy1:k41
K + 1
Pe (y
x))
f=1
Pluk x)) Pe (vv x))
E R K +
log q8.g (k: v)
(C.1)
As a reminder, the notation y pK+1 indicates that the K + 1 samples in are drawn (conditionally æ) from the K + 1 product distribution of pe(: (: x). Now, the term in the brackets is the (negative,
Note that in the integral expressions that follow, we slightly abuse the mensure-theoretic notation to imply the equivalence d ny (K+1) (dy). (d)).K41): where byy x(K+1) is the product measure.
47

---

Published in Transactions on Machine Learning Research (09/2024)
un-normalized) cross entropy between the following distributions over {l,... K + 1}: the (un-normalized) distribution p(aulz) K++ and the distribution (ko.g(k x. -
Pp (a) Jeml
Let 0. E O. By the definition of realizability, we have for a.c. (x:
plyk 2) 2)/pe(U x)
vk (L...., +1}.
2j=1 P(yj 2)/pe(Us) x)
Thus, we see by Claim C.I that the choice of 0+ O pointwise maximizes the cross-entropy term in the brackets in (C.1), and hence θ S argmaxgee (),
Now, suppose that 0 E argmaXace L)..g). Let 8, O, be arbitrary. Since L(O,E) L(Q..E) and since 0, pointwise maximizes (C.1), by Claim C.I we must have that for a.e. ():
00 (k: x, = ek x) Ak (1..... K+11 1}.
(C.2)
Note that when (C.2) holds, have:
exp(Eg(=r,))) exp(Eo, (z, ykk))
Eves exp(E)(xx,)) 4) logge (v/ x)) Ever exp(Es, (z, - logpe 2))
4k: E (1,..,,,+11) ++1
K
Because the expression on the RHS above is independent of k; then:
exp(EE(I, 9k.) Eg(x, mt)) = exp(Es, (I, 8k) Eg, (x, vk, (I....,,K Thus, defining the set A == ((x,)) (C.3) holds), have:
K K+ +1..
(C.3)
(μx x My x(K+1) (A) =1.
(C.4)
We now define the following set on X x Y2:
B = ((.,v1)) exp(Ex(2, v) - Eg(2,8') + exp(Eo, (2, (2,
Eg. (I, }))}
Notice that (x. y, v/) E if and only if (x, jr,y.1l of UYK-1 E A Thus, letting e= I d1z (here we use the finiteness of the measure space), it follows from Tonelli's theorem and (C.4):
(ux 1 (B) / 1((0, E d(ux x n39)
cc- /V 1((2,01 B) d(ux x 189) duz 11)
-c- fI((E.0n E B) d(ux x My x((--))
=c- /1((E,S) e d(ux My x(K+1)) 0.
Next, we have that μx x py measure of the set.
B' = (x, v) the cross-section (², == (v/ (I,9.v) B} is not py-null)
is also zero, due to the fact that 112 and 122 are measures, then cross-sections of v1 1/2 null sets (holding the first variable fixed) are D2 null sets, in almost everywhere (see e.g. Folland, 1999, Ch. Ex. 49). Hence
48

---

Published in Transactions on Machine Learning Research (09/2024)
for any (z, v) (B'):
exp(Ex(x,)) y))
Pz(y x))
exp(Ex (z, v²) My (dy'))
Texp(E((2, v/ Eg(o, u)) po (dy')
1
[B exp(Ee(0, 1/) Eg(x, (dy')
JB exp(E0, (x. v/) - Ea, (z. u)) by (dy')
exp(EE, (z.v) Eo. (x. v)) (dy)
pe. (y z).
This shows that po(v xx) po, (y a)) holds μox jo almost everywhere, and hence 0 O..
C.2 Asymptotic Convergence
C.2.1 Consistency
In order to prove Theorem 4.8, we require the following technical lemmaa.
Lemma C.2. Let (2(r) denote the closed Euclidean ball centered at the origin with rodius r; the dimension will be implicit from context. Let θ be compact set, and let f + IR be continuous function. Denote the marimization set.
e. = { 0 ((0) sup S(0)
ø'ee
Suppose that O, is contained in the interior of O. There erists a Eu > such that for all 0 < & Eo.
sup (e.(()(e.-B.(e))
f(0) f sup f(0), #E0
where cl(.) denoles the closure of sel.
Proof. For 0, define == cl(@) \ (0, + B2(s))). Since O is assumed be in the interior of O, there = exists an Eo 0 such that for all & & Eo. O is non-empty. It is also compact by boundedness of O. Let {0x}s k be sequence within O such that f(0,) supaeor f(0). Suppose for contradiction that supsee: f(0) fr. By compactness of θ there exists convergent subsequence {0,, h with limit # e By continuily, [(0F) = f", and thus 0= E e.. But, this means that gee,, is non-emply, contradiction. 7
We now restate and prove Theorem 4.8.
Theorem 4.8 (General Consistency). Suppose that I is realizable, O is contained in the interior of O and Assumptions 4.1 and 4.2 hold. Let 0,, argmaxage L., (0, denote arbitrary empirical risk marimizer from n samples. Then, d(Om, O.) "= almost surely-
Proof. Let 0, E O. be arbitrary. For E θ, define the function g(0) as:
g(t) sup [L(,.E) () L(0., EEE
By the continuity of L and the envelope theorem, ((0) is continuous. Also, have:
sup (0) = sup sup [T((.E) - L(O,, { = sup supsup [T(0,E) L(O,, E) = 0, Çe 0E0
0ee 0 feB
49

---

Published in Transactions on Machine Learning Research (09/2024)
by Theorem 4.4. By Lemma C.2, since θ is contained in the interior of θ there exists Eo 0 such that for all 0 < & & Eo:
E cl(@) \ (0, + B2(e))) - -> < 0.
Next, by optimality of On, L. (0n. En) 3 Ln (0,, En). Thus,
I(0,, En) L(On,Ea) = [(0,, En) In (on En) + In Ln( En) - I(ōm, En)
< L(0., Em) Ln (0., En) Ln (6n, En) L(O,,, (n)
< 2 sup AEA,EE
[L(0,) (0,
Therefore, we obtain the implication:
d(Om, .+) g(0m) <0
L(On, En) - L(... <0
2 sup [L(,E) (0, > O.
DEA,EE
Now, let En) be as indicated in Assumption 4.2, and define U(0,E) = I(O,E) I(0o, So) and similarly (,, (0, E) L,.(0: Eo). By triangle inequality:
sup (0.E) - L(0,5)) < sup 0EA,ÇEE
(0, - (0, |L, (0.: {o) - L(Go.
OEA,EE
Hence, we have the following event inclusion:
sum 0.) > >
n-000
20)
c lim sup sup n-000 Og0,
(0, - U((,5)) -(0
lim
sup sup IT, (f6, L(fo, Ea)l > 0f. R=+00
By the compactness of θ x E, continuity of (0,E) + lo (x, y), and Assumption 4.2, the first event on the RHS has measure zero by the uniform converge result of Ferguson (2017, Theorem 16(a)).
On the other, the second event on the RHS also has measure zero by the strong law of large num- bers (SLLN), which simply requires the measurability of lø.s and the existence of (fu. Hence, [im sim-sup d(O,, O.) > 2 measure zero set for every E 0 sufficiently small, establishing the re- sult.
C.2.2 Asymptotic Normality
The following technical lemma will be necessary for what follows.
Lemma C.3 (Symmctrization lemma, (cf. Ma & Collins, 2018, Lemma B.4)). Let M(x, y) be a measurable function. For any & E (1,..., + 1}, define:
Σ
go x. x) vi)
For any 0, E O,, and & E, we have:
Zn (O. E...y (M
50

---

Published in Transactions on Machine Learning Research (09/2024)
Proof. Wlog we take & = 1. By symmetry:
(0
K
K+1
K+1
plue x))
K +
En R8
glane
R++
Σ
q0,c xr M)
P xx)
K Σ
K
pluk xx)
Ex pk
ς x, F. 31)
K +
Pe (B1 x))
im1
For any E O.:
plys x)) (yi (x)
90 x,,))
Σ
K+1
k=1 plux p)/pp(uk x))
Hence:
K+1
plyA x))
(M) = K +
B₁ By +++1
Σ
Max ys)
glank
pe(u x))
K-1
Σ
E, mpt (M) (x,
K +
1=1
(Mr,)))
We prove Theorem 4.10 in two parts, starting with the following proposition.
Proposition C.4. Let 0, E O, and & E E. Pul 7. = (0,, We have that:
Var, Vol,, (z)] = -V)LL(,,
Proof. We first observe that since E. Lo (z)] by optimality, the following identity holds:
Var, [Vol. E. [Veg, (z)Vol,
Now, define the following shorthand notation (with arguments x and y implicit):
ge = Vačo, (x, yk), gk == E...e. := (k x),), Ee (l,..., 1}.
we have by symmetry and the definition of ql:
K+1
K+1
K+1
p(yk xx)
Eg Eyln.s.
Σ
9.9.9h
T
Eg E 3
pp R+1
49i9h
K +
yizi
Ps. (3) xx)
K+1
K+1
K
plus x)) Pe, Cur xx)
Eg E
PR
Σ
Ai:g qk
T
K
k=1
K+1
p(uj xx)
K+1 K+1
E. E
pk
++
9.gi qk
K
Pe. (y)
(x)
1=1
K+1
K+1
Eg E
Σ
K +11
A=
K+1
K+1
R. E.,
Σ
E Egl
Σ 9:0i
$=1,=1
51

---

Published in Transactions on Machine Learning Research (09/2024)
where 602 vvT denotes the outer product of a vector With this calculation, we have that:
Var, (Vll,(2))
( (VoEo, (1. (1, - Σ
K+1
By
go k. (i x. V)V (2,
K+1
827
VVEo, VoE,,
(x, -y
Σ
2 Vo,,
q0 go. (i
(C.5)
1=1
It remains to show that the expression (C.5) is equal to VgL(O. E.), from which we conclude the result.
To do this, we next recall the following identity:
Vkge x) q0,e (k: x, v, log ge (k x,
Hence for any 1o E {1..... K 1), letting 0o(.) denote the Jacobian of the argument,
0v(go,5(1: (A: x, V) ya))
= [o.g(k x. })
Ve mb) + VoEg( Vo log go (k y1
K+1
c0 x,
(zE + Vo
Σ
90.5 x, (o)
Recalling the following expression for the gradient VoL(O,E):
voL(0,E) = Ex,, VoEo(x, v) EyEy
Σ g0 Vo)o((()
we have that the Hessian (7.(0,E) is:
K+1
V3I(O,E) = EyRre VgEo(z, Me)
Σ go x2 (2, + Vo
K+1K+1
+ F2 E F
Σ
Σ
go x. Ei0 (2
K+1
Eg Eyng
ViEo(z, vk)
Σ
go,r x, ( Vo.5
K+1
EyEg
Σ
(
Let simplify this by evaluating the Hessian at (8, (.., Applying the symmetrization lemma (Lemma C.3) and the identity (C.5), have:
K+1 Σ
(2I(O..).) -Ez,y Eo (z. VoE, (z, 1 B,
go..e (i 2. V (x, vi)
(C.6)
im
=- (z)}
We proceed to restate and prove Theorem 4.10.
52

---

Published in Transactions on Machine Learning Research (09/2024)
Theorem 4.10 (Asymptotic Normality). Suppose that F is realizable and Assumptions 4.1, 4.2, and 4.9 hold. Let [0n] be sequence of optimizers as defined via (4.5). Denote the joint parameter 7 = (0, and joint r.v. v), and assume the following properties hold:
1. The negative proposal distribulion parameters {En} are vn-consistent about a fixed point En, i.e., vnll. {x) = Op(1).
2. The set O = {0.} is a singleton, and 7, = (0.,5.) lies in the interior of x E.
The Hessian of l is Lipschitz-contiinouut i.e.,
(z) ( < M(2)lla -ell, Ez [M(3) < co, Vn, E x E.
4. The (0, 0) block of the Hessian V2,(00. E.), denoted (2,,0,, satisfies ViL((0,, <
Then, vn(O, 0.) ~ NNO, Vo), where
Ve -VLL0..E))-- [3g (z)]
Proof. We first require the following remainder version of the mean-value theorem, which is a consequence of Taylor's theorem.
Claim C.5. By the C² smoothness and Lipschilz Hessian properlies of + En(2), for any E X Y, and 72 E θ x E, it holds:
(x) V.L,, 7) R(72 7))
where R is a remainder matriz dependent upon 71, 72, and satisfies: || < M2)]ll2 -
Applying Claim C.5 with (71, 72) = fm): where in == (On En), we obtain:
[oj Ojel
Vola, (2) V V (2) + (=) 0.) + (x) - -E.) +
Rlf 7x),
where Ia and Os are the identity and zero matrices of dimensions x respectively. Now, since 6, converges to 0, a.s., for n sufficiently large, the optimizer 0, argmaxage E also lies in the interior of O and thus, Vo 0. Applying the empirical expectation all terms above yields:
0 = (2
V (=)) + B. (3), 0.) ( Ex)
B T [R(in, Yo, 2)] - %-))
Now, notice that:
1
n
[ (R(m,, Th, =)]llop
R(ans Tap 2(0)
R(fm, 7o, =(0) llop
n
11
n
M (z) 1l - 7- = o1()).
2
40
[M(
53

---

Published in Transactions on Machine Learning Research (09/2024)
Then, the remainder term takes the form (,( 8.) + (p())(En [.). Thus,
0= E=EVVV (=)
(0,, 0.) (, (=)] + +zx(1)) (En -x)
VolVoly (z)] [ (3) 1. (2)] ()())) (0, -0.) E. E c, (2) (B. 2. (z) + (p(1) (En (.).
E.)
By the SLIN, (F.- - -..][ (=)] a.c 0 and (E_ - R.) R.) Lg (=) - and thus are both op()). Fur- thermore, notice from Theorem 4.4 that since 0, maximizes 1 + L(),E) for any fixed &, it follows that V vig ((0.. = E, I 85 G. (z) = O. Thus,
(z]]+ (0, - 0.) + op(1)(En E.) Vn (=)] vn ( Vgl,,
E. Vol. (z)]
oy))) (0.. - 0.) + vn (-(1)(m--
Now, from Theorem 4.4 and Assumption 4.9, we have that E. [=). = 0. Thus, by the CLT, (Vo,, (z)] (O, Var, (=),, Next, by the hypothesis (1) that vn(h. E.) is Op(1), the last term above is (1). Finally, by hypothesis (4), for sufficiently large n, E. (V36. (z) + +1) is invertible. Thus, applying Slutsky's theorem yields vn(0n, 0+) ~ (O,, Vo), where:
V = E. Vol,, (z) Var, [ol,, (z)] Es (z)
Applying Proposition C.4 gives the desired result.
C.3 Learning the Proposal Distribution
We now restate and prove Theorem 4.12.
Theorem 4.12. Consider the selling of Theorem 4. 10, and suppose that pe, (y x)) p(y x). Then, the asynptotic voriance matriz Vo
Vo 1 K [T xx)
Proof. From (C.6) in the proof of Proposition C.4, we have the identity:
K+1
002
(.L(0., EE, VoCo, (z, VoEa (z.z) EyEg,
Σ
q0, & (i (z,
(C.7)
With the assumption that 0, E O. and Pe, (y xx) = p(v x)), make two observations:
Pylzre = -p( s and
2. e. (i x. y) KtI for all E (1,... ...... K 1).
Hence, we have the simplifications:
K+1
827
Eg Emne.
Σ
go..s. (i x. VoEs, (x, (x.
1
ET[[(O. [)] + E.ryIVVoS (a. 82
K
From this, conclude that:
VGI(O..E+) K [ |x)] [T [x)] = =- K E. [(0,, [x)]
From Theorem 4.10, the asymptotic variance is Vo - [V. from which the claim follows. 7
54

---

Published in Transactions on Machine Learning Research (09/2024)
C.3.1 Adversarial R-NCE
To prove Proposition 4.13, we require the following technical lemmaa.
Lemma C.6 (Saddle Optimality). Fir an integer E N+ and E .l.....d.. Define the function Fj. Rd 20
+ R as:
(17 (n) (11 1)
P&(n) = log
(n, Ex)
The function Fj is convez on the domain R
Proof. A straightforward computation shows that:
F=))())= (x) ( ek
( ck 1
(() (n,
(n, ck) (n,1)
Since 7) E Rd So: have (n, ex) and hence (² > 0.
We now restate and prove Proposition 4.13.
Proposition 4.13. Suppose that f is realizable, and furthermore, the family of proposal distributions Fn is realizable. That is, there erists some non-empty E. C = s.t. for every & E E., Pe. (y x) ply 2) for a.e. (2. Then:
1
max min L(O,E) = L(O.,E.) = log K+1
(4.8)
dEB fEB
where (0., -0, x E..
Proof. Fix (1...., K 1}. From the proof of Theorem 4.4, we can write:
Po. (uu (pxll)/ppc(ulll)
L(0.,E) E, y.uKK
log
po. (ux (
sizit
Σ
K+1 j=1 Po. (V, (x) p
Define the mapping 17 (x, -
30 RK+1
as Mkk = po, (ux (p Notice that
Eg
E Ey. slaee, [1] 1.
Thus, by Lemma C.6 and Jensen's inequality:
L(O,, = E. p 9/mls
(a(1((,))) 8)) > Fx(1) = log
K +
On the other hand, since pe, (u/xr) = Ps. (y/x) = p(y(a), L(O.EE.) = log {1 chain of incqualities:
So, we have the following
K+T
1 x+1 L(O.,E)) min L(O,, < max min L(O,E) < max L(0, = L(O.,E.) = log K x1+1
log
EEE
050
K +
The last inequality above is a consequence of Theorem 4.4. Thus, we conclude that O, E, a saddle-sel of the game:
max min L(0,))- 9EB gez
Proposition 4.16. For any Stackelberg optimal pair (0, E) the population game (4.8), holds that 0 E θ. Conversely, for any O, O., there exists an E E such that (0.. {) is Stackelberg optimal for the population game (4.8).
55

---

Published in Transactions on Machine Learning Research (09/2024)
Proof. We first prove the forward direction.
By definition of Stackelberg optimality, it holds that:
inf EEE
L(0,g') > inf L(0,5)) 40 E E'EE
Suppose that By Theorem 4.4, I(O.g)) 5)= T.(0g) E) for all { e Taking the infimum on both sides, we have that:
inf L(O., inf L(O,E), EEE
EE
where the striet incquality is preserved since compact and & +> [(),E) continuous for every 0 E O so the infinums are attained. But this contradicts the Stackelberg optimality of 0.
Now for the backwards direction. Consider the chain of inequalities for any e θ:
inf L(0,C) < sup inf L(0', (') < sup EEE
L(8", 5) = E()..E) AA E,
0'E0''ES E'EE
d'ee
where the last inequality follows from optimality of 0. for any &, as per Theorem 4.4. Now, again by continuity of & L(O.,E) and compactness of =, there exists & such that infee= L(O., {) L(O., E). Hence, for any (0,E) E e
inf L(,, E') < L(e..)) < L(O.,E), e'EE
which shows that E) is Stackleberg optimal for (4.8).
Theorem 4.17 (Adversarial Consistency). Suppose that Assumption 4.1 and Assumption 4.2 hold. Let (On, E.) E θ E be a sequence of Stackelberg optimal pairs, as defined (4.10), for the finite sample game (4.9). Then, the following results hold:
Suppose that O, is contained in the interior of O. We have that d(θn, θ.) n=20 0 a.s.
2. Suppose furthermore that Assumption 4.9 holds. Let E: = E L(#., infçrez L(0., {')) for arbilrary choice of 0. E O+ (the set definition is independent of this choice), and suppose = is contained in the interior of E. We have that d(En, "
0 ">>
Proof. Part (1). We prove convergence of 0, to maxmin Stackelberg optimal point for the population game, ie., global maximizer of the function L defined as + infçe= L(0,)) From Proposition 4.16, have that L(0, is continuous. Since Egt is compact,
this is necessarily the set O.. By our assumptions, by classic envelope theorems, we have that I is continuous. By Lemma C.2, since O, is contained in the interior of θ, there exists eo > such that for all 0 < & Eo:
0 E cl(@) \ (0+ + B2(e))) - I < sup I.". 8E0
Put L as the function + infeez Ln (0, E). Since (0z. En) is a Stackleberg equilibrium, we have that supace La L... Now, observe that if 0, E θ \ (e, + Bo(e)):
L < + 0 + L. 8n > 0
sup inf L(O sup inf (0, + inf (6., EEB
E) inf L(êm.. E) 0. teE
OEEEE=
seE
Note that for any two real-valued functions (())g(() and domain Z, it holds that:
sup f(2) sup g(2)
< sup lf(x) - g(2)).
sE2
2E2
56

---

Published in Transactions on Machine Learning Research (09/2024)
Thus,
sup inf (0, sup inf [(0,E) lcB feE
< sup 0EG
inf T (0, inf T.(0g))
DEESEE
SEE
EEE
sup 0
sup - L(0,
$E
SEE
sup
(0, L(0,
(C.8)
0EA,EE
Similarly,
inf Ln Ln(On, (0n,)) - inf L( SE=
sup KEE
- LL(On: [)] -supl
Lupl-Lm (0n. E)
$EB
EEE
Ln (0n. {) L(B.. E)1
supp EEE
sup
(m, L(0,))l.
A
This yields the implication:
d(On, 0.) >
sup
[0, I(0,E) > 0.
#EA.EE
Following the proof of Theorem 4.8, the conclusion of part (1) now straightforward.
Part (2). Let 0. O, be arbitrary. The function & + L(0., E) is continuous, by Lemma C.2, since E is assumeed to be in the interior of E, there exists En > 0 such that for 0 < e & Co:
cE (E B2(e))) L(O., L inf L(O., EEE
(C.9)
Let &g denote this gap:
5. =
BE+ I(0.,E) L"
inf
By Assumption 4.9, we have that L is C on x Hence, the following constant Br, well defined and finite:
By == sup BE0,EE
(v)((,)))]
Next, let 0, E e, satisfy ||0_ - - < d(On, 0.) + 6 chain of inequalities:
and let eE E be arbitrary. Consider the following
28
(L) - L(Q., 51 < E.) L(θm, &n) [L (n) - Ln (0n. En)-
+ |L, (n) L(0., [L(,g) - (8,
< sup ((. {) - L(8,,, {) + sup
EE3
BEB,EES
+ LL (On,En) ((0.
Ti + T2 T'a.
(C.10)
We first control Tj by uniform Lipschitz contimuity; in particular, by boundedness of V and the mean-value theorem:
\oL(0,))1100 - 0.11 = =Br1 - 0mll < Bidd O.) 8
Ti = sup [L(0__)) L((n,E)) < sup 0E0,CE
SEB
57

---

Published in Transactions on Machine Learning Research (09/2024)
Next, control T3. Since En). is Stackelberg optimal for the finite sample game, we have L, (mn supage infçez Ln (). Furthermore, by Proposition 4.16, the pair (0., is Stackelberg optimal for the population game, and hence L(O,, {) = supvee infçe= L(0,E) E). By (C.8),
T3 = Ln (On. )m) - E(0.,E) = sup inf (0,5) sup inf L(0, IEe EEE
eee EC
sup 0EE,EE
[( {) - (0.5)) = T2
Combining the bounds Ti and T3 with (C.10):
[L(O,, E.) - I(Q., < 2T: + Brd(On, 0.) 6z
This bound with (C.9) yiclds the implication:
d(En. - = 8, < L(O,, En) - L'
= 6. < 2T, + Byd(dn,, Q.) + 8c
2
8g < 412 + 2B O.).
Hence, we have:
[L(0,) (0,))) > of U lim sup d(O,, 0,) 0.) 20) 12+00
lim lim B) {) n-000
lim lim sup sup n+000 8EA
Part (2) now follows, since both events on the RHS were established in part (1) as measure zero events, and since E > is arbitrarily small.
Theorem 4.18 (Adversarial Normality). Suppose that F realizable and Assumptions 4.1, 4.2, and 4.9 hold. Let In = En) E e be a sequence of Stackelberg optimal pairs, as defined in (4.10), for the Snite sample game (4.9), and define E as Theorem 4.17. Assume then that the following properlies hold:
1. The sets 0,, are singletons, i.c., O = =(-}, and E = ={}, and 7+ = (0.,E) lies in the interior of x E.
2. The Hessian of ly is Lipschittz-comtinnuus i.e.
(2), (2)loo < M(2)llm- - rell. E, [M(z)] < co, Vn, E
3. The (0,0) block of the Hessian (O,, denoted (8.( satisfies VOI(O., 0, and the (E,E) block, denoted (),, satisfies (2.(0,, ~0.
Then, vilffn 7.) N(0 Vg),, where
V = E. V (x),
Var, Vyl.( Ez (x)
(4.11)
Proof. Note first that consequence of Theorem 4.17, An "=3 % almost surely. We now proceed in manner similar to the proof for Theorem 4.10, with the following identity:
(=) (x) + (=) -w) + R(fn 77,);
where R a remainder matrix dependent upon 2, Tx. 2., in and satisfies: || & 1.M()))l - 7+ll. Ap- plying the empirical expectation all terms, and as in the proof for Theorem 4.10, upper-bounding BBAR,, in o)plop by an (1) matrix by leveraging the fact that |l >lll 2$ 0, yields:
(5n 77,) = =B.IV,L,, (E.I33E,, + (É.
(x))
E.).V (z)] + oz)( (fn 74).
58

---

Published in Transactions on Machine Learning Research (09/2024)
By SLLN, (E,- - Ex), - 0, and thus:
E. L., - (BA. + oz) (fn --7)
To complete the result, we need only establish the following: (i) [V] =0, (iii) Es [V (z)] is invertible. The result then follows by applying CLT to IV (=)] and Slutsky's theorem (see for instance, the proof for Theorem 4.10). We show (ii) and (iii) first.
(ii) (V),,,() 0, and
For (ii), notice that by the definition of Stackelberg optimality, & minimizes & ++ I((O,,)) 5) over and by assumption, is contained within the interior of E. Thus, VeL(0., O. Further, by Theorem 4.4, 0. maximizes + L( over θ for any & E, and is also assumed to be within the interior of O. Thus, VoL(0.,,) = O. Leveraging C' re-uglarrr,
conclude that Ex[V,L,, O.
To prove (iii), notice that the Hessian (0,(0,)) is block-diagonal, since by optimality of 0, for any &, it follows that T O. Further, the blocks (O..(..))) and (0,,,, assumed invertible.
Finally, to prove stationarity for the finite-sample objective at In we first require the following additional result.
Claim C.7. || In Ex) - (2I(0,,) "3300 0 a.s.
Proof. Let d == (0) + dim(f), and let gd-1 denote the unit-sphere in Rd. We first claim that:
sup 0EAREE
(..))) - E0l " a.s.
(C.11)
To see this, by the variational representation of the operator norm (Rudin, 2008, Theorem 4.4), the empirical process above is:
2 ((V.,n(0.))-V2.10..)))))))) @EE "13"1
sup
Assumption 4.9 combined with Ferguson (2017, Theorem 16(a)) yields (C.11).
Next, by Ce-regularity, (2,L(0,E) = E. (2), and therefore by Jensen's inequality and the Hessian Lips- chitz assumption:
-2 = -1 - -, (2)llop
< 1E.1V (2) -3G,, (z)llp
< (z)(i -Vlll-
Finally, then
(On, - V3, & V. V.L(O,, lop lop + - -V,L(,.,)) | L.. (0. E) 10, llop + E.M( 7all-
sup BEB
Then, (C.11) along with fn " 7. a.s. proves the stated claim.
Returning now to (i), since we only assume that (0,, Ex) is a Stackelberg optimal point, it may not satisfy the stationarity condition B. (=)] = 0 (Jin et al., 2020). However, since é. globally minimizes & + (On. E) and É. + & a.s., which is at interior point, for large enough we have VeL, (Ō,, E.) =
Furthermore, by Claim C.7 and the assumption that (LI( is non-singular, for large enough n we have that (On, E.)) is also non-singular. By the implicit function theorem, there exists a contimuously diffcrentiable functioo G. A E e E, defined for 0 in an open local neighborhood of Ô., s.t. G..(O.)) E, and (eLn(0, G))) = 0. Then, since ò, is a global interior maximum (for large enough n, since 0, - 0. a.s.) of 0 + infgez I, (0, (8, which in turn equals (0, G., (0)) for A near 0n,, stationarity dictates that (o (θ.)) O. It follows then by the chain rule that Vo (On, 0.
59

---

Published in Transactions on Machine Learning Research (09/2024)
Corollary 4.19 (Marginal Normality). Under the assumptions of Theorem 4.18, the parameters (0,, satisfy vn(0n, - 0.) NO, Vi'), where
Proof. From the proof of Theorem 4.18, the Hessian (0,, is block-diagonal since the cross-term ((0.5)) =. Thus, the (8, 0) block for Vi in (4.11) reduces to:
E. Vz (2) Var, V
Applying Proposition C.4 gives the desired result..
60


Reviews:
[{"review_id": "yMQoaAVOUR", "reviewer": "Reviewer_jagM", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n- The presented theoretical comparison between IBC and R-NCE loss is of high interest to understand the limitations of IBC and presents the clear benefit of using R-NCE instead.\n- The authors show the usefulness of Energy-Based Models as Policies if the selected training loss is proper.\n\n**Weaknesses**\n- Some of the authors claim require additional details to be clear.\n  - In page 1, 'EBM yields compact representations; model size is an important...'. Nevertheless, the authors are not only training EBM, but also a contrastive model, such as a CNF that would still be large. \n  - In page 2, authors claim as contribution 'EBMs are compatible with multiple noise resolutions'. Nevertheless, in the paper 'How to Train Your Energy-Based Models' (https://arxiv.org/abs/2101.03288) it is already suggested this possibility, that authors fail to cite.\n- Despite the authors explore the application of EBM's as policies for robotics, there are no experiments in a high dimensional robot.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors propose the use of R-NCE as a training algorithm for Energy-Based Policies. The paper claims the limitations of the training algorithm of Implicit Behavioral Cloning (IBC) and show that those limitations can be solved using R-NCE. \n\nAdditionally, the authors introduce I-R-NCE. Inspired by the multiple time steps based generative models, this work extends R-NCE to train Energy-Based Models on multiple timesteps.\n\nThe paper presents a large theoretical evaluation comparing the training loss of IBC and R-NCE to validate their claims and show empirically the application of R-NCE Energy-Based Policies in a set of Behavioral Cloning experiments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Regarding the contributions of the paper, the authors should check the highlighted weaknesses and fix them accordingly.\n\nRegarding the experimental section, I suggest the authors to evaluate their proposed algorithm in a challenging robotics experiments such as Robomimic (https://arxiv.org/abs/2108.03298).", "broader_impact_concerns": "There are no ethical implications in the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "sQPMny4Bpe", "reviewer": "Reviewer_i6Db", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "STRENGTHS\n\nThe paper presents several theoretical contributions, which allow notably to better understand previous work (IBC). Even if some of them (e.g., Theorem 4.4) are extension of existing ones, their presentation and discussion in the context of behavior cloning seem to be quite novel. Moreover, the proposition about interpolating EBM sounds promising and could lead to interesting novel developments.\n\nAlthough the paper could be more self-contained and is not always easily accessible to non-experts, I find the exposition to be well-organized and well-written, with informative side remarks and connections to other results and works. \n\n\nWEAKNESSES\n\nThe empirical validation seems to be quite limited. For instance, it is not clear to me how sensitive the experimental results are to the various chosen hyperparameters.\n\nThe paper could be written to be more accessible to non-experts.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper presents several theoretical results regarding energy-based models (EBM) in the behavior cloning setting:\n\t- The authors demonstrate that a previous method (IBC) based on InfoNCE may be not consistent, while their proposed method based on R-NCE with a learned negative sampler is.\n\t- Inspired by diffusion models, the authors propose a method to train EBM models with multiple noise resolutions.\nThe propositions are experimentally validated on three domains.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. For better legibility, I would suggest the authors to explain right away why it's useful to define the R-NCE objective as in (3.8), which looks a bit strange since the LHS doesn't depend on k, but the RHS does.\n2. In Remark 3.2, could the authors explain why Lines 11 and 12 can be performed independently? It seems that Line 12 depends on Line 11.\n3. For the path planning domain, although the authors seem to follow the previous work by Urain et al. I was wondering why c_obs (which has a square root) is not in the same scale as the other terms, which are all squared distances. Could the authors comment on that?\n4. I like the proposition of interpolating EBMs. However, could the authors discuss more clearly about the possible advantages/disadvantages of this model compared to alternative models like diffusion models or stochastic interpolants?\n5. Could you comment about the sensitivitiy of the results to the hyperparameters?\n\nSome minor issues:\nPage 3: that easier\nPage 4: to the their\nPage 5: I find the notation to define \\bar y not so great, since it mixes set and vector notations.\nPage 10: suceed\nPage 10: Please recall the definition of || . ||_op\nPage 12: Is V_R defined?\nPage 15: Stackleberg", "broader_impact_concerns": "Not applicable", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "RKcQM0GFOR", "reviewer": "Reviewer_A1FS", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n- The paper is very well written and structured. The papers start with a general explanation of why EBMs training is so challenging and highlighted shortcomings and theoretical flaws of prior training methods before introducing their own contributions. \n- The proposed algorithm and novel model class claims are supported with detailed derivations and illustrative experiments to provide strong theoretical and empirical evidence\n- Detailed experimental evaluation in different benchmarks with relevant baselines that verify the theoretical contributions with additional empirical evaluations. The propsed EBM policy does improve upon prior EBM policies by a significant margin in all experiments\n\nWeaknesses\n - given the amount of experiments and evaluations in the experiment section, I found it hard to follow the conclusion of each section. Adding some more structure and and overview for the reader would be valuable.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "- The authors provide a detailed analysis and theoretical explanation as to why Energy-Based Models (EBMs) trained  as conditional policies p(a|s) with the InfoNCE loss from Implicit Behavioral Cloning (IBC) do not converge to optimal solutions. This is an issue observed in many prior work and limited the general usage of EBMs in the context of robot learning.\n- The authors propose a novel training algorithm for EBMs called R-NCE, that uses ranked noise contrastive loss together with a negative sampling model. The loss for training conditional EBM policies improves several flaws over prior methods and achieves better convergence and performance. This approach addresses the shortcomings of traditional EBM training techniques by improving model training robustness and effectiveness.\n- A new class of interpolant-EBM models is proposed, combining elements from diffusion models with EBMs under the R-NCE training framework, that show improved performance in challenging experiments. \n- The contributions and claims made in the paper are thoroughly validated through a series of toy task experiments and extensive evaluations on robotic learning benchmarks with competitive baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- add a general overview at the beginning of the experiment section to summarize the following experiments \n- ablation experiment in one of the robotic task to ablate the negative sampling method and it’s impact to the performance and other design decision \n - A table of content for the long appendix would also improve readability \n - A reference training curve and stability would be very insightful for the readers. Chi et al., 2023 contains a figure to showcase the evaluation results over training epochs for IBC and Diffusion Policy, which compares diffusion policies and ebm rollout performance over training time. Adding a similar plot for R-NCE would further strengthen the claims regarding training stability and provide new insights. \n - The papers claims, that learning multimodal policies is crucial for imitation learning, but no experiment quantifies the diversity of learned behavior of the proposed model. Some experiment to analyze that property compared to prior policies would be interesting. Recent Imitation Learning benchmarks like [1], that contain a entropy measurement to quantify multimodal behavior, would be suited for that.\n  - Prior work such as IBC reported difficulties in using Transfomer architectures as an EBM architecture given the training instabilities of EBMs trained with InfoNCE. Thus, all prior work only use MLP models as an EBM backbone, which was another major drawback of the default IBC framework. In my opinion, this fact that R-NCE EBMs can be trained with a transformer backbone effectively is worth highlighting somewhere. \n\nAll requested changes are not cruical for my acceptance but would further strengthen the contributions of this work. \n\n[1]: Jia, Xiaogang, et al. \"Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations.\" _The Twelfth International Conference on Learning Representations_. 2023.", "broader_impact_concerns": "The work tackles improving Generative models in the context of robotics, this could be used for potential harmfull task such as:\n- improved generative models such as an EBM have the potential to be used to generate images and videos, which can be used with negative intent to create harmful content \n- Training robot policies, that act autonomously in the real world could potentially harm humans without safety measurements and controlled testing.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission investigates the use of energy based models for behavioral cloning. Based on the observation that training EBMs using the InfoNCE loss is biased (unless when using uniform negative samples, which are ineffective in practice), an alternate, unbiased training objective based on ranking noise contrastive estimation (R-NCE) with a learned negative sampling distribution is proposed. Furthermore, inspired by the success of diffusion models, the paper proposed to train an interpolating EBM (I-EBM) at different noise levels. For that variant, the EBM gets the noise level as additional input, and a continuous normalizing flow is used for negative sampling with the corresponding noise level. The proposed method is evaluated on low-dimensional toy tasks and on higher-dimensional motion planning and block pushing tasks that were used in related work. In these experiments, training EBMs with R-NCE was shown to clearly outperform the INFO-NCE loss, reaching the performance--or even slightly outperforming--diffusion models and normalizing flows.

As the reviews did not raise major concerns that would need to be addressed, I recommend acceptance as is. Yet, I encourage the authors to add the additional ablations in the block pushing environment, which they mentioned in their rebuttal, to the camera-ready.
Decision: Accept (poster)