Title: SWITCHING LINEAR DYNAMICS FOR VARIATIONAL BAYES FILTERING

Abstract: System identification 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 find 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.

1. INTRODUCTION:

Learning dynamics from raw data (also known as system identification) 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 difficult 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.

2. BACKGROUND:

We consider discretized time-series data consisting of continuous observations x t ∈ X ⊂ R nx and control inputs u t ∈ U ⊂ R nu that we would like to model by corresponding latent states z t ∈ Z ⊂ R nz . We'll denote sequences of variables by x 1:T = (x 1 , x 2 , ..., x T ).

s 1 s 2 s 3 z 1 z 2 z 3 u 1 u 2 x 1 x 2 x 3 (a) SLDS graphical model. w s 2 s 3 z 1 z 2 z 3 u 1 u 2 x 1 x 2 x 3

(b) Our generative model. (Barber, 2012) .

EQUATION

Here A ∈ R nz×nz is the state matrix, B ∈ R nz×nu control matrix, the transition noise with covariance matrix Q and η the emission/sensor noise with covariance matrix R. Finally, the observation matrix H ∈ R nx×nz defines 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:

p(x 1:T , z 1:T , s 1:T | u 1:T ) = T t=1 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 figure 1a .

2.2. STOCHASTIC:

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 ∼ N (0, 1) and applying the deterministic and differentiable transformation z = µ + σ .

2.3. THE CONCRETE DISTRIBUTION:

One simple and efficient way to obtain samples d from a k-dimensional categorical distribution with class probabilities α is the Gumbel-Max trick:

EQUATION

However, since the derivative of the argmax is 0 everywhere except at the boundary of state changes, where it is undefined, 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

d k = exp((log α k + g k )/λ) n i=1 exp((log α i + g i )/λ)

, with g 1 , . . . , g k ∼ Gumbel(0, 1)

This softmax computation approaches the discrete argmax as temperature λ → 0, for λ → ∞ it approaches a uniform distribution.

3 RELATED WORK 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 first 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).

4. PROPOSED APPROACH:

Our goal is to fit 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 figure 1b an our inference model in figure 2a .

4.1. GENERATIVE MODEL:

Our generative model for a single x t is described by

p(x t ) = s ≤t z ≤t p(x t | z t )p(z t | z t-1 , s t , u t-1 )p(s t | s t-1 , z t-1 , u t-1 )p(z t-1 , s t-1 ) (7)

which is close to the one of the original SLDS model (see figure 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 influence 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 { A (i) , B (i) , Q (i) | ∀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.

p θ (z t | z t-1 , s t , u t-1 ) = N µ, σ 2 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.2 INFERENCE 4.2.1 STRUCTURED 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 σ 2 trans 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.

EQUATION

The densities of q meas and q trans are multiplied resulting in another Gaussian density:

µ q = µ trans σ 2 meas + µ meas σ 2 trans σ 2 meas + σ 2 trans , σ 2 q = σ 2 meas σ 2 trans σ 2 meas + σ 2 trans (11)

This update scheme is highlighted in figure 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 fit. 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

While we could condition on the entire sequence, we restrict it to just the first couple of observations. (a) Inference model.

qmeas(zt | xt) qmeas(st | x≥t) × qtrans(zt | •) q(zt | x≥t, zt-1, st, ut-1) × qtrans(st | •) pθ(xt | zt) zt-1 st-1 zt-1 q(st | x≥t, zt-1, st-1, ut-1) st xt x≥t ut-1 ut-1 . . . (b) High-level overview.

Figure 2: (a) Depicts the inference model. b t is the hidden state of the backward RNN of q φ (s t | x ≥t , u ≥t ). Initial inference of w may be conditioned on the entire sequence of observations, or just a subsequence. We've omitted the arrows for sake of clarity for the rest of the graph. (b) Shows schematically how we combine the transition with the inverse measurement model in the inference network. Transitions (in blue) are (partially) shared with the generative model.

4.2.2. INFERENCE OF SWITCHING VARIABLES:

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:

EQUATION

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 sufficient, information from the inverse measurement model can flow 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 coefficients for linear systems would then be determined by a linear combination of these latent variables:

α = softmax(W s t + b) ∈ R M (14)

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.3. TRAINING:

Our objective function is the commonly used evidence lower bound for our hierarchical model.

EQUATION

We choose to factorize over time, 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 first 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.

5. EXPERIMENTS:

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) . 

5.1. MULTIPLE BOUNCING BALLS IN A MAZE:

Our first 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 confirmation of that in figure 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 figure 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 confirm 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.46. It is difficult 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. 

5.2. REACHER:

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 collision2 between lower and upper joint which is one we'd like our model to capture. We again learn a linear classifier based on latent space s to see if this is successfully encoded and reach an F1 score of 0.46.

5.3. BALL IN A BOX ON IMAGE DATA:

Finally, we evaluate our method on high-dimensional image observations using the single bouncing ball environment used by Fraccaro et al. (2017) . 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 figure 5a we compare our model to both the smoothed and generative version of the KVAE. The smoothed version receives the final 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 figure 4 .

5.4. SUSCEPTIBILITY TO THE SCALE OF TEMPORAL DISCRETIZATION:

In this section, we'd like to explore how the choice of ∆t when discretizing a system influences 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 figure 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. 

6. DISCUSSION:

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 influence z t by determining the dynamics, while previously z t influenced both the choice of transition function as well as acted inside the transition. These two opposing roles lead to conflicting 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 fit for variables determining dynamics as such variables may change very abruptly.

7. CONCLUSION:

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 fitting simple dynamics and would lead to significant improvements of this approach.

A LOWER BOUND DERIVATION For brevity we omit conditioning on control inputs u 1:T . log p(x T ) = log

z 1:T s 1:T q φ (s 1:T , z 1:T | x 1:T ) p θ (x 1:T | z 1:T )p θ (z 1:T , s 1:T ) q φ (s 1:T , z 1:T | x 1:T ) ≥ z 1:T s 1:T q φ (s 1:T , z 1:T | x 1:T ) log p θ (x 1:T | z 1:T )p θ (z 1:T , s 1:T ) q φ (s 1:T , z 1:T | x 1:T ) = T t=1 E st [E zt [p(x t | z t , s t )]] -D KL (q(z 1:T , s 1:T | x 1:T ) || p(z 1:T , s 1:T ))

A.1 FACTORIZATION 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.

D KL (q(z 1 , s 2 , . . . , s T , z T ) || p(z 1 , s 2 , . . . , s T , z T )) (Factorization of the variational approximation) = z1 s2 • • • s T z T q(z 1 )q(s 2 | z 1 ) . . . q(s T | z T -1 , s T -1 )q(z T | z T -1 , s T ) log q(z 1 )q(s 2 | z 1 ) . . . q(s T | z T -1 , s T -1 )q(z T | z T -1 , s T ) p(z 1 , s 2 , . . . , s T , z T ) (Factorization of the prior) = z1 s2 • • • s T z T q(z 1 )q(s 2 | z 1 ) . . . q(s T | z T -1 , s T -1 )q(z T | z T -1 , s T ) log q(z 1 )q(s 2 | z 1 ) . . . q(s T | z T -1 , s T -1 )q(z T | z T -1 , s T ) p(z 1 )p(s 2 | z 1 ) . . . p(s T | z T -1 , s T -1 )p(z T | z T -1 , s T ) (Expanding the logarithm by the product rule) = z1 q(z 1 ) log q(z 1 ) p(z 1 ) + z1 s1 q(z 1 )q(s 1 | z 1 ) log q(s 1 | z 1 ) p(s 1 | z 1 ) + T t=2 z1 s2 • • • s T z T q(z 1 )q(s 2 | z 1 ) . . . q(z T | z T -1 , s T ) log q(z t | z t-1 , s t ) p(z t | z t-1 , s t ) + T t=3 z1 s2 • • • s T z T q(z 1 )q(s 2 | z 1 ) . . . q(z T | z T -1 , s T ) log q(s t | z t-1 , s t-1 ) p(s t | z t-1 , s t-1 ) (Ignoring constants) = D KL (q(z 1 ) || p(z 1 )) + E z1∼q(z1) [D KL (q(s 2 | z 1 ) || p(s 2 | z 1 ))] + T -1 t=2 E st,zt-1 [D KL (q(z t | z t-1 , s t ) || p(z t | z t-1 , s t ))] + T -1 t=3 E st-1,zt-1 [D KL (q(s t | z t-1 , s t-1 ) || p(s t | z t-1 , s t-1 ))]

• 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 first 3 or 4 observations for our experiments. • q initial (s 2 ): The first switching variable in the sequence has no predecessor. We therefore require a replacement for q trans (s t | •) in the first 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 1024. 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.

B.4 HYPERPARAMETERS:

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 specific 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.

:

As an ablation study, we will compare this to modeling switching variables by a Gaussian distribution.

We roughly identify a collision to be the point where the lower joint decelerates by over a fixed value of 2.

Figures and Tables:

Figure 3: Figures (b) and (c) depict an agent's position colored by the average value of a single latent variable s marginalized over all control inputs u and velocities. Figure (d) highlights a representative activation for a single transition system for the deterministic treatment of the transition dynamics. It doesn't generalize to the entire maze and stays fairly active in proximity to the wall.

Figure 4: First row: data, second row: filtered reconstructions, third row: predictions. The first 4 steps are used to find a stable starting state, predictions start with step 5.

Fraction of incorrectly predicted pixels.

Figure 5: (a) Our dynamics model is outperforming even the smoothed KVAE for longer trajectories. (b) Modeling switching variables as Concrete random variables scales less favorably.

Figure1: (a) s t denote discrete switch variables, z t are continuous latent variables, x t continuous observed variables, u t are (optional) continuous control inputs. (b) By introducing a special latent variable w used for initial state inference, we want to make explicit that the first step is treated differently from the rest of the sequence.

Table content: <table><tr><td>2.1 SWITCHING LINEAR DYNAMICAL SYSTEMS</td></tr></table>

Mean squared error (MSE) on predicting future observations. Static refers to constantly predicting the first observation of the sequence.

Table content: <table><tr><td/><td/><td>REACHER</td><td/><td/><td>3-BALL MAZE</td><td/></tr><tr><td>PREDICTION STEPS</td><td>1</td><td>5</td><td>10</td><td>1</td><td>5</td><td>10</td></tr><tr><td>STATIC</td><td colspan="6">5.80E-02 5.36E-01 1.25E+00 1.40E-02 5.74E-01 2.65E+00</td></tr><tr><td>LSTM</td><td colspan="5">3.07E-01 7.76E-01 1.22E+00 7.20E-02 1.58E-01</td><td>2.60E-01</td></tr><tr><td>DVBF</td><td colspan="2">1.10E-01 3.08E-01</td><td>6.07E-01</td><td colspan="2">6.20E-02 1.36E-01</td><td>1.82E-01</td></tr><tr><td>DVBF FUSION</td><td colspan="2">4.90E-03 2.97E-02</td><td>8.25E-02</td><td colspan="2">4.33E-03 2.03E-02</td><td>4.88E-02</td></tr><tr><td colspan="3">OURS (CONCRETE) 1.06E-02 5.73E-02</td><td>1.56E-01</td><td colspan="2">2.28E-03 1.22E-02</td><td>3.40E-02</td></tr><tr><td>OURS (NORMAL)</td><td colspan="6">3.39E-03 1.85E-02 4.97E-02 1.30E-03 5.52E-03 1.38E-02</td></tr></table>

Overview of hyperparameters.

Table content: <table><tr><td>Multi Agent Maze Reacher Image Ball in Box</td></tr></table>

Peer review (iclr2019):
 This paper proposes a deep probabilistic model for temporal data that leverages latent variables to switch between different learned linear dynamics. The probability distributions are parameterized by deep neural networks and learning is performed end-to-end with amortized variational inference using inference networks. There has been a lot of recent research trying to combine probabilistic models and deep learning to define powerful transition models that can be learned in an unsupervised way, to be used for model-based RL. This paper fits in this research area, and presents a nice combination of several interesting ideas presented in related works (switching variables, structured inference networks, merging updates as in the Kalman filter). The novelty of this paper in terms of original ideas is limited, the novel part lies in the clever combination of known approaches. The paper reads well, but I found the explanation and notation in section 4 quite confusing (although easy to improve). The authors propose a structured variational approximation, but the factorization assumptions are not clear from the notation (I had to rely on Figure 2a to fully understand them). - In the first line of equation 7 it seems that the variational approximation q_phi for z_t only depends on x_t, but it is actually dependent also on the future x through s_t and q_meas - The first line of section 4.1.1 shows that q_phi depends on x_{1:T}. However from figure 2a it seems that it only directly depends on x_{t:T}, and that the dependence on x_{1:t-1} is modelled through the dependence on z_{t-1}. - Is there a missing s_t in q_trans in the first line of (7)? - why do you keep the dependence on future outputs in q_meas if it is not used in the experiments and not shown in figure 2a? It only makes the notation more confusing. - You use f_phi to denote all the function in 4.1.1 (with different inputs). It would be clearer to use a different letter or for example add numbers (e.g. f^1_.phi) - Despite being often done in VAE papers, it feels strange to me to introduce the inference model (4.1) before the generative model (4.2), as the inference model defines an approximation to the true posterior which is derived from the generative model. One could in principle use other type of approximate inference techniques while keeping the generative model unchanged. It is difficult for me to understand how useful are in practice the switching variables. Reading the first part of the paper it seems that the authors will use discrete random variables, but they actually use for s_t continuous relaxiations of discrete variables (concrete distribution), or gaussian variables. As described in appendix B2 by the authors, training models with such continuous relaxations is often challenging in terms of hyper-parameter tuning. One may even wonder if it is worth the effort: could you have used instead a deterministic s_t parameterized for example as a bidirectional LSTM with softmax output? This may give equivalent results and remove a lot of complexity. Also, the fact that the gaussian switching variables perform better in the experiments is an indication that this may be the case. To be able to detect walls the z variables basically need to learn to represent the position of the agent and encoding the information on the position of the walls in the connection to s_t. Would you then need to train the model from scratch for any new environment? Minor comment: - in the softmax equation (6) there are missing brackets: lambda is at the denominator both for g and the log