S WITCHING
L INEAR
D YNAMICS
FOR
V ARIATIONAL
B AYES F ILTERING
Anonymous authors
Paper under double-blind review
A BSTRACT
System identiﬁcation of complex and nonlinear systems is a central problem for model predictive control and model-based reinforcement learning.
Despite their complexity, such systems can often be approximated well by a set of linear dynamical systems if broken into appropriate subsequences.
This mechanism not only helps us ﬁnd good approximations of dynamics, but also gives us deeper insight into the underlying system.
Leveraging Bayesian inference and Variational
Autoencoders, we show how to learn a richer and more meaningful state space, e.g. encoding joint constraints and collisions with walls in a maze, from partial and high-dimensional observations.
This representation translates into a gain of accuracy of the learned dynamics which we showcase on various simulated tasks.
I NTRODUCTION
Learning dynamics from raw data (also known as system identiﬁcation) is a key component of model predictive control and model-based reinforcement learning.
Problematically, environments of interest often give rise to very complex and highly nonlinear dynamics which are seemingly difﬁcult to approximate.
However, switching linear dynamical systems (SLDS) approaches claim that those environments can often be broken down into simpler units made up of areas of equal and linear dynamics (Ackerson & Fu, 1970; Chang & Athans, 1978).
Not only are those approaches capable of good predictive performance, which often is the sole goal of learning a system’s dynamics, they also encode valuable information into so called switching variables which determine the dynamics of the next transition.
For example, when looking at the movement of an arm, one is intuitively aware of certain restrictions of possible movements, e.g. constraints to the movement due to joint constraints or obstacles.
The knowledge is present without the need to simulate; it’s explicit.
Exactly this kind of information will be encoded when successfully learning switching dynamics.
Our goal in this work will therefore entail the search for richer representations in the form of latent state space models which encode knowledge about the underlying system dynamics.
In turn, we expect this to improve the accuracy of our simulation as well.
Such a representation alone could then be used in a reinforcement learning approach that possibly only takes advantage of the learned latent features but not necessarily its learned dynamics.
To learn richer representations, we identify one common problem with prevalent recurrent Variational
Autoencoder models (Karl et al., 2017a; Krishnan et al., 2015; Chung et al., 2015; Fraccaro et al., 2016): the non-probabilistic treatment of the transition dynamics often modeled by a powerful nonlinear function approximator.
From the history of the Autoencoder to the Variational Autoencoder, we know that in order to detect features in an unsupervised manner, probabilistic treatment of the latent space is paramount.
As our starting point, we will build on previously proposed approaches by
Krishnan et al.
(2017) and Karl et al. (2017a).
The latter already made use of locally linear dynamics, but only in a deterministic fashion.
We extend their approaches by a stochastic switching LDS model and show that such treatment is vital for learning richer representations and simulation accuracy.
B ACKGROUND
We consider discretized time-series data consisting of continuous observations x t ∈ X ⊂ R n x and control inputs u t ∈ U ⊂ R n u that we would like to model by corresponding latent states z t ∈ Z ⊂ R n z .
We’ll denote sequences of variables by x 1: T = ( x 1 , x 2 , ..., x T ) .
ssszzzuuxxx(a) SLDS graphical model.
w sszzzuuxxx(b)
Our generative model.
2. S WITCHING LINEAR DYNAMICAL SYSTEMS
Switching Linear Dynamical System models (SLDS) enable us to model nonlinear time series data by splitting it into sequences of linear dynamical models.
At each time t = 1 , 2 , ..., T , a discrete switch variable s t ∈ 1 , ..., M chooses of a set LDSs a system which is to be used to transform our continuous latent state z t to the next time step (Barber, 2012).
[Equation]
Here A ∈ R n z × n z is the state matrix, B ∈ R n z × n u control matrix, (cid:15) the transition noise with covariance matrix Q and η the emission/sensor noise with covariance matrix R .
Finally, the observation matrix H ∈ R n x × n z deﬁnes a linear mapping from latent to observation space which we will replace by a nonlinear transformation parameterized by a neural net.
These equations imply the following joint distribution:
[Equation]
t = p ( x t | z t ) p ( z t | z t − 1 , u t − 1 , s t ) p ( s t | z t − 1 , u t − 1 , s t − 1 ) (2) with p ( z 1 | z 0 , u 0 , s 1 )
= p ( z 1 ) being the initial state distribution.
The corresponding graphical model is shown in ﬁgure 1a.
2. S TOCHASTIC GRADIENT VARIATIONAL B AYES
[Equation]
Given the simple graphical model in equation (3), Kingma & Welling (2014) and Rezende et al.
(2014) introduced the Variational Autoencoder (VAE) which overcomes the intractability of posterior inference of q ( z | x ) by maximizing the evidence lower bound (ELBO) of the model log-likelihood.
[Equation]
Their main innovation was to approximate the intractable posterior distribution by a recognition network q φ ( z | x ) from which they can sample via the reparameterization trick to allow for stochastic backpropagation through both the recognition and generative model at once.
Assuming that the latent state is normally distributed, a simple transformation allows us to obtain a Monte Carlo gradient estimate of E q φ ( z | x )
[ln p θ ( x | z )] w.r.t.
to φ .
Given that z ∼ N ( µ, σ 2 ) , we can generate samples by drawing from an auxiliary variable (cid:15) ∼ N (0 , 1) and applying the deterministic and differentiable transformation z
= µ + σ(cid:15) .
2. T HE CONCRETE DISTRIBUTION
One simple and efﬁcient way to obtain samples d from a k -dimensional categorical distribution with class probabilities α is the Gumbel-Max trick:
[Equation]
hot ( argmax [ g i + log α i ]) , with g 1 , . .
, g k ∼ Gumbel (0 , 1) (5)
However, since the derivative of the argmax is 0 everywhere except at the boundary of state changes, where it is undeﬁned, we can’t learn a parameterization by backpropagation.
The Gumbel-Softmax trick approximates the argmax by a softmax which gives us a probability vector (Maddison et al., 2017; Jang et al., 2017).
We can then draw samples via
[Equation]
with g
[Equation]
Gumbel
[Equation] [Equation]
[Equation]
This softmax computation approaches the discrete argmax as temperature λ
→ 0 , for λ → ∞ it approaches a uniform distribution.
R ELATED W ORK
Our model can be viewed as a Deep Kalman Filter (Krishnan et al., 2015) with structured inference (Krishnan et al., 2017).
In our case, structured inference entails another stochastic variable model with parameter sharing inspired by Karl et al.
(2017b)
and
Karl et al.
(2017a) which pointed out the importance of backpropagating the reconstruction error through the transition.
We are different to a number of stochastic sequential models like Bayer & Osendorfer (2014); Chung et al.
Shabanian et al.
; Goyal et al.
(2017) by directly transitioning the stochastic latent variable over time instead of having an RNN augmented by stochastic inputs.
Fraccaro et al.
(2016) has a transition over both a deterministic and a stochastic latent state sequence, wanting to combine the best of both worlds.
Previous models (Watter et al., 2015; Karl et al., 2017a; Fraccaro et al., 2017) have already combined locally linear models with recurrent Variational Autoencoders, however they provide a weaker structural incentive for learning latent variables determining the transition function.
Van Steenkiste et al.
(2018) approach a similar multi bouncing ball problem (see section 5.1) by ﬁrst distributing the representation of different balls into their own entities without supervision and then structurally hardwiring a transition function with interactions based on an attention mechanism.
Recurrent switching linear dynamical systems (Linderman et al., 2016) uses message passing for approximate inference, but has restricted itself to low-dimensional observations and a multi-stage training process.
Johnson et al.
(2016) propose a similar model to ours but combine message passing for discrete switching variables with a neural network encoder for observations learned by stochastic backpropagation.
Tackling the problem of propagating state uncertainty over time, various combinations of neural networks for inference and Gaussian processes for transition dynamics have been proposed (Eleftheriadis et al., 2017; Doerr et al., 2018).
However, these models have not been demonstrated to work with high-dimensional observation spaces like images.
One feature a switching
LDS model may learn are interactions which have recently been approached by employing Graph
Neural Networks (Battaglia et al., 2016; Kipf et al., 2018).
These methods are similar in that they predict edges which encode interactions between components of the state space (nodes).
P ROPOSED A PPROACH
Our goal is to ﬁt a series of continuous state z 1: T and switching variables s 2: T to a given sequence of observations x 1: T .
We assume a nonlinear mapping between observations and latent space which we generally approximate by neural networks, apart from the transition which is modeled by a locally linear function.
Our generative model is shown in ﬁgure 1b an our inference model in ﬁgure 2a.
4. G ENERATIVE M ODEL
Our generative model for a single x t is described by
[Equation]
7) which is close to the one of the original SLDS model (see ﬁgure 1a).
Latent states z t are continuous and represent the state of the system while states s t are the switching variables determining the transition.
We approximate the discrete switching variables by a continuous relaxation, namely the Concrete distribution.
1 Differently to the original model, we do not condition the likelihood of the current observation p θ ( x t | z t ) directly on the switching variables.
This limits the inﬂuence of the switching variables to choosing a proper transition dynamic for the continuous latent space.
The likelihood model is parameterized by a neural network with either a Gaussian or a Bernoulli distribution as output depending on the data.
There is both a transition on the continuous states z t and discrete latent states s t .
For the continuous state transition p ( z t | z t − 1 , s t , u t − 1 ) we follow (1) and maintain a set of M base matrices { (cid:0) A ( i ) , B ( i ) , Q ( i ) (cid:1)
| ∀ i. 0 < i < M } as our linear dynamical systems to choose from.
For the transition on discrete latent states p ( s t | s t − 1 , z t − 1 , u t − 1 ) , we usually require the learning of a Markov transition matrix.
However, since we approximate our discrete switching variables by a continuous relaxation, we can parameterize this transition by a neural network.
Therefore, our entire generative model can be learned end-to-end by (stochastic) backpropagation.
Finally, the resulting dynamics matrices are computed through a linear combination of the base matrices:
[Equation]
Both transition models – the continuous state transition p θ ( z t | z t − 1 , s t , u t − 1 ) and concrete switching variables transition p θ ( s t | s t − 1 , z t − 1 , u t − 1 ) – are shared with the inference model which is key for good performance. [Equation]
N (cid:0) µ, σ 2 (cid:1) where [ µ, σ 2 ] =
f θ ( z t − 1 , s t , u t − 1 ) p θ ( s t | s t − 1 , z t − 1 , u t − 1 ) =
Concrete ( α, λ prior ) where α = g θ
( z t − 1 , s t − 1 , u t − 1 ) (9)
4. I NFERENCE 4.2. S TRUCTURED INFERENCE OF CONTINUOUS LATENT STATE
We split our inference model q φ ( z t | z t − 1 , s t , x ≥ t , u ≥ t − 1 ) into two parts: 1) transition model q trans ( z t | z t − 1 , s t , u t − 1 ) and 2) inverse measurement model q meas ( z t | x ≥ t , u ≥ t ) as previously proposed in Karl et al. (2017b).
This split allows us to reuse our generative transition model in place of q trans ( z t | z t − 1 , s t , u t − 1 ) .
This sharing of variables is essential for good performance as it forces the reconstruction error to be backpropagated through the transition model.
For practical reasons, we only share the computation of the transition mean µ trans but not the variance σ 2trans between inference and generative model.
Both parts, q meas and q trans , will give us independent predictions about the new state z t which will be combined in a manner akin to a Bayesian update in a Kalman
Filter.
q φ ( z t | z t − 1 , s t , x ≥ t , u ≥ t − 1 )
∝ q meas ( z t | x ≥ t , u ≥ t ) × q trans ( z t | z t − 1 , s t , u t − 1 ) =
N (cid:0) µ q , σ 2 q (cid:1) q meas ( z t | x ≥ t , u ≥ t ) = N (cid:0) µ meas , σ 2meas (cid:1) where [ µ meas , σ 2meas ]
= h φ ( x ≥ t , u ≥ t ) q trans ( z t | z t − 1 , s t , u t − 1 ) =
N (cid:0) µ trans , σ 2trans (cid:1) where [ µ trans , σ 2trans ]
= f θ ( z t − 1 , s t , u t − 1 )
[Equation]
The densities of q meas and q trans are multiplied resulting in another Gaussian density:
[Equation]
This update scheme is highlighted in ﬁgure 2b.
We found empirically that conditioning the inverse measurement model q meas ( z t | x ≥ t , u ≥ t ) solely on the current observation x t instead of the entire remaining trajectory to lead to better results.
We hypothesize that the recurrent model needlessly introduces very high-dimensional and complicated dynamics which are harder to approximate with our locally linear transition model.
For the initial state z 1 we do not have a conditional prior from the transition model as in the rest of the sequence.
Other methods (Krishnan et al., 2015) have used a standard normal prior, however this is not a good ﬁt.
We therefore decided that instead of predicting z 1 directly to predict an auxiliary variable w that is then mapped deterministically to a starting state z 1 .
A standard Gaussian prior is then applied to w .
Alternatively, we could specify a more complex or learned prior for the initial state like the VampPrior (Tomczak & Welling, 2017).
Empirically, this has lead to worse results.
[Equation]
[ µ [Equation] [Equation] While we could condition on the entire sequence, we restrict it to just the ﬁrst couple of observations.
bbw sszzzuuxxx(
[Figure]
model.
q meas ( z t | x t ) q meas ( s t | x ≥ t ) × q trans ( z t | · ) q ( z t | x ≥ t , z t − 1 , s t , u t − 1 ) × q trans ( s t | · )
p θ ( x t | z t )
z t −s t −z t − 1 q ( s t | x ≥ t , z t − 1 , s t − 1 , u t − 1 ) s t x t x ≥ t u t − 1 u t −
[Figure]
[Figure]
4.2. I NFERENCE OF S WITCHING V ARIABLES
Following Maddison et al.
(2017) and Jang et al.
(2017), we can reparameterize a discrete latent variable with the Gumbel-softmax trick.
Again, we split our inference network q φ ( s t | s t − 1 , z t − 1 , x ≥ t , u ≥ t − 1 ) in an identical fashion into two components: 1) Transition model q trans ( s t | s t − 1 , z t − 1 , u t − 1 ) and 2) inverse measurement model q meas ( s t | x ≥ t , u ≥ t ) .
The transition model is again shared with the generative model and is implemented via a neural network as we potentially require quick changes to chosen dynamics. The inverse measurement model is parametrized by a backward LSTM.
However, for the case of concrete variables, we cannot do the same Gauss multiplication as in the previous case.
Therefore, we let each network predict the logits of a Concrete distribution and our inverse measurement model q φ ( s t | x ≥ t , u ≥ t ) produces an additional vector γ , which determines the value of a gate deciding how the two predictions are to be weighted: q φ ( s t | s t − 1 , z t − 1 , x ≥ t , u ≥ t − 1 ) =
Concrete ( α, λ posterior ) with α = γα trans + (1 − γ ) α meas q meas ( s t | x ≥ t , u ≥ t ) = Concrete ( α meas , λ posterior ) where [ α meas , γ ] = k φ ( x ≥ t , u ≥ t ) q trans ( s t | s t − 1 , z t − 1 , u t − 1 ) =
Concrete ( α trans , λ prior ) where α = g θ
( z t − 1 , s t − 1 , u t − 1 ) (13)
The temperatures λ posterior and λ prior are set as a hyperparameter and can be set differently for the prior and approximate posterior. The gating mechanism gives the model the option to balance between prior and approximate posterior.
If the prior is good enough to explain the next observation, γ will be pushed to 1 which ignores the measurement and minimizes the KL between prior and posterior by only propagating the prior.
If the prior is not sufﬁcient, information from the inverse measurement model can ﬂow by decreasing γ and incurring a KL penalty.
Since the concrete distribution is a relaxation of the categorical, our sample will not be a one-hot vector, but a vector whose elements sum up to 1.
We face two options here: we could take a categorical sample by choosing the linear system corresponding to the highest value in the sample (hard forward pass) and only use the relaxation for our backward pass.
This, however, means that we will follow a biased gradient.
Alternatively, we can use the relaxed version for our forward pass and aggregate the linear systems based on their corresponding weighting (see (8)).
Here, we lose the discrete switching of linear systems, but maintain a valid lower bound.
We note that the hard forward pass has led to worse results and focus on the soft forward pass for this paper.
Lastly, we could go further away from the theory and instead treat the switching variables also as normally distributed.
If this worked better than the approach with Concrete variables, it would highlight still existing optimization problems of discrete random variables.
As such, it will act as an ablation study for our model.
The mixing coefﬁcients for linear systems would then be determined by a linear combination of these latent variables:
[Equation]
Our inference scheme for normally distributed switching variables is then identical to the one described in the previous section.
We compare both approaches throughout our experimental section.
4. T RAINING
Our objective function is the commonly used evidence lower bound for our hierarchical model.
[Equation]
We choose to factorize over time
[Equation]
so the loss for a single observation x t becomes:
[Equation]
The full derivation can be found in appendix A.
We learn the parameters of our model by backpropagation through time and we (generally) approximate the expectations with one sample by using the reparametrization trick.
The exception is the KL between two Concrete random variables in which case we take 10 samples for the approximation.
For the KL on the switching variables, we further introduce a scaling factor β < 1 (as ﬁrst suggested in Higgins et al.
(2016), although they suggested increasing the KL term) to down weigh its importance.
More details on the training procedure can be found in appendix B.2.
E XPERIMENTS
In this section, we evaluate our approach on a diverse set of physics and robotics simulations based on partially observable system states or high-dimensional images as observations.
We show that our model outperforms previous models and that our switching variables learn meaningful representations.
Models we compare to are Deep Variational Bayes Filter (DVBF)
(Karl et al., 2017a), DVBF Fusion (Karl et al., 2017b) (called fusion as they do the same Gauss multiplication in the inference network) which is closest to our model but doesn’t have a stochastic treatment of the transition, the Kalman
VAE (KVAE)
(Fraccaro et al., 2017) and a LSTM (Hochreiter & Schmidhuber, 1997).
[Figure]
5. M ULTIPLE B OUNCING B ALLS IN A M AZE
Our ﬁrst experiment is a custom 3-agent maze environment simulated with Box2D. Each agent is fully described by its x and y coordinates and its current velocity and has the capability to accelerate in either direction.
We learn in a partially observable setting and limit the observations to the agents’ positions, therefore x ∈ R 6 while the true state space is in R 12 and u ∈ R 6 .
First, we train a linear regression model on the latent space z to see if we have recovered a linear encoding of the unobserved velocities.
We achieve an R2 score of 0.92 averaged over all agents and velocity directions.
Our focus shifts now to our switching variables which we expect to encode interactions with walls.
We provide a visual conﬁrmation of that in ﬁgure 3 where we see switching variables encoding all space where there is no interaction in the next time step, and variables which encode walls, distinguishing between vertical and horizontal ones.
In ﬁgure 3d one can see show that if the choice of locally linear transition is treated deterministically, we don’t learn global features of the same kind.
To conﬁrm our visual inspection, we train a simple decision tree based on latent space s in order to predict interaction with a wall.
Here, we achieve an F1 score of 0 .
It is difﬁcult to say what a good value should look like as collisions with low velocity are virtually indistinguishable from no collision.
We compare our prediction quality to several other methods in table 1 where we outperform all of our chosen baselines.
Also, modeling switching variables by a Normal distribution outperforms the
Concrete distribution in all of our experiments.
Aside from known practical issues with training a discrete variable via backpropagation, we explore one reason why that may be in section 5.4, which is the greater susceptibility to the scale of temporal discretization.
We provide plots of predicted trajectories in appendix D.
Transitioning multiple agents with a single transition matrix comes with scalability issues with regards to switching dynamics which we explore further in appendix C. refers to constantly predicting the ﬁrst observation of the sequence.
[Table]
We then evaluate our model on the Roboschool reacher environment.
To make things more interesting, we learn only on partial observations, removing time derivative information (velocities), leaving us with just the positions or angles of various joints as observations.
Table 1 shows a comparison of various methods on predicting the next couple of time steps.
One critical point is the possible collision 2 between lower and upper joint which is one we’d like our model to capture.
We again learn a linear classiﬁer based on latent space s to see if this is successfully encoded and reach an F1 score of 0 . 46 .
5. B ALL IN A B OX ON I MAGE D ATA
Finally [Table] we evaluate our method on high-dimensional image observations using the single bouncing ball environment used by Fraccaro et al.
They simulated 5000 sequences of 20 time steps each of a ball moving in a two-dimensional box, where each video frame is a 32 × 32 binary image.
There are no forces applied to the ball, except for the fully elastic collisions with the walls.
Initial position and velocity are randomly sampled.
In ﬁgure 5a we compare our model to both the smoothed and generative version of the KVAE.
The smoothed version receives the ﬁnal state of the trajectory after the n predicted steps which is fed into the smoothing capability of the KVAE.
One can see that our model learns a better transition model, even outperforming the smoothed KVAE for longer sequences.
For short sequences, KVAE performs better which highlights the value of it disentangling the latent space into separate object and dynamics representation.
A sample trajectory is plotted in ﬁgure 4.
5. S USCEPTIBILITY TO THE SCALE OF TEMPORAL DISCRETIZATION
In this section, we’d like to explore how the choice of ∆ t when discretizing a system inﬂuences our results.
In particular, we’d expect our model with discrete (concrete) switching latent variables to be more susceptible to it than when modeled by a continuous distribution.
This is because in the latter case the switching variables can scale the various matrices more freely, while in the former scaling up one system necessitates scaling down another.
For empirical comparison, we go back to our custom maze environment (this time with only one agent as this is not pertinent to our question at hand) and learn the dynamics on various discretization scales.
Then we compare the absolute error’s growth for both approaches in ﬁgure 5b which supports our hypothesis.
While the discrete approximation even outperforms for small ∆ t , there is a point where it rapidly becomes worse and gets overtaken by the continuous approximation.
This suggests that ∆ t was simply chosen to be too large in both the reacher and the ball in a box with image observations experiment.
[Figure]
D ISCUSSION
We want to emphasize some subtle differences to previously proposed architectures that make an empirical difference, in particular for the case when s t is chosen to be continuous.
In Watter et al.
(2015) and Karl et al. (2017a), the latent space is already used to draw transition matrices, however they do not extract features such as walls or joint constraints.
There are a few key differences from our approach.
First, our latent switching variables s t are only involved in predicting the current observation x t through the transition selection process.
The likelihood model therefore doesn’t need to learn to ignore some input dimensions which are only helpful for reconstructing future observations but not the current one.
There is also a clearer restriction on how s t and z t may interact: s t may now only inﬂuence z t by determining the dynamics, while previously z t inﬂuenced both the choice of transition function as well as acted inside the transition.
These two opposing roles lead to conﬂicting gradients as to what should be improved.
Furthermore, the learning signal for s t is rather weak so that scaling down the KL-regularization was necessary to detect good features.
Lastly, a (locally) linear transition may not be a good ﬁt for variables determining dynamics as such variables may change very abruptly.
C ONCLUSION
We have shown that our construction of using switching variables encourages learning a richer and more interpretable latent space.
In turn, the richer representation led to an improvement of simulation accuracy in various tasks.
In the future, we’d like to look at other ways to approximate the discrete switching variables and exploit this approach for model-based control on real hardware systems.
Furthermore, addressing the open problem of disentangling latent spaces is essential to ﬁtting simple dynamics and would lead to signiﬁcant improvements of this approach
A L OWER BOUND DERIVATION
For brevity we omit conditioning on control inputs u 1: T .
[Equation]
A. F ACTORIZATION OF THE KL DIVERGENCE
The dependencies on data x T and u T as well as parameters φ and θ are omitted in the following for convenience.
[Equation]
Factorization of the variational approximation
[Equation]
(Factorization of the prior)
[Equation]
(Expanding the logarithm by the product rule)
[Equation]
[Table]
B D ETAILS OF THE E XPERIMENTAL S ETUP B. E NVIRONMENTS
[Table]
To generate data, we follow a Uniform distribution U ∼ [ − 1 , 1] as the exploration policy.
Before we record data, we take 20 warm-up steps in the environment to randomize our starting state.
We take the data as is without any other preprocessing.
[Table]
[Table]
M ULTI A GENT M AZE
Observations are normalized to be in [ − 1 , 1] .
Both position and velocity is randomized for the starting state.
We again follow a Uniform distribution U ∼ [ − 1 , 1] as the exploration policy.
B. T RAINING
Overall, training the Concrete distribution has given us the biggest challenge as it was very susceptible to various hyperparameters.
We made use of the fact that we can use a different temperature for the prior and approximate posterior (Maddison et al., 2017) and we do independent hyperparameter search over both.
For us, the best values were 0 .
75 for the posterior and 2 for the prior.
Additionally, we employ an exponential annealing scheme for the temperature hyperparameter of the Concrete distribution.
This leads to a more uniform combination of base matrices early in training which has two desirable effects.
First, all matrices are scaled to a similar magnitude, making initialization less critical.
Second, the model initially tries to ﬁt a globally linear model, leading to a good starting state for optimization.
We also tried increasing the number of samples taken (up to 100 ) to approximate the KL between the Concrete distributions, however we have not observed an improvement of performance.
We therefore restrict ourselves to 10 samples for all experiments.
In all experiments, we train everything end-to-end with the ADAM optimizer.(Kingma & Ba, 2015)
We start with learning rate of 5e − 4 and use an exponential decay schedule with rate 0 .
97 everyiterations.
B. N ETWORK ARCHITECTURE
For most networks, we use MLPs implemented as residual nets (He et al., 2016) with ReLU activations.
Networks used for the reacher and maze experiments.
• q meas ( z t | · ) : MLP consisting of two residual blocks with 256 neurons each.
We only condition on the current observation x t although we could condition on the entire sequence.
This decision was taken based on empirical results.
• q trans ( z t | · ) :
In the case of Concrete random variables, we just combine the base matrices and apply the transition dynamics to z t − 1 .
For the Normal case, the combination of matrices is preceded by a linear combination with softmax activation.
(see equation 14) • q meas ( s t | · ) : is implemented by a backward LSTM with 256 hidden units.
We reuse the preprocessing of q meas ( z t | x t ) and take the last hidden layer of that network as the input to the LSTM.
• q trans ( s t | · ) : MLP consisting of one residual block with 256 neurons.
• q initial ( w | · ) : MLP consisting of two residual block with 256 neurons optionally followed by a backward LSTM.
We only condition on the ﬁrst 3 or 4 observations for our experiments.
• q initial ( s 2 ) :
The ﬁrst switching variable in the sequence has no predecessor.
We therefore require a replacement for q trans ( s t | · ) in the ﬁrst time step, which we achieve by independently parameterizing another MLP.
• p ( x t | z t ) : MLP consisting of two residual block with 256 neurons.
• p ( z t | · ) :
Shared parameters with q trans ( z t | · ) .
• p
( s t | · ) :
Shared parameters with q trans ( s t | · ) .
We use the same architecture for the image ball in a box experiment, however we increase number of neurons of q meas ( z t | · ) to
[Table]
B. H YPERPARAMETERS
[Table]
C O N S CALING I SSUES OF SWITCHING LINEAR DYNAMICAL SYSTEMS
Let’s consider a simple representation of a ball in a rectangular box where its state is represented by its position and velocity.
Given a small enough ∆ t , we can approximate the dynamics decently by just 3 systems: no interaction with the wall, interaction with a vertical or horizontal wall (ignoring the corner case of interacting with two walls at the same time).
Now consider the growth of required base systems if we increase the number of balls in the box (even if these balls cannot interact with each other).
We would require a system for all combinations of a single ball’s possible states: 3 2 .
This will grow exponentially with the number of balls in the environment.
One way to alleviate this problem that requires only a linear growth in base systems is to independently turn individual systems on and off and let the resulting system the sum of all activated systems.
A base system may then represent solely the transition for a single ball being in speciﬁc state, while the complete system is then a combination of N such systems where N is the number of balls.
Practically, this can be achieved by replacing the softmax by a sigmoid activation function or by replacing the categorical variable s of dimension M by M Bernoulli variables indicating whether a single system is active or not.
We do this for our multiple agents in a maze environment.
Theoretically, a preferred approach would be to disentangle multiple systems (like balls, joints) and apply transitions only to their respective states.
This, however, would require a proper and unsupervised separation of (mostly) independent components.
We defer this to future work.
[Figure]
[Figure] [Figure]
I MAGE B
[Figure]
IN A B
[Figure]