Variational Inference MPC using Normalizing Flows and Out-of-Distribution Projection

The connection between control and inference is long established [11, 31, 30]. Attias [2] first framed planning as an inference problem, and proposed a tractable inference algorithm for discrete state and action spaces. Further work has used inference techniques for planning [20, 21] and Stochastic Optimal Control (SOC) [33, 27, 35]. Two widely used sampling based MPC techniques, MPPI [36] and CEM [15], use importance sampling to generate low-cost control sequences, and have strong connections to the inference formulation of SOC which was explored in [34]. Several recent papers have considered the SOC problem as Bayesian inference, and proposed methods of performing Variational Inference (VI) to approximate a posterior over low-cost control sequences [17, 34, 23, 3]. These methods differ in how they represent the variational posterior. VI methods often use an independent Gaussian posterior, known as the mean-field approximation [4]. Okada and Taniguchi [23] represent the control sequence as a Gaussian mixture, and Lambert et al. [17] use a particle representation, extended to handle parameter uncertainty in [3]. These representations allow for greater flexibility in representing complex posteriors. We will similarly use a flexible class of distributions to represent the posterior, but will further make the posterior dependent on the start, goal, and environment. To our knowledge our approach is the first to amortize the cost of computing this posterior by learning a conditional control sequence posterior from a dataset.

Our work is related to work learning sampling distributions from data for motion planning. Zhang et al. [39] proposed learning a sampling distribution that is trained across multiple environments, but is independent of the environment. Others have proposed learning a sampling distribution which is dependent on the environment, start and goal [10, 26]. These methods were restricted to geometric planning, but Li et al. [18] proposed an approach for kinodynamic planning which learns a generator and discriminator which are used to sample states that are consistent with the dynamics. Recent work by Lai et al. [16] uses a diffeomorphism to learn the sampling distribution; a model that is similar to a normalizing flow. The model we propose will also learn to generate samples conditioned on the start, goal and environment, though in this work we are considering online MPC and not offline planning. Loew et al. [19] uses probabilistic movement primitives (ProMPs) learned from data as the sampling distribution for sample-based trajectory optimization, however the representation of these ProMPs only allows for uni-modal distributions and the sampling distribution is not dependent on the environment. Adaptive and learned importance samplers have been used for sample-based MPC [12, 6], but these methods only consider a single control problem and the learned samplers do not generalize to different goals & environments.

We can reformulate SOC as an inference problem (as in [27, 32, 23, 17]). First, we introduce a binary ‘optimality’ random variable $o$ for a trajectory such that

We place a prior $p(U)$ on $U$, resulting in a prior on $\tau$, $p(\tau)=p(X|U)p(U)$ and aim to find posterior distribution $p(\tau|o=1)\propto p(o=1|\tau)p(\tau)$. In general, this posterior is intractable, so we use variational inference to approximate it with a tractable distribution $q(\tau)$ which minimizes the KL-divergence $\mathcal{KL}(q(\tau)||p(\tau|o=1))$ [4]. Since we define the trajectory by selecting the controls, the variational posterior factorizes as $p(X|U)q(U)$. Thus, we must compute an approximate posterior over control sequences. The quantity to be minimized is

Simplifying and omitting terms that do not depend on $\tau$ yields the variational free energy

Where $\mathcal{H}(q(U)$ is the entropy of $q(U)$. Intuitively, we can understand that the first term promotes low-cost trajectories, the second is a regularization on the control, and the entropy term prevents the variational posterior collapsing to a maximum a posteriori (MAP) solution. Note that $\log p(U)$ can be appropriately combined with the cost, i.e. a Gaussian prior can be incorporated as a squared cost on the control, so will be omitted for the rest of the paper.

Normalizing flows are bijective transformations that can be used to transform a random variable from some base distribution (i.e. a Gaussian) to a more complex distribution [28, 8, 14]. Consider a random variable $z\in\mathbb{R}^{d}$ and with known pdf $p(z)$. Let us define a bijective function $f:\mathbb{R}^{d}\to\mathbb{R}^{d}$ and a random variable $y$ such that $y=f(z)$ and $z=f^{-1}(y)$. According to the change of variable formula, we can define $p(y)$ in terms of $p(z)$ as follows:

Normalizing flows can be used as a parameterization of the variational posterior [28]. By selecting a base PDF $p(z)$ and a family of parameterized functions $f_{\theta}$, we specify a potentially complex set of possible densities $q_{\theta}(y)$. Suppose that we want to approximate some distribution $p(y)$ with some distribution $q_{\theta}(y)$. The variational objective is to minimize $\mathcal{KL}(q_{\theta}(y)||p(y))$. This is equivalent to:

Thus we can optimize the parameters $\theta$ of the bijective transform $f_{\theta}$ in order to minimize the variational objective. We will use a normalizing flow to represent the control sequence posterior in our method.

The control sequence posterior introduced in section IV-A is specific to each MPC problem. Our approach is to use dataset $\mathcal{D}$ to learn a conditional control sequence posterior $q(U|x_{0},x_{G},E)$. We will use a Conditional Normalizing Flow [37] to represent this conditional posterior as $q_{\zeta}(U|C)$. $C$ is the context vector which we compute as follows: First, we input $E$ into the encoder of a Variational Autoencoder (VAE) [13] to produce a distribution over environment embeddings $h$. We then sample from this distribution to produce an $h$. A neural network $g_{\omega}$ then produces $C$ from $(x_{0},x_{G},h)$ (Figure 2). Essentially $C$ is a representation of what is important about the start, goal, and environment for generating low-cost trajectories.

The above models are trained on the dataset $\mathcal{D}$, which consists of randomly sampled starts, goals and simulated environments. To train the system we iteratively generate samples from the control sequence posterior, weigh them by their cost, and perform a gradient step on the parameters of our models to maximize the likelihood of low-cost trajectories.

At inference time, we simply compute $C$ and generate control sequence samples from $q_{\zeta}(U|C)$. Below we describe each component of the method to learn $q_{\zeta}(U|C)$ in detail.

As discussed, our dataset $\mathcal{D}$ consists of environments, starts and goals. The details of the dataset generation for each task can be found in section VI-A. Since the environment is a high dimensional SDF, we must first compress it to make it computationally tractable to train the control sequence posterior. To encode the environment, we use a VAE with environment embedding $h$. The VAE consists of an encoder $q_{\theta}(h|E)$, which is a Convolutional Neural Network (CNN) that outputs the parameters of a Gaussian. The decoder is a transposed CNN which produces the reconstructed SDF $\hat{E}$ from $h$. The decoder log-likelihood $p_{\psi}(E|h)$ is $||\hat{E}-E||_{2}$, where $\psi$ are the parameters of the decoder CNN. Chen et al. [7] showed that learning a latent prior can improve VAE performance, so we parameterize the latent prior $p_{\phi}(h)$ as a normalizing flow and learn the prior during training. The loss for the VAE is as follows:

We then use a Multilayer Perceptron (MLP) network $g_{\omega}$ to generate a context vector $C$ to use in the normalizing flow, via $C=g_{\omega}(x_{0},x_{G},h)$, which has parameters $\omega$.

We use a conditional normalizing flow parameterized by $\zeta$ to define the conditional variational posterior, i.e. $q_{\zeta}(U|C)$ is defined by $U=f_{\zeta}(Z,C)$ for $Z\sim p(Z)=\mathcal{N}(0,I)$. The variational free energy 3 then becomes:

We can then optimize $\zeta$ to minimize the free energy.

By using a conditional normalizing flow, we are amortizing the cost of computing the posterior across environments. The conditional normalizing flow $U=f_{\zeta}(Z,C)$ is invertible with respect to $Z$, i.e. $Z=f^{-1}(U,C)$. For our conditional Normalizing Flow we use an architecture based on Real-NVP [8] architecture with conditional coupling layers [37], the structure is specified in section VI-C.

Minimizing eq. (8) via gradient descent requires the cost and dynamics to be differentiable. To avoid this, we estimate gradients, using the method in [23]: At each iteration, we sample $R$ control sequences $U_{1..R}$ from $q_{\zeta}(U|C)$ and compute weights

We present a method for using the learned control sequence posterior for a control task. Given a $C$ computed from $(x_{0},x_{G},E)$, the control sequence posterior $q_{\zeta}(U|C)$ can be used as a sampling distribution in a sampling-based MPC approach. We propose a method for using the control sequence posterior with MPPI [36], which we term FlowMPPI (Algorithm 1).

MPPI iteratively perturbs a nominal control sequence with Gaussian noise and performs a weighted sum of the perturbations to find a new control sequence. Empirically, we found that standard MPPI is good at performing local optimization on an already-feasible nominal trajectory. On the other hand, the control sequence posterior is able to sample collision-free goal-directed trajectories, but locally improving trajectories with samples is difficult as small changes in the control sequence posterior latent space $Z$ often lead to large differences in the resulting control sequence. As a result, we observed that we obtained trajectories which reached the goal and avoided obstacles with very few samples, however the cost of the best trajectory did not improve much with more iterations of MPPI.

FlowMPPI combines sampling in the latent space $Z$, and sampling perturbations to trajectories to get the advantages of both. For a given sampling budget $K$, we generate half of the samples from perturbing the nominal trajectory as in MPPI, and the other half from sampling from the control sequence posterior. These samples will be combined as in standard MPPI. Since the control sequence posterior is invertible w.r.t $U$, a given nominal trajectory $U$ can be transformed to a latent state $Z$. For the samples from the control sequence posterior, we apply a perturbation cost on the distance of the sampled trajectory from the nominal in latent space. This cost mirrors a similar cost in standard MPPI which penalizes perturbations based on distance to the nominal in the control space.

A novel environment can be OOD for the control sequence posterior and result in poor performance. We present an approach where we project the OOD environment embedding $h$ in-distribution in order to produce low-cost trajectories when it is used as part of the input to $f_{\zeta}$. The intuition behind this approach is that our goal is to sample low-cost trajectories in the current environment. Given that $f_{\zeta}$ will have been trained over a diverse set of environments, if we can find an in-distribution environment that would elicit similar low-cost trajectories, then we can use this environment as a proxy for the actual environment when sampling from the flow. Thus we avoid the problem of samples from the control sequence posterior being unreliable when the input is OOD.

In order to do this projection, we first need to quantify how far out-of-distribution a given environment is. Once we have such an OOD score, we will find a proxy environment embedding $\hat{h}$ by optimizing the score, while also regularizing to encourage low-cost trajectories. For the OOD score, we use the VAE we have discussed in section V-B. VAEs and other deep latent variable models have been used to detect OOD data in prior work [9, 38, 22], however these methods are typically based on evaluating the likelihood of an input, in our case $p(E)$. For VAEs this requires reconstruction. We would like to avoid using reconstruction in our OOD score for two reasons. First, reconstruction, particularly of a 3D SDF, adds additional computation cost and we would like to evaluate the OOD score in an online control loop. Second, optimizing an OOD score based on reconstruction would drive us to find an environment embedding proxy which is able to approximately reconstruct the entire environment. This makes the problem more difficult than is necessary, as we do not need $\hat{h}$ to accurately represent the entire environment, only to elicit low-cost trajectories from the control sequence posterior.

To determine how close $h$ is to being in-distribution, we use the following OOD score:

We can perform gradient descent on $\mathcal{L}_{OOD}$ to find $\hat{h}$, thus projecting the environment to be in-distribution. Note that without regularization this process will converge to a nearby maximum likelihood solution, which may lose key features of the current environment. Since our aim is to sample low-cost trajectories from the control sequence posterior, we use $\mathcal{L}_{flow}$ as a regularizer for this gradient descent. Our intuition here is that in order to generate low-cost trajectories in the true environment, the projected environment embedding should preserve important features of the environment relevant for that particular planning query. The new environment embedding is then given by

In this section we will introduce the systems and the environments we use for evaluation. For all systems and environments, a task is considered a failure if there is a collision or if the system does not reach the goal region within a timeout of 100 timesteps. The cost function for both systems is given by $J(\tau)=100d_{G}(x_{T})+\sum^{T-1}_{t=1}10d_{G}(x_{t})_{2}+\sum^{T}_{t=1}10000D(x_{t})$, where $T$ is the MPC horizon, $d_{G}$ is a distance to goal function, and $D$ is an indicator function which is 1 if $x_{t}$ is in collision and 0 otherwise. For all of our experiments, the control horizon $T=40$. We use a Gaussian prior over controls is $p(U)=\mathcal{N}(0,\sigma^{2}I)$ which induces a cost on the squared magnitude of the actions. For all of our experiments the dynamics are deterministic. Further details of the generation of training data can be found in appendix -C.

To confirm the efficacy of our OOD score, we computed this score for the training and OOD environments for each system above. Figure 3 shows that this score is clearly able to distinguish in-distribution environment embeddings from OOD ones. To show the necessity of using both the OOD score and the regularization in projection, we perform an ablation on these two components in appendix -F for the quadrotor system.

For both the control sequence posterior flow $f_{\zeta}$ and the VAE prior $p_{\phi}(h)$ we use an architecture based on Real-NVP [8]. For the VAE prior $p_{\phi}(h)$ we use a flow depth of 4, while for the control sampling flow $f_{\zeta}$ we use a flow depth of 10. For the control sampling flow we use the conditional coupling layers from [37]. For the VAE encoder we use four CNN layers with a kernel of 3 and a stride of 2, followed by a fully connected layer. For the VAE decoder we used a fully connected layer followed by four transposed CNN layers. For the 3D case we use 3D convolutions. The dimensionality of both $h$ and $C$ was $64$ for the planar navigation environments $256$ for 3D 12DoF quadrotor environments. $g_{\omega}$ was defined as an MLP with a single hidden layer of size $256$. For nonlinear activations we used ReLU throughout.

For training, we use $10000$ randomly generated environments for planar navigation task, and $20000$ for the 3D 12DoF quadrotor task. At each epoch, for each environment, we randomly select one of $100$ start and goal pairs. We train the control sequence posterior flow $f_{\zeta}$, the VAE parameters ($\theta,\phi,\psi$) and the context MLP $g_{\omega}$ end-to-end using Adam for $1000$ epochs with a learning rate of $1e^{-3}$, with a decay rate of $0.9$ every 50 epochs. After 100 epochs, we freeze the VAE and do not continue training with $\mathcal{L_{VAE}}$. This is primarily because the VAE converges quickly and training proceeds more quickly without reconstruction. When training the VAE we divide the loss by the total dimensionality of the SDF and use $a=5$. For every environment for the flow training, hyperparameters we use $\beta=1$ and we linearly increase $\alpha$ from $1$ to $500$ during training. Empirically we found that low initial $\alpha$ was required for the flow to learn to generate goal-directed trajectories early on during training, and that increasing $\alpha$ later during training increases the diversity of the flow sampling distribution. A more details list of training hyperparameters can be found in appendix 6

For our baselines we use several state-of-the-art sampling-based MPC methods: MPPI [36], Stein Variational MPC (SVMPC) [17] and improved CEM (iCEM) [24]. MPPI uses a Gaussian distribution as the sampling distribution, iCEM uses colored noise, and SVMPC uses a mixture of Gaussians. For each baseline, we tune the hyperparameters to get the best performance based on the training environments, and maintain these hyperparameters when switching to the out-of-distribution environments. When evaluating our two proposed methods and the baselines, each method is given the same sampling budget per timestep. This means that for methods that require multiple iterations per timestep, the sampling budget is distributed across the iterations. A more detailed list of the hyperparameters for each controller can be found in appendix -D. Evaluating $\mathcal{L}_{flow}$ during projection requires sampling and evaluating control sequences. When we consider the sampling budget of different algorithms in evaluation, we will include these samples. FlowMPPIProject uses half of the allowed sampling budget during the project step, and the other half for the FlowMPPI control algorithm. While it does take longer to sample from the flow than from the distributions in the baselines, we observe that the cost of evaluating control sequences dominates over the cost of sampling. For example, for the 3D 12DoF quadrotor system, sampling 1024 control sequences from the flow and evaluating the cost of these control sequences takes on average $9$ms and $80$ms, respectively on an i7-8700K CPU and Nvidia 1080 Ti GPU.

The results comparing our MPC methods to baselines are shown in Tables I and II. For the planar navigation case, we see that FlowMPPI and FlowMPPIProject are competitive with the baselines on the training environments. Our method reaches the goal region more often, while attaining slightly higher average cost. For the out-of-distribution environments, our method reaches the goal significantly more often. For example, with a sampling budget of $256$, the success rates for FlowMPPIProject is $0.65$ and increases to $0.87$ for a sampling budget of 1024. The next closest baseline, iCEM, has successes rates of $0.46$ and $0.62$ for sampling budgets of $256$ and $1024$, respectively. The projection process for the planar navigation task is visualized for an OOD environment in Figure 5.

We observed during this experiment that when iCEM and SVMPC are able to generate a trajectory which reaches the goal region, they are able to locally optimize this trajectory better than FlowMPPI variants, while FlowMPPI is better able to generate sub-optimal trajectories to the goal region.

For the quadrotor system, FlowMPPIProject outperforms all other methods in both cost and success rate across all environments and sampling budgets. With a sampling budget of $256$, FlowMPPIProject attains a $71\%$ success rate compared to $35\%$ by iCEM and $11\%$ by SVMPC for OOD environments. For a sampling budget of $1024$ the success rate of FlowMPPIproject rises to $93\%$ vs. $63\%$ for iCEM. The dynamics of this task make it much more difficult, particularly because stabilizing around the goal is non-trivial. We found that the baselines struggled to find trajectories which both reached and stabilized to the goal, and thus were more susceptible to becoming stuck in local minima.

Table II shows the results when evaluating our method in simulation in two environments generated from real-world data. FlowMPPIProject outperforms all baselines in cost & success rate, despite only being trained on simulated environments consisting of large spherical obstacles. For the challenging stairway environment, FlowMPPIProject achieves $85\%$ success, while the next closest baseline, iCEM, has $58\%$ success. FlowMPPI achieves only $50\%$ success rate for this task, highlighting the importance of projection for real-world environments.