On Uncertainty in Deep State Space Models for Model-Based Reinforcement Learning
Authors: Philipp Becker
Gerhard Neumann
Venue: Accepted by TMLR ()

Abstract: Improved state space models, such as Recurrent State Space Models (RSSMs), are a key factor behind recent advances in model-based reinforcement learning (RL). 
Yet, despite their empirical success, many of the underlying design choices are not well understood. 
We show that RSSMs use a suboptimal inference scheme and that models trained using this inference overestimate the aleatoric uncertainty of the ground truth system. 
We find this overestimation implicitly regularizes RSSMs and allows them to succeed in model-based RL. 
We postulate that this implicit regularization fulfills the same functionality as explicitly modeling epistemic uncertainty, which is crucial for many other model-based RL approaches.
Yet, overestimating aleatoric uncertainty can also impair performance in cases where accurately estimating it matters, e.g., when we have to deal with occlusions, missing observations, or fusing sensor modalities at different frequencies.
Moreover, the implicit regularization is a side-effect of the inference scheme and not the result of a rigorous, principled formulation, which renders analyzing or improving RSSMs difficult.
Thus, we propose an alternative approach building on well-understood components for modeling aleatoric and epistemic uncertainty, dubbed Variational Recurrent Kalman Network (VRKN). 
This approach uses Kalman updates for exact smoothing inference in a latent space and Monte Carlo Dropout to model epistemic uncertainty. 
Due to the Kalman updates, the VRKN can naturally handle missing observations or sensor fusion problems with varying numbers of observations per time step.
Our experiments show that using the VRKN instead of the RSSM improves performance in tasks where appropriately capturing aleatoric uncertainty is crucial while matching it in the deterministic standard benchmarks.

Published in Transactions on Machine Learning Research (10/2022)
On Uncertainty in Deep State Space Models for Model- Based Reinforcement Learning
Philipp Becker
philipp. becker Gkit. edu
Autonomous Learning Robots Lob, Korlsruhe Institute of Technology
Gerhard Neumann
gerhard. Gkit. edu
Autonomous Learning Robots Lab, Karlsruhe Institute of Technology
Reviewed on OpenReview: https: openreview. net/ mett forun? id- UQX&QyoRZh
Abstract
Improved state space models, such as Recurrent State Space Models (RSSMs), are a key fac- behind recent advances in model-based reinforcement learning (RL). Yet, despite their empirical success, many of the underlying design choices not well understood. We show That RSSMs use suboptimal inference scheme and that models trained using this inference overestimate the aleatoric uncertainty of the ground truth system. We find this overesti- mation implicitly regularizes RSSMs and allows them to succeed in model-based RL. We postulate that this implicit regularization fulfills the same functionality as explicitly model- ing epistemic uncertainty, which crucial for many other model-basod RI. approaches. Yet, overestimating alcatoric uncertainty can also impair performance in cases where accurately estimating it malters, e-g, when we have to deal with occlusions, missing observations, or fusing sensor modalities at different frequencies. Moreover, the implicit regularization is side-cffect of the inference scheme and not the result of a rigorous, principled formulation, which renders analyzing or improving RSSMs difficult. Thus, we propose an alternative approach building on well-understood components for modeling aleatoric and epistemic un- certainty, dubbed Variational Recurrent Kalman Network (VRKN). This approach uses Kalman updates for exact smoothing inference a latent space and Monte Carlo Dropout to model epistemic uncertainty. Due to the Kalman updates, the VRKN handle missing observations sensor fusion problems with varying numbers of observations per time step. Our experiments show that using the VRKN instead of the RSSM improves performance in tasks where appropriately capturing aleatoric uncertainty is crucial while matching in the deterministic standard benchmarks!
naturally
1 Introduction
Accurate system models are crucial for model-based control and reinforcement learning (RL) in au- tonomous systems applications under partial observability. Practitioners commonly use state space models (SSMs) (Murphy, 2012) to formalize such systems. SSMs consist of a dynamics model, describing how one state relates to the next, and an observation model, which describes how system states generate observations. Yet, dynamics and observation models are unknown for most relevant problems, and exact. inference in the resulting SSM is usually intractable. Researchers have proposed numerous approaches to learn the models from data and approximate the inference to solve those issues.
Recurrent State Space Models (RSSMs) (Hafner et al., 2019) are of particular interest here. Using RSSMs as the backbone for their Deep Planning Network (PlaNet), Hafner et al. (2019) showed that variational latent dynamics learning can succeed in image-based RL for complex control tasks. Combined with simple
Code available https://githan com/pbbeckr993/vvkn

---

Published in Transactions on Machine Learning Research (10/2022)
planning, RSSMs can match the performance of model-free RL while requiring significantly fewer environ- ment interactions. The authors later improved upon their original model, including a parametric policy trained on imagined trajectories (Dreamer) (Hafner et al., 2020). In general, approaches based on RSSMs have found considcrable interest in the model-based RL community. Yet, while RSSMs draw inspiration from classical SSMs, they use a simplified inference scheme. During inference, they assune the belief is independent of future observations instead of using the correct smoothing assumptions (Murphy, 2012) to obtain the belief. We formalize this observation in Section 2 and discuss how these simplified assumptions result in a theoretically looser variational lower bound. Further, on model learning and find they cause an overestimation of aleatoric uncertainty. Such aleatoric uncertainty stems from partial observability or inherent stochasticity of the system (Hillermeier & Waegeman, 2021). It plays an important role in many realistic lasks, e.g., in the form of occlusions, missing observations, or observations arriving in different modalities and frequencies. Consider, for example, low-frequency camera images providing external information and high-frequency proprioceptive mcasurements of the robot's inter- nal state. Given such sensory inputs, we require an appropriate estimation of aleatorie uncertainty to fuse all information optimally and form accurate belief states. Our experiments show that RSSMs performance sub-optimally in such cases, which we attribute to the wrong estimation of aleatoric uncertainty.
analyze the effects of these assumptions
Furthermore, we argue that the RSSM's 's inference assumptions are a double-edged sword for model-based RL. On the one hand, the overestimated aleatoric uncertainty can be beneficial as it. leads to dynamics models that generalize better and are more robust to objective mismatch (Luo et al., 2019; Lambert et al., 2020). Such objective mismatch arises because the model aims to maximize the data likelihood while the metric we care about is the agent's reward. Additionally, as the agent explores during training, it will encounter states and actions the model has not seen, causing a distribution shift between training and data collection. While many approaches rely on explicit epistemic uncertainty to lackle this issue (Chua et al., 2018; Janner et al., 2019), RSSMs succeed without capturing epistemic uncertainty. Such epistemic uncertainty (Hüllermeier & Wacgeman, 2021) stems from the lack of training data and plays an important role in model-based RL, where the data is not given but collected during the training process. On the other hand, this heuristic approach to address objective mismalch complicates the design and analysis of the approach as the robustness is a side-producct of the inference scheme and does not follow well-motivated principles. For example, purely stochastic versions of the RSSM gives unsatisfactory results, and it is unclear why this is the case. Arguably, explanation could be the overestimated aleatoric uncertainty causing the purely stochastic model to be unstable. As a remedy, Hafner et al. (2019) introduce deterministic path, ie.. they combine stochastic and deterministic features to form the belief state.
We show that removing this issue for RSSMs by implementing a naive approach to smoothing yields un- satisfactory results, even if the model uses explicit epistemic uncertainty to compensate for the missing overestimation of the aleatoric uncertainty. To make smoothing inference work, we redesign the model using well-understood and theoretically founded components to model aleatoric and epistemic uncertainty. We dub the resulting model Variational Recurrent Kalman Network (VRKN). It uses a linear Gaussian State Space Model embedded in a latent space, which allows closed-form smoothing inference and proper estimation of aleatoric uncertainty. Furthermore, as the inference builds on Bayes rule, it provides a natural treatment of missing observations and sensor fusion settings. Additionally, the VRKN uses Monte Carlo Dropout (Gal & Ghahramani, 2016) for a Bayesian treatment of its transition model's paramcters, explicitly modeling epistemic uncertainly.
We first show that the resulting architecture allows model-based agents to perform comparably to RSSM- based agents on standard benchmarks from the Deep Mind Control Suite (Tassa et al., 2018). Those enviromments are almost deterministic, i.e., the aleatoric uncertainty is low. Subsequently, we modify tasks from the Deep Mind Control Suite to create three different scenarios which exhibit high aleatoric uncertainty, i.e., (i) we introduce occlusions, (ii) missing observations, and (iii) sensor fusions problems of external camera images and internal proprioceptive feedback which are available at different frequencies. Using those, we show that the VRKN improves the agents' performance due to its accurate estimation of aleatoric uncertainty. Our approach builds well-motivated components for alcatoric and epistemic uncertainty. Additionally, it does not require heuristic model components such as deterministtcc path. As such, it may serve as basis for future research into improving state space models for model-based RL.
2

---

Published in Transactions on Machine Learning Research (10/2022)
To summarize our contributions:
We analyze the assumptions underlying the RSSM and show that they are not only suboptimal but also have subtle effects on model learning. They lead to an overestimation of the aleatoric uncertainty (Section 2.1 and Section 2.2).
2. We argue that, counterintuitively, this suboptimal inference is beneficial for the RSSMs performance as it addresses objective mismatch in heuristic way (Section 2.3). We evaluate several modifications to the RSSM to provide empirical evidence to support this hypothesis (Section 4.1).
We introduce the VRKN, providing a more principled inductive bias for smoothing inference than the RSSM (Section 3). We again show that a smoothing inference without additional measures results in suboptimal performance, yet, when combined with epistemic uncertainty, the VRKN's improved inductive bias allows it to close the performance gap on standard benchmarks (Section 4.1).
4. We show that VRKN-based agents improve performance in tasks where correct uncertainty estima- tion matters. Here we consider tasks with partial observability, missing information, or tasks that require sensor fusion. (Section 4.2)
2 Inference and Learning in State Space Models
State Space Models (SSMs) (Murphy, 2012) assume that a sequence of observations ogT = (or is generated by a sequence of latent state variables ZST = (zet=o--Tr given a sequence of actions aşt = (ar)t=o...r In SSMs, each observation Oz is assumed to only depend on the current latent state 2, via observation model p(o,/z,). Further, they assume the latent states are Markovian, ie., each latent state only deponds on its direct predecessor and the correspond- ing action via a dynamics model p(ze+1 /2,, at). Finally, the initial state is distributed according a distribu- tion p(zo). Figure la shows the corresponding graphical model. Typically, when inferring latent states from ob- servations, consider three different beliefs for each 2,- ie., the belief
a1
a2
a
Z1
Z
Z,2
Z3
01
02
01
02
03
(a) State Space Model (b) RSSM Inference Model
Figure 1: (a) State space model, serving as the generative model throughout this work. (b) Graphical Model underlying the RSSM (Hafner el al., 2019) inference scheme. In contrast to the generative SSM, the direction between observa- tions and latent states is inverted. These indepen- dence assumptions result in a simplified inference and subtle effects on model learning.
Those are the prior Sieelorri,,gr-- before observing or, the posterior, os,, age- i.c., the belief after observing Ot, as well as the smoothed belief (ogT, agr) which is conditioned on all future observations and actions, until the last time step T. We refer to those estimates guish them from dynamics distributions, which are con- ditioned on the previous state Z;-1 such as the prior dynamics p(zz/zt-1 a1-1), the posterior dynamics p(z, 21-1, Ot, 1) and the smoothed dynamics p(zz/ze-1, O21, a21-1). To get from a dynamics distribution to the corresponding state belief we need to marginalize out the previous state, e.g., (1 = I -p
beliefs to distin-
for the prior.
Traditionally, the independence assumptions of the generative model, shown in Figure 1a, are also used for inference. Yet, RSSMs (Hafner et al., 2019) work with different set of assumptions during inference, shown in Figure 1b?. In particular, in this graphical model the state 2, is conditionally independent of all observations Oy and actions a2: given 2,-1, a;-1, and o,, which is not the case for the standard SSM. We
2The full model of Hafner ct al. (2019) also includes dotermiinistic-patt which of concern regarding the discussion here. Thus, omit it for brevity. Furthermore, Hafner al. (2019) compare their RSSM bascline they abbreviate SSM (stochastic state model), which also builds on the simplified assumptions. this work, we refer all approaches based the simplified assumptions RSSM.
3

---

Published in Transactions on Machine Learning Research (10/2022)
discuss the effects of those assumptions on inference and model learning. We refer to Appendix A for more detailed derivations of the following identities.
2.1 Variational Inference for State Space Models
Inferring beliefs over the latent states given observations and actions is intractable for most models. Thus, we usually use approximate inference methods. this work, we focus on variational inference. Such variational methods introduce auxiliary distribution q over the latent variables given the observable variables. They decompose the data log- ligllikel into variational lower bound and Kullback-Leible divergence (KL) term, mcasuring how close the bound is to the data log- oogliiclihood. For State Space Models (SSMs) this decomposition is given as logp(o agr)
E B((esrlosrarr) log p(ogr, (gr) q OST
KKL (olzsrlogr,ast
ll p(zer los, agr)]
(1)
KL term(20)
Lower Bound
Variational inference methods try to find the optimal by maximizing the lower bound minimiz- ing the KL term. If can be found such that g(zgrlogr agr) = p(zerlogi, the KL term is O, and the bound is said to be tight. While the docomposition is valid for arbitrary distributions (ikeroorraaa agr), we need to pose a set of independence assumptions to obtaín tractable variational models. If we again use the independence assumptions of the generative model, shown in Figure la, we can obtain the inference model by explicitly inverting the generative direction using Bayes rule, q(zgr Ogr, agr) p(oo))
I 1 p(z)) II p(z//z.11, II (z.)) For certain model parametrizations, this variational distribution can be computed analytically, e.g., by Kalman smoothing the model is linear and Gaussian. In this casc, the KL term vanishes, and the bound is tight.
Yet, RSSMs, introduced in (Hafner et al., 2019); assume the variational distribution factorizes
T
g(zerlogr. agr) = ((z) Il a1-1))
This assumption results in a simplified inference procedure, as the belief over z, is assumed to be independent of all future observations O>t, given Z;-1, O₁, and a;-1-
Inserting this assumptions into the KT. term yields KT, (gzzgrlogr, p(ZerlosT,
T Eg(zr ploc--loog-ngg--- JKL (1 a.-1
(2)
For general distributions p, the individual terms in this sum will not be 0, as the right-hand side distributions receive information, i.e., future observations and actions, which are ignored by the variational distribution. Thus, this bound can in general not be tight, not even for linear Gaussian models as the resulting variational distribution can in general not represent p(zg:logn agr). Typically, tight variational bounds are preferable as they allow for faster optimization of the marginal log-likelihood.
The above discussion may be hypothetical, as all considered architectures do not provide tight lower bounds due to the use of deep neural networks, which prevents analytic solutions for inference. Still, as tight lower bound does not even theoretically exist for the RSSM assunptions, we believe this is already an indication of the misspecification of its inference distribution.
2.2 Model Learning under Different Inference Assumptions
In the model-basod RL setting considered in this work, the generative model is usually assumed be unknown and we jointly learn a parametric generative model pe and an inference model gΦ using an auto- encoding variational Bayes approach (Kingma & Welling, 2013; Sohn et. al., 2015) by maximizing the lower
4

---

Published in Transactions on Machine Learning Research (10/2022)
10a
Log State Probability
Model Error (Frob.)
10"
Cum (FF)
Transition Variance
Ground Truth a
Ceom (NN)
Lisam
10
10
-1
GT Smoothed
GT Posterior
10
10
100
200
0
100
200
0
100
200
Train Epochs
Train Epochs
Train Epochs
(a) Inference Tracking Error
(b) Model Learning Error
(c) Transition Variance
Figure 2: Comparison of inference and model learning results on a simple linear- Gaussian state space model without actions, under both RSSM and SSM inference assumptions. We report average results over 10 seeds and error bars indicating 95% bootstrapped confidence intervals. Note that the error bars are too small to be visible. We consider both a closed-form (CF) version based on Kalman Smoothing and a version with a neural network (NN) as an inference model for the SSM-based approaches and find that only the objective matters, not the parametrization the inference model. (a) Log- probbabbbyly of the ground truth states under the learned models. We compare against. the quality of the ground truth smoothed (GT Smoothed) and posterior (GT Posterior) beliefs, computed using a Kalman Smoother and ground truth generative model. While the SSM objective reaches the quality of the smoothed belief, the RSSM-based inference fails to attain the quality of even the posterior belief. (b): Distance between ground truth transition matrix and learned transition matrix, measured using the Frobenius norm. Here, the SSM inference yiclds model that is an order magnilude closer to the ground-truth model than that learned by the RSSM-bound. (c): Transition variance øI. With the SSM bound we recover the ground-truth aleatoric uncertainty, yet with the RSSM bound the alcatoric uncertainty significantly overestimated.
bound part of Equation 1. Inserting the RSSM-asee aasM--ssuupptonss into the general lower bound gives the objective introduced by Hafner et al. (2019), Crssm (OsT, açr)
T (xglarog po [z]] - Eg (ew
[KL (4 (Ze ar-1,
(3)
If one instcad uses the same factorization assumptions as the generative model, the inference distribution lactorizes as gu (25rloos (gr)) = Ye (Zo (2olo agr) II (2,2//22110 a2(-1, Inserting this factorization into the general lower bound leads to Lesm (OST, agr) [ =
T
Σ E Beslerlorraagr po (o: [2:)] - E 10 [KL (2//2e , agr-1,
|| po (o0(z:/z111 ar]] (4)
Opposed to the RSSM-objective, this objective uses the smoothed belief-state g (z agr) and the smoothed dynamics Qup (Zr 1122,/z1111 2 (21-1, The variational distribution considers all future information until t t= T and thus can theoretically yield a tight bound.
These differences have interesting effects on the learned transition model belief over past states can, by definition, not change due to additional observations. Thus, any discrepancy between this past belief and the current observation must be explained by the transition model. In contrast, smoothing inference can also explain the discrepancy by propagating information from observations to past beliefs. In Equation 3, this observation is reflected in the expected KL-term, where the expectation does not consider o, or other future observations. Thus, the transition model has to explain the transitions from a given 2 to 2,, even 2 would be rendered implausible by a future observation.
For the RSSM, the
For a thought experiment illustrating these effects, consider the following scenario. You meet person holding box, and they tell you there hamster inside. As your prior experience that people are
5

---

Published in Transactions on Machine Learning Research (10/2022)
usually trustworthy, you chose to believe them. Next, the box opens, and a cat jumps out. As you trust your eyes, you now believe it is a cat. Yet, under the RSSM-assumpiion ions, you cannot revise your belief of the first time step and thus still believe it originally was a hamster. When updating your model based on this interaction, you would learn that hamsters can turn into cats, as you cannot capture the arguably more likely explanation that the person lied. Learning under these assumptions requires you to model unlikely events as more likely than they are. Thus, you overestimate the alcatoric uncertainty in the world.
More formally, we can demonstrate this effect using simple linear- Shear-Gaussian State Space Model without actions. We will use a state dimension of and the ground-truti truth generative model given by
p(zo) = -( p(o, = (N,,, 0.025T), p( 1 (z:) = (2 Azz, 0.01T)
where I denotes the identity matrix. The transition matrix A induces a slightly damped, oscillating behavior. The complete matrix A, together with further details regarding the exact setup of this experiment, can be found in Appendix B. Using this generative model, we generate 1, 000 sequences of length 50. Even in this simple setting, computing the optimal inference distribution for the RSSM-bound (Equation 3) is impossible, and we thus resort to numerical methods. We parameterize Fo (Zr|z,-1, o,) as locally linear- Gaussian distributions and learn the paramcters using a ncural network. For the SSM-bound (Equation 4 we can either compute the optimal inference in closed form, or again parameterize qu ( Iz--1, 02)) as a neural network. To condition on future observations, we use a GRU (Cho et al., 2014) which runs backward over the observation sequence. For the generative model, we learn transition matrix A and an isotropic covariance øI jointly with the parameters of the inference model, i.e., 0 = {A, We fix the remaining parts of the generative model to the ground-truth values. Figure 2 summarizes the results and demonstrates that the RSSM-bound leads to a suboptimal inference and consequently to learning a wrong model. In particular, we can see that for the RSSM, the transition variance ā is much larger than the ground-truth value a= = 0.01 as all unexpected observations have to be explained by the transitions instead of correcting the beliefs of past time steps.
2.3 The Interplay of Policy Optimization and Regularization
Despite the theoretical considerations in the previous section, RSSMs work well for model-based RL. It is well known that model-based RL suffers from an objective mismatch (Luo al., 2019; Lambert et al., 2020) issue. This issue arises because the model aims to maximize the ELBO (Equation 1) but is evaluated based the agent's reward. The effect is further amplified by the distribution shift between training and data collection, as data collection is typically performed only after a policy improvement step. In RL, we explicitly want the agent to explore unseen parts of the state-action space, encountering observations the model has not seen before. Thus, training the underlying model requires careful regularization such that wrong predictions do not prevent the agent from exploring relevant parts of the state space. Many model- based RL approaches (Chua el al., 2018; Janner el al., 2019) handle this issue by explicitly modeling the epistemic uncertainty of the model, which is not required by the RSSM Instead, we argue that RSSMS rely the overestimated alcatoric uncertainty caused by suboptimal inference to address objective mismatch in a more heuristic manner. The overestimation implicitly regularizes the RSSM as it forces the transition model to model unlikely events with higher probability. This regularization thus implicitly prevents overconfident model predictions due to overfitting. Yet, while it alleviates the objective mismatch issue, there are also drawbacks to this heuristic solution. First, it complicates the model design, analysis, and improvement of RSSMs. As already observed by Hafner et al. (2019), a fully stochastic model based on the RSSM- assumptions underperforms without additional measures as fails at reliably propagating information for multiple time steps. As a remedy, Hafner et al. (2019) introduce delerminnistic-pahh, i.e., a Gated Recurrent Unit (Cho et al., 2014), and base the belief update on this instead of the stochastic belief. Second, we show in Section 4.2, there are settings where appropriately capturing the alcatoric uncertainty is important, and failing to do so can hurt performance.
In a first attempt to address those issues, we minimally adapt the RSSM to be capable of smoothing. This Smoothing RSSM uses a GRU which added before the actual RSSM and runs backwards over the representations extracted by the encoder, effectively accumulating all future information. The RSSM then receives the output of this GRU instead of the original observation input. This model is an extension of the
6

---

Published in Transactions on Machine Learning Research (10/2022)
a1
a2
A D,
=0(1 a.))
Linear + ReLU
Linear + ReLU
Linear Output
A,
%
72
Gated Unit
b
W1
W2
W3
dyn
02
(a) VRKN Overview
(b) Overview of the VRKN's Dynamics Network
Figure 3: (a) We use an encoder (blue) to extracts an intermediate representation w. and the observation noise o from the original observation o,. Both, Wt and o,", then used to update the state estimate 2, using the closed-form Kalman belief update (Bayes rule for Gaussians) (orange). A decoder (red) reconstructs the observation given a state sample. We propagate the latent state estimate to the next time step using the transition model (0 a:-1) (green). The parameters of this distribution, At, b and of dyn are computed using the network Φ(u ar), shown in (b). The network receives the current posterior mean ju," and the action a as inputs. Thus, the resulting dynamics are lincarizod around the posterior mean, which allows closed-form belief propagation, similar an extended Kalman Filter. For stability during training, the network includes gated unit, ie.,, GRU cell, which receives the posterior mean μt," at the memory input. For Bayesian treatment of the transition model's parameters, we include Monte Carlo Dropout layers at the positions indicated by the purple dashed lines.
neural network based inference model for Casm, introduced in the previous section ( Csam (CF) in Figure 1). We compare this approach to the original RSSM and present the results Scction 4.1, in particular Figure Those results show that the performance decreases compared the original RSSM. We argue that with proper inference, the overestimated aleatoric uncertainty longer regularizes the model, making it more prone to objective mismatch. Following other methods, we try to improve the results by modeling epistemic uncertainty. To this end, we use Monte Carlo Dropout (MCD) but find that it does not help to improve the Smoothing RSSM's performance. From these results, again presented in Figure 4, we conclude that solely addressing the suboptimal inference assumption is insufficient, but we also need to rethink the model's parameterization. We postulate that the additional GRU for the backward pass a poor inductive bias and that we require smoothing approach that adds as little complexity possible to the model. In the next section, will introduce architecture that allows for parameter-free smoothing.
3 Variational Recurrent Kalman Networks
To provide a thcoretically better-grounded alternative to the RSSM, require model which allows tractable inference while still scaling to complex image-based control tasks. Further, the architecture should allow efficient computation of smoothed and posterior state beliefs and dynamics. While we need smoothed beliefs for training, we require (filtered) posteriors for online control. In particular, as only the seoothed beliefs are explicitly used for training, the model needs to provide a strong inductive bias that enables it to still produce reasonable posterior estimmates.
To meet these criteria, we introduce a new parametrization of the latent dynamics based on linear- Gaussian state space model (LGSSM) (Murphy, 2012) embedded in latent space. The Gaussinnn assumptions allow rigorous treatment of uncertainties while working in a learned latent space allows for modeling high-dimensional and non-lincar systems. Inference in this LGSSM amounts to (extended) Kalman filtering and smoothing, enabling smoothing inference without introducing additional parameters besides the latent observation and dynamics model. Due to this property, the architecture can perform proper smoothing inference for training and also produce reasonable posterior belicfs for online control. As an additional benefit, the formulation can naturally handle missing observations and the fusion of multiple sensors, making it amenable to many realistic problems. To include epistemic uncertainty in our approach, we Monte
7

---

Published in Transactions on Machine Learning Research (10/2022)
Carlo Dropout (Gal & Ghahramani, 2016) for Bayesian treatment of the LGSSM's transition model. We name the resulting approach Varialional Recurrent. Kalman Network (VRKN). Figure 3 shows schematic overview.
In the following, we first introduce the VRKN's dynamics model, which is shared between the inference and generative parts of the model. Next, we introduce the parameterization of the inference model dhe and the generative model pg. We describe how to train the model and use it for control and conclude by elaborating the natural fusion mechanism.
3.1 The VRKN's Dynamics Model
Both, the VRKN's inference model dh and the generative model Po; use the same dynamics model. Such parameter sharing is common for variational time-series models (Karl et al., 2016; Fraccaro et al., 2017; Hafner al., 2019; Klushyn et al., 2021), it simplifies the architecture, reduces the number of parameters, and simplifies training.
We model the latent dynamics
po (0(2++1/20 [ee a,)= =N -N(2+11As(m1 az)z + bo ar), ot (M. ai)).
where pe denotes the mean of the posterior state estimate qu (w(z) provided by the inference model. Thus, we learn a lincarized dynamics around the current posterior mean. This parametrization builds on the work of Shaj et al. (2020), who propose using a model that is locally linear the state while depending non-lincarly on the action. Building on the assumption that there no uncertainty the actions, is maximally flexible while still allowing the usage of extended Kalman filtering (Jazwinski, 1970). Here, the linearization around the current posterior mean pu," enables closed-form propagation of state beliefs.
We model A, to be diagonal matrix which is emitted together with the offsct term b, and the transition noise O dyn by single neural network (A,, by, of v = o(p at). We carefully design this network to prevent the state estimates and gradients from growing indefinitely during training. First, as A, is diagonal, its values are its eigenvalues, and constraining them in appropriate range ensures stable dynamics. To this end, we use an activation of the form f(x) = s sigmoid (2 b) + m where we choose s, b, and m such that it saturates at 0.1 and 0.99 while ((0) = 0.9. Here, the intuition is that we want plausible and stable dynamics, which we initialize as slightly dampened system. Second, we employ a gating mechanism to mitigate problems with vanishing and exploding gradients. To this end, we use a standard GRU cell implementation but feed the posterior mean fe into the memory input, i.e., of a.) (eRGUU(01 (011(( at), (e),),)).. The resulting model is still fully stochastic and linear in 2,- In contrast to the RSSM, it does not use determinstic- as the GRU cell does not introduce additional deterministic memory and is used solely for addressing problems arising from unstable dynamics and gradients. We empirically found this helpful and took inspiration from standard recurrent architectures (Hochreiter & Schmidhuber, 1997; Cho et al., 2014), where gating mechanisms are a well-known approach to mitigate instabilitics. Figure 3b provides an overview of the dynamics model architecture.
Modeling Epistemic Uncortainty. As discussed in Soction 2.3, the overestimated alcatoric uncortainty of the RSSM's transition model avoids overconfident estimates due to overfiuting and compensates for the lack of explicit epistemic uncertainty. Thus, our approach captures the aleatoric uncertainty correctly, we need explicitly consider epistemic uncertainty obtain a model that is also useable for policy optimization. We use Monte Carlo Dropout (Gal & Ghahramani, 2016) due to its simplicity and include corresponding layers at appropriate points in the transition model. Those point are shown in Figure 3b.
3.2 Inference Model
We aim for model that can work with high-dimensional observations which depend non-linearly on the latent state while still allowing efficient inference of filtered and smoothed beliefs. Thus, we introduce an auxiliary, intermediate representation w, for the original observations o,, Such an intermediate representation allows
8

---

Published in Transactions on Machine Learning Research (10/2022)
to capture the highly complex relations between state and observations in the mapping from o, to w, while using a simple observation model 2e (w.lzz) in the latent space. This approach is common for variational latent state space models (Karl et al., 2016; Fraccaro et al., 2017; Klushyn et al., 2021) as it enables efficient inference in latent space. These approaches model We as a latent Gaussian random variable and learn it's mean and variance based on Og- They then sample we given these parameters. Yet, this assumption complicates inference and training, as now an additional latent variable, Wt, has to be addressed. Further, it makes uncertainly propagation from observations to states harder (Becker et al., 2019; Volpp et al., 2021), as explicit uncertainty information is lost during sampling. Thus, following (Haarnoja et al., 2016; Becker et al., 2019; Shaj et al., 2020; Volpp et al., 2021), we instead formulate the observation model as ge (w,/2)) = W (welze, ))), where w, = (o,) and the diagonal covariance o" enc,(o;) are given by an encoder network laking the original observation. In practice, the encoder consists of a single neural network with two output heads. This formulation does not suffer from the previously mentioned issues. First, as it models w, as a direct mapping of the obscrvation and not a random variable, we do not nced to infer w, but can assume is observable. This assumption results in a simplified inference and training objective as we do not have to account for unobserved latent variables other than the latent state itself. Second, the encoder can directly extract the uncertainty in the observation, which tells the inference procedure how much it can rely on W.. For example, when estimating w, from images, some images might contain certain information, e.g. the positions of an object, while others do not. The former case would result in a low uncertainty and the latter in a high one. Given We and ox we can update belief states using the Kalman update, i.e., Bayes rule for Gaussian distributions,
1 go (2 ag( = z 90 (cnco(o:) ze 9w q (g--))
Here, of contributes to the computation of the Kalman gain, which highlights how the model uses the encoder's uncertainty to trade-off information from the prior belief and observation
The inference observation model is combined with the shared locally linear dynamics model d (Z++1 (e, at) = po (Z(+1 [2, ar) to obtain a latent LGSSM for inference. Given this dynamics model, we can forward beliefs in time using closed-form Gaussian marginalization,
gep (x (0(2)+100gr ast) = / (2 (2 a
Repeating the described forms of forwarding beliefs in time and updating them based on observations amounts to standard (extended) Kalman filtering (Kalman, 1960; Jazwinski, 1970). Similarly, given these models, we can compute smoothed belief using the Mauch-Tung-Strine In particular, the smoothing only requires beliefs and transition models computed during filtering and does not introduce additional learnable parameters. Note that both the filtering and smoothing equations simplify under our assumptions of diagonal covariances and are thus efficient and scalable. Thus, the smoothing in- ference only slightly increases the computational load during training. Especially in the image-based settings considered here, the computational cost, largely stems from encoding and reconst reconstructing observations.
(Rauch et al., 1965) equations.
3.3 Generative Model
Recall that we assume generative State Space Model consisting of three parts. Those are the dynam- ics model, the observation model, and the initial state distribution. First, we have the dynamics model po pe(241 |Z,, (z) ax) described in Section 3.1. Second, assume Gaussian generative observation model po (2)) parameterized by neural network (decoder). Following Hafner et al. (2019), we parameterize the mean by a neural network and assume a fixed variance. This parameterization is often empirically beneficial, especially in the image-based setting considered in our experiments, but not a theoretical constraint. We could also parameterize the variance by the network if required. Third, we have an initial state distribution Po (zo). Here we use Gaussian with zero mean and learned diagonal variance which we initialize with the identity matrix.
9

---

Published in Transactions on Machine Learning Research (10/2022)
3.4 Stochastic Gradient Variational Bayes for Model Learning
We Irain our model using a version of stochastic gradient variational Bayes (Kingma & Welling, 2013) and maximize Sam. (Equation 4). Like most approaches building on (Kingma & Welling, 2013) we approximate all expectations in Equation 4 using Monte Carlo estimation with a single sample from the latent variable and jointly optimize the parameters of qe and pe. For training, we infer the required belief states using the Kalman-basee smoothing procedure described in Section 3.2. To compute the KT, term in Equation we agz,
additionally need the smoothed dynamics Fy (Zc+1 &2,, 020+1) which we obtain with minimal overhead by extending the RTS equations detailed in subsection A.2.
The VRKN's formulation tightly couples posterior and smoothed beliefs by fixed (not learned), deter- ministic, and well-motivated procedure. Smoothing using the Rauch-TTug-Strieeee equations ensures that reasonable smoothed beliefs can only stem from reasonable posterior beliefs. Thus, while the posterior belief is not explicitly part of the objective, it is still learned during training as a prerequisite of the smoothed beliefs.
3.5 Using the Model for Online Control and Reinforcement Learning
When using the model for online control, we cannot smooth but have to act based on samples from the posterior belief (z( agg, Ogt). We can rely on Kalman filtering, as described in Section 3.2, to obtain this posterior belief and omit the smoothing step, future observations are unavailable. For the VRKN, the tight coupling between smoothed and posterior beliefs ensures the posteriors are reasonable and usable for control.
To act based on the latent. state beliefs, we add a decoder network to predict rewards from latent states. Following (Hafner et al., 2019), this decoder is trained by adding a reconstruction loss term to the objective. Given the predicted reward, various forms of control are applicable to act optimally. Yel, in this work, we focus on the underlying state space model and thus reuse two previously introduced methods building on the RSSM (Hafner et al., 2019; 2020). The PlaNet (Hafner et al., 2019) approach plans actions using the cross entropy method by rolling out trajectories the model. The Dreamer (Hafner et al., 2020) approach learns a parametric policy and value function based on latent imagination. Such latent imagination uses the model as differentiable simulator and optimizes the policy based on the estimated values of predicted future states.
3.6 Sensor Fusion
Given the possibility of using the Kalman update for incorporating observations, we can use the VRKN for a simple but principled approach to sensor fusion. Formally, we assume the observation o, factorizes Po(o, (k) lz.).
into K different observations 0!*h), i.c., po(0) P (z.). Those obscrvations can be of various modalities and be available at different. frequencies, e.g., high-frequency velocity information from an inertial measurement unit and low-freen wow-freuency camera images of the surroundings. In this scenario, we have K encoders and K decoders, one for each of Similar to more traditional sensor fusion approaches (Gustafsson, 2010), we then accumulate the intermodiate obscrvation representation w(()) by repeatedly applying the Kalman update. This approach reflects the invariance to permutations of all observations for a single time step.
Furthermore, it enables the model to omit the update if some of the K observations are unavailable for time step.
4 Evaluation
We compare the VRKN, the original RSSM,add the modified RSSM version introduced in Scction 2.3 on image-based continuous control tasks using the DeepMind Control Suite (Tassa et al., 2020). Prior works (Lambert et al., 2020; Lutter et al., 2021) concluded that the model's predictive performance is of- ten uninformative about the quality of the model-based agent. We concluded the same after preliminary experiments and want to study the effects of the different assumptions and parametrizations on the perfor- mance in model-based RL setting. Thus, we evaluate the state space models as backbones for model-based
10

---

Published in Transactions on Machine Learning Research (10/2022)
Aggregated PlaNet Agents
Aggregated Dreamer Agents
1,000
Expected Return
800
600
400
RSSM MCD-RSSM
Smooth RSSM
200
Smooth MCD-RSSM
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Eavironment Steps (x109)
Environment Steps (x109)
Figure 4: Comparison of the RSSM, the simple smoothing extension (Smooth RSSM) well as versions of both models using Monte Carlo Dropout (MCD) to capture epistemic uncertainty in the dynamics (MCD- RSSM and Smooth MCD-RSSM). Note that we use different tasks for PlaNet and Dreamer-based agents. Thus the left and right plots are not directly comparable. For both types of agents, we find that proper inference by smoothing deteriorates performance. The additional epistemic uncertainty does not compensate for this decrease in performance.
agents and directly consider the achieved reward. As mentioned, use both the PlaNet and the Dreamer approaches for control and thus closely follow the experiment setup described by Hafner et al. (2019; 2020). Appendix C gives further details about the experimental setup and used basclines.
For our experiments, we run a varying number of tasks from the DeepMind Control Suite (Tassa et al., 2020) and report aggregated results over these tasks. Such aggregation is possible as the rewards in the DeepMind Control Suite are normalized, and all sequences are of equal length (1,000 steps). Thus the returns are normalized, and aggregation yields better performance estimates (Agarwal al., 2021). Again following the suggestions of Agarwal et al. (2021), report interquartile means while indicating 95% stratified bootstrap confidence intervals by shaded areas. These metrics are computed using the provided library". We base our conclusions on those aggregated results. Unless noted otherwise, we use 10 seeds for each agent-environment pair, train for million environment steps, and approximate the expected return using rollouts. Appendix D provides reward curves for the individual tasks and additional quantitative results, such box plots of their final performance.
4.1 Evaluation of the Effect of Epistemic Uncertainty on Different Smoothing Architectures
We start our evaluation by comparing the original RSSM with its extended smoothing version using a GRU with and without. Monte Carlo Dropout (MCD), described Section 2.3. For completeness, we also include a version of the SigginllRSSSM. with MCD model epistemic uncertainty. For the PlaNet- agents, we evaluate the 6 lasks originally used in (Hafner et al., 2019). For the Dreamer-agents tasks. Those Cheetah Run, Walker Walk, Cartpole Swingup, Cup Catch, Reacher Easy, Hopper Hop, Pendulum Swingup, and Walker Run. As indicated by the results in Figure 4, the proper smoothing inference significantly deteriorates performance. Adding epistemic uncertainty in the form of MCD does, on average, neither affect the performance of the original RSSM nor the smoothing RSSM. These results underpin the argument that naive smoothing with the RSSM gives suboptimal performance. Additionally, they show that epistemic uncertainty in the form of Monte Carlo Dropout cannot serve as a remedy for the RSSM.
use 8
Next, we compare the VRKN to the RSSM using the same environments and also include a version of the VRKN without Monte Carlo Dropout (MCD), i.e., without epistemic uncertainty, dubbed VRKN (no MCD). Figure 5 shows the aggregated results for both PlaNet and Dreamer agents. Those results demonstrate that matching the RSSM's performance is possible with a principled smoothing inference by explicitly modeling epistemic uncertainty and providing a more appropriate inductive bias. Additionally, we find the VRKN tends to reach a given performance with fewer environment interactions, especially for the Dreamer-agents.
8rliable: wtt.ghtwww.oooooogoeeerrrrr/aaaa
11

---

Published in Transactions on Machine Learning Research (10/2022)
Aggregated PlaNet Agents
Aggregated Dreamer Agents
1,000
Expected Return
800
600
VRKN
400
VRKN MCD)
200
RSSM
0.2 0.4 0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Eavironment Steps (x109)
Environment Steps (x109)
Figure 5: Comparison of VRKN and RSSM based PlaNet and Dreamer agents standard DeepMind Control Suite tasks. Note that we use different tasks for PlaNet and Dreamer-based agents. Thus the left and right plots are not directly comparable. The VRKN without epistemic uncertainty ( (VRKN (no MCD)) cannot compete with the RSSM-agents, especially using the more involved Dreamer-based agents. Yet, the VRKN profits from the epistemic uncertainty, and closes the performance gap to the RSSM and even shows tendency be more sample efficient.
Together, the results from Figure 4 and Figure 5 emphasize the importance of regularization for model-based RI. This regularization can be either implicitly by suboptimal inference or explicitly by capturing epistemic uncertainty. Additionally, they indicate that epistemic uncertainty alone is insufficient for approaches using a correct smoothing inference. Yel, the VRKN's appropriate inductive bias makes the more principled approach of combining proper inference with explicit epistemic uncertainty feasible.
4.2 Evaluation on Tasks where Aleatoric Uncertainty Matters
The standard versions of the Deep Mind Control Suite Tasks have a deterministic simulation and rendering process. Thus, their alcatoric uncertainty is low. To better analyze the approaches capabilitics to capture and handle aleatoric uncertainty, we design tasks where doing so necessary. To this end, we modify the Cheetah Run, Walker Walk, Cartpole Swingup, and Cup Catch tasks. First, we introduce transition noise by adding Gaussian noise to the actions before execution. We added noise with a standard deviation of 0.2 for Cheetah Run and Walker Walk and 0.3 for Cartpole Swingup and Cup Catch. Note that valid actions are between -1 and for all tasks. Second, we modify the observations to contain only partial information, are missing for several time steps, or are available in different modalities at different frequencies. For details, we refer to the individual experiments below.
We use these tasks to study the effects of appropriately capturing aleatoric uncertainty tasks where this form of uncertainty matters. We compare VRKNN-aase based and the RSSM-based Dreamer agents and find that precise estimation of the aleatoric uncertainly makes difference. Contrary to the original tasks, where RSSM and -RKNdbased agents performed similarly, agents building on the VRKN outperform their RSSM-based counterparts in all considered scenarios.
4.2.1 Partial Observability through Occlusions
In a first approach to include observation uncertainty, we render two types of occlusions over the images, i.e., discs and walls. See Figure 6 for an explanation and some examples. Due to these occlusions, the individual observations have varying amounts of relevant information. Thus, the models need to correctly capture uncertainties in the system, allowing them to trade off information from the prior belief and current observation. For training, we use masked reconstruction, i.e., only non-occluded pixels contribute to the roconstruction loss4. We also consider a bascline where we train the model solely to reconstruct the reward to show that the approaches can still extract information from the highly occluded observations. Figure
We want emphasize that we do consider the availability of such loss masks a realistic scenario but see the task reasonable benchmark evaluate the models' capabilities cope with uncertainties.
12

---

Published in Transactions on Machine Learning Research (10/2022)
= 0 f ==
t= 10
t= 15
t= = 20
t= 0
t= 5
t= 10
f= 15
t= 20
(a) Disk Occlusions
(b) Wall Occlusions
We introduce two types occlusions to induce partial observability. (a) Disk Occlusion: Slow- moving disks float through the image and bounce off its walls. (b) Occlusions: Walls slide over the image from right to left at a constant speed. Their width is sampled randomly between half the image width and the image width. In both cases, the models have to accurately estimate which parts of the system are visible and which are not. Yet, the basl-based occlusions more correlated over time and require memorization and prediction over longer time periods.
Figure 6:
Disk Occlusions
Wall Occlusions
1,000
Expected Return
VRKN
RSSM
800
Reward Only
600
400
200
0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Environment Steps (x109)
Environment Steps (x109)
Figure 7: Aggregated results for the partial observability experiments with both occlusion types. For the experiments, except for the reward-only baseline, we used 20 seeds per task. Both approaches clearly perform better than the reward-only baseline, showing that they can extract useful information from the occluded images. Yet, especially in the wall-occlusion task, the VRKN-based agents outperform the RSSM-baset agents. They do not only achieve higher reward but also converge faster. Especially in the wall occlusion task, it is insufficient to only extract information from the occluded images, but reasonable belief also needs to be sustained over time.
shows the results of the comparison. Under the heavy occlusions considered here, the VRKN performs significantly better than the RSSM. In particular for the wall occlusions, the VRKN-based agents achieve a higher reward using fewer environment interactions. The wall occlusions are more correlated, i.e., occluded parts in one image are more likely also occluded the previous and next images. Thus, they test the approaches' capability to not only form reasonable belief but also to propagate it for multiple time steps.
Additionally, qualitatively compare images reconstructed from posterior beliefs of both approaches to gain further insights into the quality of the belicf state. Figure 8 shows some of those reconstructions for the Cup Catch task, further images can be found in Appendix D.6. From these images, appears that the VRKN better captures the actual system state and uncertainty and thus allows the model-based agent to achieve a higher reward. Yet, while the reconstructions are visually much better, we only observe a small increase in performance, which we attribute to objective missmatch.
4.2.2 Dealing with Missing Observations
Another common setting where estimating alcatoric uncertainty is important is missing observations. It requires the models to propagate a reasonable belief without information and update it if an observation is available. Additionally, if no observations are available for multiple time steps, the belief state gradually drifts away from the real state due noise and model inaccuracics. Once a new obscrvation arrives, the discrepancy between this observation and the belief has to be explained. Recall that, due to its inference assumptions, the RSSM must explain the discrepancy by the last transition and thus will overestimate the
13

---

Published in Transactions on Machine Learning Research (10/2022)
Cound Thatk
Figure 8: Exemplary sub-sequence of reconstructions, based on the model's posterior beliefs. The first row the noise-free ground truth image, which the models never The second row is the model input, followed by the VRKN and the RSSM reconstructions. Even though the ball is partly visible in most images, the RSSM fails to reconstruct its position. The VRKN manages to do so and even provides a reasonable es- limate for cup position. These results indicate the VRKN's improved ability capture the system state in noisy scenarios.
Inpan Tosert
Without Transition Noise
With Transition Noise
1,000
Expected Return
VRKN
VRKN (Cat)
800
RSSM
600
400
200
0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Environment Steps (x109)
Environment Steps (x109)
Figure 9: Aggregated results in the missing observations setting with and without transition noise. First, we see that the VRKN works equally well if we provide the information about which information are valid explicitly (VRKN) or we provide it by concatenation (VRKN (Cal)). Second, the RSSM performs significantly worse then both both VRKN version. Especially if the system subject to transition noise it fails to give good results.
aleatoric uncertainty even further. On the other hand, the VRKN with its smoothing inference can capture the drift of the belief state and evenly attribute the discrepancy to all transitions since the last observation. Thus the VRKN should be better equipped to solve such task.
We consider a setting where we only provide every n-th image, where is sampled uniformly between 4 and 8, We assume knowledge about whether the observation for a given time step is valid and feed that knowledge to the models as an additional input flag. For valid observation, the model receives the observation and a 1. For invalid observations, it receives a default value and The RSSM handles this flag by appending it the encoder output or the default value. The VRKN allows for two ways of handling the flag. First, can again concatenate it to the observation before computing the latent observation and corresponding uncertainty estimate (VRKN (Cal)). Second, the VRKNs design allows us to provide the information more explicitly and the Kalman update step for invalid observations during inference. We evaluate with and without additional transition noise and show the results in Figure 9. Those results underpin the initial hypotheses that the VRKN is better equipped to learn in this setting and significantly outperforms the RSSM. The improvement independent of how we provide information about missing information to the VRKN and, expected, more significant with transition noise.
4.2.3 Fusing Information from Multiple Sensors Different Frequencies
We extend the missing observations setting using proprioceptive information which is available at every time step. The exact form the proprioceptive information is task-dependent. For example, for Cup Catch, we define the cup position as proprioceptive, but the ball position, which has to be inferred from
For the RSSM refer the convolution network processing the inputs as encoder, see Appendix C.3 for details.
14

---

Published in Transactions on Machine Learning Research (10/2022)
Without Transition Noise
With Transition Noise
1,000
Expectod Return
800
600
400
VRKN RSSM
200
0
0.2
0.4
0.6
0.8
1.0
0.2
0.4
0.6
0.8
1.0
Environment Steps (x109)
Environment Steps (x109)
Figure 10: Aggregated results in the fusion sctting with and without transition noisc. When comparing to Figure 9 we see that both approaches are able to exploit the additional proprioceplive information and improve their performance. Yet, the VRKN seems to better exploit and accumulate the available information, as the resulting agents significantly outperform their RSSM-based counterparts.
images. Appendix C.4 provides an overview for all considered tasks. This experiment mimics a common robotics scenario where we have proprioceptive information about the robot at high frequencies but need to estimate the environment's 's state based on lower-frequency images. It tests the approaches abilities to form reasonable state estimates from observations that arrive in different modalities and at varying frequencies. Here, the models need to trade off information encoded in the prior belief with the information available both sensor sources. This trade off requires accurate estimates about the aleatoric uncertainty in both state and observations. For the fusion lask, the VRKN can rely on the natural fusion mechanism described in Section 3. For the RSSM, we also multiple encoder networks, decoders, and loss terms and concatenate the outputs of the encoders before forming the posterior belief. For details, see Appendix C.3.
As shown in Figure 10, the VRKN achieves a higher reward than the RSSM, especially in the setting with transition noise. This result again emphasizes the VRKNs capability to exploit all available information, appropriately capture the system's 's uncertainty, and form belief states that yicld good performance.
5 Related Work
Recurrent State Space Models. While earlier works (Wahlström et al., 2015; Watter et al., 2015; Ban- ijamali et al., 2018; Ha & Schmidhuber, 2018; Buesing et al., 2018) discussed and showed the feasibility of control using Icarned latent state space models, the work originally proposing the RSSM (Hafner et al., 2019) was the first to show that such approaches can achieve similar performance to model-free RI on pixel- based complex continuous control tasks while using significantly fewer environment interactions. Since then, Hafner et al. (2020) improved their approach using a parametric policy learned on imagined trajectories and categorical latent spaces (Hafner et al., 2021). These approaches gained interest in the model-basee RLL community and are empirically successful, yet, little attention has been paid to the underlying state space model itself, the assumptions it builds upon, and its parametrization.
State Space Models. The Machine Learning community extensively studied and used state space models (SSMs). Besides classical approaches using linear models (Shumway & Stoffer, 1982) and works using Gaussian Processes (Eleftheriadis et al., 2017; Doerr et al., 2018), most recent methods build on Neural Networks (NNs). The first class of NN-based models of particular relevance for this work embeds linear- Gaussian SSMs (LGSSM) into latent spaces (Watter et al., 2015; Karl et al., 2016; Fraccaro et al., 2017; Banijamali ct al., 2018; Becker-Ehmck et al., 2019; Klushyn et al., 2021). These approaches assume actuated systems and learn using stochastic gradient variational Bayes (Kingma & Welling, 2013). Yet, non of these approaches were used to model or even control systems of the complexity considered by (Hafner et al., 2019) and here. They are not directly applicable to these sccnarios for various reasons. First, they use full transition matrices and covariances, which prevents them from scaling to sufficiently high-dimensional latent spaces. (Karl et al., 2016; Becker-Fhhndd et al., 2019) do not allow smoothing. (Fraccaro et al., 2017;
15

---

Published in Transactions on Machine Learning Research (10/2022)
Klushyn et al., 2021) model the latent observations random variables which are inferred jointly with the latent states and use constant observation uncertainty for the filtering in the latent space. This choice complicates inference and training and prevents principled usage of the observation uncertainty for filtering, Our parameterization of the LGSSM alleviates these issucs by building on factorization assumptions which yield scalable architecture. Further, it allows smoothing and principled usage of observation uncertainly during filtering by modcling the observations in latent space as deterministic. Finally, non of these approaches considered modeling epistemic uncertainty.
Another class of approaches directly uses NN-based, nonlinear parametrization for SSMs (Archer et al., 2015; Krishnan et al., 2015; Gu et al., 2015; Zheng et al., 2017; Krishnan et al., 2017; Yingzhen & Mandt, 2018; Schmidt & Hofmann, 2018; Naesseth et al., 2018; Moretti et al., 2019). Out of this class, Structured Inference Networks (SINs) (Krishnan et al., 2017) are the most relevant for our work. SINs build on the same variational objective as VRKN, yet without conditioning on actions. The smoothing- RSSM baseline introduced in Section 2.3 can be considered an instance of SIN. Yet while builds on the same loss and fundamental ideas, the underlying NN architecture is different.
Kalman Updates in Deep Latent Space. Haarnoja et al. (2016) first proposed using an encoder to extract uncertainty estimates from high-dimensional observations for filtering. They only Iearned the encoder while assuming the transition dynamics to be known. Becker et al. (2019) proposed an efficient factorization to additionally learn a high-dimensional, latent, locally-linear dynamics model. Shaj et al. (2020) extended this approach by introducing a principled form of action conditioning. While the VRKN builds on many of their design choices, there are also considerable differences. Those differences mainly concern the parametrization of the dynamics model and further simplifying the factorization assumptions. These changes are necessary to make the approaches scale to the complex control tasks considered in this work. Additionally, Haarnoja et al. (2016); Becker et al. (2019); Shaj et al. (2020) train using regression and do not learn a full generative model. Thus, they cannot produce the reasonable latent trajectories needed for model-based RL.
Epistemic Uncertainty for Model-Based RL. Ample work emphasises the importance of modeling epistemic uncertainty for model-based RL (Deisenroth & Rasmussen, 2011; Chua et al., 2018; Janner et al., 2019) and several authors equipped RSSMs with epistemic uncertainty. Okada et al. (2020) use an ensemble of RSSMs and showed improved results on modified versions of the Deep Mind Control Suite (Tassa et al., 2020) benchmarks. Sekar et al. (2020) also combine an ensemble with the RSSM but focus on exploration and generalization to unseen tasks. Yet, neither of these works questioned the assumptions underlying the RSSM or analyzed their effects on the learned models.
6 Conclusion
We analyzed the independence assumptions underlying Recurrent Stale Spuce Models (RSSMs) and found they are theoretically suboptimal. Yet, they implicitly regularize the model by causing an overestimated alcatoric uncertainty and are crucial to the RSSMs success in model-based RL. When trying to avoid this heuristic approach and use the correct assumptions while replacing the implicit regularization with a more explicit approach using epistemic uncertainty, we found a simple extension of the RSSM architecture is insufficient. Thus, we redesigned the model using well-understood and established components, providing a more appropriate inductive bias for smoothing. The resulting Variational Recurrent Kalman Network (VRKN) uses a latent linear Gaussian State Space Model (LGSSM) to address aleatoric uncertainty and Carlo Dropout to model epistemic uncertainty explicitly. Building on an LGSSM allows exact in- ference in the latent space using Kalman filtering and smoothing to obtain both smoothed and posterior belief states efficiently. While agents based on the VRKN and the RSSM perform similar on the standard DeepMind Control Suite (Tassa et al., 2020) benchmarks, the VRKN-based agents significantly outperform those using the RSSM on tasks where capturing uncertainties is more relevant. Additionally, the VRKN provides a natural approach to sensor fusion and outperforms the RSSM on tasks that require fusing sensor observations from several sensors at different frequencies.
Limitations. We showed that designing a state space model out of well-founded components that matches improves the RSSMs performance is possible. This insight opens a path to improve them individually.
16

---

Published in Transactions on Machine Learning Research (10/2022)
Yet, here we used simple instances of these components and have not yet further investigated how to improve them. Further, we have not investigated the interplay between the models and the controllers used on top of them but used the control approaches proposed in (Hafner et al., 2019) and (Hafner et al., 2020) with default parameters. Due to the intricate interplay between model learning and using the resulting controller for data collection, it reasonable to rethink the design of the controller when changing the model.
Acknowledgments
The authors acknowledge support by the state of Baden- Württemberg through bwHPC, as well as the HoreKa supercomputer funded by the Ministry of Science, Research and the Arts Baden-Württeni Württemberg and by the German Federal Ministry of Education and Research.
References
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304 29320, 2021.
Evan Archer, I1 Memming Park, Lars Buesing, John Cunningham, and Liam Paninski. Black box variational inference for state space models. arXiv preprint crXio:1511.07367, 2015.
Ershad Banijamali, Rui Shu, Hung Bui, Ali Ghodsi, et al. Robust locally- Menely-Iinear controllable embedding. In International Conference on Artificial Intelligence and Statistics, pp. 1751-1759. PMLR, 2018.
Philipp Becker, Harit Pandya, Gregor Gebhardt, Cheng Zhao, C James Taylor, and Gerhard Neumann. Recurrent kalman networks: Factorized inference high-dimensional deep feature spaces. In International Conference on Machine Learning, pp. 544 544-552, 2019.
Philip Becker- -hknck, Jan Peters, and Patrick Van Der Smagt. Switching linear dynamics for variational Bayes filtering. In Kamalika Chaudhuri and Ruslan Salakhutdinor (eds.), Proceedings of the 36th Inter- national Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 553 562. PMLR, 09 15 Jun 2019.
Lars Buesing, Theophane Weber, Sébastien Racaniere, SM Eslami, Danilo Rezende, David P Reichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, et al. Learning and querying fast generative models for reinforcement learning. arXiv preprint arXi::1802.0300 03006, 2018
Kyunghyun Cho, Bart Van Merrienboer, Dzmitry Bahdanau, and Yoshua Bengio. On the properties of neural machine translation: Encoder-decoder approaches. arXiv preprint 1xXiv:1409. 1259, 2014.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in Neural Information Processing Systems. 31, 2018.
Mare Deiscnroth and Carl E Rasmussen. Pilco: A model-based and data-cfficient approach to policy search. In Proceedings of the 28th International Conference machine learning (ICML-11), pp. 465 472. Citeseer, 2011.
Andreas Doerr, Christian Daniel, Martin Schiegg, Nguyen-Tuong Duy, Stefan Schaal, Mare Toussaint, and Trimpe Sebastian. Probabilistic recurrent state-space models. In International Conference on Machine Learning, pp. 1280-1289. PMLR, 2018.
Stefanos Eleftheriadis, Tom Nicholson, Mare Peter Deisenroth, and James Hensman. Identification of gaus- sian process state space models. In NIPS, pp. 5309-5319, 2017.
Marco Fraccaro, Simon Kamronn, Ulrich Paquct, and Ole Winther. A disentangled recognition and nonlincar dynamics model for unsupervised learning. In Advonces in Neural Information Processing Systems, pp. 3601-3610 2017.
17

---

Published in Transactions on Machine Learning Research (10/2022)
Yarin Gal and Zoubin Ghahramani. Dropout bayesian approximation: Representing model uncertainty in deep learning. In international conference. on machine learning, pp. 1050-1059. PMLR, 2016.
S Gu, 7 Ghahramani, and RE Turner. Neural adaptive sequential monte carlo. Advonces in Neural Infor-
mation Processing Systems, 2015:2629 2637, 2015.
Fredrik Gustafsson. Statistical sensor fusion. Studentlittcratur, 2010.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint ar arxiv:1803.1000 2018.
Thomas Haarnoja, Anurag Ajay, Sergey Levine, and Picter Abbeel. Backprop kf: learning discriminative deterministic state estimators. In Proceedings of the 30th International Conference on Neurul Information Processing Systems, pp.
2016.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565 PMLR, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models. In International Conference on Learning Representations, 2021. URL https: //openreview htte:/oormmiddd0oaayy200000
Sepp Hochreiter and Jürgen Schmidluber. Long short-term memory. Neural computation, 988):1735- 1997.
Eyke Hüllermeier and Willem Waegeman. Aleatoric and epistemic uncertainty in machine learning: An introduction to concepts and mothods. Machine Learning, 110(3):457-006, 2021.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When trust your model: Model- based policy optimization. Advances in Neural Information Processing Systems, 32:12519-128 2019.
AH Jazwinski. Stochastic processes and filtering theory. ACADEMIC PRESS, INC.., 1970.
Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Journal of basic Engineering, 82(1):35-455 1960.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. arXiv preprint arXiv:1605.064 06432, 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In ICLR (Poster), 2015.
Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.61 6111, 2013.
Alexej Klushyn, Richard Kurle, Maximilian Soelch, Botond Cseke, and Patrick van der Smagt. Latent matters: Learning deop state-space models. Advances in Neural Information Processing Systems, 34, 2021.
Rahul Krishman, Uri Shalit, and David Sontag. Structured inference networks for nonlinear state space models. In Proceedings of the AAAT Conference on Artificial Intelligence, volume 31, 2017.
Rahul G Krishnan, Uri Shalit, and David Sontag. Deep kalman filters. arXiv preprint arXiv: 151:1011..5111. 2015.
Nathan Lambert, Brandon Amos, Omry Yadan, and Roberto Calandra. Objective mismatch in reinforcement learning. Proceedings of Machine Learning Research vol, 120:1-15, 2020.
18

---

Published in Transactions on Machine Learning Research (10/2022)
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic frame- work for model-based deep reinforcement learning with theoretical guarantees. In International Conference on Learning Representations, 2019.
Michael Lutter, Leonard Hasenclever, Arunkumar Byravan, Gabriel Dulac-Arnold, Piotr Trochim, Nicolas Heess, Josh Merel, and Yuval Tassa. Learning dynamics models for model predictive agents. arXiv preprint AxXiv2209.143111 2021.
Antonio Moretti, Zizhao Wang, Luhuan Wu, and Itsik Peler. Smoothing nonlincar variational objectives with sequential monte carlo, 2019.
Kevin P Murphy. Machine learnóng: probabilisttc perspective. MIT press, 2012.
Christian Naesseth, Scott. Tinderman, Rajesh Ranganath, and David Blei. Variational sequential monte carlo. In International Conference Artificial Intelligncce and Statistics, pp. 968-977. PMLR, 2018.
Masashi Okada, Norio Kosaka, and Tadahiro Taniguchi. Planet of the bayesians: Reconsidering and im- proving deep planning network by incorporating bayesian inference. arXiv preprint arXiv:2003.0077 003700 2020.
Herbert E Rauch, F Tung, and Charlotte T Striebel. Maximum likelihood estimates of linear dynamic systems. AIAA journal, 5(8):1445-1455) 1965.
Florian Schmidt and Thomas Hofmann. Deep state space models for unconditional word generation. Advances in Neural Informalion Processing Systems 31, 31:6158 6168, 2018.
Ramanan Sckar, Olch Rybkin, Kostas Danillidis, Picter Abboel, Danijar Hafner, and Decpak Pathak. Plan- ning to explore via self-supervised world models. In International Conference on Machine Learning, pp. PMLR, 2020.
Vaisakh Shaj. Philipp Becker, Dieter Buchler, Harit Pandya, Niels van Duijkeren, C James Taylor, Marc Hanheide, and Gerhard Neumann. Action-conditional recurrent kalman networks for forward and inverse dynamics learning. Conference Robot Learning, 2020.
Robert H Shumway and David S Stoffer. An approach to time series smoothing and forecasting using the cm algorithm. Journal of time series analysis, ((4):253-244) 1982.
Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep condi- tional generative models. Advances in neural information processing systems, 28:3483 3491, 2015.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq. et al. Deepmind control suite. arXiv preprint wxXiv:1801.00690 2018.
Yuval Tassa, Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, and Nicolas Heess, dm control: c Software and tasks for continuous control, 2020.
Michael Volpp, Fabian Flürenbrock, Lukas Grossberger, Christian Daniel, and Gerhard Neumann. Bayesian context aggregation for neural processes. In International Conference Learning Representations, 2021.
Niklas Wahlström, Thomas B Schön, and Mare Peter Deisenroth. From pixels to torques: Policy learning with deep dynamical models. arXie preprint (xXi::500.0225:: 2015.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally lincar latent dynamics model for control from raw images. In Advances in neural information processing systems, pp. 2746 2754, 2015.
Li Yingzhen and Stephan Mandt. Discntangled sequential autoencoder. In Jennifer Dy and Andreas Krause (eds.); Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 5670-560000 5679. PMLR, 10-15 Jul 2018.
19

---

Published in Transactions on Machine Learning Research (10/2022)
Xun Zheng, Manzil Zaheer, Amr Ahmed, Yuan Wang, Eric P Xing, and Alexander J Smola. State space Istm models with particle meme inference. arviv preprint ar 7711.11111 11111119,, 2017.
20


Reviews:
[{"review_id": "SAS6bnh2ok", "reviewer": "Reviewer_uvqA", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. Investigates widely used state space models in model-based RL approaches (Dreamer, PlaNet, and its successors) and formulates a core assumption and its weakness thereafter. This is done in two parts,  \n- the misspecification of the inference distribution is shown by spelling out the variational lower bound decomposition to the log-likelihood of the observations given the actions, to indicate that the assumption of RSSMs leads to no tightness in the lower bound which is a strong limitation.\n- the effect of inference assumption on model learning has been demonstrated using a simple linear-gaussian state space model, with a state dimensionality of 4 and a ground truth generative model. Figure 2 shows the RSSM/SSM model learning results where evidence is provided that RSSM indeed has higher transition variance, higher model error, and underestimates the log state probability.\n\n2. The paper identifies the issue in RSSM arises because the model aims to maximize the ELBO while being evaluated based on the agent’s reward, to address this naively the authors propose Monte Carlo Dropout to called smoothing RSSM. Claims to be corroborated empirically later on in the paper somewhere.\n\n3. The proposed method VRKN (Variational Recurrent Kalman Networks) seems theoretically grounded, which allows for \n- parameter-free smoothing using a locally linear state space model in a latent space, while still scaling to image-based control tasks.\n- one network architecture that can compute both smoothed distributions for training and (filtered) posteriors for online control, to avoid an over-parametrization, and naturally integrate multiple sensors that emit data at different frequencies into a common latent state representation.\n\n4. A diverse set of experiments that probe for specific questions/weaknesses of the proposed method\n \n**Weaknesses**\n\n**Overall** I find the paper is not very well written and it is not very clear when it comes to stating its contributions. The empirical results are not so convincing either, which I guess leads the authors to kind of position the paper like this approach is more justified than the current SOTA and achieves the same performance. This positioning kind of hinges on establishing the theoretical justification first and foremost, making it the main contribution. However, I fear this also may not be a particularly strong point of the paper. Section 2.1 is the most theoretical part, but the soundness and significance of the analysis are not that clear to me. The main conclusions are made in sections 2.2 and 2.3, but the authors seem to be more hand-wavy/marginal empirical justifications than actual formal theoretical results/solid empirical analysis with rigorous ablations. Here are more detailed comments:\n\n1. Figure 1, it is not clear to me why there are **no concrete error bars** here considering there are 1000 sequences of length 50 each, ideally there should be multiple runs here. Maybe I am missing something. Can you please clarify? From what I understand so far, you are plotting the median. But even that requires error bars here. For instance in col 2, orange there is a lot of variances noticed. \n\n2. The paper is **not very clearly written and organized**. The introduction is very dense and it is hard to read. The authors introduce Smoothing RSSM in Section 2.3 and then claim there that “it does not help to improve the Smoothing RSSM’s performance. From these results, we conclude that solely addressing the sub-optimal inference assumption is insufficient, but we also need to rethink the model’s parameterization.” Where is the evidence for this provided? There is no reference or result here until I go to further sections which are not even referenced here. This is hard to parse.  \n\n3. Considering the paper is very much ablation-focused, **only 5 seeds** is not at all sufficient given that the performance gains are marginal especially when the performance remains at par with the baseline. It would be important to test experiments for at least 10 seeds or more to solidify these claims.\n\n4. In Figure 6, we would expect to see a greater difference in performance as the main point of the paper is that the overestimation of the aleatoric uncertainty is not precise enough in RSSMs. However, the difference in performance is pretty narrow. Would be nice to know a bit more about why the method proposed by VRKN does not seem to do so well.\n\n5. Another weakness is the assumption of a Gaussian generative distribution p(ot|zt) with fixed variance. What happens when this is not true in more practical settings?\n\n6. The presentation of the VRKN method is not clear enough. Many design choices, such as introducing latent observations, are hardly motivated and are not well-situated enough with respect to other parts of the method (Figure 3b).\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper investigates the inference scheme used by Recurrent State Space Models (RSSMs) mainly pointing out the flaws in the core assumption of RSSMs is that the belief is independent of future observations instead of the correct smoothing assumptions. **Key contributions** are claimed to be as follows:\n\n1. The authors discuss the shortcomings in the inference scheme of RSSMs, in that they formalize this shortcoming and analyze the assumptions’ effect on model learning leading to overestimating the aleatoric uncertainty. This is shown empirically that overestimating the aleatoric uncertainty can lead to poor performance in settings where “correctly estimating it is important”. \n\n2. They implement a naive approach of smoothing to show that even this does not help cope with the over-estimation of the aleatoric uncertainty. \n\n3. The authors present VRKN as a solution; it uses a latent linear Gaussian parameterization, combining components for both aleatoric and epistemic \n\n4. Introduce a bayesian treatment of the transition model’s parameters where the systems epistemic uncertainty is modelled using Monte Carlo Dropout. This results in: \n- performance is comparable to RSSM-based agents in a deterministic environment\n- performance is improved when the tasks required aleatoric uncertainty estimation.\n", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "With the paper currently as it stands, I am unfortunately not able to make a recommendation for acceptance. Here are specific adjustments which I believe will strengthen the work:\n\n- Please add clearly and precisely what the key contributions of this work are. At present, the introduction has dense paragraphs and it is left to the reader to enlist these. \n\n- Please cross-refer to each contribution which section is covering what. At present, it is hard to tell which section is covering formalism, etc. \n\n- 5 seeds are not at all sufficient given that the performance gains are marginal especially when the performance remains at par with the baseline. It would be important to repeat these experiments for at least 10 seeds or more to solidify these claims.\n\n- Please add error bars wherever missing. \n\n- What do the error margins indicate where they are actually drawn? Can you please elaborate on this, especially given there is both aleatoric and epistemic uncertainty being dealt with and claimed about here? \n\n- It is not exactly clear how experiments in section 4.3 relate to the general idea/message of the paper. The section tries to an argument to relate it to uncertainty estimation, but it seems a bit tenuous.\n\n- The presentation of the VRKN method is not clear enough. Many design choices, such as introducing latent observations, are hardly motivated and are not well-situated enough with respect to other parts of the method (Figure 3b). I would encourage the authors to carefully attend to this section, by highlighting differences/similarities with previous work.\n\n- When the authors say in conclusion that “we analyzed the independence assumptions underlying Recurrent State Space Models (RSSMs) and found they are theoretically sub-optimal. Yet, they implicitly regularize the model by causing an overestimated aleatoric uncertainty and are crucial to the RSSMs success in model-based RL.” It is not at all convincing to me that the current set of experiments demonstrates this sufficiently. Alternatively, is this formally proven somewhere? If yes, I encourage the authors can present this in a formal theoretical statement. It is unclear to me to what degree the authors position this as a contribution or rather more as motivation for the specific approach they construct. \n\n- Further, the authors say “When trying to avoid this heuristic approach and use the correct assumptions while replacing the implicit regularization with a more explicit approach using epistemic uncertainty, we found a simple extension of the RSSM architecture is insufficient. Thus, we redesigned the model from the first principles, using an inductive bias, which is appropriate for smoothing. As a result, we propose the Variational Recurrent Kalman Network (VRKN) which builds on well-understood tools for modelling the aleatoric and epistemic uncertainties. It uses extended Kalman smoothing for exact inference in a latent space to capture aleatoric uncertainty and explicit models the epistemic uncertainty using Monte Carlo Dropout.”  I really struggle following this part of the story. If I understand this correctly, the authors mean here that “redesigned model from first principles” includes as its main features “Kalman smoothing” and “Monte Carlo dropout”. To me, these are like well-known techniques that perform well as heuristics for helping estimate uncertainty, and the fact that they would be positioned here as being part of “first principles” is quite confusing to me. \n\n- While the sensor fusion section is interesting, and I believe in itself could be the main positioning. But I believe that is not the case. If this is the main focus of the paper, it should have been discussed during the formulation. \n\n- It seems that a lot is happening in the paper but it definitely could be both written and presented with significantly more clarity than the current standing of the work. \n", "broader_impact_concerns": "To the best of my knowledge, there are not any immediate broader impact concerns. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "iPM02HrluY", "reviewer": "Reviewer_5H8i", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**:\n\nThe paper is pleasant to read. Overall, it’s written clearly and it’s easy to follow as a result. I appreciated the background sections and the explanations around the effect of the approx inference process in Sec 2.2. The paper addresses potential shortcomings in a popular RL strand of work, the work combines many areas of expertise across approximate probabilistic inference, deep learning and model-based RL, and this is done in a way that is accessible to many. \n\nThe experiments are clear and illustrative, they generally support what is being described and the different hypotheses at play. \n\n**Weaknesses**:\n\n* Probably the main weakness of the paper is the limited gain that is demonstrated using the proposed model in the MuJoCo experiment. I think this is fine given the other contributions.\n\n* The paper mostly focuses on the model architecture and inference, but there is little discussion of how this model is integrated in a control setting (the new proposed model is used as a drop-in replacement for the existing RSSM model).  \n\n* The authors have tried a simple change to integrate smoothing in the RSSM model (“smoothing RSSM”), which didn’t work well in practice. They then introduce the VRKN model, which is a big change from the existing model, and I wish there were more discussion of the design choices and consequences here (beyond just the smoothing aspect). For example, what is the additional cost to doing smoothing? Is it still a suitable model for RL / planning? e.g. Can simulated latent trajectories when planning be decoded to sensible observation trajectories? Is it more/less sensitive to action-conditioning?\n\n* Sensor fusion: If I understood correctly, missing observations in the VRKN under the “natural” fusion mechanism means repeated observations. It seems unfair to compare it to RSSM with default values for unavailable observations. If the main positive effect here is through the observation repeat, then it’s not a terribly convincing demonstration of the fusion capability.\n\n---\nOther comment:\n\nAnother paper to consider citing in the space of SSMs for model-based RL: “Learning and Querying Generative Models for RL”\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies a popular type of state space models (SSM) used in recent model-based RL work: RSSMs. The authors point out an important limitation in the inference process (namely the lack of smoothing). The authors speculate this limitation has a beneficial effect on final agent performance through a regularization effect (namely overestimating aleatoric uncertainty), which means simply “fixing it” naively, in the current model architecture, degrades performance. The authors instead propose a different model architecture based on a Linear-Gaussian SSM where smoothing doesn’t require additional parameters (through a Kalman backward algorithm), that can still deal with high-res inputs, and where the beneficial regularization can be separately controlled with Monte Carlo Dropout. \n\nThe different approaches are tested with PlaNet and Dreamer agents on 3 Mujoco tasks. Some further analysis is done (e.g.  to check how the model deals with occlusion).", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "If my comment about sensor fusion above is justified, maybe soften the claims about sensor fusion.\n", "broader_impact_concerns": "No concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "vs5h5qZWcv", "reviewer": "Reviewer_TTvU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "I would like to first hear the authors' thoughts about the following question. Suppose we can find p and q such that inequality (1) becomes equality, when the agent wants to query the model for planning, it can not estimate the smoothing belief, which requires future information. Given that only the posterior is available anyway at decision time, why do we use the smoothing belief when learning the model? Could it be even worse than using the posterior for model learning (what RSSM does)? To think about the above question, maybe it is good to consider the following two questions first. What is the uncertainty in your state estimation? Where does it come from?\n\nStrengths:\nThe problem that the paper considers is important novel and nontrival.\n\nThe authors' perspective in analyzing this problem is novel and promising to me.\n\nThe paper shows that the problem indeed happens in practice (Figure 7).\n\nThe solution method proposed by the paper seems to resolve the problem (Figure 7) well.\n\nThe limitations of the paper were acknowledged and discussed.\n\nRelevant papers were cited.\n\nWeakness:\n\n1. Many math notations used to describe the proposed algorithm are used in a confusing way. I think the authors mix the notations of quantities to be estimated and their estimates. For example, p(z_{t+1} | z_t, a_t) is defined to be the transition probability of the latent states in Section 2 but is also used as a parameterized function modeling this transition probability in Section 3.2. Similarly, in Section 3.2, the authors define the latent observation model p(w_t | z_t), which is a definition of the generation process of observations from latent states. But the definition involves \\sigma_t^w, which is the encoder's uncertainty estimator, a solution parameter instead of a problem parameter. Similarly, the \"posterior state estimate p(zt|o≤t, a≤t)\". Similarly, z1, z2, z3 in figure 3 are estimates instead of actual latent states.\n\n2. Overall, I find the description of the proposed algorithm not complete and is thus difficult to follow. For example, it is not clear how the algorithm produces the estimate of the posterior of the latent state without future information. It is also not clear how learning works. The paper only says \"We train all parts of the model jointly by maximizing Equation 3\". But Equation 3 involves q and the algorithm description does not mention how it deals with q (e.g., how does the algorithm parameterize it). Also, Equation 3 involves expectations and the paper does not mention how to deal with them. It is also not clear how data were generated and how they were used in training.\n\n3. Some important statements are not justified: \n\nWhen presenting the requirements of the desired algorithm, the paper says\n\"In addition, such an architecture can naturally integrate multiple sensors that emit data at different frequencies\ninto a common latent state representation.\"\nThis sentence looks strange because up to this point, the paper has not discussed why integrating data from multiple sensors is important and why RSSM can not achieve this goal.\n\nSimilarly, the paper says in the requirements:\n\"Both should be computable by the same network architecture to avoid an over-parametrization and overfitting as a consequence.\"\nThis is again unjustified.\n\n4. Some important terminologies/statements need clear definition:\nSection 2 is called \"Inference and Learning in State Space Models\" but I don't see a definition/description of what \"inference\" and \"learning\" mean in your problem.\n\nThe paper needs to either define or cite other works for aleatoric and epistemic uncertainty.\n\n\"Traditionally, the independence assumptions of the generative model are also used for inference.\" What do the independence assumptions mean?\n\n\"As exact inference is intractable for most models of interest, we usually use approximate methods such as\nvariational inference. For a single sequence, ...\" It seems that you are about to explain the variational inference method here, but you didn't say it. And this terminology is not used anymore in the paper.\n\n5. A comparison between MCD-VRKN and MCD-RSSM in the experiments producing Figures 5 and 6 is important but is missing.\n\n6. I don't find a discussion in the main text referring to Figure 4.\n\n7. It is unclear how Smoothing-RSSM does inference when future information is not available.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper studies a popular model (RSSM) and its associated learning algorithm used in many recent successful model-based reinforcement learning algorithms. It shows that the model learning algorithm uses a sub-optimal inference scheme leading to an over-estimation of aleatoric uncertainty. The main goal of the paper is to seek a new variational model learning method and a new model that involve an optimal inference scheme.\n\nIt empirically shows that, by performing inference in an appropriate way to deal with the over-estimation problem (I can not understand how this is done though), the performance of PlaNet and Dreamer, two popular algorithms using RSSM, becomes worse, which is surprising. The paper argues that this is because of the lack of modeling the epistemic uncertainty. It then modifies the algorithms (PlaNet and Dreamer with appropriate inference) to add epistemic uncertainty following the Monte Carlo Dropout (MCD) approach but still obtains worse performance than the original performance of PlaNet and Dreamer.\n\nThe paper then postulates that the parameterization of the RSSM model is not suitable for smoothing (the appropriate inference technique), and proposed a new parameterization of the model using the linear-Gaussian model, VRKN. It shows that with VRKN and smoothing, in deterministic environments, both PlaNet and Dreamer match the performance achieved by the same algorithms with RSSM (original algorithms). In partial observable environments, VRKN (can not tell if MCD is used or not) is shown to be better than RSSM.", "claims_and_evidence_yn": "", "claims_explanation": "", "audience_yn": "", "audience_explanation": "", "requested_changes": "I would like the authors to think about my question and also discuss it in the paper (maybe some new experiments are needed to support your hypothesis) and also address the 7 weakness points mentioned above.", "broader_impact_concerns": "Not applied.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Three knowledgeable reviews all agree on acceptance. The initial submission had issues with the empirical setup and the authors went to great lengths to address these issue along with addressing the clarity concerns raised by the reviewers.


Decision: Accept (poster)