LEO: Learning Energy-based Models in Factor Graph Optimization
Abstract:
We address the problem of learning observation models end-to-end for estimation.
Robots operating in partially observable environments must infer latent states from multiple sensory inputs using observation models that capture the joint distribution between latent states and observations.
This inference problem can be formulated as an objective over a graph that optimizes for the most likely sequence of states using all previous measurements.
Prior work uses observation models that are either known a-priori or trained on surrogate losses independent of the graph optimizer.
In this paper, we propose a method to directly optimize end-to-end tracking performance by learning observation models with the graph optimizer in the loop.
This direct approach may appear, however, to require the inference algorithm to be fully differentiable, which many state-of-the-art graph optimizers are not.
Our key insight is to instead formulate the problem as that of energy-based learning.
We propose a novel approach, LEO, for learning observation models end-to-end with graph optimizers that may be non-differentiable.
LEO alternates between sampling trajectories from the graph posterior and updating the model to match these samples to ground truth trajectories.
We propose a way to generate such samples efﬁciently using incremental Gauss-Newton solvers.
We compare
LEO against baselines on datasets drawn from two distinct tasks: navigation and real-world planar pushing.
We show that LEO is able to learn complex observation models with lower errors and fewer samples.
Introduction
We focus on the problem of learning observation models end-to-end for estimation.
Consider a robot hand manipulating an object with only tactile feedback.
It must reason over a sequence of touch observations over time to collapse uncertainty about the latent object pose.
A common way to solve this is as an inference over a factor graph which relies on observation models that can map between states and observations [1, 2, 3].
However, in many domains, sensors that produce observations are complex and difﬁcult to model.
Can we instead learn observation models from data?
Given a batch of ground truth trajectories and corresponding measurements, how should we learn observations models?
One approach would be to learn a direct mapping from measurement to state, for example as a regression or classiﬁcation problem [4, 5, 6, 7].
However, while easy to optimize given a direct supervised loss, this only minimizes a surrogate loss independent of the graph optimizer and is not guaranteed to minimize the ﬁnal tracking errors that we care about.
The other option would be to directly minimize ﬁnal tracking errors between optimized and ground truth trajectories [8, 9].
However, many state-of-the-art factor graph optimizers, e.g. iSAM2 [10], are not natively differentiable due to operations such as dynamic re-linearizations 1 .
In such cases, we are limited to black-box search for learning parameters which is very sample inefﬁcient [11, 12, 13].
Instead of differentiating through the optimization process, we note that what we ultimately care about is the ﬁnal solution from the optimizer, which depends only on the shape of the optimized cost function.
We would like a cost function that has low cost around the observed ground truth trajectories and high cost elsewhere.
This is precisely what energy-based models aim to do by
Code and supplementary material can be found on https://psodhi.github.io/leo 5th Conference on Robot Learning (CoRL 2021), London, UK.
[Table]
x gt ̂ x trajectories trajectories
[Figure]
shaping an “energy” function to be low around observed data and high elsewhere [14, 15, 16].
We cast our problem of learning observation models for estimation as energy-based learning.
We propose a novel approach, L earning E nergy-based Models in Factor Graph O ptimization (LEO), for learning observation models end-to-end with graph optimizers that may be non-differentiable.
LEO alternates between sampling trajectories from the graph posterior and updating the model to match the distribution of samples to ground truth trajectories.
We show that we can generate such samples efﬁciently using incremental Gauss-Newton graph solvers that create a Gaussian approximation for the Boltzmann distribution deﬁned over a continuous space of trajectories.
To the best of our knowledge, this is the ﬁrst paper to show that even non-differentiable factor graph models can be learned efﬁciently.
Our main contributions are:
A novel formulation of learning observation models in estimation as energy-based learning.
An algorithm, LEO, to learn observation models that minimize end-to-end tracking errors in factor graph optimizers that may be non-differentiable.
Empirical evaluation on two distinct tasks: synthetic navigation and real-world planar pushing.
Related Work
Filtering and smoothing based estimation.
Observation models represent joint or conditional distributions between states and measurements.
Early estimation methods utilized these models within ﬁltering contexts such as Kalman ﬁlters and EKFs/UKFs [17, 18, 19].
Recent methods have also looked at making these ﬁlters differentiable for end-to-end learning [20, 21, 22].
However, ﬁltering can be inconsistent for nonlinear estimation problems due to linearization based on past, marginalized states that cannot be undone [23].
Instead, a wide consensus amongst the simultaneous localization and mapping (SLAM) community is to solve these problems as smoothing or nonlinear optimization objectives [3, 2].
Unlike ﬁltering, smoothing allows dynamic relinearization of past states as well as exploits the inherent sparsity resulting in efﬁcient and more accurate solutions.
Factor graphs are an increasingly popular way for solving such smoothing problems [1, 2, 24, 4, 5, 25, 26, 27, 28], offering a ﬂexible modeling framework while being efﬁcient to optimize.
We use the factor graph based smoothing framework in this paper.
Learning for smoothing.
Typically, observation models used in factor graphs are analytic models deﬁned a-priori [29, 30, 31].
Recent work has looked at using learned models [4, 5, 6] or learned measurement representations [25, 32] either to be used independently or to be plugged into a graph optimizer.
However, these model parameters are learned on surrogate losses independent of the graph optimizer, and hence do not directly attempt to minimize ﬁnal tracking errors.
An alternative approach is to view observation model parameters as hyper-parameters of the graph and perform a hyper-parameter search.
This is usually done as sampling and derivative-free approaches using black-box solvers [11, 12, 13].
However, the high computational burden associated with sampling and evaluation can be prohibitive for large parameter spaces.
Differentiable optimization.
Instead of a search, one may learn these model parameters via a continuous gradient-based optimizer matching graph optimizer output against a desired solution.
Minimizing this loss would, however, require differentiating through the argmin optimization process.
One way to do this is via unrolling where the optimization is represented as a series of differentiable, gradient based updates [33, 16, 34].
Since smoothing is a non-linear least-squares optimization, it too can be unrolled in a similar fashion via differentiable Levenberg-Marquardt updates [35, 9] or differentiable Gauss-Newton updates [36, 37, 8].
However, vanilla unrolling suffers from some drawbacks, notably the learned cost function can be sensitive to the speciﬁc optimization procedure, e.g. number of unrolling steps [38].
Moreover, this approach cannot be readily applied to commonly used optimizers relying on non-differentiable heuristics such as those in GTSAM, SNOPT [39, 40].
An alternate to unrolling is to employ the implicit function theorem which states that an optimal energy function must satisfy ﬁrst-order necessary conditions [41, 42].
While this is promising for convex problems, for non-convex problems, one may recover an energy function with multiple minima such that the optimizer at test time may ﬁnd itself in a different basin from ground truth.
Energy-based learning.
Instead of differentiating through the optimization process, we note what we ultimately care about is the ﬁnal solution from the optimizer, which depends only on the shape of the cost function.
This is what energy-based models aim to do by shaping an “energy” function to be low around observed data and high elsewhere [14].
Energy-based modeling is a family of unsupervised learning methods popular in many structured prediction tasks such as inverse reinforcement learning [43, 44] and computer vision [45, 46].
Recently, a number of approaches have looked at training deep energy-based models [47, 48, 16].
The challenge with such methods is generating samples from the energy based model.
Methods either rely on MCMC sampling [49, 50, 51] which does not scale well with dimension or resort to learning a separate generator [52, 53, 15].
We propose a way to generate such samples efﬁciently using incremental Gauss-Newton solvers.
Problem Formulation
We formalize our problem of learning observation models under the general theoretical framework of energy based models.
We begin by describing the observation model, how it is optimized at inference time and ﬁnally how its parameters are learned from training data (Fig. 1).
3.1 Model
We deﬁne an energy-based model that captures dependencies between variables in a graph by associating a scalar energy E ( · ) to each conﬁguration of the graph.
Our goal is to model the likelihood of a sequence of latent states x given a sequence of measurements z .
We adopt a factor graph framework (Fig. 1(a)) for expressing this likelihood as,
[Equation]
where, x := { x 1 . . .
x T } is the sequence of latent states, and z :=
{ z 1 . . .
z T } is the sequence of measurements.
|| f k
θ, x ; z ) || 2 are the factor costs with f k ( θ, x ; z ) being the local observation model mapping a subset of states x and measurements z to a cost, and θ the learnable model parameters[-
-]
Z ( θ ; z ) = (cid:82) x exp( − E ( θ, x ; z ))
dx is the normalization, also referred to as the partition function.
Note that Eq.
1 is a special case of the Boltzmann distribution where the energy E ( · ) is a sum-ofsquares.
Notably, when f k ( θ, x ; z ) is linearized, E ( · ) is quadratic in x and subsequently P θ ( x | z ) becomes a Gaussian distribution.
This is the structure that Gauss-Newton solvers employ, and we will leverage this fact to efﬁciently sample trajectories as described later in Section 4.3.
3.2 Inference
At inference time, given a sequence of measurements z , we wish to solve for the most likely sequence of latent states x .
We formulate this as maximizing the log-likelihood of the model,
[Equation]
Substituting the model expression from Eq. 1, and dropping constants Z ( θ ; z ) and 1 / 2 , results in a nonlinear least-squares minimization of the form,
[Equation]
We use an efﬁcient online Gauss-Newton graph solver for this objective.
This enables both real-time inference as well as efﬁcient sampling for learning model parameters during training (Section 4.3).
3.3 Learning
At train time, our goal is to learn a model that explains training data of pairs ( x gt , z ) .
We express this as minimizing a loss L ( θ ) over the training data
[Equation]
where, { x igt , z i } ∈ D is a training dataset of ground truth trajectories and measurements.
There are various choices of losses L ( · ) for learning model parameters θ ,
Loss 1:
Minimize ﬁnal tracking loss.
A straightforward choice for L ( · ) is to directly minimize ﬁnal tracking errors between optimized ˆ x i and ground truth x igt trajectories, i.e. L ( θ ; x gt , z )
= 1 / |D| (cid:80) i ||
ˆ x i (cid:9) x igt ||
While this loss directly minimizes ﬁnal tracking errors, it may not be differ- entiable since the graph inference path from θ to ˆ x is not differentiable.
In such cases, we are limited to black-box hyperparameter search which is very sample inefﬁcient.
Loss 2: Minimize surrogate loss.
An alternate choice for the loss L ( · ) is to directly map measurements to ground truth states, for example as a regression or classiﬁcation problem.
Such a surrogate loss is independent of graph inference, and unlike Loss 1, is differentiable and easy to optimize.
However, this does not minimize the ﬁnal tracking errors that we care about.
Loss 3:
Minimize energy-based loss.
Finally, we consider a loss that is a function of the energy, i.e. L ( E ( θ, · ); x igt , z i ) .
Intuitively, this assigns a low loss to well-behaved energy functions, i.e. functions that give the lowest energy to training data of ground truth trajectories (correct answers) and higher energy to unseen data (incorrect answers).
This energy-based loss, unlike Loss 2, is highly correlated with the ﬁnal tracking loss as it shapes the energy (or cost) landscape so as to make the inference step return trajectories closer to the ground truth.
Moreover, unlike Loss 1, this loss is only a function of energies which is differentiable irrespective of inference .
Hence, we use Loss 3 and detail our approach in the next section.
Approach
We instantiate the broad framework of energy-based learning within the context of factor graph estimation in 4.1 and 4.2.
In 4.3 we show how incremental Gauss Newton approaches enable efﬁcient sampling from these energy based models.
4.1 Normalized Negative Log-Likelihood Loss
We deﬁne the overall loss as the negative log-likelihood of the data under the energy model 3 :
[Equation]
Substituting P θ ( x | z ) as the Boltzman distribution expression from Eq.
1 we have,
[Equation]
[Equation]
⋮ ̂ μ i , ̂Σ i θ k
Graph optimizer oeeeoô x i ̂ x i x igt
Match trajectories { x igt , z i } (parallel runs) ∇ igt ∇ iopt ,
Sampler ∇ 1 g t ∇ 1 o p t , ∇ | D | g t ∇ | D | o p t , { x 1 g t , z 1 } { x | D | g t , z | D | }
Training Dataset θ + ← θ − η 1 | D | | D | ∑ i =1 { ∇ igt −∇ iopt } ∇ igt ∇ iopt ⋮ ⋮ (Eq. 7)
Input Initial θ , Training data { x igt , z i } |D| i =1 while <!
converged> do for each training example i do ˆ µ i , ˆΣ i = argmin
[Figure]
end for
[Figure]
end while
The expression above has an intuitive explanation.
The ﬁrst term penalizes high energy over ground truth training samples.
The second term penalizes low energy everywhere else, thus acting as a “contrastive” term.
This has a desirable effect of globally shaping the energy function so that the minima lies at the training samples.
The contrastive term prevents the energy surface from becoming ﬂat and thus degenerate [14].
Unfortunately, we cannot directly plug in Eq.
6 as a standard machine learning loss (e.g. in PyTorch) to optimize.
This is because the second log partition term is over a continuous space of trajectories x , which renders exact computation intractable for arbitrary nonlinear cost functions [55].
Instead, let us take the gradient of this loss, and substitute P θ ( x | z i ) expression from Eq. 1,
[Equation]
We see the log partition term transforms into an expectation over the learned distribution.
The gradient expression above has a similar interpretation as the loss.
The ﬁrst term pushes down the energy of ground truth training samples.
Interestingly, the second term now pulls up the energy of samples drawn from the learned Boltzmann distribution.
If one is able to generate such samples, computing the gradient is straightforward.
Generally, sampling from the continuous Boltzmann distribution is intractable, relying typically on
MCMC methods.
But as we will detail in Section 4.3, we can indeed generate such samples efﬁciently since our graph optimizer maintains a Gaussian approximation of this distribution.
Additionally in practice, we found that introducing a temperature term T inside the Boltzmann distribution, i.e. P θ ( x | z ) ∝ exp( − 1 / T E ( θ, x ; z )) , was critical for controlling the variance of the samples.
Setting T = 0 corresponds to only using the mode of the distribution which in this case is the mean trajectory.
This corresponds to minimizing a generalized perceptron loss [14].
4.2 LEO Algorithm
Algorithm 1 summarizes the LEO algorithm, also illustrated in Fig.
It takes as inputs the training data containing ground truth state and measurement pairs, i.e. { x igt , z i } |D| i =1 , along with initial values for the learnable observation model parameters θ .
The goal is to converge to a θ that minimizes the log-likelihood loss in Eq.
5 over the given training set i = 1 . . .
|D| .
LEO invokes a black-box, non-differentiable graph optimizer iSAM2 in the loop to generate samples.
Within each LEO iteration, for every training example i , we invoke the graph optimizer to obtain a Gaussian approximation N (ˆ µ
i , ˆΣ i ) to the original Boltzmann distribution P θ ( x | z ) about its mode.
We then sample trajectories ˆ x i ∼ N (ˆ µ i , ˆΣ i ) from this distribution.
The next step is to compute gradients of the energy, i.e. ∇ θ E ( θ ; x igt , z i ) and ∇ θ E ( θ ;
ˆ x i , z i ) , at the ground truth x igt and the samples ˆ x i , respectively.
Note that the gradients here are only w.r.t.
θ , and not x , and hence do not require the graph optimizer to be differentiable.
Once we have the two gradients, we use the difference between them to update the observation model parameters θ .
These updated parameters are then used in the next LEO iteration.
4.3 Efﬁcient Inference via Incremental Gauss-Newton
The only thing that remains to be speciﬁed in Algorithm 1 is — How do we generate samples efﬁciently from the graph inference?
In Section 3.2 we noted that our inference algorithm must satisfy two requirements (a) real-time inference and (b) sampling from the Boltzmann distribution.
We achieve this using via incremental Gauss-Newton to solve the nonlinear least squares in Eq. 3.
Gauss-Newton proceeds by linearizing observation model functions f k ( · ) about a point x 0 ,
[Equation]
(8) where, x 0 is the linearization point, F k = ∂f k
[Equation] x [Equation] ∂x (cid:12)(cid:12)(cid:12) x 0 is the measurement Jacobian, and δx = x − x 0 the state update vector.
Measurements z and learnable parameters θ are treated as constant within the graph optimization step.
Substituting the Taylor expansion in Eq.
8 back into Eq. 3,
[Equation]
where, A , b are constructed by concatenating all F k , f k ( θ, x 0 ; z ) together into a single matrix and vector respectively.
We then iterate x 0 ← x 0 + δx ∗ until convergence.
We can solve for the objective in Eq.
9 efﬁciently by exploiting two properties of the problem: sparsity and online measurements.
For SLAM problems, both A and A T A are large but sparse matrices, and can hence be factorized efﬁciently using sparse QR or Cholesky factorization respectively [56].
However, this still does not guarantee real-time inference during test time since A grows with measurements over time, making the method increasingly slower.
Instead, subsequent incremental solvers [10, 57, 58] reuse matrix factorizations from previous time steps.
In this work, we use iSAM2 [10], an incremental Gauss-Newton graph solver whose resulting MAP solution can be expressed as a Gaussian probability distribution with mean and covariance,
[Equation]
This is the Gaussian approximation of the original Boltzmann distribution P θ ( x | z i ) which we can now sample from efﬁciently.
The graph posterior covariance ˆΣ for such a Gauss-Newton solver is obtained using just Jacobians A without needing to compute a Hessian [2].
Note that sampling ˆ x ∼
N (ˆ µ, ˆΣ) is done at the tangent space of the linearized graph, i.e. ˆ x ← ˆ µ ⊕ δx where δx ∼ N (0 , ˆΣ) and ⊕ is a retraction [59].
The Gauss-Newton method for inference is a well-known technique for sampling from intractable models.
A recent overview of the Laplace approximation with GaussNewton method can be found in [60].
Here, we build a speciﬁc Laplace approximation using the incremental solver iSAM2 that scales to sampling efﬁciently in real-time from very large graphs.
Results and Evaluation
We evaluate LEO on synthetic navigation and real-world planar pushing.
We compare against baselines on metrics like ﬁnal tracking error and sample efﬁciency.
We implement LEO in PyTorch [61] and interface it with the GTSAM C++ library [39], which contains several non-differentiable factor graph optimizers.
We use the iSAM2 [10] optimizer for real-time, online optimization.
[Table]
[Figure]
[Table]
<
[Table]
/JRo4FSwr9RLDYkLHZMi6lkoSMeOns6szfGyVAQ6VtiUBz9TfEymJjJlEge2MCIzMojcV//O6CYRXfsplnACTdL4oTAQGhacR4AHXjIKYWEKo5vZWTEdEEwo2qJINwVt8eZm0zmreRc29Pa/WK3kcRXSEKugEeegS1dENaqAmokijZ/SK3pxH58V5dz7mrQUnnzlEf+B8/gDPTpKb</latexit> o t - <latexit sha1_base64="tvrC7Sk82mSfkqV/Zwwa9D09RkQ=">AAAB9XicbVBNS8NAEN3Ur1q/qh69hBbBiyURRY8FLx4r2A9oY9lsJ+3SzSbsTtQS8j+8eFDEq//Fm//GbZuDtj4YeLw3w8w8PxZco+N8W4WV1bX1jeJmaWt7Z3evvH/Q0lGiGDRZJCLV8akGwSU0kaOATqyAhr6Atj++nvrtB1CaR/IOJzF4IR1KHnBG0Uj30E+xh/CE6WnmZv1y1ak5M9jLxM1JleRo9MtfvUHEkhAkMkG17rpOjF5KFXImICv1Eg0xZWM6hK6hkoagvXR2dWYfG2VgB5EyJdGeqb8nUhpqPQl90xlSHOlFbyr+53UTDK68lMs4QZBsvihIhI2RPY3AHnAFDMXEEMoUN7fabEQVZWiCKpkQ3MWXl0nrrOZe1Jzb82q9ksdRJEekQk6ISy5JndyQBmkSRhR5Jq/kzXq0Xqx362PeWrDymUPyB9bnD7+kkpE=</latexit> e tFactor Graph
[Figure]
5.1 Overview
Experimental setup: For synthetic navigation (Datasets N1–N4), the goal is to estimate latent robot poses x t ∈ SE (2) from relative odometry and absolute GPS measurements.
For real-world planar pushing (Datasets P1–P3), the goal is to estimate latent object poses o t ∈ SE (2) from tactile image measurements using a DIGIT tactile sensor [62] as shown in the setup in Fig.
Both these problems are modeled as factor graphs with poses as variable nodes and measurements as factor nodes.
Observation model parameters: Recall factor costs || f k ( θ, x ; z ) || 2 are deﬁned by the local observation model f k ( · ) and learned model parameters θ .
This can be expanded out as || f k ( θ, x ; z ) || 2 :=
||
µ k ( x ) − z
|| 2Σ k , where θ := { µ k , Σ k } are the mean and covariance parameters.
For the navigation task, we learn (a) ﬁxed covariances θ :=
{ Σ odom , Σ gps } (N1, N2), and (b) covariances as a function of measurements θ := { Σ odom ( z ) , Σ gps ( z ) } (N3, N4).
For the planar pushing task, we learn (a) ﬁxed covariances θ := { Σ tac , Σ qs } (P1, P2), and (b) means as a function of tactile image measurements θ := { µ tac ( z ) } (P3).
For all experiments, we use a diagonal covariance model.
Baselines: We compare against a learned sequence model such as LSTM, and black-box hyperparameter search methods such as CMAES [11] and Nelder-Mead, all of which optimize Loss 1.
We also compare against a surrogate supervised learning method [4] which optimizes Loss 2.
Appendix D expands further on the setup, factor graph models, and additional results.
5.2 End-to-end Tracking Errors
Fig.
4 shows ﬁnal tracking results for navigation datasets.
We initialize LEO with random θ such that the graph optimizer returns trajectories far from the ground truth.
At the ﬁnal iteration, LEO converges to a θ that matches the optimized trajectories to ground truth.
We summarize quantitative performance against baselines minimizing ﬁnal tracking loss (Loss 1).
LEO consistently outperforms baselines on ﬁnal rmse tracking errors.
The difference is more pronounced for datasets N3,
N4 where the varying covariance model has a higher dimensional parameter space.
Fig.
5 shows ﬁnal tracking results for real-world planar pushing datasets.
Similar to navigation, we show LEO converging to ground truth trajectory for two different objects, a rectangle and ellipse, with differing local contact patches.
We report quantitative performance on datasets P1, P2.
LEO outperforms different hand-tuned choices of covariances as used in prior work [4].
[Figure]
t
[Figure]
s
[Figure]
t Nf e v a l s t r a n s l a t i o n a l e rr o r s ( m ) r o t a t i o n a l e rr o r s ( r a d )
[Figure]
LEO
[Figure]
[Table]
[Figure]
L E O
L o ss
[Figure]
LEO iter
[Figure] [Figure]
LOGO
[Figure]
LOGO ﬁnal iteration
[Figure]
[Table] [Table]
[Table] [Table]
– LOGO
[Figure]
5.3 Sample Efﬁciency and Convergence
Black-box hyper-parameter search is sample inefﬁcient and does not scale well with increasing dimensionality.
LEO, on the other hand, is efﬁcient as it uses gradients.
This can be seen in Fig.
4, where LEO consistently has lower graph optimizer calls (fevals).
Moreover, LEO can also be used to learn models with large parameter spaces, e.g. Fig.
7 learns parameters for a tactile factor neural network mapping image features to transformations.
Finally, Fig.
6 shows LEO consistently converging for varying θ initializations on navigation dataset
N1.
We further show how the LEO samples evolve over iterations.
Initially, when the cost is incorrect, the samples are more spread out.
Over iterations, as the cost converges, the samples concentrate about the ground truth trajectory.
5.4 Surrogate Loss Analysis
We further analyze why a surrogate loss (Loss 2), independent of the graph optimizer, may not minimize end-to-end tracking errors.
We train a network on surrogate supervised learning loss to directly map a pair of tactile image features to ground truth poses similar to prior work [4].
Fig. 7(b) shows the per time step surrogate loss on two trajectories — an initial trajectory that deviates from ground truth, and a ﬁnal converged trajectory.
While the tracking errors of ﬁnal is much less than initial, the surrogate loss does not strictly improve.
This is further reﬂected in the table in Fig.
7(c) where the surrogate loss baseline has higher ﬁnal tracking errors over LEO for dataset P1.
Conclusion
We presented an approach, LEO, for learning observation models end-to-end with graph optimizers that may be non-differentiable.
We compared LEO against baselines for different observation model types on two distinct tasks of navigation and real-world planar pushing.
Overall, LEO is able to learn observation models with lower errors and fewer samples.
While we investigated LEO with non-differentiable optimizers in this paper, as future work, we would like to apply LEO to differentiable optimizers and compare against alternate solutions such as unrolling (see Appendix A for a more detailed discussion).
We would also like to scale to more complex observation models with applications to 3D object tracking using tactile sensors.