Title: Learning and Policy Search in Stochastic Dynamical Systems with Bayesian Neural Networks
Authors: Stefan Depeweg
José Miguel Hernández-Lobato
Finale Doshi-Velez
Steffen Udluft
Venue: ICLR 2017
Keywords: Deep learning
Reinforcement Learning

Abstract: We present an algorithm for policy search in stochastic dynamical systems using
model-based reinforcement learning. The system dynamics are described with
Bayesian neural networks (BNNs) that include stochastic input variables.  These
input variables allow us to capture complex statistical
patterns in the transition dynamics (e.g. multi-modality and
heteroskedasticity), which are usually missed by alternative modeling approaches. After
learning the dynamics, our BNNs are then fed into an algorithm that performs
random roll-outs and uses stochastic optimization for policy learning. We train
our BNNs by minimizing $\alpha$-divergences with $\alpha = 0.5$, which usually produces better
results than other techniques such as variational Bayes. We illustrate the performance of our method by
solving a challenging problem where model-based approaches usually fail and by
obtaining promising results in real-world scenarios including the control of a
gas turbine and an industrial benchmark.


Review discussion:
Dear Reviewers,
we uploaded a revision of the paper yesterday (27th of November):
In the original version, at training time, we sampled the latent
variables z from the prior. In the revision, we learn a posterior
approximation q(z) for each z.
That is, for each data point (x_i,y_i)
\in D we infer a posterior approximation z_i ~ N(\mu_i,\sigma_i).This
is very similar to how the EM algorithm works. After training, the
model is used in the same way as before by sampling z from the prior
N(0,\gamma).
Because of this changes we repeated the experiments (also in a more
sparse data setting) and improved a bit the presentation, simplified
when it was too complicated and made the paper more compact. We also
readjusted the focus of this work a bit, emphasising more its strength
over standard neural networks and gaussian processes.
1. [Related work] "[...] popular modes for transition functions [...]". It might be relevant to mention the work on Predictive State Representation which played an important role in RL in the early 2000s and subsequently applied to real world problems with Boots' TPSRs (2011).
2. The claim that the proposed method allows for "previously impossible" problems to be learned might have to backed up with more extensive experimental comparaisons. I'm thinking for example of Guided Policy Search, which also builds a model, and seems to be rather sample efficient. If the goal the paper is to improve within the class of Bayesian methods (which on its own, is already good !), then it should be stated more clearly. To make a fully convincing point, I would also like to see wall clock times and number of samples in table 3.
3. I've been re-reading your explanation on why $\epsilon_n$ is needed and I'm still not 100% sure that I fully understand. From what I understand:  $z_n$ is really part of your  problem/system/environment while $\epsilon_n$ has to do with your choice of function approximator (BNN), on the side of the agent.  Since you already estimate $z_n$, why can't you model the entire stochastic effects through it ? A lack of representational power ? Couldn't you reparametrize your model and push everything in $z_n$ ?
4. Line 14 of Algorithm 1: [UNFOLD(s_n)], do you assume that you can restart your system in the same state ?
1. Thanks for pointing this out. We will extend the related work section for
the final version.
2. We agree: the point of this paper is not that we can now solve RL problems
that were impossible to solve before. We have two main contributions:
a) In the model-based framework we propose a Bayesian model architecture (BNN
with stochastic inputs trained with \alpha-divergences) that allows us to
describe systems with complex stochastic transitions.
Yes, "we try to improve within the class of Bayesian methods" but also more
generally within the class of model-based RL methods.Our motivation was to come up with a model that is flexible & scalable like a
NN, is fully Bayesian like a Gaussian Process/BNN and can model complex
stochastic patterns (e.g bi-modality, heteroskedasticity, etc.).
b) An algorithm that trains a parameterized policy on a batch of data using a
BNN. Algorithm 1 in the paper shows the interplay between model and policy and
describes the sampling process. The roll-outs give particle-like trajectories
from which we obtain gradient information. We consider this to be interesting
and novel, and for industry setting the applicability to 'batch settings' is an
important criteria.
We would also like to point out that in our experience model-free approaches
have issues with convergence and stability in these batch settings.
Often times policy performance is unstable and at some point in training starts to detoriate. A good
model is very useful here for evaluation.
3. We want to make two clarifications first:
3.1 We assume there exists a latent variable $z$ with unknown realizatons  $z_1,\cdots,z_N$
for the state transitions \{(x_1,y_1),\cdots,(x_N,y_N)\}$. Our models tries to infer these realizations by learning a
variatonal distribution \{(\mu_z_1,\sigma_z_1),\cdots,(\mu_z_N,\sigma_z_N)\}.   This  inference is possible
because we see both x and y.
When we use the model after training to make predictions,  we will sample z from the prior N(0,\gamma), because  we only have a prior belief over what z might be (since y is unknown).
3.2  $\epsilon_1,\cdots,\epsilon_d$ are the variances of the gaussian output noise with zero mean. We have 1 parameter per output
dimension. In regression one normally assumes the data is corrupted by zero-mean gaussian noise:
y = f(x,W_true) + N(0,\epsilon)
Our model is:
y = f(x,W_true,z) + N(0,\epsilon)
In this framework sampling W and z will give a  mixture of a diagonal
d-dimensional Gaussians (with each mixture component having the same variance).
If we would not assume  gaussian output noise we instead would have a mixture of delta
distributions. This would make optimization difficult as  the gradients
for the energy function would be either infinitely large or zero.
4. We are a bit unsure what is meant with this question, so we give two
related answers:
a) We are doing batch RL, the result of our optimization is a policy that maps
from a state to an action; during optimization no interaction with the 'target
system' is required, all inference is done on the batch and a given reward
function.
So the property that the target system can be reset to different
states is not relevant for the training process.
b) When we train our policy we initialize algorithm 1 to a starting state s_0
and simulate for $T$ steps the interaction between model, policy and reward function.
When the policy is deployed it usually interacts
constantly with the target system.
If the horizon $T$  that we chose is too short, the policy may behave
suboptimal in deployment. We are aware of this;
but it is fundamental to a lot
of model-based RL approaches, such like the 'receding horizon' approach in
model predictive control. (or the discount factor and number of update steps in the updates
to the Q-function)
So 'no', at test/deployment time we don't assume the system is repeatedly
initialized to a particular state, we usually have continuous autonomous
interaction. We are aware that this requires the horizon $T$ to be chosen to be
long enough.
This paper considers the problem of model-based policy search. The authors
consider the use of Bayesian Neural Networks to learn a model of the environment
and advocate for the $\alpha$-divergence minimization rather than the more usual
variational Bayes.
The ability of alpha-divergence to capture bi-modality however
comes at a price and most of the paper is devoted to finding tractable approximations.
The authors therefore use the approach of Hernandez-Lobato
et al. (2016) as proxy to the alpha-divergence .
The environment/system dynamics is clearly defined as a well as the policy parametrization
(section 3) and would constitute a useful reference point for other researchers.
Simulated roll-outs, using the learned model, then provide samples of the expected
return. Since a model of the environment is available, stochastic gradient descent
can be performed in the usual way, without policy gradient estimators, via automatic
differentiation tools.
The experiments demonstrate that alpha-divergence is capable of capturing multi-model
structure which competing methods (variational Bayes and GP) would otherwise
struggle with. The proposed approach also compares favorably in a real-world
batch setting.
The paper is well-written, technically rich and combines many recent tools
into a coherent algorithm.
However, the repeated use of approximations to original
quantities seems to somehow defeat the benefits of the original problem formulation.
The scalability and computational effectiveness of this approach is also questionable
and I am uncertain if many problem would warrant such complexity in their solution.
As with other Bayesian methods, the proposed approach would probably shine in low-samples
regime and in this case might be preferable to other methods in the same class (VB, GP).
We want to thank the reviewer for the review.
The authors propose a novel way of using Bayesian NNs for policy search in stochastic dynamical systems.
Specifically, the authors minimize alpha-divergence with alpha=0.5 as opposed to standard VB. The authors claim that their method is the first model-based system to solve a 20 year old benchmark problem; I'm not very familiar with this literature, so it's difficult for me to assess this claim.
The paper seems technically sound. I feel the writing could be improved. The notation in sections 2-3 feels a bit dense and there are a lot of terminology / approximations introduced, which makes it hard to follow. The writing could be better structured to distinguish between novel contributions vs review of prior work. If I understand section 2.3 correctly, it's mostly a review of black box alpha divergence minimization. If so, it would probably make sense to move this to the appendix.
There was a paper at NIPS 2016 showing promising results using SGHMC for Bayesian optimization: "Bayesian optimization with robust Bayesian neural networks" by Springenberg et al. Could you comment on applicability of stochastic gradient MCMC (SGLD / SGHMC) for your setup?
Can you comment on the computational complexity of the different approaches?
Section 4.2.1: why can't you use the original data? in what sense is it fair to simulate data using another neural network? can you evaluate PSO-P on this problem?
We want to thank the reviewer for the review.  We would like to give answers to the reviewer's questions:
1. Could you comment on applicability of stochastic gradient MCMC (SGLD / SGHMC) for your setup?
SGLD / SGHMC are alternative methods for approximate inference in the proposed neural networks with stochastic inputs.
Instead of learning a parametric Gaussian approximation q(w) to the posterior, we could have as well used such methods
to  produce a collection of samples, whose empirical distribution would approximate the exact posterior.  For training
the policy we would simply exchange line 4 of algorithm 1 to draw samples from the empirical distribution.
We chose to use deterministic methods because of 1) their simplicity, 2) lower computational cost and 3) ability to
compactly represent and to easily update the posterior approximation.
Note also that while SGLD / SGHMC have already been applied to Bayesian neural networks (BNNs) for learning a posterior
on the weights, these methods have never been applied to BNNs with random inputs, where one is also required to learn
a posterior over the random inputs that were used to generate the training data.
There are different trade-offs to consider when choosing whether to use sampling based methods (SGLD / SGHMC) or
deterministic approximation methods based on optimization (alpha-divergence minimization, variational Bayes)
for approximate inference:
1 - Sampling-based methods are asymptotically unbiased and can therefore produce more accurate posterior approximations than
deterministic approaches providing enough computation time is available. Deterministic methods are by contrast biased.
2 - Sampling-based methods have higher computational cost. Requiring a significantly higher number of passes through the
data than deterministic approaches.
3 - Sampling-based method usually include many hyper-parameters that are highly data and model dependent and that require tuning.
By contrast, deterministic approaches often include less hyper-parameters (in our case, we only tune the learning rate).
4 - It is difficult to determine when a sampling-based method has converged and is drawing samples from the correct stationary
distribution.
By contrast it is easy to determine when a deterministic approach has converged.
For example, when
the objective function that is being optimized does not improve anymore beyond a specific threshold.
5 - In deterministic methods the posterior approximation is compactly represented by a collection of parameters. These parameters
can be easily updated when new data is available. Updating the samples generated by sampling-based methods is by
contrast very challenging.
2.  "Section 4.2.1: why can't you use the original data? in what sense is it fair to simulate data using another neural network? can you evaluate PSO-P on this problem?"
We use original data recorded from a gas turbine to train the world model. We do this because  we need to transform the
recordings into a model that an agent can interact with (We cannot compare the different methods on the real turbine directly,
for obvious safety concerns). To that end we use a neural network as an approximator for the turbine dynamics and sample
data from it in a partial observable scenario.  This process is repeated in each of the 5 repetitions of the
experiment.
Regarding fairness:  The artificial batch only has about 10% of original features visible.
Because of this partial
observability, we expect that there should be not much advantage from the fact that the MLP and BNNs share the same
base function class with the world model as compared to the GPs.
Indeed, the results in table 1 show that the policies
of the  BNNs and the GP outperform the MLP although the MLP and the world model do their stochastic transition in a
similar fashion. We mentioned the ‘fairness’ aspect here because we saw the MLP as a simple baseline and used
this class as the lowest common denominator--- it would be arguably more unfair had we used a BNN or GP as function
class for the world model.
Yes, we will include PSO-P results in the final version. We excluded it in the current version because it is computationally
demanding.
3. "Can you comment on the computational complexity of the different approaches?"
We would like to comment on the complexity of training the model (usually faster) and training the policy
separately (usually much slower):
All models were trained using theano and a single GPU. Training  the standard neural network is the fastest (SGD + the
early stopping process that slows things a bit down). The training time for this method was between 5 - 20 minutes,
depending on data set size and dimensionality of the benchmark.
In theano, the computational graph of the BNNs is similar to that of an ensemble of standard neural networks. The training time
for the BNNs varied between 30 minutes to 5 hours depending on data size and dimensionality of benchmark.
The sparse Gaussian Process was optimized using an expectation propagation algorithm and after training, it was approximated
with a Bayesian linear model with fixed basis functions whose weights are initialized randomly (see Appendix B).
We choose the inducing points in the GPs and the number of training epochs for these models so that the resulting
training time was comparable to that of the BNNs.
For policy training we used a single CPU. All methods are of similar complexity as they are all trained using Algorithm 1.
Depending on the horizon, data set size and network topology, training took between 20 minutes  (Wet-Chicken,T=5),
3-4 hours  (Turbine, T=20) and 14-16  hours (industrial benchmark, T=75).
Thanks for your detailed response. Yes, please include the PSO-P results in the final version. It would be useful to include the timing discussion as well so that readers get a sense of the computational complexity.
Dear reviewer, in the light of our responses and the other reviewers' comments, did your initial assessment of our work change?
I didn't have any pre-review question. Have posted my review titled "Policy search using Bayesian NNs" now.
This paper introduces an approach for model-based control of stochastic dynamical systems with policy search, based on (1) learning the stochastic dynamics of the underlying system with a Bayesian deep neural network (BNN) that allows some of its inputs to be stochastic, and (2) a policy optimization method based on simulated rollouts from the learned dynamics. BNN training is carried out using \alpha-divergence minimization, the specific form of which was introduced in previous work by the authors. Validation and comparison of the approach is undertaken on a simulated domain, as well as real-world scenarios.
The paper is tightly written, and easy to follow. Its approach to fitting Bayesian neural networks with \alpha divergence is interesting and appears novel in this context. The resulting application to model-based control appears to have significant practical impact, particularly in light of the explainability that a system model can bring to specific decisions made by the policy.
As such, I think that the paper brings a valuable contribution to the literature.
That said, I have a few questions and suggestions:
1) In section 2.2, it should be explained how the random z_n input is used by the neural network: is it just concatenated to the other inputs and used as-is, or is there a special treatment?
2) Moreover, much case is made for the need to have stochastic inputs, but only a scalar input seems to be provided throughout. Is this enough? How computationally difficult would providing stochastic inputs of higher dimensionality be?
3) How important is the normality assumption in z_n? How is the variance \gamma established?
4) It is mentioned that the hidden layers of the neural network are made of rectifiers, but no further utilization of this fact is made in the paper. Is this assumption somehow important in the optimization of the alpha-divergence (beyond what we know about rectifiers to mitigate the vanishing gradient problem) ?
5) Equation (3), denominator \mathbf{y} should be \mathbf{Y} ?
6) Section 2.3: it would be helpful to have an overview or discussion of the computational complexity of training BNNs, to understand whether and when they can practicably be used.
7) Between eq (12) and (13), a citation to the statement of the time embedding theorem would be helpful, as well as an indication of how the embedding dimension should be chosen.
8) Figure 1: the subplots should have the letters by which they are referenced in the text on p. 7.
9) In section 4.2.1, it is not clear if the gas turbine data is publicly available, and if so where.
In addition more details should be provided, such as the dimensionality of the variables E_t, N_t and A_t.
10) Perhaps the comparisons with Gaussian processes should include variants that support stochastic inputs, such as Girard et al. (2003), to provide some of the same modelling capabilities as what’s made use of here.
At least, this strand of work should be mentioned in Section 5.
References:
Girard, A., Rasmussen, C. E., Quiñonero Candela, J., & Murray Smith, R. (2003). Gaussian process priors with uncertain inputs-application to multiple-step ahead time series forecasting. Advances in Neural Information Processing Systems, 545-552.
We want to thank the reviewer for the review.  We would like to give answers to the reviewer's questions:
1) In section 2.2, it should be explained how the random z_n input is used by the neural network: is it just concatenated
to the other inputs and used as-is, or is there a special treatment?
Yes, the random input z_n is an additional input  variable of  dimension 1, concatenated to the other inputs.
In contrast to  the other inputs  however,  we optimize two  free parameters (mu_z_n,\sigma_z_n) for each z_1,\cdots,z_N
during training:
We assume there exists a latent variable $z$ with unknown realizations  $z_1,\cdots,z_N$ for the state transitions
\{(x_1,y_1),\cdots,(x_N,y_N)\}$. Our models tries to infer these realizations by learning a  variational distribution
\{(\mu_z_1,\sigma_z_1),\cdots,(\mu_z_N,\sigma_z_N)\}.
2) Moreover, much case is made for the need to have stochastic inputs, but only a scalar input seems to be provided throughout.
Is this enough? How computationally difficult would providing stochastic inputs of higher dimensionality be?
A central motivation for the stochastic inputs is so that the Bayesian Model has an explicit place where model uncertainty is
(in the uncertainty over the W) and an explicit place where noise uncertainty is (the transformation of the input noise z sampled
from the prior). These two components give principled Bayesian inference over stochastic functions.
Higher dimensional stochastic input noise can be used without any restriction, note however, that by default a BNN will blend
together all input sources x and z in each layer.
Thus, while the stochastic input  was just 1-dimensional at the input layer,
a transformation to a  high-dimensional and complex random variable will occur as it propagates through the network.
In this sense, we do not believe that using only 1-dimensional input noise is a limitation in our method.
3) How important is the normality assumption in z_n? How is the variance \gamma established?
The normality assumption in z_n is not important at all. We could have as well used any other sampling distribution.
Any arbitrary distribution could be obtained by transforming Gaussian samples through non-linearities, e.g. by
applying the Gaussian cdf and then the inverse cdf of any desired distribution. As the random samples of z are propagated
through the neural network the could be transformed to have the same effect as if they had been sampled from
any other arbitrary distribution.
We define the prior over  the latent variable z to be a Gaussian with \mu_z_n=0, \sigma_z_n = \gamma = \sqrt(d)
where d is the dimensionality of x.  This is an attempt to keep the effect of z  on the network activation independent
of the input dimensionality (compare the Wet-Chicken benchmark with only 4 inputs to the industrial benchmark with around
50 input dimensions). By doing this, the effect of the input noise is not vanishing in high dimensional problems.
In our experiments, the results obtained are not very sensitive to the particular choice of \gamma that we make.
4) It is mentioned that the hidden layers of the neural network are made of rectifiers, but no further utilization of this
fact is made in the paper. Is this assumption somehow important in the optimization of the alpha-divergence
(beyond what we know about rectifiers to mitigate the vanishing gradient problem) ?
The choice for the activation function is unrelated to alpha-divergence minimization.  We choose rectifiers
because these are standardly used in literature. As with neural networks any differentiable activation  function can be used.
5) Equation (3), denominator \mathbf{y} should be \mathbf{Y} ?
This is correct. Thanks for pointing this out!
6) Section 2.3: it would be helpful to have an overview or discussion of the computational complexity of training BNNs,
to understand whether and when they can practicably be used.
Thanks for the feedback. We would like to refer the reviewer to our answer given to question 3 to  AnonReviewer2.
Short answer: In theano, the computational graph of the BNNs is similar to that of an ensemble of standard deterministic
neural networks (with ensemble size equal to the number of samples from the posterior approximating that are used to
approximate the expectations in the objective function for black-box alpha)
7) Between eq (12) and (13), a citation to the statement of the time embedding theorem would be helpful,
as well as an indication of how the embedding dimension should be chosen.
Thanks for the feedback. A citation will be added to the final version.
8) Figure 1: the subplots should have the letters by which they are referenced in the text on p. 7.
Thanks for the feedback. We will include this in the final version.
9) In section 4.2.1, it is not clear if the gas turbine data is publicly available, and if so where.
In addition more details should be provided, such as the dimensionality of the variables E_t, N_t and A_t.
The gas turbine data unfortunately is not publicly available at this point. Similar experiments can be found in [1].
For the final version we will be more specific about the dimensionalities of these variables.
10) Perhaps the comparisons with Gaussian processes should include variants that support stochastic inputs,
such as Girard et al. (2003), to provide some of the same modelling capabilities as what’s made use of here.
At least, this strand of work should be mentioned in Section 5.
Thanks for the reference, we will be include it in the final  version. We chose the GP to resemble the
PILCO approach[2], it is to our knowledge the most direct competitor to the method we propose.
We agree that comparing to GPs with stochastic inputs would be interesting. Extending PILCO to learn a
variational distribution on z_1,\cdots,z_N is not trivial, to our knowledge such an  architecture has
never been used in the context of model-based RL or policy search. It  may be of special interest for a
separate, independent contribution.
[1] Schaefer, Anton Maximilian, et al. "A neural reinforcement learning approach to gas turbine control."
2007 International Joint Conference on Neural Networks. IEEE, 2007.
[2] Deisenroth, Marc, and Carl E. Rasmussen. "PILCO: A model-based and data-efficient approach to policy search."
Proceedings of the 28th International Conference on machine learning (ICML-11). 2011.

Meta-review:
role: PC
data: comment: Despite it's initial emphasis on policy search, this paper is really about a learning method for Bayesian neural networks, which it then uses in a policy search setting. Specifically, the authors advocate modeling a stochastic system using a BNN trained to minimize alpha-divergence with alpha=0.5 (this involves a great deal of approximation to make computationally tractable). They then use this in the policy search setting.
 
 The paper is quite clear, and proposes a nice approach to learning BNNs. The algorithmic impact honestly seems fairly minor (the idea of using different alpha-divergences instead of KL divergence has been considered many times in the content of general variational approximations), but combining this with the policy search setting, and reasonable examples of industrial control, together these all make this a fairly strong paper.
 
 Pros:
 + Nice derivation of alternative variational formulation (I'll still call it variational even though it uses alpha-divergence with alpha=0.5)
 + Good integration into policy search setting
 + Nice application to industrial control systems
 
 Cons:
 - Advance from the algorithmic standpoint seems fairly straightforward, if complicated to make tractable (variational inference using alpha-divergence is not a new idea)
 - The RL components here aren't particularly novel, really the novelty is in the learning

Decision: Accept (Poster)