# Learning Sampling Distributions for Model Predictive Control

Jacob Sacks
University of Washington
<EMAIL_ADDRESS>&Byron Boots
University of Washington
<EMAIL_ADDRESS>
###### Abstract

Sampling-based methods have become a cornerstone of contemporary approaches to
Model Predictive Control (MPC), as they make no restrictions on the
differentiability of the dynamics or cost function and are straightforward to
parallelize. However, their efficacy is highly dependent on the quality of the
_sampling distribution itself_ , which is often assumed to be simple, like a
Gaussian. This restriction can result in samples which are far from optimal,
leading to poor performance. Recent work has explored improving the
performance of MPC by sampling in a learned latent space of controls. However,
these methods ultimately perform all MPC parameter updates and warm-starting
between time steps in the control space. This requires us to rely on a number
of heuristics for generating samples and updating the distribution and may
lead to sub-optimal performance. Instead, we propose to carry out all
operations in the latent space, allowing us to take full advantage of the
learned distribution. Specifically, we frame the learning problem as bi-level
optimization and show how to train the controller with backpropagation-
through-time. By using a normalizing flow parameterization of the
distribution, we can leverage its tractable density to avoid requiring
differentiability of the dynamics and cost function. Finally, we evaluate the
proposed approach on simulated robotics tasks and demonstrate its ability to
surpass the performance of prior methods and scale better with a reduced
number of samples.

> Keywords: Model Predictive Control, Normalizing Flows

##  1 Introduction

Sequential decision making under uncertainty is a fundamental problem in
machine learning and robotics. Recently, model predictive control (MPC) has
emerged as a powerful paradigm to tackle such problems on real-world robotic
systems. In particular, MPC has been successfully applied to helicopter
aerobatics [1], robot manipulation [2, 3, 4], humanoid robot locomotion [5],
robot-assisted dressing [6], and aggressive off-road driving [7, 8, 9].
Sampling-based approaches to MPC are becoming particularly popular due to
their simplicity and ability to handle non-differentiable dynamics and cost
functions. These methods work by sampling controls from a policy distribution
and rolling out an approximate system model using the sampled control
sequences. They then use the resulting trajectories to compute an approximate
gradient of the cost function and update the policy. Next, the controller
samples an action from this distribution and applies it to the system. It
repeats the process from the resulting next state, creating a feedback
controller. Between time steps, it warm-starts the optimization process using
a modification of the solution at the previous time step.

An important design decision is the form of the sampling distribution, which
is often simple, e.g. a Gaussian, such that we can efficiently sample and
tractably update its parameters. However, this also has drawbacks: without
much control over the distribution form, samples often lie in high-cost
regions, hindering performance. This can be particularly problematic in
complex environments with sparse costs or rewards, as a poorly parameterized
distribution may hinder efficient exploration, leading the system into bad
local optima. A side effect is that we often require many samples to
accomplish the objective, increasing computational requirements. There have
been extensions which target more complex distributions, such as Gaussian
mixture models [10] and a particle method based on Stein variational gradient
descent (SVGD) [11]. However, there is a large amount of structure in the
environment that these methods fail to exploit. Instead, an alternative
approach is to _learn_ a sampling distribution which can leverage the
environmental structure to draw more optimal samples.

Prior work on learning MPC sampling distributions generally requires
differentiability of the dynamics and cost function [12, 13]. Power and
Berenson [14, 15] circumvent this issue by leveraging normalizing flows (NFs)
[16, 17, 18], which have a tractable log-likelihood. This property allows them
to learn flexible distributions by directly optimizing the MPC cost without
requiring differentiability via the likelihood-ratio gradient. However, a
limitation of their approach is that all online updates to the distribution
and warm-starting between time steps occur entirely in the control space,
leaving the latent distribution fixed. This forces them to apply heuristics to
generate samples by combining those from the learned distribution with
Gaussian perturbations to the current control-space mean. These restrictions
prevent us from fully taking advantage of the learned distribution and
potentially throws away useful information. Additionally, their approach does
not allow for the incorporation of control constraints directly into the
sampling distribution, which is important for real-world robots.

Instead, we propose to alter the optimization machinery to operate entirely in
the latent space. As the NF latent space follows a simple distribution, it
remains feasible to perform MPC updates in this learned space and update the
latent distribution online. Specifically, during an episode, the parameters of
the latent distribution are updated with MPC while those of the NF remain
fixed. Then during training, after each episode, the parameters of the NF are
updated. We can frame this setup as a bi-level optimization problem [19] and
derive a method for computing an approximate gradient through the latent MPC
update. This involves treating MPC as a recurrent network, where the control
distribution acts as a form of memory, and unrolling the computation to train
with backpropagation-through-time (BPTT). However, it is no longer clear how
to warm-start between time steps because there is no clear delineation of time
in the latent space. Moreover, the usual method of warm-starting, which simply
shifts the current plan forward in time, may be sub-optimal. Therefore, we
propose to learn a shift model, which performs all warm-starting operations in
the latent space. Finally, we show how to alter the NF architecture to
incorporate box constraints on the sampled controls.

Contributions: In this work, we build upon recent efforts to learn sampling
distributions for MPC with NFs by moving all online parameter updates and
warm-start operations into the latent space. We accomplish this by framing the
learning problem as bi-level optimization and derive an approximate gradient
through the MPC update of the latent distribution to train the network with
BPTT. Additionally, we show how to parameterize the flow architecture such
that we can incorporate box constraints on the controls. Finally, we
empirically evaluate our proposed approach on simulated navigation and
manipulation tasks. We demonstrate its ability to improve performance over the
baselines by taking full advantage of the learned latent space. We find that
the performance of the controllers with our learned sampling distributions
scales more gracefully with a reduction in the number of samples.

##  2 Sampling-Based Model Predictive Control

We consider the problem of controlling a discrete-time stochastic dynamical
system, which is in state $x_{t}\in{}^{N}$ at time step $t$. Upon the
application of control $u_{t}\in{}^{M}$, the system incurs the instantaneous
cost $c(x_{t},u_{t})$ and transitions to $x_{t+1}\sim f(x_{t},u_{t})$. We wish
to design a policy $u_{t}\sim\pi(\cdot|x_{t})$ such that the system achieves
good performance over $T$ steps. Instead of finding a single, globally optimal
policy, MPC re-optimizes a local policy at each time step by predicting the
system’s behavior over a finite horizon $H<T$ using an approximate model
$\hat{f}$. Specifically, in sampling-based MPC, we sample a control sequence
$\bm{\hat{u}}_{t}\sim\pi_{\theta}(\cdot)$, where
$\bm{\hat{u}}_{t}\triangleq(\hat{u}_{t},\hat{u}_{t+1},\cdots,\hat{u}_{t+H-1})$,
and our policy is parameterized by $\theta\in\Theta$. We rollout our model
starting at $x_{t}$ using these sampled controls to get our predicted state
sequence
$\bm{\hat{x}}_{t}\triangleq(\hat{x}_{t},\hat{x}_{t+1},\cdots,\hat{x}_{t+H})$,
with $\hat{x}_{t}=x_{t}$. The total trajectory cost is

$C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})=\sum_{h=0}^{H-1}c(\hat{x}_{t+h},\hat{u}_{t+h})+c_{term}(\hat{x}_{t+H}),$
(1)

where $c_{term}(\cdot)$ is a terminal cost function. We then construct a
statistic $\hat{J}(\theta;x_{t})$ defined on cost
$C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})$ and use the rollouts to solve
$\theta_{t}\leftarrow\operatorname*{arg\,min}_{\theta\in\Theta}\hat{J}(\theta;x_{t})$.
After optimizing our parameters, we sample the control sequence
$\bm{\hat{u}}_{t}\sim\pi_{\theta_{t}}(\cdot)$, apply the first control to the
real system (i.e. $u_{t}=\hat{u}_{t}$), and repeat the process. Because each
parameter $\theta_{t}$ depends on the current state, MPC effectively yields a
state-feedback policy, even though the individual distributions give us an
open-loop sequence.

In this paper, we consider a popular sampling-based MPC algorithm known as
Model Predictive Path Integral (MPPI) control [7, 8]. MPPI optimizes the
exponential utility or risk-seeking objective:

$\hat{J}(\theta;x_{t})=-\log\mathbb{E}_{\pi_{\theta},\hat{f}}\Big{[}\exp\Big{(}-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})\Big{)}\Big{]},$
(2)

where $\beta>0$ is a scaling parameter, known as the temperature. As we do not
assume that the dynamics or cost function are differentiable, we compute the
gradients via the likelihood-ratio derivative:

$\nabla\hat{J}(\theta;x_{t})=-\frac{\mathbb{E}_{\pi_{\theta},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\nabla_{\theta}\log\pi_{\theta}(\bm{\hat{u}}_{t})\Big{]}}{\mathbb{E}_{\pi_{\theta},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}}.$
(3)

In MPPI, the policy is assumed to be a factorized Gaussian of the form

$\pi_{\theta}(\bm{\hat{u}})=\prod_{h=0}^{H-1}\pi_{\theta_{h}}(\hat{u}_{t+h})=\prod_{h=0}^{H-1}\mathcal{N}(\hat{u}_{t+h};\mu_{t+h},\Sigma_{t+h}).$
(4)

Previous work by Wagener et al. [9] has shown that optimizing this objective
with dynamic mirror descent (DMD) [20] and approximating with Monte Carlo
estimates gives us the MPPI update rule:

$\mu_{t+h}=(1-\gamma_{t})\tilde{\mu}_{t+h}+\gamma_{t}\sum_{i=1}^{N}w_{i}\hat{u}_{t+h}^{(i)},\quad
w_{i}=\frac{e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t}^{(i)},\bm{\hat{u}}_{t}^{(i)})}}{\sum_{j=1}^{N}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t}^{(j)},\bm{\hat{u}}_{t}^{(j)})}}$
(5)

where $\tilde{\mu}_{t+h}$ is the current mean for each time step and
$\gamma_{t}$ is the step size. Between time steps of DMD, we get
$\tilde{\mu}_{t+h}$ from our previous solution $\mu_{t+h}$ by using a shift
model $\tilde{\mu}_{t+h}=\Phi(\mu_{t+h})$. This shift model aims to predict
the optimal decision at the next time step given the previous solution. In the
context of MPC, it allows us to warm-start the optimization problem to speed
up convergence, as we can only approximately solve the optimization problem
due to real-time constraints.

##  3 Learning the Sampling Distribution of MPC

###  3.1 Representation of the Learned Distribution

Instead of using uninformed sampling distributions, learned distributions can
potentially exploit structure in the environment to draw samples which are
more likely to be collision-free and close to optimal. However, such learned
distributions must be sufficiently expressive in order to better capture near-
optimal, potentially multimodal, behavior. They must also be parameterized
such that it is tractable to sample from and update online. If the
distribution has a large number of parameters, the number of samples required
to efficiently update them online may be computationally infeasible. And
ideally, the form of our distribution would be such that we could find a
closed-form update.

One path towards meeting these criteria is to maintain a simple latent
distribution from which we can sample, and then learn a transformation of the
samples which maps them to a more complex distribution. During training, we
learn the parameters of this transformation, which can be conditioned on
problem-specific information, such as the starting and goal configurations of
the robot and obstacle placements. However, when executing the policy during
an episode, the parameters of this learned transformation remain fixed, and
instead, we update the parameters of the latent distribution. Concretely, we
consider learning a distribution $\pi_{\theta,\lambda}$ which is defined
implicitly as follows:

$\bm{\hat{z}}_{t}\sim
p_{\theta}(\cdot),\quad\bm{\hat{u}}_{t}=h_{\lambda}(\bm{\hat{z}}_{t};c)$ (6)

where
$\bm{\hat{z}}_{t}\triangleq(\hat{z}_{t},\hat{z}_{t+1},\cdots,\hat{z}_{t+H-1})$,
$c$ is a context variable describing the relevant information of the
environment, $p_{\theta}$ is the latent distribution with parameters $\theta$,
and $h_{\lambda}$ is the learned conditional transformation with parameters
$\lambda$. Moving forward, we assume that both $\bm{\hat{z}}_{t}$ and
$\bm{\hat{u}}_{t}$ are stacked as vectors in MH. If $p_{\theta}$ is a Gaussian
factorized as in Equation 4 and we assume that $h_{\lambda}$ is invertible, we
prove in Section A.6 that the corresponding DMD update to the latent mean is
simply Equation 5, except that we replace the controls in the weighted sum
with the latent samples.

###  3.2 Formulating the Learning Problem

Learning the distribution $\pi_{\theta,\lambda}$ amounts to solving a bi-level
optimization problem [19], in which one optimization problem is nested in
another. The lower-level optimization problem involves updating the latent
distribution parameters at each time step, $\theta_{t}$, by minimizing the
expected cost with DMD. The upper-level optimization problem consists of
learning $\lambda$ such that MPC performs well across a number of different
environments. To formalize this, first consider that we have some distribution
of environments $c\sim\mathcal{C}(\cdot)$ over which we wish MPC to perform
well. For each environment, our system has some conditional initial state
distribution $x_{0}\sim\rho(\cdot|c)$. The objective we wish to minimize is
then

$\ell(\bm{\theta},\lambda;c)=\mathbb{E}_{\bm{\pi}_{\bm{\theta},\lambda},\rho,f}\Bigg{[}\sum_{t=0}^{T-1}\hat{J}(\theta_{t},\lambda;x_{t},c)\Bigg{]}$
(7)

where $\bm{\theta}=(\theta_{0},\theta_{1},\cdots,\theta_{T-1})$ and our cost
statistic, $\hat{J}$, now depends on $\lambda$ and $c$ as well. This objective
measures the expected performance of the intermediate plans produced by MPC
along the $T$ steps of the episode. Our desired bi-level optimization problem
can be formulated as:

$\displaystyle\min_{\lambda}\
\mathbb{E}_{\mathcal{C}}\Big{[}\ell(\bm{\theta}(\lambda),\lambda;c)\Big{]}\quad\text{s.t.}\
\bm{\theta}(\lambda)\approx_{\lambda}\operatorname*{arg\,min}_{\bm{\theta}}\ell(\bm{\theta},\lambda;c)$
(8)

where $\approx_{\lambda}$ indicates that we approximate the solution of the
optimization problem with an iterative algorithm that may also be
parameterized by $\lambda$, as the exact minimizer is not available in closed
form. Moreover, the notation $\bm{\theta}(\lambda)$ indicates the dependence
of the lower-level solution on the upper-level parameters. In our case, we
solve the lower-level problem with DMD, where we also parameterize the shift
model, $\Phi_{\lambda}(\cdot;c)$, making it a learnable component and
conditioned on $c$.

The normal shift model in MPC simply shifts the control sequence forward one
time step and appends a zero or random control at the end. However, because we
are performing this update in the latent space, there is no clear delineation
between time steps of the latent controls, as they are coupled according to
the learned transformation. Therefore, there is no way to easily perform the
equivalent shift operation in the latent space. As such, we instead learn this
shift model along with the transformation. Besides, the standard approach
described above may not be optimal. By learning it, we may be able to further
improve performance. This is especially true because the performance hinges
greatly on the quality of the shift model since we only run one iteration of
the DMD update.

###  3.3 Parameterizing with Normalizing Flows

In order to optimize the upper-level objective in Equation 8 with respect to
$\lambda$, we need to be able to compute the density $\pi_{\theta,\lambda}$
directly. Therefore, we choose to represent $h_{\lambda}$ with a normalizing
flow (NF) [16, 18, 17], which explicitly learns the density by defining an
invertible transformation that maps latent variables to observed data.
Generally, we compose a series of component flows together, i.e.
$h_{\lambda}=h_{\lambda_{K}}\circ h_{\lambda_{K-1}}\circ\dots\circ
h_{\lambda_{1}}$, which define a series of intermediate variables
$\bm{\hat{y}}_{0},\dots,\bm{\hat{y}}_{K-1},\bm{\hat{y}}_{K}$, with
$\bm{\hat{y}}_{0}=\bm{\hat{z}}$ and $\bm{\hat{y}}_{K}=\bm{\hat{u}}$. The log-
likelihood of the composed flow is given by:

$\log\pi_{\theta,\lambda}(\bm{\hat{u}}|c)=\log
p_{\theta}(\bm{\hat{z}})-\sum_{i=1}^{K}\log\Big{|}\text{det}\
\frac{\partial\bm{\hat{y}}_{i}}{\partial\bm{\hat{y}}_{i-1}}\Big{|}.$ (9)

In this work, we make use of the affine coupling layer proposed by Dinh et al.
[18] as part of the real non-volume-preserving (RealNVP) flow. The core idea
is to split the input $\bm{\hat{u}}$ into two partitions
$\bm{\hat{u}}=(\bm{\hat{u}}_{I_{1}},\bm{\hat{u}}_{I_{2}})$, where $I_{1}$ and
$I_{2}$ are a partition of $[1,MH]$, and apply

$\displaystyle\bm{\hat{y}}_{I_{1}}=\bm{\hat{u}}_{I_{1}},\quad\bm{\hat{y}}_{I_{2}}=\bm{\hat{u}}_{I_{2}}\odot\exp{s_{\lambda}(\bm{\hat{u}}_{I_{1}},c)}+t_{\lambda}(\bm{\hat{u}}_{I_{1}},c),$
(10)

where $s_{\lambda}$ and $t_{\lambda}$ are the scale and translation terms,
which are represented with arbitrary neural networks, and $\odot$ is the
Hadamard product. This makes computing the log-determinant term in Equation 9
and inverting the flow fast and efficient. Now, in robotics, we often have
lower and upper limits on the controls. These are usually enforced in
sampling-based MPC by either clamping the control samples or passing them
through a scaled sigmoid. However, instead of enforcing the constraints
heuristically after sampling, we learn a constrained sampling distribution
directly. Since the sigmoid function is invertible and has a tractable log-
determinant (shown in Section A.8), we can simply append one after
$h_{\lambda_{K}}$ in the NF and scale it by the control limits. This ensures
that control constraints are satisfied by design and taken into account while
learning the distribution.

###  3.4 Training the Sampling Distribution

Computing gradients through the upper-level objective is not straightforward,
as both the expectation and the inner terms of Equation 7 depend on $\lambda$.
Therefore, the state distribution depends on the NF and latent shift model.
One way around this issue is to consider a modified objective at each batch
$d$:

$\ell_{d}(\bm{\theta},\lambda;c)=\mathbb{E}_{\bm{\pi}_{\bm{\theta},\lambda_{d}},\rho,f}\Bigg{[}\sum_{t=0}^{T-1}\hat{J}(\theta_{t},\lambda;x_{t},c)\Bigg{]},$
(11)

which fixes the outer expectation to be with respect to the current policy.
Intuitively, this choice trains the NF to optimize the MPC cost function under
the state distribution resulting from the current policy
$\bm{\pi}_{\bm{\theta},\lambda_{d}}$. This effectively ignores the dependency
of the action chosen to interact with the true environment and the learned
parameters. We then update the outer expectation distribution at each episode,
overcoming the covariate shift problem that would otherwise arise.

Now, we only have to focus on computing the gradient
$\nabla_{\lambda}\hat{J}(\theta_{t}(\lambda),\lambda;x_{t},c)\rvert_{\lambda=\lambda_{d}}$
for each time step, which can be computed similar to Equation 3 and
approximated with Monte Carlo sampling:

$\nabla\hat{J}(\theta_{t}(\lambda),\lambda;x_{t},c)\approx-\sum_{i=1}^{N}w_{i}\nabla_{\lambda}\log\pi_{\theta_{t}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c),$
(12)

where the weights $w_{i}$ are defined according to Equation 5. The log-
likelihood is given by Equation 9, the gradient of which involves computing
the backwards pass through the network $h_{\lambda}$. However, we also have to
consider the dependence of the latent distribution parameters
$\bm{\theta}(\lambda)$ on $\lambda$. Therefore, we must backpropagate through
the MPC update:

$\bm{\mu}_{t}(\lambda)=(1-\gamma_{t})\bm{\tilde{\mu}}_{t}(\lambda)+\gamma_{t}\Delta\bm{\mu}_{t},\quad\Delta\bm{\mu}_{t}=\frac{\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}h_{\lambda}(\bm{\hat{u}}_{t};c)\Big{]}}{\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}}$
(13)

where the previous shifted mean $\bm{\tilde{\mu}}_{t}(\lambda)$ is given by
the learned latent shift model. Note that we have rewritten the expectations
in terms of the control distribution, rather than the latent distribution.
This is necessary in order to derive the following approximate gradient
without requiring differentiability. To compute the gradient of Equation 13,
we must approximate the gradient of $\Delta\bm{\mu}_{t}$ with respect to
$\lambda$, which we can compute as
$\frac{\partial\Delta\bm{\mu}_{t}}{\partial\lambda}\approx M_{1}-M_{2}M_{3}$,
where we define:

$\begin{gathered}M_{1}=\sum_{i=1}^{N}w_{i}\Big{[}\nabla_{\lambda}h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)+h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)\Big{]},\\\
M_{2}=\sum_{i=1}^{N}w_{i}h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c),\quad
M_{3}=\sum_{i=1}^{N}w_{i}\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c).\end{gathered}$
(14)

The derivation of this approximate gradient can be found in Section A.7. Note
that computing the gradients
$\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)$
will also require us to backpropagate through the shift model due to the
dependence of $\tilde{\theta}$ on $\lambda$. Therefore, even when the step
size is set to one, i.e. $\gamma_{t}=1$, we introduce a form of recurrence
between time steps. Please see Section A.1 for additional visualizations of
the computational graph generated by an episode and further descriptions of
the overall algorithm.

##  4 Related Work

There have been many recent attempts at combining machine learning with MPC,
which mostly center around learning or fine-tuning a good dynamics model [6,
8, 21, 22, 23, 24, 25], potentially from high-dimensional observations [3, 26,
27, 28, 29, 30, 31]. Alternatively, some methods target the cost function,
learning terminal value functions [32, 33, 34, 35, 36] or cost-shaping terms
[37] to improve performance with reduced prediction horizons. Instead of
learning dynamics and cost separately, other methods propose learning the
entire optimal controller [12, 38, 39, 40, 41, 42, 37] or planner [43, 44, 45,
46] end-to-end. In contrast, recent work has explored leaving the dynamics and
cost function fixed and instead focused on improving the optimization process.
For example, Sacks and Boots [47] propose to use learning to replace the
update rule employed to solve the optimization problem at each time step with
a neural network. They showed that this can substantially reduce the number of
samples required to achieve a target task performance with sampling-based MPC.

Using alternative sampling distributions is another viable strategy for
improving the optimization process in MPC. Multiple works have considered
sampling distributions beyond simple Gaussians, such as Gaussian mixture
models [10] and a particle method based on Stein variational gradient descent
(SVGD) [11]. Additionally, most implementations use heuristics to modify
samples and squeeze out additional performance gains [4, 48]. In terms of
learning the distribution, Amos and Yarats [12] learn a latent action space
for their proposed differentiable cross-entropy method (CEM) controller.
However, they require all components of the pipeline to be differentiable and
do not consider learning a shift model. Agarwal et al. [13] learn a
normalizing flow in the latent space of a variational autoencoder (VAE). Yet,
they also require differentiability, use expert demonstrations to learn the
latent space of the VAE, and have no means of warm-starting between time
steps. Wang and Ba [49] propose to use a learned feedback policy to warm-start
MPC, but still rely on a Gaussian perturbations to the proposed action
sequence. The authors also explore performing online planning in the network’s
parameter space, which results in a massive action space that may be hard to
scale.

Power and Berenson [14, 15] also train a normalizing flow to use as the
sampling distribution for MPC, but they do not learn a latent shift model and
perform all operations in the control space. Moreover, they mix the latent
samples with Gaussian perturbations to the current control-space mean, as they
do not update the latent distribution directly. This prevents them from fully
taking advantage of the learned distribution and throws away useful
information which could potentially improve performance. In fact, we show in
Section A.4 that the learned shift model contributes significantly to the
performance gains. Additionally, a primary focus of their work is on how to
handle out-of-distribution (OOD) environments by learning a posterior over
environment context variables. We could combine their approach with ours by
conditioning the learned shift model on the inferred environment context for
improved generalization. Finally, normalizing flows have also been used to
improve exploration in RL [50, 51, 52] by providing a more flexible, and
potentially multimodal, distribution. They have also been employed in
sampling-based motion planning [53, 54] to provide good proposal
configurations to speed up convergence.

##  5 Experimental Results

In all experiments, we denote our proposed approach as NFMPC, the baseline
MPPI implementation as MPPI, and the method by Power and Berenson [14, 15] as
FlowMPPI. Details about the hyperparameters, implementation, tasks, and
training can be found in Section A.2. We evaluate on a fixed set of
environments, which includes start states, goal locations, and obstacle
placements, and run 32 rollouts for each sample amount. Our primary metrics
for comparison are the success rate, defined as the percentage of times the
task goal was achieved, and the average cost of trajectories which
successfully completed the task. Additionally, we cannot ignore the overhead
introduced by the NF in the control pipeline. Therefore, we measured the
average change in wall clock time across different amounts of samples for
NFMPC and FlowMPPI in Section A.5.

###  5.1 Planar Robot Navigation

We begin by applying NFMPC to a planar robot navigation task in which a 2D
holonomic point-robot must reach a goal position while avoiding eight dynamic
obstacles, which move around randomly at each time step (PNRandDyn). The
point-robot has double integrator dynamics with stochasticity on the
acceleration commands to create a mismatch between the predictive model used
by MPC and the true environment. Each obstacle’s current position is perturbed
with Gaussian noise and clipped to be within map bounds, and the starting and
goal locations of the robot are randomized in each episode. The NF for both
NFMPC and FlowMPPI is conditioned on the obstacle locations, current state,
and goal position. Note that we do not condition the shift model, as this
consistently hurt performance.

Figure 1: Success rate and cost distribution on the PNRandDyn environment
across a different number of samples.

NFMPC

FlowMPPI

MPPI

Figure 2: Visualization of a trajectory and top samples from (top) NFMPC,
(middle) FlowMPPI, and (bottom) MPPI on the PNRandDyn task.

We quantitatively compare all controllers in Figure 1. The trajectory cost box
plots represent the median and quartiles of the distribution. We find that
NFMPC consistently outperforms the both MPPI and FlowMPPI. While all
controllers reach a 100% success rate at 1024 samples, NFMPC achieves a 20%
and 8% lower median cost over MPPI and FlowMPPI, respectively. We also find
that NFMPC scales more gracefully overall as the number of samples is reduced.
For instance, it is able to withstand a $32\times$ decrease in the number of
samples (1024 to 32) while reducing success rate by only 3% and increasing
median cost by 8%, nearly matching FlowMPPI at 1024 samples. In comparison,
MPPI reduces success rate by 6% and increases median cost by 43%. Similarly,
NFMPC outperforms FlowMPPI at all sample amounts in terms of median cost,
although at 16 samples FlowMPPI achieves a slightly higher success rate than
NFMPC. We visualize trajectories and top samples in Figure 2. The green and
red dots are the starting and goal locations, respectively, and the blue dot
is the current position of the robot at the given time step. The thick black
line is the resulting path taken by the controller, while the gray lines are
the top samples generated at the current state. In this example, FlowMPPI
nearly collides with an obstacle, while NFMPC and MPPI are able to safely
reach the goal. FlowMPPI over commits to a narrow corridor and is unable to
reroute in time to account for the new obstacle location. NFMPC takes a
similar trajectory to FlowMPPI, however, it is able to pause until the
obstacle moves out of the way to proceed towards the goal. We also provide
additional experiments with static environments and comparing unconditional
and conditional models in Section A.3.

Figure 3: Success rate and cost distribution on the FrankaObstacles
environment across a different number of samples.

NFMPC

FlowMPPI

MPPI

Figure 4: Visualization of a trajectory and top samples from (top) NFMPC,
(middle) FlowMPPI, and (bottom) MPPI on the FrankaObstacles task.

###  5.2 Franka Panda Arm

Next, we apply NFMPC to the Franka task, which involves controlling a 7
degree-of-freedom (DOF) Franka Panda robot arm and steering it towards a
randomly placed target goal from a fixed starting pose while avoiding a single
pole obstacle. The NF for both NFMPC and FlowMPPI is conditioned on the
obstacle locations, initial state, and goal position. It is also important to
note that no controller achieves a $100\%$ success rate, as not every randomly
generated environment is feasible. As shown in Figure 3, at 1024 samples,
NFMPC achieves a success rate of $84\%$, while FlowMPPI and MPPI only succeed
$81\%$ and $78\%$ of the time, respectively. Again, NFMPC scales better with a
reduced number of samples. With a $64\times$ decrease in the number of samples
(1024 to 16), NFMPC only drops in success rate by $29\%$. Meanwhile, FlowMPPI
decreases by nearly $70\%$ to a success rate of $25\%$. However, both fair
better than MPPI, which drops by $80\%$ to a success rate of $16\%$. These
results support the hypothesis that learning to perform MPC updates in the
latent space of the NF and training the controller as a recurrent network
improves performance.

To more clearly understand what NFMPC is doing differently, we visualize the
performance of all three controllers on a held-out validation environment in
Figure 4. The green and red markers indicate the end-effector and goal
positions, respectively, while the single pole is the obstacle which must be
avoided. Additionally, the green trajectories represent the top samples from
the controller at each time step. While MPPI collides with the obstacle, both
NFMPC and FlowMPPI learn to take a different path which is collision-free.
FlowMPPI generates better initial trajectories than NFMPC, which are more
pointed towards the goal location and achieve a greater velocity. However,
NFMPC is able to better adapt the sampling distribution throughout the episode
and reach the goal more quickly. Finally, in Section A.4, we present results
for the environment with no obstacles, a comparison of conditional and
unconditional models, and a breakdown of the individual cost terms for all
models to gain insight into the learned sampling distributions. We also
perform an ablation which removes the learned shift model from NFMPC to
illustrate that it is a crucial component.

##  6 Limitations

A major limitation of NFMPC is that the learned distribution and shift model
are only valid for a fixed horizon and control dimensionality. Therefore,
these components cannot be directly transferred to new robots or for alternate
horizons without being retrained. However, this could potentially be remedied
by novel architectural innovations and training distributions across both
environments and robots. Moreover, the learned distribution is specific to the
environmental distributions on which it was trained. Therefore, it does not
always perform as well when transferred to out-of-distribution environments.
However, this is always going to be a challenge for any learning-based method
and addressing it is an open problem for future research. For instance, our
approach could be combined with the method developed by Power and Berenson
[14, 15], which learned a generative model of environments and used this
learned distribution to perform a projection step on the conditional NF.
However, this approach would only be valid for conditional models, and
addressing transfer of unconditional models is an unresolved question.
Finally, both our approach and FlowMPPI introduce an additional overhead to
due to running the NF which cannot be ignored. This increase in wall clock
time may be worth the additional performance gains. Additionally, we scale
better with a reduction in samples compared to all baselines. Therefore, we
can potentially alleviate some of the introduced overhead by reducing the
number of samples while still meeting the application demands.

##  7 Conclusion

We presented a method for learning MPC sampling distributions with normalizing
flows (NFs) which moves all online parameter updates and warm-starting
operations into the latent space. We show how to frame the problem as bi-level
optimization and derive an approximate gradient through the MPC update to
train the distributions. Additionally, we illustrate how to incorporate
control box constraints directly into the NF architecture. Through our
empirical evaluations in both simulated navigation and manipulation problems,
we demonstrate that our approach is able to surpass the performance of all
baselines. Moreover, we find that controllers which move all operations into
the latent space are often able to scale more gracefully with a reduction in
the number of samples. These results indicate the importance of leveraging the
latent space in learned sampling distributions for MPC. Finally, because we
learn all components of the controller through episodic interactions with the
environment, they can potentially be trained to account for the modeling
errors in the MPC controller.

## References

  * Abbeel et al. [2010] P. Abbeel, A. Coates, and A. Y. Ng.  Autonomous Helicopter Aerobatics through Apprenticeship Learning.  _The International Journal of Robotics Research (IJRR)_ , 29(13):1608–1639, 2010.
  * Kumar et al. [2016] V. Kumar, E. Todorov, and S. Levine.  Optimal Control with Learned Local Models: Application to Dexterous Manipulation.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 378–383. IEEE, 2016.
  * Finn and Levine [2017] C. Finn and S. Levine.  Deep Visual Foresight for Planning Robot Motion.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 2786–2793. IEEE, 2017.
  * Bhardwaj et al. [2021] M. Bhardwaj, B. Sundaralingam, A. Mousavian, N. Ratliff, D. Fox, F. Ramos, and B. Boots.  STORM: An Integrated Framework for Fast Joint-Space Model-Predictive Control for Reactive Manipulation.  _arXiv preprint arXiv:2104.13542_ , 2021.
  * Erez et al. [2013] T. Erez, K. Lowrey, Y. Tassa, V. Kumar, S. Kolev, and E. Todorov.  An integrated system for real-time Model Predictive Control of humanoid robots.  In _IEEE-RAS International Conference on Humanoid Robots (Humanoids)_ , pages 292–299. IEEE, 2013.
  * Erickson et al. [2018] Z. Erickson, H. M. Clever, G. Turk, C. K. Liu, and C. C. Kemp.  Deep Haptic Model Predictive Control for Robot-Assisted Dressing.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 4437–4444. IEEE, 2018.
  * Williams et al. [2016] G. Williams, P. Drews, B. Goldfain, J. M. Rehg, and E. A. Theodorou.  Aggressive driving with model predictive path integral control.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 1433–1440. IEEE, 2016.
  * Williams et al. [2017] G. Williams, N. Wagener, B. Goldfain, P. Drews, J. M. Rehg, B. Boots, and E. A. Theodorou.  Information theoretic MPC for model-based reinforcement learning.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 1714–1721. IEEE, 2017.
  * Wagener et al. [2019] N. Wagener, C.-A. Cheng, J. Sacks, and B. Boots.  An Online Learning Approach to Model Predictive Control.  _arXiv preprint arXiv:1902.08967_ , 2019.
  * Okada and Taniguchi [2020] M. Okada and T. Taniguchi.  Variational Inference MPC for Bayesian Model-based Reinforcement Learning.  In _Conference on Robot Learning (CoRL)_ , pages 258–272. PMLR, 2020.
  * Lambert et al. [2020] A. Lambert, A. Fishman, D. Fox, B. Boots, and F. Ramos.  Stein Variational Model Predictive Control.  _arXiv preprint arXiv:2011.07641_ , 2020.
  * Amos and Yarats [2020] B. Amos and D. Yarats.  The Differentiable Cross-Entropy Method.  In _International Conference on Machine Learning (ICML)_ , pages 291–302. PMLR, 2020.
  * Agarwal et al. [2020] S. Agarwal, H. Sikchi, C. Gulino, and E. Wilkinson.  Imitative Planning using Conditional Normalizing Flow.  _arXiv preprint arXiv:2007.16162_ , 2020.
  * Power and Berenson [2021] T. Power and D. Berenson.  Variational Inference MPC for Robot Motion with Normalizing Flows.  _Advances in Neural Information Processing Systems (NeurIPS) Workshop on Robot Learning: Self-Supervised and Lifelong Learning_ , 2021.
  * Power and Berenson [2022] T. Power and D. Berenson.  Variational Inference MPC using Normalizing Flows and Out-of-Distribution Projection.  _arXiv preprint arXiv:2205.04667_ , 2022.
  * Dinh et al. [2014] L. Dinh, D. Krueger, and Y. Bengio.  Nice: Non-linear Independent Components Estimation.  _arXiv preprint arXiv:1410.8516_ , 2014.
  * Rezende and Mohamed [2015] D. Rezende and S. Mohamed.  Variational Inference with Normalizing Flows.  In _International Conference on Machine Learning (ICML)_ , pages 1530–1538. PMLR, 2015.
  * Dinh et al. [2016] L. Dinh, J. Sohl-Dickstein, and S. Bengio.  Density estimation using Real NVP.  _arXiv preprint arXiv:1605.08803_ , 2016.
  * Bard [2013] J. F. Bard.  _Practical Bilevel Optimization: Algorithms and Applications_ , volume 30.  Springer Science & Business Media, 2013.
  * Hall and Willett [2013] E. Hall and R. Willett.  Dynamical Models and Tracking Regret in Online Convex Programming.  In _International Conference on Machine Learning (ICML)_ , pages 579–587. PMLR, 2013.
  * Lenz et al. [2015] I. Lenz, R. A. Knepper, and A. Saxena.  DeepMPC: Learning Deep Latent Features for Model Predictive Control.  In _Robotics: Science and Systems (R:SS)_. Rome, Italy, 2015.
  * Kocijan et al. [2004] J. Kocijan, R. Murray-Smith, C. E. Rasmussen, and A. Girard.  Gaussian process model based predictive control.  In _American Control Conference (ACC)_ , volume 3, pages 2214–2219. IEEE, 2004.
  * Fu et al. [2016] J. Fu, S. Levine, and P. Abbeel.  One-Shot Learning of Manipulation Skills with Online Dynamics Adaptation and Neural Network Priors.  In _IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_ , pages 4019–4026. IEEE, 2016.
  * Chua et al. [2018] K. Chua, R. Calandra, R. McAllister, and S. Levine.  Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models.  _arXiv preprint arXiv:1805.12114_ , 2018.
  * Nagabandi et al. [2018] A. Nagabandi, G. Kahn, R. S. Fearing, and S. Levine.  Neural Network Dynamics for Model-Based Deep Reinforcement Learning with Model-Free Fine-Tuning.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 7559–7566. IEEE, 2018.
  * Wahlström et al. [2015] N. Wahlström, T. B. Schön, and M. P. Deisenroth.  From Pixels to Torques: Policy Learning with Deep Dynamical Models.  _arXiv preprint arXiv:1502.02251_ , 2015.
  * Watter et al. [2015] M. Watter, J. T. Springenberg, J. Boedecker, and M. Riedmiller.  Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images.  _arXiv preprint arXiv:1506.07365_ , 2015.
  * Banijamali et al. [2018] E. Banijamali, R. Shu, H. Bui, and A. Ghodsi.  Robust Locally-Linear Controllable Embedding.  In _International Conference on Artificial Intelligence and Statistics (AISTATS)_ , pages 1751–1759. PMLR, 2018.
  * Ebert et al. [2018] F. Ebert, C. Finn, S. Dasari, A. Xie, A. Lee, and S. Levine.  Visual Foresight: Model-Based Deep Reinforcement Learning for Vision-Based Robotic Control.  _arXiv preprint arXiv:1812.00568_ , 2018.
  * Ha et al. [2018] J.-S. Ha, Y.-J. Park, H.-J. Chae, S.-S. Park, and H.-L. Choi.  Adaptive Path-Integral Autoencoder: Representation Learning and Planning for Dynamical Systems.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 2018\.
  * Hafner et al. [2019] D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson.  Learning Latent Dynamics for Planning from Pixels.  In _International Conference on Machine Learning (ICML)_ , pages 2555–2565. PMLR, 2019.
  * Zhong et al. [2013] M. Zhong, M. Johnson, Y. Tassa, T. Erez, and E. Todorov.  Value Function Approximation and Model Predictive Control.  In _IEEE Symposium on Adaptive Dynamic Programming and Reinforcement Learning (ADPRL)_ , pages 100–107. IEEE, 2013.
  * Rosolia and Borrelli [2017] U. Rosolia and F. Borrelli.  Learning Model Predictive Control for Iterative Tasks. A Data-Driven Control Framework.  _IEEE Transactions on Automatic Control_ , 63(7):1883–1896, 2017.
  * Lowrey et al. [2018] K. Lowrey, A. Rajeswaran, S. Kakade, E. Todorov, and I. Mordatch.  Plan Online, Learn Offline: Efficient Learning and Exploration via Model-Based Control.  _arXiv preprint arXiv:1811.01848_ , 2018.
  * Bhardwaj et al. [2020a] M. Bhardwaj, S. Choudhury, and B. Boots.  Blending MPC & Value Function Approximation for Efficient Reinforcement Learning.  _arXiv preprint arXiv:2012.05909_ , 2020a.
  * Bhardwaj et al. [2020b] M. Bhardwaj, A. Handa, D. Fox, and B. Boots.  Information Theoretic Model Predictive Q-Learning.  In _Learning for Dynamics & Control (L4DC)_, pages 840–850. PMLR, 2020b.
  * Tamar et al. [2017] A. Tamar, G. Thomas, T. Zhang, S. Levine, and P. Abbeel.  Learning from the Hindsight Plan — Episodic MPC Improvement.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 336–343. IEEE, 2017.
  * Karkus et al. [2017] P. Karkus, D. Hsu, and W. S. Lee.  QMDP-Net: Deep Learning for Planning under Partial Observability.  _arXiv preprint arXiv:1703.06692_ , 2017.
  * Okada et al. [2017] M. Okada, L. Rigazio, and T. Aoshima.  Path Integral Networks: End-to-End Differentiable Optimal Control.  _arXiv preprint arXiv:1706.09597_ , 2017.
  * Amos et al. [2018] B. Amos, I. D. J. Rodriguez, J. Sacks, B. Boots, and J. Z. Kolter.  Differentiable MPC for End-to-end Planning and Control.  _arXiv preprint arXiv:1810.13400_ , 2018.
  * Okada and Taniguchi [2018] M. Okada and T. Taniguchi.  Acceleration of Gradient-based Path Integral Method for Efficient Optimal and Inverse Optimal Control.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 3013–3020. IEEE, 2018.
  * Pereira et al. [2018] M. Pereira, D. D. Fan, G. N. An, and E. Theodorou.  MPC-Inspired Neural Network Policies for Sequential Decision Making.  _arXiv preprint arXiv:1802.05803_ , 2018.
  * Tamar et al. [2016] A. Tamar, Y. Wu, G. Thomas, S. Levine, and P. Abbeel.  Value Iteration Networks.  _arXiv preprint arXiv:1602.02867_ , 2016.
  * Srinivas et al. [2018] A. Srinivas, A. Jabri, P. Abbeel, S. Levine, and C. Finn.  Universal Planning Networks: Learning Generalizable Representations for Visuomotor Control.  In _International Conference on Machine Learning (ICML)_ , pages 4732–4741. PMLR, 2018.
  * Yu et al. [2019] T. Yu, G. Shevchuk, D. Sadigh, and C. Finn.  Unsupervised Visuomotor Control through Distributional Planning Networks.  _arXiv preprint arXiv:1902.05542_ , 2019.
  * Bhardwaj et al. [2020] M. Bhardwaj, B. Boots, and M. Mukadam.  Differentiable Gaussian Process Motion Planning.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 10598–10604. IEEE, 2020.
  * Sacks and Boots [2022] J. Sacks and B. Boots.  Learning to Optimize in Model Predictive Control.  In _IEEE International Conference on Robotics and Automation (ICRA)_ , pages 10549–10556. IEEE, 2022.
  * Pinneri et al. [2020] C. Pinneri, S. Sawant, S. Blaes, J. Achterhold, J. Stueckler, M. Rolinek, and G. Martius.  Sample-efficient Cross-Entropy Method for Real-time Planning.  _arXiv preprint arXiv:2008.06389_ , 2020.
  * Wang and Ba [2019] T. Wang and J. Ba.  Exploring Model-Based Planning with Policy Networks.  _arXiv preprint arXiv:1906.08649_ , 2019.
  * Tang and Agrawal [2018a] Y. Tang and S. Agrawal.  Boosting Trust Region Policy Optimization by Normalizing Flows Policy.  _arXiv preprint arXiv:1809.10326_ , 2018a.
  * Tang and Agrawal [2018b] Y. Tang and S. Agrawal.  Implicit Policy for Reinforcement Learning.  _arXiv preprint arXiv:1806.06798_ , 2018b.
  * Mazoure et al. [2020] B. Mazoure, T. Doan, A. Durand, J. Pineau, and R. D. Hjelm.  Leveraging exploration in off-policy algorithms via normalizing flows.  In _Conference on Robot Learning (CoRL)_ , pages 430–444. PMLR, 2020.
  * Lai and Ramos [2020] T. Lai and F. Ramos.  Learning to Plan Optimally with Flow-based Motion Planner.  _arXiv preprint arXiv:2010.11323_ , 2020.
  * Lai et al. [2021] T. Lai, W. Zhi, T. Hermans, and F. Ramos.  Parallelised Diffeomorphic Sampling-based Motion Planning.  _arXiv preprint arXiv:2108.11775_ , 2021.
  * Paszke et al. [2019] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al.  PyTorch: An Imperative Style, High-Performance Deep Learning Library.  _Advances in Neural Information Processing Systems (NeurIPS)_ , 32:8026–8037, 2019.
  * Halton [1964] J. H. Halton.  Algorithm 247: Radical-inverse quasi-random point sequence.  _Communications of the ACM_ , 7(12):701–702, 1964\.
  * Makoviychuk et al. [2021] V. Makoviychuk, L. Wawrzyniak, Y. Guo, M. Lu, K. Storey, M. Macklin, D. Hoeller, N. Rudin, A. Allshire, A. Handa, et al.  Isaac Gym: High Performance GPU-Based Physics Simulation For Robot Learning.  _arXiv preprint arXiv:2108.10470_ , 2021.
  * Ba et al. [2016] J. L. Ba, J. R. Kiros, and G. E. Hinton.  Layer Normalization.  _arXiv preprint arXiv:1607.06450_ , 2016.
  * Kingma and Ba [2014] D. P. Kingma and J. Ba.  Adam: A Method for Stochastic Optimization.  _arXiv preprint arXiv:1412.6980_ , 2014.

##  Appendix A Appendix

###  A.1 Computational Graph Visualization

Figure 5: Computational graph of an episode which illustrates the interaction
between the normalizing flow (NF), learned latent shift model (Shift), the
MPPI update (MPPI) of the latent mean, the simulator (Sim), and the
environment (Env).

We visualize the computational graph of a single episode in Figure 5. The
normalizing flow (NF), shift operator (Shift), and loss computation (Loss) are
colored blue, indicating that they are differentiable modules. The MPPI update
(MPPI) is colored in purple, to indicate that we are approximating the
gradient as above, which can be implemented as a custom backwards pass in
autodifferentiation software. All of the paths of the graph through which the
gradients flow during the backwards pass are colored in purple. The simulator
(Sim), colored in green, is solely used to compute the weights for the MPPI
update, which are reused in the backwards pass. At each time step, the updated
mean $\bm{\mu}_{t}$ defines a latent Gaussian, which is used to generate
controls applied to the actual environment (Env), denoted by the yellow box.
Note that both the latent variables and controls are being passed to the MPPI
module. This is because during the forward pass, we can simply take the
weighted sum of latent variables. However, during the backward pass, we have
to re-run the network with the controls to compute the approximate gradients,
as discussed in Section 3.4.

###  A.2 Experimental Details

#### Controller Details.

Table 1: Controller Hyperparameters |  | Environment
---|---|---
Parameter | PNGrid | PNRand | Franka
Horizon (H) | 32 | 64 | 32
Temperature ($\beta$) | $10^{-32}$ | $10^{-32}$ | $10^{-4}$
MPPI | Init. cov. ($\sigma^{2}$) | 10 | 100 | 0.1
Step size ($\gamma$) | 0.7 | 1 | 1
Spline knots ($n$) | None | None | 4
FlowMPPI | Init. cov. ($\sigma^{2}$) | 10 | 10 | 0.1
Latent cov. | 1 | 1 | 0.1
Latent mean penalty ($\lambda$) | $10^{-4}$ | $10^{-3}$ | 1
NFMPC | Latent cov. | 1 | 1 | 0.1

We use a modified version of the MPPI implementation by Bhardwaj et al. [4],
which is implemented in PyTorch [55]. For the MPPI baseline, we perform
covariance adaptation of the full covariance matrix across the horizon and
control dimensions. The initial covariance matrix is always an identity matrix
scaled by an initial covariance scalar hyperparameter $\sigma^{2}$.
Additionally, this implementation uses Halton sequences [56] for generating
control sequence samples and smooths the sampled trajectories with $n$-degree
B-splines in some tasks. When B-splines are used, we sample the Halton
sequence once at the beginning of a rollout and then transform it using the
mean and covariance of the Gaussian distribution. Additionally, we ensure that
the mean of the Gaussian is always in the set of samples. All hyperparameters
of the controllers were chosen via a grid search and the final choices for
each task are listed in Table 1. When it improves performance, we warm-start
the controllers prior to the first time step by running the MPC update for 100
iterations to ensure convergence. Additionally, we normalize the total
trajectory costs prior to computing the softmax weights, as discused by Okada
and Taniguchi [10].

In general, we use the same settings for NFMPC and FlowMPPI that were found in
this grid search. However, we do not performance covariance adaptation on the
latent Gaussian and assume the flow learns how to adjust sample spread as
needed. Additionally, we take the previous mean in the control space, shift it
forward, and add it to the set of control samples at the next time steps.
While the learned latent shift model handles this well in most cases, we found
adding this sample sped up training and slightly improved performance. We
always use Halton sequences to sample from the latent Gaussian, but never use
B-splines. For FlowMPPI, we always use half the samples for sampling from the
NF and half for the Gaussian perturbations on the current control-space mean.
We do perform covariance adaptation on the perturbation Gaussian and re-tune
its initial covariance. Finally, we have the additional $\lambda$ parameter,
which penalizes the latent Gaussian samples from deviating too much from the
projection of the current control-space mean into the latent space.

#### Planar Robot Navigation.

The planar navigation environment has a state space of $x_{t}\in{}^{4}$, which
consists of the robot’s 2D position, $(p_{x},p_{y})$, and velocity,
$(v_{x},v_{y})$, and a control space of $u_{t}\in{}^{2}$, which are the
robot’s 2D acceleration commands. The robot has double-integrator dynamics
with additive Gaussian noise on the controls, as described by the following
equations:

$\begin{bmatrix}p_{x}\\\ p_{y}\\\ v_{x}\\\
v_{y}\end{bmatrix}_{t+1}=\begin{bmatrix}1&0&\Delta t&0\\\ 0&1&0&\Delta t\\\
0&0&1&0\\\ 0&0&0&1\end{bmatrix}\begin{bmatrix}p_{x}\\\ p_{y}\\\ v_{x}\\\
v_{y}\end{bmatrix}_{t}+\begin{bmatrix}0&0\\\ 0&0\\\ \Delta t&0\\\ 0&\Delta
t\end{bmatrix}(u_{t}+w_{t}),\quad w_{t}\sim\mathcal{N}(0,\sigma\textrm{I}),$
(15)

where we set $\Delta t=0.1$ and $\sigma=1$. Additionally, we added
acceleration limits of
$\stackunder[1.2pt]{$u$}{\rule{3.44444pt}{0.32289pt}}=-10$ and $\bar{u}=10$
for both directions. The cost function consists of the Euclidean distance to
the goal, a signed-distance field representation of the obstacles, and a term
which encourages the robot to stay within the bounds of the map, and a
quadratic control penalty:

$c(x,u)=w_{goal}||x-x_{g}||_{2}^{2}+w_{bound}c_{bound}(x_{t})+w_{coll}\text{SDF}(p_{x},p_{y})+w_{ctrl}||u||_{2}^{2},$
(16)

where we define the map bound cost as:

$c_{bound}(x)=\sum_{i\in(x,y)}\mathbb{I}[{(p_{i}>\bar{p}_{i})\ ||\
(p_{i}<\stackunder[1.2pt]{p}{\rule{3.44444pt}{0.32289pt}}_{i})}]\min{\big{(}(p_{i}-\bar{p}_{i})^{2},(p_{i}-\stackunder[1.2pt]{p}{\rule{3.44444pt}{0.32289pt}}_{i})^{2}\big{)}}.$
(17)

In the above equations, $x_{g}$ is the goal state, $\bar{p}_{i}$ and
$\stackunder[1.2pt]{$p$}{\rule{3.44444pt}{0.32289pt}}_{i}$ are the upper and
lower bounds of the map for each coordinate, and $\text{SDF}(\cdot,\cdot)$
indexes an image which represents the signed-distance field. We mainly focus
on the PNRandDyn task, which involves steering the robot towards a goal
position while avoiding eight dynamic obstacles. While the obstacles drift
randomly in the environment with Gaussian steps, their positions are clipped
to be within map bounds. We do not update their position if the perturbation
would bring it too close to the robot or goal location to prevent collisions
which the robot cannot react to in time to avoid. An episode lasts for 200
time steps and is considered successful if the agent reaches the goal without
colliding into any obstacles. In Section A.3, we also consider a static
version of this environment (PNRand), which simply places the eight obstacles
randomly in the environment, and a version which arranges the obstacles in a
fixed grid (PNGrid).

#### Franka Panda Arm.

The Franka Panda arm environment defines the robot state in joint space with
$x_{t}\in{}^{21}$, consisting of each joint’s angle $\theta_{i}$, angular
velocity $\dot{\theta}_{i}$, and angular acceleration $\ddot{\theta}_{i}$. Its
control space is $u_{t}\in{}^{7}$, which are the angular acceleration commands
for each joint. The dynamics are deterministic and implemented by the Nvidia
Isaac Gym simulator [57]. However, the MPC controllers use a simpler kinematic
model defined by Bhardwaj et al. [4], which is implemented in a batch fashion
by leveraging its linearity:

$\bm{\ddot{\Theta}}=\bm{u},\quad\bm{\dot{\Theta}}=\dot{\Theta}_{t}+S_{l}(1)\text{diag}(\Delta
t)\bm{\ddot{\Theta}},\quad\bm{\Theta}=\Theta_{t}+S_{l}(1)\text{diag}(\Delta
t)\bm{\dot{\Theta}},$ (18)

where the bold symbols indicate that they consist of values along the entire
horizon, $\bm{\Theta}$ includes the angles for all joints, $S_{l}(1)$ is a
lower triangular matrix filled with 1, and $\Delta t$ is a vector of time
steps across the horizon. We use smaller time steps earlier along the horizon
and larger ones for later time steps. By implementing the dynamics in batch,
we avoid iteratively unrolling the dynamics, speeding up controller
computation significantly. For computing cost, we also require the Cartesian
poses $X$, velocities $\dot{X}$, and accelerations $\ddot{X}$ of the end-
effector. These are obtained via:

$X=\text{FK}(\Theta),\quad\dot{X}=J(\Theta)\dot{\Theta},\quad\ddot{X}=\dot{J}(\Theta)\dot{\Theta}+J(\Theta)\ddot{\Theta}$
(19)

where $\text{FK}(\Theta)$ are the forward kinematics and $J(\Theta)$ is the
kinematic Jacobian. The cost function is a weighted sum of a number of terms,
as defined by Bhardwaj et al. [4], which includes: distance of end-effector to
the goal pose $c_{pose}$, a time varying velocity limit cost $c_{stop}$ that
enables stopping within the specified horizon, a joint limit cost $c_{joint}$,
a manipulability cost $c_{manip}$ which encourages the arm to avoid singular
configurations, a self-collision cost $c_{self-coll}$, and a obstacle
collision cost $c_{coll}$. The overall final cost function is then:

$c(x,u)=w_{p}c_{pose}(x)+w_{s}c_{stop}(x)+w_{j}c_{joint}(x)+w_{m}c_{manip}(x)+w_{c}(c_{self-
coll}(x)+c_{coll}(x)).$ (20)

The self-collision cost is implemented with a neural network that predicts the
closest distance between the links of the robot given a configuration. The
collision cost is a binary cost which uses a learned collision checking
function that operates directly on raw point cloud data and classifies if a
robot link is in collision. See Bhardwaj et al. [4] for further details about
each of the individual cost terms. For the FrankaObstacles task, we control
the 7 degree-of-freedom (DOF) Franka Panda robot arm and steer it towards a
target goal from a fixed starting pose while avoiding a single pole obstacle.
The obstacle and goal positions are randomized at the beginning of each
episode, which lasts for 600 time steps. An episode is considered successful
if the end effector reaches the target position under the time constraints
while avoiding the obstacle. In Section A.4, we also consider a simplified
version which has no obstacles (Franka).

#### Architectural and Training Details.

All NFs for both NFMPC and FlowMPPI were implemented in PyTorch and contain
affine coupling layers which use multilayer perceptrons (MLPs) for both the
scale and translation terms. The scale and translation networks use Tanh and
ReLU activations, respectively. We also employ layer normalization [58] in
these networks to help prevent overfitting. Interestingly, we found that
adding batch normalization between each layer, as proposed by Dinh et al.
[18], actually hurt performance and was therefore excluded. In all
environments, we use 5 RealNVP blocks for the NF, and each MLP has a hidden
dimensionality of 128 neurons. For the PNGrid, Franka, and FrankaObstacles
tasks, the shift model is also an MLP with a single hidden layer of 128
neurons and a ReLU activation function. In PNRand and PNRandDyn the shift
model is implemented as an LSTM with a hidden dimensionality of 128 neurons.
For a fair comparison, the same architecture was used for both FlowMPPI and
NFMPC. We trained all networks with the Adam optimizer [59] using a learning
rate of $10^{-4}$.

Input: Environment dist. $\mathcal{C}$, initial state dist. $\rho$, initial
param. $\tilde{\theta}_{0}$, $\lambda_{0}$

Parameters: # episodes $D$, episode length $T$, # samples $N$

for  _$d=1,2,\dots,D$_ do

Sample environment $c\sim\mathcal{C}(\cdot)$

Sample initial state $x_{0}\sim\rho(\cdot|c)$

Initialize episode loss $l_{d}$ $\leftarrow$ $0$

for  _$t=0,1,\dots,T_{k}-1$_ do

$(\bm{\hat{u}}_{t},\bm{\hat{z}}_{t},w_{t})^{(1:N)}$ $\leftarrow$
Rollout($x_{t}$, $c$, $\tilde{\theta}_{t}$, $\lambda_{d}$)

Update $\tilde{\theta}_{t}$ to $\theta_{t}$ using latent MPPI update

Sample $u_{t}\sim\pi_{\theta_{t},\lambda_{d}}$ or $u_{t}\leftarrow
h_{\lambda_{d}}^{-1}(\mu_{t};c)$

Apply control to system $x_{t+1}\sim f(x_{t},u_{t})$

Shift parameters $\tilde{\theta}_{t+1}$ $\leftarrow$
$\Phi_{\lambda_{d}}(\theta_{t},c)$

Compute loss $\hat{J}_{t}(\theta_{t},\lambda_{d};x_{t},c)$

Accumulate loss $\ell_{d}\leftarrow\ell_{d}+\hat{J}_{t}$

end for

Update $\lambda_{d}$ to $\lambda_{d+1}$ with $\nabla_{\lambda}\ell_{d}$ using
SGD

end for

Algorithm 1 Training Loop

To train the NFMPC variants, we following the training procedure in Algorithm
1, which is carried out training over $D$ episodes. In each episode $d$, we
sample an environment from $\mathcal{C}$ and initial state from $\rho$. We
then perform our rollouts by sampling latent controls
$\bm{\hat{z}}_{t}^{(1:N)}$, passing them through the normalizing flow to get
controls $\bm{\hat{u}}_{t}^{(1:N)}$, and then applying them to our approximate
dynamics model and cost function to get weights $w_{t}^{(1:N)}$. These
variables are used to update the latent distribution of the policy to
$\theta_{t}$. Next, we can either sample a control from the policy or use a
control corresponding to the latent mean. We apply this control to the true
system and shift the latent distribution parameters forward with the shift
model. Finally, we compute the loss for the current time step, accumulate the
loss to our running sum, and repeat for $T$ time steps. Once the episode is
complete, we update $\lambda_{d}$ using the gradient of the loss and carry on
to the next episode. Every 100 environments, we test the controller on 10 held
out environments and save the current model if it outperforms the previous
best on this validation set. While this introduces a variable number of total
episodes used to train the models, we generally find convergence between 2000
and 7000 episodes.

In contrast, to train FlowMPPI, we do not actually run an episode. Instead, we
generate the environment and take a gradient step on the initial distribution
of the flow conditioned on the environmental information. We train FlowMPPI
over 10000 randomly generated environments for each task. We achieved the best
performance by initializing the flow with the FlowMPPI solution, and then
refining the flow and learning the shift model jointly as above. When we
condition the NF on additional information, this is simply appended to the
current input of each shift and translation network directly. This conditional
information includes start and goal locations for PNGrid and Franka and start,
goal, and obstacle locations for PNRand and FrankaObstacles. For PNRandDyn we
use the current state instead of the initial location, which was found to
improve performance for all controllers. For the obstacles, the conditional
information is the Cartesian coordinates of each object of interest stacked
together in a vector.

###  A.3 Additional Planar Navigation Experiments

Figure 6: Success rate and cost distribution on the PNGrid task across a
different number of samples.

NFMPC

FlowMPPI

MPPI

Figure 7: Visualization of trajectories in the PNGrid task across multiple
random seeds for a fixed environmental layout.

NFMPC

FlowMPPI

MPPI

Figure 8: Visualization of a trajectory and top samples from (top) NFMPC,
(middle) FlowMPPI, and (bottom) MPPI on the PNGrid task.

####  PNGrid.

We consider a variant of the planar navigation task that involves static
obstacles arranged in a grid (PNGrid), described in Section A.2. All other
task details are the same as in the PNRandDyn task, except that we use an
unconditional NFMPC controller. We quantitatively compare all controllers in
Figure 6 and find that NFMPC consistently matches or outperforms both MPPI and
FlowMPPI at each sample quantity in terms of both success rate and average
trajectory cost of successful trajectories. We also find that NFMPC scales
more gracefully than MPPI as the number of samples is reduced. In fact, we
found that while FlowMPPI improves over MPPI at higher sample counts, it
actually performs significantly worse with fewer samples. This is in contrast
to the results on more complex environments considered in the main paper, in
which FlowMPPI generally outperforms MPPI at lower sample counts as well. This
is potentially because in the standard implementation of FlowMPPI, half of the
samples come from the NF and the other half are Gaussian perturbations of the
current control-space mean. Initially, the samples coming from the NF provide
a good initialization for the control distribution mean. However, as the
latent Gaussian distribution used by the NF is never updated, half of our
samples are always coming from this same distribution. When the environment or
task is more complex, the conditioning information provided to the NF is
enough to transform the samples in a useful way. However, in this simple
environment, our hypothesis is that as the robot moves in the environment,
these samples may cease to be as useful or informative. We then have to rely
on the other half of samples coming from Gaussian perturbations of the
control-space mean to do most of the work. As such, we effectively have half
the budget of samples to work with than MPPI would, as the samples from the NF
potentially do not provide much useful information. Meanwhile, because NFMPC
is trained recurrently and updates are performed in the latent space, it can
better exploit structure in the environment to transform samples.

To better understand what NFMPC is doing differently, we superimpose 32
different trajectories with fixed start and goal positions using each
controller in Figure 7. We find that both MPPI and FlowMPPI always select the
same path through the environment. Meanwhile, NFMPC is able to discover
different paths through the environment, allowing it to better react to the
stochastic perturbations that knock it off the current plan and improve
performance. Additionally, we visualize the resulting trajectories and top
samples drawn from the distributions for all controllers on one of the
validation environments in Figure 8. As we would expect, the initial
trajectory from FlowMPPI is better than those of the other two controllers,
which basically lay in straight lines in front of the robot. However, NFMPC is
able to discover a slightly faster route to the goal as it proceeds in the
environment. The baseline MPPI controller takes a similar, but slightly
longer, route to the goal. Additionally, it takes a longer time to ramp up its
velocity compared to the other two controllers, contributing to its sub-
optimal performance.

Figure 9: Success rate and cost distribution on the PNRand environment across
a different number of samples.

NFMPC

FlowMPPI

MPPI

Figure 10: Visualization of a trajectory and top samples from (top) NFMPC,
(middle) FlowMPPI, and (bottom) MPPI on the PNRand task.

####  PNRand.

Next, we consider a variant of PNRandDyn in which the eight obstacles are
static (PNRand). We consider the case where we condition the NF on obstacle
locations, initial state, and goal position. However, it is important to note
that we do not condition the shift model, as we found this consistently hurt
performance. In Figure 9, we display the quantitative results and find that
NFMPC performed about on par with FlowMPPI, which outperformed MPPI. Unlike on
the PNGrid task, both NFMPC and FlowMPPI scale similarly with a reduction of
samples and better than MPPI. This can be partly attributed to the fact that
the obstacles are more spaced out and there are more ”holes” in the
environment than in the grid. Therefore, it is easier to avoid collisions,
possibly contributing to the higher success rates when the controller has
access to fewer samples. Moreover, conditioning on the obstacle locations
provides the NF more information, which can be exploited without updating the
latent distribution.

We again visualize the trajectories and top samples for all controllers on a
validation environment in Figure 10. First, we note that the samples in
FlowMPPI appear to be better spread around in the environment to search for
good paths towards the goal. Meanwhile, NFMPC seems to have all top samples
directed in one direction. All models seem to find the same path, with MPPI
oscillating more near the goal and reaching the goal more slowly than NFMPC
and FlowMPPI. The similar performance of NFMPC and FlowMPPI can be partially
attributed to the fact that all we may really need to succeed in this
environment is a good initial trajectory which steers us towards the goal.
Whether we refine this trajectory with a learned latent shift model or with
Gaussian perturbations in the control space does not appear to make much
difference in performance. However, in PNRandDyn, we demonstrated that there
is a significant advantage to our approach, indicating that dynamic
environments may be a good application for NFMPC.

Figure 11: Comparison of unconditional and conditional models on the PNRand
environment across a different number of samples.

Finally, in order to evaluate the benefit of conditioning the NF, we compare
the performance of NFMPC with and without conditioning the flow on PNRand in
Figure 11. We find that the conditional model consistently outperforms the
unconditional model in terms of median cost, with the gap growing at reduced
sample counts. This may be because the dynamics in PNRand are rather simple,
and there may not be much general structure for the unconditional model to
exploit since the obstacle locations are entirely random. Therefore, while the
unconditional model does fairly well, conditioning the flow, and not the shift
model, seems to enable further gains, even in this simple task.

Figure 12: Comparison of unconditional and conditional models on the PNRandDyn
environment across a different number of samples when trained in an
environment with no obstacles (PNRand) and retrained with obstacles present.

####  PNRandDyn.

In addition to the experiments in the main paper, we also performed additional
ablation studies. Specifically, we also considered training an unconditional
model, as we did before in the PNRand task. Moreover, we explored transferring
controllers trained in the PNRand task to this dynamic version of the
environment. We plot the quantitative results from these ablation studies in
Figure 12. Again, we find that conditional models consistently outperform
unconditional ones, with the gap in performance more pronounced. Not only is
there a reduction in median cost for unconditional controllers, but it
sometimes fails in environments in which its conditional counterpart succeeds.
Transferring the models trained in PNRand works surprisingly well at 1024
samples. However, at lower sample counts, there is a more pronounced
difference in the transferred controllers to those specifically trained in the
dynamic environment. Therefore, it appears that controllers are more efficient
at exploring the environments they are trained on, and unsurprisingly, using
more samples can partially overcome this gap.

###  A.4 Additional Franka Experiments

Figure 13: Success rate and cost distribution on the Franka environment across
a different number of samples.

####  Franka.

We consider a variant of the FrankaObstacles task which involves no obstacles,
just a target goal, which we call Franka. We plot our quantitative results in
Figure 13 and find that NFMPC again consistently matches or outperforms MPPI
and FlowMPPI at each sample amount. In fact, NFMPC always achieves a $100\%$
success rate at all sample counts, while both MPPI and FlowMPPI significantly
drop in performance at lower amounts of samples. Moreover, FlowMPPI sometimes
actually performed worse than MPPI, indicating that conditioning on just goal
location does not help as much in this more complex scenario.

Figure 14: Comparison of an unconditional model trained with and without
obstacles to a conditional model in the FrankaObstacles environment across a
different number of samples.

####  FrankaObstacles.

In addition to the results in the main paper, we perform an additional
ablation study in which we again compare an unconditional and conditional
model on the FrankaObstacles environment in Figure 14. We also show the
performance of transferring the unconditional model trained on the Franka
environment without obstacles to an environment which contains the single pole
obstacle. At 1024 samples, the conditional model performs the best in terms of
success rate. Surprisingly, the unconditional model trained without obstacles
performs best in terms of median cost and scales better to fewer samples.
Therefore, while the conditional model more often finds a feasible path to the
goal, the unconditional model is better able to exploit structure across
environments to find lower cost trajectories. One possible explanation is that
because the unconditional model is trained without knowing the specific
obstacle locations, it has to be more robust to variation in the environment.
This also shows that transferring the learned distribution to novel
environments is possible. However, since we only have a single static
obstacle, it is not clear if these findings would generalize to more complex
environments.

#### Breakdown of Trajectory Cost.

We would like to better understand how the learned controllers improve upon
the baseline on the FrankaObstacles task. As discussed in Section A.2, the
cost function for the Franka tasks is composed of multiple terms. By
inspecting the averages for each term, we hope to gain insight into the
learned sampling distributions. Briefly, we consider a manipulability cost
(Manip), which encourages the arm to avoid singular configurations, a self-
collision cost (Self), an obstacle collision cost (Obstacle), and a distance-
to-goal cost (Goal).

Figure 15: Breakdown of different cost terms for each controller executing the
FrankaObstacles task.

The cost breakdown for these terms in shown in Figure 15. The baseline MPPI
controller achieves the lowest manipulability cost. Since this term is not
weighted as highly, it makes sense that training the normalizing flow would
focus on minimizing terms which were more heavily weighted. Meanwhile,
FlowMPPI achieves the lowest self-collision cost, which may be due to the fact
that it starts off with a better initial plan. However, NFMPC trained with
obstacles (NFMPC (Obs)) achieves the lowest obstacle collision and distance-
to-goal cost. One possible explanation for this improvement is that, because
we train NFMPC with BPTT, we are potentially able to account for errors which
arise due to the inaccurate model. Additionally, learning the shift model may
be an important component to this improvement, which we explore below.
Finally, we note that NFMPC (Obs) is better than NFMPC trained without
obstacles (NFMPC (No Obs)) in all cost terms except the manipulability cost,
as it is in distribution for the task.

Figure 16: Success rate and cost distribution on the Franka task with (NFMPC)
and without (NFMPC (No Shift) and NFMPC (No Shift, Retrained) the learned
shift model.

#### Shift Model Ablation.

One of our main contributions which sets us apart from prior work is that we
learn a latent shift model and train the controller as a recurrent network. To
determine the impact of this choice, we consider an alternate scenario where
the control sequence is instead shifted in control space. That is, we shift
the mean in control space, as in the baseline MPPI controller, and then run
the flow backwards to infer the corresponding latent mean, which is used to
bootstrap the next iteration. We consider two cases: 1) taking a pre-trained
controller and removing the shift model (NFMPC (No Shift)) and 2) retraining
the controller entirely from scratch without the shift model (NFMPC (No Shift,
Retrained)). In Figure 16, we show the results of this ablation study on the
Franka task. Removing the shift model strictly hurts performance, even when we
retrain the normalizing flow from scratch. Moreover, retraining the flow
actually results in worse performance than simply removing the shift model
from the pre-trained controller. This implies that training the entire
controller with BPTT allows it to discover lower-cost trajectories to the
goal.

###  A.5 Performance Overhead of Normalizing Flow

We measured the average change in wall clock time across different amounts of
samples for NFMPC and FlowMPPI compared to the baseline MPPI implementation on
our NVIDIA Titan V GPU. For the Planar Robot Navigation and Franka Panda Arm
experiments, the average change is $1.61\times$ and $1.91\times$,
respectively. Moreover, compared to MPPI with 1024 samples, the change in wall
clock time for NFMPC and FlowMPPI with 16 samples is approximately
$1.01\times$ and $1.14\times$, respectively. Therefore, the introduction of
the normalizing flow (NF) has a notable impact on wall clock time for both
methods. However, this overhead is expected and does not prohibit either
method’s utility in the real world. And there are still clear performance
advantages of NFMPC over the baselines in terms of success rate and average
trajectory cost for a given sample amount.

Furthermore, the performance of MPPI reported in the paper and the above
timing comparisons are when the optimization is run for a single iteration per
time step. However, the performance of MPPI generally improves with an
increased number of iterations at the cost of an increased runtime. For
instance, we compare the performance of MPPI run for 3 iterations per time
step with NFMPC run for a single iteration, both using 1024 samples. In the
Franka task, we can reduce the average trajectory cost of MPPI to be only
$12\%$ worse than NFMPC, rather than the previous $19\%$. When we introduce
obstacles for the FrankaObstaclestask, MPPI with 3 iterations can actually
match the success rate of NFMPC, albeit with a worse average trajectory cost.
However, in this case, NFMPC actually results in an average reduction of wall
clock time by $24.8\%$. Similarly, for the PNGrid task, MPPI with 3 iterations
reduces the average trajectory cost to be only $10\%$ worse than NFMPC, rather
than the previous $40\%$. However, this again comes at the cost of increased
runtime, as NFMPC reduces the average wall clock time by $26.8\%$. Therefore,
NFMPC allows us to surpass the performance of MPPI run with more iterations
while reducing the required runtime.

Additionally, NFMPC has substantial runtime benefits over FlowMPPI due to its
improved scaling. For the PNGridtask, NFMPC with 128 samples outperforms
FlowMPPI with 1024 samples while reducing average runtime by $22\%$.
Similarly, for the Frankatask, NFMPC with 64 samples outperforms FlowMPPI with
1024 samples while reducing average runtime by $43\%$. And for both Franka
tasks (Franka and FrankaObstacles), NFMPC with 16 samples outperforms FlowMPPI
with 128 samples while reducing average runtime by $8\%$. As such, there are
clear runtime benefits for NFMPC over both FlowMPPI and MPPI run for
additional iterations. Finally, it is also important to note that we did not
perform a hyperparameter sweep on the NF, and it may be possible to
significantly reduce the size of the network while retaining performance
benefits.

###  A.6 Proof for DMD Update of Latent Parameters

######  Theorem A.1.

Consider the optimization problem

$\theta_{t}=\operatorname*{arg\,min}_{\theta\in\Theta}\langle\gamma_{t}g_{t},\theta\rangle+D_{KL}(\pi_{\theta,\lambda}||\pi_{\tilde{\theta}_{t},\lambda})$
(21)

where we define

$\pi_{\theta,\lambda}(\bm{\hat{u}}|c)=p_{\theta}(h_{\lambda}(\bm{\hat{u}};c))\Big{|}\text{det}\
\frac{\partial h_{\lambda}}{\partial\bm{\hat{u}}}\Big{|}$ (22)

and $h_{\lambda}$ is an invertible, deterministic transformation, $p_{\theta}$
a Gaussian factorized as in Equation 4, and $\theta$ is the natural parameters
of the Gaussian. Then the corresponding update to the mean of the latent
Gaussian is given by:

$\bm{\mu}_{t}=(1-\gamma_{t}^{\mu})\bm{\tilde{\mu}}_{t}+\gamma_{t}^{\mu}\frac{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}h_{\lambda}(\bm{\hat{u}}_{t};c)\Big{]}}{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}}=(1-\gamma_{t}^{\mu})\bm{\tilde{\mu}}_{t}+\gamma_{t}^{\mu}\frac{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c))}\bm{\hat{z}}_{t}\Big{]}}{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c))}\Big{]}}$
(23)

###### Proof.

First, we note that

$g_{t}=\nabla\hat{J}(\tilde{\theta};x_{t})=-\frac{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\nabla_{\tilde{\theta}}\log\pi_{\tilde{\theta},\lambda}(\bm{\hat{u}}_{t})\Big{]}}{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}},$
(24)

and that

$\log\pi_{\tilde{\theta},\lambda}(\bm{\hat{u}}|c)=\log
p_{\tilde{\theta}}(h_{\lambda}(\bm{\hat{u}};c))+\log\Big{(}\Big{|}\text{det}\
\frac{\partial h_{\lambda}}{\partial\bm{\hat{u}}}\Big{|}\Big{)}.$ (25)

Therefore, when computing the gradient of Equation 25 with respect to
$\tilde{\theta}$, we can drop the log-determinant term as it does not depend
on $\tilde{\theta}$. As such, we are left with the gradient of the latent
Gaussian with respect to its natural parameters, or
$\nabla_{\tilde{\theta}}\log\pi_{\tilde{\theta},\lambda}(\bm{\hat{u}}|c)=\nabla_{\tilde{\theta}}\log
p_{\tilde{\theta}}(h_{\lambda}(\bm{\hat{u}};c))$. Taking this gradient of the
log-likelihood term with respect to the natural parameters, we have:

$g_{t}=-\frac{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\big{(}\phi(h_{\lambda}(\bm{\hat{u}}_{t};c))-\tilde{\phi}_{t}\big{)}\Big{]}}{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}},$
(26)

where $\tilde{\phi}_{t}$ is the expectation parameter corresponding to natural
parameter $\tilde{\theta}_{t}$ and $\phi(\cdot)$ is the sufficient statistics
of the latent distribution. We can rewrite the expectations in terms of
$p_{\tilde{\theta}}(\cdot)$:

$g_{t}=-\frac{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c)))}\big{(}\phi(\bm{\hat{z}}_{t})-\tilde{\phi}_{t}\big{)}\Big{]}}{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c))}\Big{]}},$
(27)

Next, looking at the KL divergence term, we can write:

$\displaystyle D_{KL}(\pi_{\theta,\lambda}||\pi_{\tilde{\theta}_{t},\lambda})$
$\displaystyle=\mathbb{E}_{\pi_{\theta,\lambda}}\Bigg{[}\log\frac{\pi_{\theta,\lambda}(\bm{\hat{u}})}{\pi_{\tilde{\theta}_{t},\lambda}(\bm{\hat{u}})}\Bigg{]}$
(28)
$\displaystyle=\mathbb{E}_{\pi_{\theta,\lambda}}\Bigg{[}\log\frac{p_{\theta}(h_{\lambda}(\bm{\hat{u}};c))}{p_{\tilde{\theta}}(h_{\lambda}(\bm{\hat{u}};c))}\cancel{\frac{\Big{|}\text{det}\
\frac{\partial h_{\lambda}}{\partial\bm{\hat{u}}}\Big{|}}{\Big{|}\text{det}\
\frac{\partial h_{\lambda}}{\partial\bm{\hat{u}}}\Big{|}}}\Bigg{]}$
$\displaystyle=\mathbb{E}_{p_{\theta}}\Bigg{[}\log\frac{p_{\theta}(\bm{\hat{z}})}{p_{\tilde{\theta}}(\bm{\hat{z}})}\Bigg{]}$
$\displaystyle=D_{KL}(p_{\theta}||p_{\tilde{\theta}_{t}}),$

Then, using Proposition 1 from Wagener et al. [9], we can write the update
rule as

$\phi_{t}=(1-\gamma_{t})\tilde{\phi}_{t}+\gamma_{t}\frac{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c)))}\phi(\bm{\hat{z}}_{t})\Big{]}}{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c))}\Big{]}}$
(29)

And when the sufficient statistic is
$\phi(\bm{\hat{z}}_{t})=(\bm{\hat{z}}_{t},\bm{\hat{z}}_{t}\bm{\hat{z}}_{t}^{T})$,
then we arrive at the usual MPPI update for the mean, but now defined in terms
of the latent samples:

$\bm{\mu}_{t}=(1-\gamma_{t}^{\mu})\bm{\tilde{\mu}}_{t}+\gamma_{t}^{\mu}\frac{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c))}\bm{\hat{z}}_{t}\Big{]}}{\mathbb{E}_{p_{\tilde{\theta}},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},h_{\lambda}^{-1}(\bm{\hat{z}}_{t};c))}\Big{]}}$
(30)

which we can equivalently rewrite in terms of $\pi_{\tilde{\theta},\lambda}$
as:

$\bm{\mu}_{t}=(1-\gamma_{t}^{\mu})\bm{\tilde{\mu}}_{t}+\gamma_{t}^{\mu}\frac{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}h_{\lambda}(\bm{\hat{u}}_{t};c)\Big{]}}{\mathbb{E}_{\pi_{\tilde{\theta},\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}}$
(31)

∎

###  A.7 Proof of Approximate Gradient Through MPPI Update

######  Theorem A.2.

Let the MPPI update of the latent mean be given by

$\bm{\mu}_{t}(\lambda)=(1-\gamma_{t}^{\mu})\bm{\tilde{\mu}}_{t}(\lambda)+\gamma_{t}^{\mu}\Delta\bm{\mu}_{t},\qquad\Delta\bm{\mu}_{t}=\frac{\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}h_{\lambda}(\bm{\hat{u}}_{t};c)\Big{]}}{\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}}.$
(32)

Then the gradient of $\Delta\bm{\mu}_{t}$ with respect to $\lambda$ can be
approximated as

$\frac{\partial\Delta\bm{\mu}_{t}}{\partial\lambda}\approx M_{1}-M_{2}M_{3}$
(33)

where

$\begin{gathered}M_{1}=\sum_{i=1}^{N}w_{i}\Big{[}\nabla_{\lambda}h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)+h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)\Big{]},\\\
M_{2}=\sum_{i=1}^{N}w_{i}h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c),\quad
M_{3}=\sum_{i=1}^{N}w_{i}\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c).\end{gathered}$
(34)

###### Proof.

First, we rewrite $\Delta\bm{\mu}_{t}=\frac{N(\lambda)}{D(\lambda)}$, where

$N(\lambda)=\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}h_{\lambda}(\bm{\hat{u}}_{t};c)\Big{]}\qquad
D(\lambda)=\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{]}.$
(35)

Then by the quotient rule of calculus, we have

$\frac{\partial\Delta\bm{\mu}_{t}}{\partial\lambda}=\frac{\nabla_{\lambda}N(\lambda)}{D(\lambda)}-\frac{N(\lambda)}{D(\lambda)}\frac{\nabla_{\lambda}D(\lambda)}{D(\lambda)}.$
(36)

We can compute each of these individual terms using the likelihood ratio
gradients

$\nabla_{\lambda}N(\lambda)=\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\Big{(}\nabla_{\lambda}h_{\lambda}(\bm{\hat{u}}_{t};c)+h_{\lambda}(\bm{\hat{u}}_{t};c)\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)\Big{)}\Big{]}$
(37)
$\nabla_{\lambda}D(\lambda)=\mathbb{E}_{\pi_{\tilde{\theta}_{t}(\lambda),\lambda},\hat{f}}\Big{[}e^{-\frac{1}{\beta}C(\bm{\hat{x}}_{t},\bm{\hat{u}}_{t})}\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)\Big{]}$
(38)

Since each of the terms that make up our gradient in Equation 36 are divided
by $D(\lambda)$, when we approximate them with Monte Carlo sampling, we can
actually write them in terms of the same weights used by MPPI in the forward
pass:

$\displaystyle\frac{\nabla_{\lambda}N(\lambda)}{D(\lambda)}$
$\displaystyle=\sum_{i=1}^{N}w_{i}\Big{[}\nabla_{\lambda}h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)+h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)\Big{]}=M_{1}$
(39a) $\displaystyle\frac{N(\lambda)}{D(\lambda)}$
$\displaystyle=\sum_{i=1}^{N}w_{i}h_{\lambda}(\bm{\hat{u}}_{t}^{(i)};c)=M_{2}$
(39b) $\displaystyle\frac{\nabla_{\lambda}D(\lambda)}{D(\lambda)}$
$\displaystyle=\sum_{i=1}^{N}w_{i}\nabla_{\lambda}\log\pi_{\tilde{\theta}(\lambda),\lambda}(\bm{\hat{u}}_{t}^{(i)}|c)=M_{3}$
(39c)

∎

###  A.8 Sigmoid Flow Layer

We wish to use a sigmoid layer in our normalizing flow to constrain our
control sample $\bm{\hat{u}}$ such that each control along the horizon is
between $\stackunder[1.2pt]{$u$}{\rule{3.44444pt}{0.32289pt}}$ and $\bar{u}$.
Since the sigmoid function is invertible, if we append a sigmoid layer at the
end of our flow, we have that

$\bm{\hat{u}}=w\sigma(\bm{\hat{y}}_{K-1})+b\quad\Longleftrightarrow\quad\bm{\hat{y}}_{K-1}=\sigma^{-1}\Big{(}\frac{\bm{\hat{u}}-b}{w}\Big{)}=\log\Big{(}\frac{\bm{\hat{u}}-b}{w-\bm{\hat{u}}+b}\Big{)},$
(40)

where $w=\bar{u}-\stackunder[1.2pt]{$u$}{\rule{3.44444pt}{0.32289pt}}$ and
$b=\stackunder[1.2pt]{$u$}{\rule{3.44444pt}{0.32289pt}}$, the sigmoid and
logit functions are applied element-wise, and the scaling and translation are
broadcasted to each element of the vector $\bm{\hat{u}}$. The derivative of
the forward transformation is given by:

$\displaystyle\frac{\partial}{\partial x}\Big{(}w\sigma(x)+b\Big{)}$
$\displaystyle=w\sigma(x)(1-\sigma(x))$ (41)
$\displaystyle=\frac{w}{1+e^{-x}}\Big{(}1-\frac{1}{1+e^{-x}}\Big{)}$
$\displaystyle=\frac{w}{(1+e^{-x})(1+e^{x})}.$

Since the sigmoid is applied element-wise, it has a diagonal Jacobian, the
log-determinant of which is simply the sum of the log of its diagonal terms:

$\log\Big{|}\text{det}\
\frac{\partial\bm{\hat{u}}}{\partial\bm{\hat{y}}_{K-1}}\Big{|}=\sum_{i=1}^{MH}\log(w)-\log(1+e^{-\hat{u}_{i}})-\log(1+e^{\hat{u}_{i}}),$
(42)

where we can implement the last two terms with the Softplus activation
function. In the reverse direction, we have that:

$\displaystyle\frac{\partial}{\partial
x}\Bigg{(}\sigma^{-1}\Big{(}\frac{x-b}{w}\Big{)}\Bigg{)}$
$\displaystyle=\frac{1}{x-b}+\frac{1}{w-x-b}$ (43)

Therefore, the log-determinant is given by:

$\log\Big{|}\text{det}\
\frac{\partial\bm{\hat{y}}_{K-1}}{\partial\bm{\hat{u}}}\Big{|}=-\sum_{i=1}^{MH}\Big{(}\log(\hat{u}_{i}-b)+\log(w-\hat{u}_{i}-b)\Big{)},$
(44)

As such, computing the inverse and the log-determinant terms of the Jacobian
is efficient and fast in both directions, adding minimal overhead to the flow.