PIPPS: Flexible Model-Based Policy Search Robust to the Curse of Chaos

Introduction

We were motivated by Probabilistic Inference for Learning Control (PILCO) (Deisenroth & Rasmussen, 2011), a model-based reinforcement learning (RL) algorithm which showed impressive results by learning continuous control tasks using several orders of magnitude less data than model-free alternatives. The keys to its success were a principled approach to model uncertainty, and analytic moment-matching (MM) based Gaussian approximations of trajectory distributions.

Unfortunately the MM framework is computationally inflexible. For example, it cannot be used with neural network models. Thus, this work searches for an alternative flexible method for evaluating trajectory distributions and gradients.

Particle sampling methods are a general scheme, which can be applied in practically any setting. Indeed, model-free algorithms have previously been successfully used with particle trajectories from the same types of models as used in PILCO (Kupcsik et al., 2014), (Chatzilygeroudis et al., 2017). Aiming for better performance, model-based gradients can be evaluated using the reparameterization (RP) trick to differentiate through stochasticities. This approach was previously attempted in the context of PILCO, but surprisingly, it did not work, with the poor performance attributed to local minima (McHutchon, 2014).

Recently, several works have attempted a similar method using Bayesian neural network dynamics models and the RP trick. Depeweg et al. (2016) successfully used this approach to solve non-standard problems. Gal et al. (2016), on the other hand, found that the direct approach with RP did not work on the standard cart-pole swing-up task. It is difficult to compare the new approaches to PILCO, because they simultaneously change multiple aspects of the algorithm: they switch the model to a more expressive one, then modify the policy search framework to accommodate this change. In our work, we perform a shorter step – we keep everything about PILCO the same, only changing the framework used for prediction. This approach allows better explaining how MM helps with learning. We find that reducing local minima was in fact not the main reason for its success over particle-based methods. The primary issue was a hopelessly large gradient variance when using particles and the RP trick.

We show that the large variance is due to chaos-like sensitive dependence on initial conditions – a common property in calculations involving long chains of nonlinear mappings. We refer to this problem as ”the curse of chaos”. Our work suggests that RP gradients and backpropagation alone are not enough. One either needs methods to prevent the curse from occurring, or other types of gradient estimators.

We derive new flexbile gradient estimators, which combine model-based gradients with the likelihood ratio (LR) trick (Glynn, 1990), also called REINFORCE (Williams, 1992) or the score function estimator. Our use of LR differs from the typical use in model-free RL – instead of sampling with a stochastic policy in the action space, we use a deterministic policy, but sample with a stochastic model in the state space. We also develop an importance sampling scheme for use within a batch of particles. Our estimators obtain accurate gradients, and allow surpassing the performance of PILCO.

Our results – LR gradients perform better than RP with backpropagation – are contradictory to recent work in stochastic variational inference, which suggest that even a single sample point yields a good gradient estimate using the RP trick (Kingma & Welling, 2014; Rezende et al., 2014; Ruiz et al., 2016). In our work, not only is a single sample not enough, even millions of particles would not suffice! In contrast, our new estimators achieve accurate gradients with a few hundred particles. As LR gradients are also not perfect, we further invent the total propagation algorithm, which efficiently combines the best of LR and RP gradients.

Background

Learning alternates between executing the policy on the system, then updating $\theta$ to improve the performance on the following attempts. Policy gradient methods directly estimate the gradient of the objective function $\frac{\text{d}~{}}{\text{d}\theta}J(\theta)$ and use it for optimization. Some model-based policy search methods use all of the data to learn a model of $f$ denoted by $\hat{f}$, and use it for ”mental rehearsal” between trials to optimize the policy. Hundreds of simulated trials can be performed per real trial, greatly increasing data-efficiency. We utilize the fact that $\hat{f}$ is differentiable to obtain better gradient estimators over model-free algorithms. Importantly, our models are probabilistic, and predict state distributions.

2 Stochastic Gradient Estimation

Reparameterization gradient (RP): Consider sampling from a univariate Gaussian distribution. One approach first samples with zero mean and unit variance $\epsilon\sim\mathcal{N}(0,1)$, then maps this point to replicate a sample from the desired distribution $x=\mu+\sigma\epsilon$. Now it is straight-forward to differentiate the output w.r.t. the distribution parameters, namely $\frac{\text{d}x}{\text{d}\mu}=1$ and $\frac{\text{d}x}{\text{d}\sigma}=\epsilon$. Averaging samples of $\frac{\text{d}\phi}{\text{d}x}\frac{\text{d}x}{\text{d}\theta}$ gives an unbiased estimate of the gradient of the expectation. This is the RP gradient for a normal distribution. For a multivariate Gaussian, the Cholesky factor ($L,\text{s.t.}\Sigma=LL^{T}$) of the covariance matrix can be used instead of $\sigma$. See (Rezende et al., 2014) for non-Gaussian distributions.

3 Trajectory Gradient Estimation

The probability density $p(\tau)=p({\bf x}_{0},{\bf u}_{0},{\bf x}_{1},{\bf u}_{1},...,{\bf x}_{T})$ of observing a particular trajectory can be written as $p({\bf x}_{0})\pi({\bf u}_{0}|{\bf x}_{0})p({\bf x}_{1}|{\bf x}_{0},{\bf u}_{0})...p({\bf x}_{T}|{\bf x}_{T-1},{\bf u}_{T-1})$.

To use RP gradients, one must know or estimate the dynamics $p({\bf x}_{t+1}|{\bf x}_{t},{\bf u}_{t})$ – in other words, RP is not applicable to the model-free case. With a model, a predicted trajectory can be differentiated by using the chain rule.

4 PILCO

The higher level view of PILCO follows Section 2.1 and the policy gradient evaluation is detailed in Algorithm 1.

4.2 Moment matching Prediction

In general, when a Gaussian distribution is mapped through a nonlinear function, the output is intractable and non-Gaussian; however, in some cases one can analytically evaluate the moments of the output distribution. Moment-matching (MM) approximates the output distribution as Gaussian by matching the mean and variance with the true moments. Note that even though the state-dimensions are modelled with separate functions $\hat{f}_{a}$, MM is performed jointly, and the state distributions can include covariances.

Particle Model-Based Policy Search

In general, particle trajectory predictions are simple – predict at all particle locations, sample from the output distributions, repeat. However, we also compare to a scheme based on Gaussian resampling (GR), used by Gal et al. (2016) to apply PILCO with neural network dynamics models.

Gaussian resampling (GR): MM can be stochastically replicated. At each time step, the mean $\hat{\mu}=\sum_{i=1}^{P}{\bf x}_{i}/P$ and covariance $\hat{\Sigma}=\sum_{i=1}^{P}({\bf x}_{i}-\hat{\mu})({\bf x}_{i}-\hat{\mu})^{T}/(P-1)$ of the particles are estimated. Then the particles are resampled from the fitted distribution $x^{\prime}_{i}\sim\hat{\mu}+L{\bf z}_{i}~{}|~{}{\bf z}_{i}\sim\mathcal{N}(0,I)$, where $L$ is the Cholesky factor of $\hat{\Sigma}$. One can differentiate this resampling operation by using RP. Obtaining the gradient $\text{d}{L}/\text{d}{\hat{\Sigma}}$ is non-trivial, but (Murray, 2016) presents an overview. We use the provided symbolic expression.

2 Hybrid Gradient Estimation Techniques

In our case, RP gradients can be used. However, surprisingly they were hopelessly inaccurate (see Figure 2(d)). To solve the problem, we derived new gradient estimators which combine model derivatives with LR gradients. In particular, our approach allowed for within batch importance sampling to increase sample efficiency.

Batch Importance Weighted LR (BIW-LR): We use parallel computation, and sample multiple particles simultaneously. The state distribution is represented as a mixture distribution $q({\bf x}_{t+1})=\sum_{i=1}^{P}p({\bf x}_{t+1}|{\bf x}_{i,t};\theta)/P$. Analogously to the derivation of LR in Section 2.2, one can derive a lower variance estimator with importance sampling within the batch for each time step: $\sum_{i=1}^{P}\sum_{j=1}^{P}\left(\frac{\text{d}{p({\bf x}_{j,t+1}|{\bf x}_{i,t};\theta)}/\text{d}{\theta}}{\sum_{k=1}^{P}p({\bf x}_{j,t+1}|{\bf x}_{k,t})}(G_{t+1}({\bf x}_{j,t+1})-b_{i,t+1})\right)/P$

We choose to estimate a leave-one-out mean of the returns by normalized importance sampling with the equation: $b_{i,t+1}=\left(\sum_{j\neq i}^{P}c_{j,t+1}G_{t+1}({\bf x}_{j,t+1})\right)/\sum_{j\neq i}^{P}c_{j,t+1}$, where $c_{j,t+1}=p({\bf x}_{j,t+1}|{\bf x}_{i,t})/\sum_{k=1}^{P}p({\bf x}_{j,t+1}|{\bf x}_{k,t})$. Without normalizing, a large variance of the baseline estimation leads to poor LR gradients. Note that we compute $P$ baselines for each time-step, whereas there are $P^{2}$ components in the gradient estimator. To obtain a true unbiased gradient, one should compute $P^{2}$ leave-one-out baselines – one for each particle for each mixture component of the distribution. The paper contains evaluations only with the baseline presented here – we found that it already removes most of the bias.

RP/LR weighted average: The bulk of the computation is spent on the $\text{d}p({\bf x}_{t+1}|{\bf x_{t}};\theta)/\text{d}\theta$ terms. These terms are needed for both LR and RP gradients, so there is no penalty to combining both estimators. A well known statistics result states that for independent estimators, an optimal weighted average estimate is achieved if the weights are proportional to the inverse variance, i.e. $\mu=\mu_{LR}k_{LR}+\mu_{RP}k_{RP}$, where $k_{LR}=\hat{\sigma}_{LR}^{-2}/(\hat{\sigma}_{LR}^{-2}+\hat{\sigma}_{RP}^{-2})$ and $k_{RP}=1-k_{LR}$.

A naive combination scheme would compute the gradient separately for the whole trajectory for both estimators, then combine them; however, this approach neglects the opportunity to use reparameterization gradients through shorter sections of the trajectory to obtain better gradient estimates. Our new total propagation algorithm (TP) goes beyond the naive method. TP uses a single backward pass to compute a union over all possible RP depths, automatically giving greater weight to estimators with lower variance.

3 Policy Optimization

Deterministic optimizer: The random number seed can be fixed to turn a stochastic problem deterministic, also known as the PEGASUS trick (Ng & Jordan, 2000). With a fixed seed, the RP gradient is an exact gradient of the objective, and quasi-Newton optimizers, such as BFGS can be used.

Experiments

We performed experiments with two purposes: 1. To explain why RP gradients are not sufficient (Section 4.1). 2. To show that our newly developed methods can match up to PILCO in terms of learning efficiency (Section 4.2).

We perturb the policy parameters $\theta$ in a randomly chosen fixed direction, and plot both the objective function, and the magnitude of the projected gradient as a function of $\Delta\theta$. The results of this experiment are perhaps the most striking component of our paper, and motivated the term ”the curse of chaos”. Refer to Figure 2 for the results, and to Section 4.1.1 for a full explanation.

The plots were generated in the nonlinear cart-pole task, using similar settings as explained in Section 4.2. We used 1000 particles, and while keeping the random number seed fixed to demonstrate that the high variance in 2(d) is not caused by randomness, but by a chaos-like property of the system. The confidence intervals were estimated by $Var/P$ where $Var$ is the sample variance, and $P$ is the number of particles. In Section 4.1.2 we plot the dependence of the variance on $P$ using a more principled approach.

Figure 2(d) contains a peculiar result where the RP gradient behaves well for some regions of the parameters $\theta$, but when $\theta$ is perturbed, a phase-transition-like change causes the variance to explode. The variance at $\Delta\theta=1.5$ is $\sim$$4\times 10^{5}$ times larger than at $\Delta\theta=0$, meaning that $\sim$$4\times 10^{8}$ particles would be necessary for the RP gradient to become accurate in that region. For practical purposes, optimizing with the RP gradient would lead to a simple random walk.

Since the seed was fixed, the RP gradient in 2(d) is an exact gradient of the value in 2(a). Therefore, there is an infinitesimal deterministic ”noise” at the right of 2(a). The value averaged across 1000 particles is though not the true objective – that would require averaging an infinite number of particles. When averaging an infinite amount of particles, is there still a ”noise”, or does the function become smooth?

Our new gradient estimators in Figures 2(e) and 2(f) suggest that the true objective is indeed smooth. To provide more evidence, we estimated the magnitude of the gradient from finite differences of the value in 2(a) using a sufficiently large perturbation in $\theta$, such that the ”noise” is ignored. The fact that two separate approaches agree – one which varies the policy parameters $\theta$, and another which keeps $\theta$ fixed, but estimates the gradient from the trajectories – provides convincing evidence that the true objective is smooth.

Figures 2(b) and 2(c) explain the reason for the explosion of the variance when using RP gradients (2(d)). Figure 2(b) corresponds to the leftmost parameter setting ($\Delta\theta=0$); Figure 2(c) corresponds to the rightmost parameter setting ($\Delta\theta=1.5$). The plots show how the value $V({\bf x};\theta)$ (the remaining cumulative cost) varies as a function of the state position ${\bf x}$. Note that because the random number seed is fixed, the value $V$ is the same as the remaining return $G$. This definition differs from the typical value function, which averages the return over an infinite amount of particles. The figures were created by predicting the trajectory at each point for 4 particles with different fixed seeds, then averaging the costs of the trajectories. We chose to predict 4 particles after trying 1 particle, for which the value appeared to include a step-like part, but was otherwise less interesting than the current figure. As the average value of the 4 particles is erratic, at least one of the 4 particles must have a highly erratic value in the shown region.

The boxes (2(b), 2(c)) are centered at the mean prediction from the center of the initial state distribution (if unclear, consider Figure 1 with $p({\bf x}_{0})$ as a point mass, then $p({\bf x}_{1})$ depicts the location of the box). The axes on the boxes are slightly different, because when $\theta$ is changed, the predicted location $p({\bf x}_{1};\theta)$ changes. The side lengths correspond to 4 standard deviations of the Gaussian distributions $p({\bf x}_{1};\theta)$. The velocities were kept fixed at the mean values.

RP estimates $\frac{\text{d}}{\text{d}\theta}\int p({\bf x}_{1};\theta)V({\bf x}_{1})\text{d}{\bf x}$. It samples points inside the box, computes the gradient $\frac{\partial V}{\partial\theta}=\frac{\partial V}{\partial x}\frac{\text{d}x}{\text{d}\theta}$ and averages the samples togetherNote that the same evaluation of the value gradient has to be performed at subsequent time-steps, and in practice the sum is evaluated simultaneously using backpropagation, but we ignore this for the purpose of the explanation.. In Figure 2(c) finding the gradient of the expectation by differentiating $V$ is completely hopeless. In contrast, the LR gradient (2(e)) only uses the value $V$, not its derivative, and does not suffer from this problem.

Finally, even though we do not show the plotted value and gradient for the Gaussian resampling case, both of these were smooth functions for a fixed random seed. Thus, resampling also beats the curse of chaos.

1.2 Policy Gradient Variance Evaluation

In Figure 3, we plot how the variance of the gradient estimators at $\Delta\theta=0$ and $\Delta\theta=1.5$ depends on the number of particles $P$. The variance was computed by repeatedly sampling the estimator for a large number of times and calculating the variance from the set of evaluations. We compare RP, TP as well as LR gradients both with and without batch importance weighting (BIW) to show that our importance sampling scheme reduces the variance. We used the importance sampled baseline; in practice the regular LR gradient would use a simpler baseline, and have even higher variance. The RP gradient is omitted from 3(b), because the variance was between $10^{8}$-$10^{15}$. The TP gradient combined the BIW-LR and RP gradients.

The results confirm that BIW significantly reduces the variance. Moreover, our TP algorithm was the best. Importantly, in 3(b), even though the variance of the RP gradient for the full trajectory is over $10^{6}$ larger than the other estimators, TP utilizes shorter path-length RP gradients to obtain 10-50% reduction in variance for 250 particles and fewer.

2 Learning Experiments

We compare PILCO in episodic learning tasks to the following particle-based methods: RP, RP with a fixed seed (RP${}_{\text{FS}}$), Gaussian resampling (GR), GR with a fixed seed (GR${}_{\text{FS}}$), model-based batch importance weighted likelihood ratio (LR) and total propagation (TP). Moreover, we evaluate two variations of the particle predictions: 1. TP while ignoring model uncertainty, and adding only the noise at each time step ($\text{TP}-\sigma_{f}$). 2. TP with increased prediction noise ($\text{TP}+\sigma_{n}$). We used 300 particles in all cases.

We performed learning tasks from a recent PILCO paper (Deisenroth et al., 2015): cart-pole swing-up and balancing, and unicycle balancing. The simulation dynamics were set to be the same, and other aspects were similar to the original PILCO. The results are in Tables 2 and 1, and in Figure 4.

Discussion

PILCO performs well in scenarios with no noise, but with noise added the results deteriorate. This deterioration is most likely caused by an accumulation of errors in the MM approximations, previously observed by Vinogradska et al. (2016), who used quadrature for predictions. Particles do not suffer from this issue, and using TP gradients consistently outperforms PILCO with high noise.

On the other hand, at low noise levels, the performance of TP as well as LR reduces. If all of the particles are sampled from a small region, it becomes difficult to estimate the gradient from changes in the return – in the limit of a delta distribution an LR gradient could not even be evaluated. The TP gradient is less susceptible to this problem, because it incorporates information from RP. Finally, if the uncertainty in predictions is very low (as in $k=10^{-2}$), one can consider model noise as a parameter that affects learning, and increase it to acquire more accurate gradients: see $\text{TP}+\sigma_{n}$, where the model noise variance was multiplied by 100.

Notably, approaches which use MM, such as PILCO, and GR outperform the others when using the Tip Cost. The reason may be the multi-modality of the objective – with the Tip Cost, the pendulum may be swung up from either direction to solve the task; with the Angle Cost there is only one correct direction. Performing MM forces the algorithm along a unimodal path, whereas the particle approach could attempt a bimodal swing-up where some particles go from one side, and the rest from the other side. Thus, MM may be performing a kind of ”distributional reward shaping”, simplifying the optimization problem. Such an explanation was previously provided by Gal et al. (2016).

Finally, we point to the surprising $\text{TP}-\sigma_{f}$ experiment. Even though the predictions ignore model uncertainty, the method achieves 93% success rate. It is difficult to explain why learning still worked, but we hypothesize that the success may be related to the 0 prior mean of the GP. In regions where there is no data, the mean of the GP dynamics model goes to 0, meaning that the input control signal has no effect on the particle. Therefore, for the policy optimization to be successful, the particles would have to be controlled to stay in regions where there exists data. Note that a similar result was found by Chatzilygeroudis et al. (2017) who used an evolutionary algorithm and achieved 85-90% success rate at the cart-pole task even when ignoring model uncertainty.

2 The Curse of Chaos in Deep Learning

Most machine learning problems involve optimizing the expectation of an objective function $J({\bf x};\theta)$ over some data generating distribution $p_{Data}({\bf x})$, where this distribution can only be accessed through sample data points $\{{\bf x}_{i}\}$. Our predictive framework is analogous to a deep model: $p({\bf x}_{0})$ is the data generating distribution, $p({\bf x}_{t};\theta)$ are obtained by pushing $p_{Data}({\bf x})$ through the model layers. The most common method of optimization is SGD with pathwise derivatives computed by backpropagation. Our results suggest that in some situations – particularly with very deep or recurrent models – this approach could degenerate into a random walk due to an exploding gradient variance.

While we have yet to computationally confirm our deep hypothesis, several works have investigated chaos in neural networks (Kolen & Pollack, 1991; Sompolinsky et al., 1988), although we believe we are the first to suggest that chaos may cause gradients to degenerate when computed using backpropagation. Notably, Poole et al. (2016) suggested that such properties lead to ”exponential expressivity”, but we believe that this phenomenon may instead be a curse.

Conclusions & Future Work

We may have described a limitation of optimizing expectations using pathwise derivatives, such as those computed by backpropagation. Moreover, we show a method to counteract this curse by injecting noise into computations, and using the likelihood ratio trick. Our total propagation algorithm provides an efficient method to combine reparameterization gradients on arbitrary stochastic computation graphs with any amount of other gradient estimators – even gradients computed using a value function could be used. There are countless ways to expand our work: better optimization, incorporate natural gradients, etc. The flexible nature of our method should make it easy to extend.

Acknowledgements

We thank Chris Reinke and the anonymous reviewers for comments about clarity. This work was supported by JSPS KAKENHI Grant Number JP16H06563 and JP16K21738.

References