Flow-based Recurrent Belief State Learning for POMDPs
Authors: Xiaoyu Chen
Yao Mu
Ping Luo
Shengbo Eben Li
Jianyu Chen
Venue: ICLR 2022 Submitted (2022)
Keywords: Partially Observable Markov Decision Process
POMDP
Model-based Reinforcement Learning
Visual Control Task

Abstract: Partially Observable Markov Decision Process (POMDP) provides a principled and generic framework to model real world sequential decision making processes but yet remains unsolved, especially for high dimensional continuous space and unknown models. The main challenge lies in how to accurately obtain the belief state, which is the probability distribution over the unobservable environment states given historical information. Accurately calculating this belief state is a precondition for obtaining an optimal policy of POMDPs. Recent advances in deep learning techniques show great potential to learn good belief states, but they assume the belief states follow certain types of simple distributions such as diagonal Gaussian, which imposes strong restrictions to precisely capture the real belief states. In this paper, we introduce the \textbf{F}l\textbf{O}w-based \textbf{R}ecurrent \textbf{BE}lief \textbf{S}tate model (FORBES), which incorporates normalizing flows into the variational inference to learn general continuous belief states for POMDPs. Furthermore, we show that the learned belief states can be plugged into downstream RL algorithms to improve performance. In experiments, we show that our methods successfully capture the complex belief states that enable multi-modal predictions as well as high quality reconstructions, and results on challenging visual-motor control tasks show that our method achieves superior performance and sample efficiency.

Under review as a conference paper at ICLR 2022
FLOW-BASED RECURRENT BELIEF STATE LEARNING FOR POMDPS
Anonymous authors Paper under double-blind review
ABSTRACT
Partially Observable Markov Decision Process (POMDP) provides a principled and generic framework to model real world sequential decision making processes but yet remains unsolved, especially for high dimensional continuous space and unknown models. The main challenge lies in how to accurately obtain the be- lief state, which the probability distribution over the unobservable environment states given historical information. Accurately calculating this belief state is a precondition for obtaining an optimal policy POMDPs. Recent advances deep learning techniques show great potential to learn good belief states, but they assume the belief states follow certain types of simple distributions such as di- agonal Gaussian, which imposes strong restrictions precisely capture the real belief states. In this paper, we introduce the F1Ow-based Recurrent BElief State model (FORBES), which incorporates normalizing flows into the variational in- ference to learn general continuous belief states for POMDPs. Furthermore, we show that the learned belief states can be plugged into downstream RL algorithms to improve performance. In experiments, we show that our methods successfully capture the complex belief states that enable multi-modal predictions as well as high quality reconstructions, and results on challenging visual-motor control tasks show that our method achieves superior performance and sample efficiency.
INTRODUCTION
Partially Observable Markov Decision Process (POMDP) (Aström 1965) provides principled and generic framework to model real world sequential decision making processes. Unlike Markov De- cision Process (MDP), the observations of POMDP are generally mon-Markovian. Therefore, to make optimal decisions, the agent needs to consider all historical information, which is usually intractable. One effective solution is obtain the belief state. The belief state is defined as the probability distribution of the unobservable environment state conditioned on the past observations and actions (Kaelbling et al.) 1998). Such belief state accurately summarizes the history. Traditional methods of calculating belief states (Smallwood & Sondik 1973 Sondik 1971 Kaelbling et al., 1998) assume finite discrete space with a known model. many real world problems, however, the underlying model remains unknown, and the state space large and even continuous.
With the recent advances of deep learning technologies, a branch of works have been proposed to leam the belief states of POMDPs with unknown model and continuous state space (Krishnan al., 2015 Gregor et al.) 2019 Lee et al. 20200 Hafner et al. 2019b a 2021). These works solve the be- lief state learning problem by sequentially maximizing the observation probability each timestep using the variational inference and achieve the state-offthe-ar performance on many visual-motor control tasks (Hafner al.) 2019a| Zhu et al., 20200) Okada et al. [020 Ma et al., (020a)) However, they still cannot capture general belief states due to the intractability of complex distributions in high-dimensional continuous space and instead approximate the belief states with diagonal Gaus- sians. This approximation imposes strong restrictions and is problematic. As shown in Figure [] the blue area denotes the unobservable state space of the POMDP. Given the past information T, the agent maintains prior distribution of the state s, denoted as p(s]T) (the distribution in white). Each colored distribution corresponds to the belief state after receiving a different new observation denoted as the posterior distribution g)) Consider example of the true beliefs as shown in Figure [bb,, with their Gaussian approximations shown in Figure T[a). The approximation error of Gaussian distributions will easily result in problems of intersecting belief which leads to mixed-up

---

Under review as a conference paper at ICLR 2022
state (e.g., the white triangle), and empty belief, which leads to meaningless state (e.g.. the grey tri- angle). This also explains the poor reconstruction problems interactive environments observed by Okada & Taniguchi (202T). Furthermore, mentioned in Hafner al. (202T), the Gaussian approx- imation of belief states also makes it difficult to predict multi-modal future behaviours. Therefore, it is preferable relax the Gaussian assumptions and
flexible family of distributions
learn accurate belief states as shown in Figure
A
5
S
(a) Approximated Gaussian Bellef
(b) True Belief
p(|rr)
q(s(r,
q(str, 0)
qsstr, 0)
Figure 1: Difference between (a) spherical Gaussian belief states and (b) true belief states (Better viewed in color).
In this paper, we propose new method called FlOw-based Recurrent BElief State model (FORBES) that able learn general continuous belief states POMDPs. FORBES incorpo- rates Normalizing Flows (Tabak & Turner, 2013; Rezende & Mohamed] 2015 Dinh et al. 2017) into the variational inference step to construct flexible belicf states. In experiments, we show that FORBES allows the agent to maintain flexible belief states, which result in multi-modal and precise predictions as well as higher quality reconstructions. We also demonstrate the results combining FORBES with downstream RL algorithms challenging visual-moo motor control tasks (DeepMind Control Suite, Tassa et al. )2018). ) The results show the efficacy of FORBES in terms of improving both performance and sample efficiency.
Our contributions can be summarized as follows:
We propose FORBES, the first flow-based belief state learning algorithm that is capable of learning general continuous belief states for POMDPs.
We propose a POMDP RL framework based on FORBES for visual-motor control tasks, which uses the learned belief states from FORBES as the inputs to the downstream RL algorithms.
Empirically, we show that FORBES allows the agent learn flexible belief states that enable multi-modal predictions as well as high quality reconstructions and help improve both performance and sample efficiency for challenging visual-motor control tasks.
2 PRELIMINARIES
2.1 PARTIALLY OBSERVABLE MARKOV DECISION PROCESS
Formally, Partially Observable Markov Decision Process (POMDP) is a 7-tuple (S, A, T, R, 7). where a set of states, .A is of ac- tions, 7 is set of conditional transition probabilities between states, R the reward function, observation probabilities, and is the discount factor.
O-1
set of observations, O is set of conditional
At each timestep t 1, the state the environment is 5t-1 S. The agent takes an action d E A, which causes the environment to transit to state s, with probability T (s; st-1, a4-1). The agent then receives an observation or Ω which depends on the
a
Figure 2: The PGM of POMDP
state of the environ- ment s, with probability O (o, s)). Finally, the agent receives a reward rt-1 equal to R(s-1)). The agent's goal is to maximize the the expected
2

---

Under review as a conference paper at ICLR 2022
sum of discounted rewards E \ Such POMDP model can also be described proba- bilistic graphical model (PGM) as shown in Figure 2 After having taken action az-1 and observing o,, agent needs to update its belief state, which defined the probability distribution of the environment state conditioned on all historical information:
b(8)) p( 0t)
where = {01, a1, 1,-1, az
2.2 NORMALIZING FLOW
Instead of using the Gaussian family to approximate the prior and posterior belief distributions, we believe it is more desirable to use family of distributions that is highly flexible and preferably flexible enough to describe all possible true belief states. Therefore, we use Normalizing Flows Tabak & Tumnr 2013 Rezende & Mohamed 2015) to parameterize those distributions.
Rather than directly parameterizing statistics of the distribution itself, Normalizing Flows model the transformations, the "flow" progress, needed to derive such distribution. More specifically, it describes a sequence of invertible mappings that gradually transform a relatively simple probability density more flexible and complex one.
Let fo R" + R be an invertible and differentiable mapping in state space parameterized by 8. Given random variable Rº with probability distribution p(x), of the transformed random variable z = fo(x) by applying the change of variable formula:
derive the probability
0fg
p(z) p(x) det
(2)
dz
0
pgz)(z) log p(x) log det
Öz
To construct a highly flexible family of distributions, we can propagate the random variable at be- ginning 2o through sequence of K mappings and get 2 = sox Jox ....o So, (zo) with the probability
Σ log det 0ʃ..
log pk (Rx) log p(z))
0z--1
Given a relatively simple distribution Zu. say, Gaussian distribution, by iteratively applying the transformations, the flow is capable of representing a highly complex distribution with the probabil- ity that remains tractable. The parameters .1....,,p, 0x determine the transformations of the flow.
An effective transformation that is widely accepted is affine coupling layer (Dinh et al. 2017; Kingma Dhariwal, 2018 Kingma et al. 2017). Given the input for scale and translation functions which are usually parameterized by neural networks, where 8,6 & + D k < D. The output, y, can be viewed as a concatenation of its first k di- mensions y 1.k and the remaining part y
R.D. let and stand
Yi.k = Xisk: Yk+1.D = Xx+1.D O exp(s(x1x)) + t(x11))
(5)
where O denotes the element-wise product (see details about affine coupling layer in Appendix A.1
3 FLOW-BASED RECURRENT BELIEF STATE LEARNING
3.1 FLOW-BASED RECURRENT BELIEF STATE MODEL
We propose the FlOw-base Recurrent BElief State model (FORBES) which learns general con- tinuous belief states via normalizing flows under the variational inference framework. Specifically, the FORBES model consists of components needed to construct the PGM of POMDP as shown in Figure
State transition model ((50/81-1,441)
Observation model: p((r(s))
(6)
Reward model: p(re(st)
3

---

Under review as a conference paper at ICLR 2022
In addition, have belief inference model q(sxlTe: 0z) to approximate the true posterior distri- bution p(5t Fr. oz). where 7 = {01, a1...... Ot-1, ) is the past information. The above com- ponents of FORBES can be optimized jointly by maximizing the Evidence Lower BOund (ELBO) (Jordan et al. 1999) or more generally the variational information bottleneck (Tishby et al. 2000
Alemi al. 2016):
log p(01.T, P1r((11))
2 aar--1)
Σ (1n((or(s:)+ + (mp(r()s)) 1e
[np(r=(s:))-
Eq(o; als 2.91 2.11 1)
= [Mode]
here we use the factorization of ((81:T O1.T, a1 Il ( (1--1,
Detailed derivations be found in Appendix In practice, the state transition model, observa- tion model, reward model, and belief inference model can be represented by stochastic deep neural networks parameterized by th:
Pe (s(_)))--1,
po(or)s)), (relss))), (u(Se)T,, oz)
where their outputs usually follow simple distributions such as diagonal Gaussians. The parameter- ized belief inference model gw (s, acts as encoder that encodes the historical information using a combination of convolutional neural networks and recurrent neural networks. Note that below is obtained through the factorization of ((81:1/01::: (1:7-1) II. 144
In FORBES we provide special treatments for the belief inference model and the state transition model to represent more complex and flexible posterior and prior distributions. As shown in Fig- ure 3 a), the input images Ouæ and actions d1 are encoded with qu (SelT,, oz) (the blue and the path). Then our final inferred belief is obtained by propagating (T, o1) through set of normalizing flow mappings denoted fex ..... fe, to get a representative posterior distribution Ty.o (sx/Tt, For convenience, we denote go du and ęK = fv.0- the other hand, 01:-1 and 01:1-2 are encoded with go (8t-1 17 0t-1) (the blue path), then the state transition model is used to obtain the prior guess of the state Pa: | T₂) =
(P &t-1, (t-1;((-1) (the green path). Then our final prior obtained by propagating (s: rr)) through another set of nor- malizing flow mappings denoted fwx o...o Sw, to get representative prior distribution (8 (T.). For convenience, we denote Po = Pe and Pк Pe... Then as shown Figure Bbb, we can sample the initial state 8 (the yellow and purple triangles) from the belief states qK ($t Te,00). For each sampled initial state, we can use the state transition model to predict the future states Seth given the future actions Ou:+hh-1, and then use the observation model reconstruct the observations ôeth. where h is the prediction horizon.
With the above settings, Equation 7 with Normalizing Flow:
substitute the density probability inside the KL-divergence term in
DKL(qx (se ITe (x) (=)
(gx(ssalreoo)) - [log qK (Sa/Te, o.) loggx((s(TT)
K
Ofo= log + log det f Isk
= -Eyx(sul(Tee))
log (o
Σ
log det Ōsc
8t-1, 0t-1) = (st Tt) given the sampled 8t-1 from qK (s (Te, Te, or). 8t, is the state variable 8, transformed by & layers of normalizing flows, and st,0 St.
where (st |
To further demonstrate the properties of FORBES, we provide the following theorems.
Theorem
The approximation error of the lig-tikelihood when maximizing the JModal (the derived
ELBO) defined in Equation] is:
DKa(q(SulTe (:) Te, ail)
po Purllarr)
JModel Egggce (rr)
(10)
where p(s: Te, Ot) denotes the true belief states.
4

---

Under review as a conference paper at ICLR 2022
6
O:47
Po(s,lt,)
Pr(s,)),
St 0x(5,1t,,))
0r11
Dc+2
P1
90(111 4x(Selte))))
(a) Belief state inference
(b) Predictions beginning from different samples
Figure 3: The algorithm framework of FORBES. Figure 3a shows how to calculate prior and poste- rior belief distribution given previous information. The blue arrows bring in historical observations and actions, and the green path shows the evolution of prior belief distribution. The red path takes an additional o and shows the evolution of posterior belief distribution. Figure shows prediction of future trajectories given the future actions.
Detailed proofs can be found in Appendix A.9 Theorem suggests that, when the learning algo- rithm maximizes the JModel (the derived ELBO), then the DKL terms in the right-hand side are minimized, which indicate the KL-divergence between the learned belief states ((sx Trl Te. o,) and the true belief states p(*+ Tr. 0²). Clearly, if p(*, Te, 0²) is a complex distribution and g(&t TT, oc) is chosen from restricted distribution class such as diagonal Gaussian, then when the algorithm maximizes the Model (the derived ELBO), there will still be a potentially large KL-divergence be- tween the learned and the true belief states. Therefore, naturally there raises the problem that is normalizing flow a universal distributional approximator that is capable of accurately representing arbitrarily complex belief states, the KL-divergence terms in the right-hand side of Equation (IO) can be minimized to approach zero? The answer is yes for a wide range of normalizing flows. To be specific, Teshima al. (2020) provides theoretical results for the family of the flow used in FORBES.
In fact, there always exists diffeomorphism that Besides the aforementioned affine coupling flow, many works show the distributional universality of other flows (Kong& Chaudhuri, 2020 Huang al. 2018). Ideally, the universal approximation property of the flow model 9K (st
well-behaved distribution to another.
allows us to approximate the true posterior p(s= Te, or) with arbitrary accuracy. Thus, compared to previous methods, FORBES helps close the gap between the log-likelihood and the ELBO to obtain a more accurate belief state. Though we usually cannot achieve the ideal zero KL-divergence in practice, our method can get a smaller approximation error, equally higher ELBO than previous works. We verify this statement in section 4.1
3.2 POMDP RL FRAMEWORK BASED ON FORBES
To show the advantage of the belief states inferred by the FORBES model compared the existing belief inference method in visual-motor control tasks, we design flow-based belief reinforcement learning algorithm for learning the optimal policy POMDPs. The algorithm follows an actor-critic framework. The critic estimates the accumulated future rewards, and the actor chooses actions to maximize the estimated cumulated rewards. Both the actor and critic operate on top of the samples of learned belief states, and thus benefit from the accurate representations learned by the FORBES model. Note that this an approximation of the true value on belief, which avoids the intractable integration through observation model.
The critic eg (s7) aims to predict the discounted sum of future rewards that the actor can achieve given an initial state s,, known the state value ( "r-"r=) where & denote the parameters of the critic network and // is the prediction horizon. We leverage temporal-difference to learn this value, where the critic is trained towards a value target that is constructed from the intermediate reward and the critic output for the next step's 's state. In order to trade-off the bias and the variance of the state value estimation, we use the more general TD(A) target (Sutton & Barto 2018). which is weighted average of n-step returns for different horizons and is defined as follows:
5

---

Under review as a conference paper at ICLR 2022
(1 -11 ()) +VA. T+1 if T<t+ H, if T=f+
= +
(11)
((g(5L+H)
To better utilize the flexibility belief states from FORBES, we run the sampling method multiple times to capture the diverse predictions. Specifically, sample N states from the belief state given by FORBES and then rollout trajectories future states and rewards using the state transition model and the reward model. Finally, we train the critic regress the TD(A) target return using mean squared error loss:
E (E) 51,0"9K SRy 1i, Pay
Σ N Eit (1e((8,,7) -gg((
(12)
where sg(.) is the stop gradient operation. The actor a, qo (a, sT) aims to output actions that maximize the prediction of long-term future rewards made by critic and trained directly by backpropagating the value gradients through the sequence of sampled states and actions, i.e., maximize:
N Σ
I+H
Jactor 98 do Si.r
(13)
1,7
We jointly optimize the model loss Model with respect to the model parameters v, 0 and w, the critic loss JCritic with respect to the critic parameters and the actor JActor loss with respect to the actor parameters & using the Adam optimizer with different learning rates:
min $.E.0,8
JFORBES anJCCrite on Jxcoor O2JMMade (tb, 0 w)
(14)
where 0g. 01, 02 are coefficients for different components, and we summarize the whole framework of optimizing in Algorithm]
Algorithm FORBES Algorithm
Input: buffer B, imagination horizon H, interacting step T, batch size B, batch length L, number trajetories N.
Draw B data sequences {(or, a, (,, &++ from B Infer belief state qx (x(TT)T,,))).
for i 1, .., .N do
Rollout imaginary trajectories ((iirr a1,z) ) with belief transition model.
end for
For each 8i,r, predict rewards pe (ri ( and values Uo (S Update 0, w, E. Φ, &b using Equation 7/. actor
Calculate imaginary returns
and 14 b Optimize the model, critic and
4 EXPERIMENTS
Our experiments evaluate FORBES on two image-based tasks. We first demonstrate the belief learn- ing capacity digit writing task in (ection[.11 and show that FORBES captures beliefs that allow for multi-modal yet precise long-term predictions as well as higher ELBO. For large-scale experi- ments, we test the proposed POMDP RL framework based FORBES in Section 4.2 The results of multiple challenging visual- visual-motor motor control tasks from Deep.ind Control Suite (Tassa et al. 2018) show that FORBES outperforms baselines terms of performance and sample efficiency. Sec- tion 4.3, we further provide ablation studies of the multiple imagined trajectories technique used in our method.
4.1 DIGIT WRITING TASKS
In this experiment, validate the capacity of FORBES by modelling the partially observable
quence with visual inputs. We adopt the MNIST Sequence Dataset (D. De Jong 2016) that consists of sequences of handwriting MNIST digit stokes. This problem be viewed a special case

---

Under review as a conference paper at ICLR 2022
Ground Truth
Model Inputs
RSSM predictions: 030
FORBES predictions: 030
5
5
5
S
S S
S
S
5
4
5
S
S
S
$
10
015
5
K
5
5
5
6
6
5
6
6 6
2
5
3
5
5
5
6
6
6
6
0g
01
05
010
15
7
7
7
7
7
3 3
3
3
3 8
7
7
7
7
7
7
7
7
7
7
7
0g 01 0s0
Figure 4: Predictions on sequential MNIST of two models.
of POMDP, whose action space is 0 and rewards remain Such a problem setting separates the belief learning and policy optimizing problem and allows to concentrate this section. We convert the digit stroke to a sequence of images of size 28 x 28 to simulate the writing process. At time step the agent observe or that has already written / pixels, and we train the agent maximizing MModel in Equation 7 except for the reward reconstruction term. As shown in Figure we randomly select three digits examples (see Appendix A. 10 for more results) and show the inputs as well as the prediction outputs of our model and the RSSM (Hafner et al., 2019b) baseline, which is the previous state-of-th thetthe-rrt method for learning continuous belief states of POMDPs. The leftmost column is the ground truth of the fully written digits. During the testing, we feed the initial 15 frames {01, 02, 015} to the model, and the columns in grey exhibit part of the inputs. Then sample several states from the inferred belief state and rollout via the learned state transition model (Equation (6)) for 15 steps and show the recon- struction results of the predictions. As shown in the blue and green columns on the right of Figure though RSSM can also predict the future strokes in general, the reconstructions are relatively blurred and mix different digits up. It also fails to give diverse predictions. ing However, FORBES can make precise yet diverse predictions. Each prediction clear and distinct from other digits. Given the beginning of the digit 7, FORBES suc- cessfully predicts both 7 and 3 since they have similar beginning. The results can be partially explained via the mixed-up belief and the empty belief as shown in Figure [ which support the claim that FORBES can better capture the complex belief states.
the former in
-724.0
-724.5
ELBO
-72..0
-725.5
FORBES RSSM
-726.0
0.5 1.0
1.
2.0 2.5
Training Steps (M)
Figure 5: ELBO on digit writ-
We also provide the quantitative results in Figure 5 which is the ELBO on test digits sequence set that is never seen during training. The results show that FORBES can achieve tighter ELBO, which verifies the theoretical results in 3.1 Details of the network can be found in Appendix A.2
4.2 VISUAR.-MOTOR CONTROL TASKS
We experimentally evaluate the performance of FORBES on Reinforcement I Learning on a variety of visual-motor control tasks from the DeepMind Control Suite (Tassa al. 2018). illustrated in Figure Across all the tasks, the observations are 64 x 64 x 3 images. These environment provide different challenges. The Sartpole-Swinuup task requires long planning horizon and memorizing the state of the cart when it is out of view; Finger-Spinning includes contact dynamics between the finger and the object; Cheetah-Run exhibits high-dimensional state and action spaces; the Walker- Walk and Walker-Run are challenging because the robot has to learn to first stand up and then walk; Hopper Stand is based on a single-legged robot, which sensitive to the reaction force on the ground and thus needs more accurate control. As for basclines,
include the scores for A3C Mnih et al. (2016) with state inputs (Ie9 steps), D4PG Barth-Maron et al. (2018) (le9 steps), PlaNet (Hafner et al. 2019b) (1e6 steps) and Dreamer Hafner et al. (2019a) with pixel inputs. All the scores of baselines are aligned with the ones reported in Hafner et al. (2019a) (see details in Appendix A.3).
The details of the implementations can be found in Appendix A.2
Our experiments empirically show that FORBES achieves superior performance and sample effi- ciency challenging visual-motor control tasks. As illustrated in Figure 6 FORBES achieves

---

Under review as a conference paper at ICLR 2022
cartpole swingup
cheetah
finger spin
Eploode Rettim
000
200
100
320
360
44
590
140
24
330
400
530
330
300
593
hopper stand
walker run
walker walk
633
Epicode Ratum
530
193
200
250
TEs
Environment Steps (K)
Environment Steps IK]
Environment Steps (K)
FORBES
D4PG stops)
PlaNot steps)
Dreamer
A3C (1e9 steps, proprio)
Figure 6: Performance DeepMind Control Suite. The shaded areas show the standard deviation across 3 seeds. FORBES achieves better performance and sample efficiency in various challenging tasks.
higher scores than Dreamer (Hafner et al. 2019a) in most of the tasks and achieves better perfor- mance than PlaNet (Hafner et al.) 2019b) within much fewer environment steps. We provide some insights into the results. As shown in Section 4.1 baselines with Gaussian assumptions may suffer from the mixed-up belief and empty belief issues, while FORBES can better capture the general be- lief states. Furthermore, multiple imagined trajectories can better utilize the diversity in the rollout. Therefore, the inner coherency within the model components allows the agent better performance. We further discuss the role of multiple imagined trajectories in the next section.
4.3
ABLATION STUDY
cartpole swingup
cheetah
finger
Eploode Returr
533
613-
230
hopper: stand
walker run
walker walk
Epicode Return
1.0
390
Environment Stops Stops (K)
Environment Staps (K)
Enviranment Steps
FORBES
Dreamer
Dreamer Multiple Imagined Trajectories
Figure 7: Comparison the performance between FORBES and Dreamer with multiple imagined trajectories.
In order to verify that the outperformance of FORBES is not simply due to increasing the number of imagined trajectories, we conducted an ablation study in this section. We compare FORBES with the "Dreamer + multiple imagined trajectories" baseline by increasing the number of imagined trajectories in Dreamer to the same as in FORBES. As shown in Figure no consistent and obvious gain can be observed after increasing the number of trajectories to Dreamer. The agent gains slight improvements in two environments and suffers from slight performance loss on other tasks. This result indicates that increasing the number of imagined trajectories may only be effective when the agent can make diverse predictions as FORBES. The Gaussian assumptions lead to the lack of trajectory diversity, that increasing the number of imagined trajectories will not effectively help.
8

---

Under review as a conference paper at ICLR 2022
5 RELATED WORK
POMDP: POMDP solving approaches can be divided into two categories based on whether their state, action and observation spaces are discrete or continuous. Discrete space POMDP solvers, in general, either upproximate the value function using biint-based methods (Kurniawati et al., 20088 Shani et al. 2013) or using Monte-Carlo sampling in the belief space (Silver Veness 2010 Kur- niawati & Yadav 2016) to make the POMDP problem tractable. Continuous space POMDP solvers often approximate the belief states as a distribution with few parameters (typically Gaussian) and solve the problem analytically either using gradients ( (Van Den Berg et al. [0022] Indelman et al,,
2013) or using random sampling in the belief space (Agha-Mohammadi et al., (014; Hollinger & Sukhatme 2014). However, the classical POMDP methods mentioned above are all based an accurately known dynamic model, which is a restricted assumption in many real world tasks.
MBRL for visual-motor control: Recent researches in model-based reinforcement learning (MBRL) for visual-motor control provides promising methods to solve POMDPs with high- dimensional continuous space and unknown models since visual-motor control tasks can be naturally modelled as POMDP problems. Learning effective latent dynamics models to solve challenging visual-motor control problems is becoming feasible through advances in deep generative modeling and latent variable models (Krishnan et al., 2015 Karl et al., 2016 Doerr et al., 2018 Buesing et al., 2018 Ha & Schmidhuber 2018 Hafner al. 2019b a). Among which, the recurrent state-space
model (RSSM) based methods (Hafner et al. 2019b a provide principled way to learn continuous latent belief states for POMDPs by variational inference and learns behaviours based on the belief states using model-based reinforcement learning, which achieves high performance visual-motor control tasks. However, they assume the belief states obey diagonal Gaussian distributions. Such assumptions impose strong restrictions to belief inference and lead to limitations in practice, includ- ing mode collapse, posterior collapse and object vanishing in reconstruction/ coanstructionBBomaan Bowmmnn et al., [016)] Salimans et al., 2015, 2015/Okada Okada & Taniguchi 2020). A few works propose particle filter based methods that use samples to approximate the belief states (Ma et al. 2020b [g] et al. 2018). However, they suffer from insufficient sample efficiency and performance.
Normalizing Flows: Normalizing Flows (NF) are a family of generative models which produce tractable distributions with analytical density. For a transformation f R.D R.D, the compu- tational time cost of the log determinant is ( Thus most previous works choose to make the computation more tractable. Rezende& & Mohamed (0135: : van den Berg et al. (2019) propose to use restricted functional form of f. Another choice is to force the Jacobian of to be lower triangular by using an autoregressive model (Kingma et al., 2016 Papamakarios et al. 2018). These models usually excel at density estimation, but the inverse computation can be time-consuming. Dinh et al.) (015 2017); Kingma & Dhariwal (2018) propose the coupling method to make the Jacobian trian- gular and ensure the forward and inverse can be computed with single pass. The applications of NF include image generation (Ho et al.) / & Dhariwal 2018). video generation (Kumar
ct al. 2019) and reinforcement learning (Mazoure et al. 2020 Ward et al. 2019 Touati et al. 2020).
6 CONCLUSION
General continuous belief states inference is a crucial yet challenging problem in high-dimensional Partially Observable Markov Decision Process (POMDP) problems. In this paper, we propose the FlOw-based Recurrent BElief State model (FORBES) that can learn general continuous belief states by incorporating normalizing flows into the variational inference framework and then effectively uti- lize the learned belief states in downstream RL tasks. We show that theoretically, our method can accurately learn the true belief states and we verify the effectiveness of our method in terms of both the quality of learned belief states and the final performance of our extended POMDP RL frame- work on two visual input environments. The digit writing tasks demonstrate that our method can learn general belief states that enable precise and multi-modal predictions and high-quality recon- structions. General belief inference plays a vital role in solving the POMDP, and our method paves way towards it. In the future, will explore further approaches to improve the accuracy of belief states inference and information secking, such as combining contrastive learning and using advanced network architectures such as transformers build normalizing flows.
9

---

Under review as a conference paper at ICLR 2022
7 REPRODUCIBILITY
We implement FORBES based on the Pytorch version of Dreamer, provided by |nttps://githb.co httpyy/ssuukuuaaaaa/arreeeeeeeee pytorch and our code will be released after the author notification in "https://anp http::::///onnoouuuu.oo 4open. science/r/ Lmpplementariiooof--Pooe -0-000---5000 We describe the implementation details as well as the hyperparameters in Appendix A.2 We train our FORBES model on NVIDIA Tesla V100 with Pytorch 1.5.0.
REFERENCES
Ali-Akbar
Suman Chakravorty, and Nancy M Amato. Firm: Sampling-basee feedback motion-planning under motion uncertainty and imperfect measurements. The Interna- tional Journal of Robotics Research, 33(2)2668-00 2014.
Alexander A Alemi, lan Fischer, Joshua V Dillon, and Kevin Murphy. Deep variational information bottleneck. arXiv preprint /612.0/6/..004///// 2016.
Gabriel Barth-Maron, Matthew W. Hoffman, David Budden, Will Dabney, Dan Horgan, Dhruva TB, Alistair Muldal, Nicolas Heess, and Timothy Lillicrap. Distributed distributional deterministic policy gradients, 2018
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Ben- gio. Generating sentences from a continuous space. CONLL, 2016.
Lars Buesing, Theophane Weber, Sébastien Racaniere, SM Eslami, Danilo Rezende, David P Re- ichert, Fabio Viola, Frederic Besse, Karol Gregor, Demis Hassabis, et al. Learning and querying 2018.
fast generative models for reinforcement learning, arXiv preprint
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre. Dzmitry Bahdanau, Fethi Bougares, Hol- ger Schwenk, and Yoshua Bengio. Learning phrase representations using encoder-decoder for statistical machine translation. arXiv preprint urXiv::4000 1078, 2014.
Edwin D. De Jong. The mnist sequence dataset. :ttps://edvirr de- jong.github. blog/mni tet-ssqqueceeeeeee
2016. Accessed: 2019-07-07.
Laurent Dinh, David Krueger, and Yoshua Bengio. Nice: Non-linear independent components esti- mation. In ICLR, 2015.
Laurent Dinh, Jascha Sohl- Dickssits,,,, and Samy Bengio. Density estimation using real nvp. 2017.
Andreas Doerr, Christian Daniel, Martin Schiegg. Nguyen-Tuong Duy, Stefan Schaal, Mare Tous- saint, and Trimpe Sebastian. Probabilistic recurrent state-space models. International Confer- on Machine Learning, pp. 1280-1289. PMLR, 2018.
Karol Gregor, George Papamakarios, Frederic Besse, Lars Buesing, and Theophane Weber. Tempo- ral difference variational auto-encoder, 2019.
David Ha and Jürgen Schmidhuber. World models. arXiv preprint arXiv:1803.1022 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01600 2019a.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels, 2019b.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with dis- crete world models, 2021.
Jonathan Ho. Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbeel. Flow++: Improving flow- based generative models with variational dequantization and architecture design. In International Conference Machine Learning, pp. 272--273.. PMLR, 2019.
10

---

Under review as a conference paper at ICLR 2022
Geoffrey A Hollinger and Gaurav S Sukhatme. Sampling- robotic information gathering algo- rithms. The International Journal of Robotics Research, 33(9):1271-122) 2014.
Chin-Wei Huang, David Krueger, Alexandre Lacoste, and Aaron Courville. Neural autoregressive flows. In ICML, 2018.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps, 2018.
Vadim Indelman, Luca Carlone, and Frank Dellaert. Planning in the continuous domain: A gener- alized belief space approach for autonomous navigation in unknown environments. The Interna- tional Journal of Robotics Research, 34(7):849-882, 2015.
Michael Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 37(2):183-233 1999.
Leslie Pack Kaelbling, Michael L. Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence. 1998.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt. Deep varia- tional bayes filters: Unsupervised learning state space models from raw data. arXiv preprint 1rXx:::605.06443.. 2016.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.698) 2014.
Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improving variational inference with inverse autoregressive flow, 2017.
Durk P Kingma and Prafulla Dhariwal. Glow: Generative flow with invertible 1x1 convolutions. In NeurIPS, 2018.
Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Im- proved variational inference with inverse autoregressive flow. In NIPS, 2016.
Zhifeng Kong and Kamalika Chaudhuri. The expressive power of class of normalizing flow mod- els, 2020.
Rahul G Krishnan, Uri Shalit, and David Sontag. Deep kalman filters. arXiv preprint arXiv: 151.055111 2015.
Manoj Kumar, Mohammad Babaeizadeh, Dumitru Erhan, Chelsea Finn, Sergey Levine, Laurent Dinh, and Durk Kingma. Videoflow: A flow-based generative model for video. arXiv preprint 1933v:1903.01434 2019.
2019.
Hanna Kurniawati and Vinay Yadav. An online pomdp solver for uncertainty planning in dynamic environment. In Robotics Research, pp. 61 611-629. Springer, 2016.
Hanna Kurniawati, David Hsu, and Wee Sun Lee. Sarsop: Efficient paint-based pomdp planning by approximating optimally reachable belief spaces. In Robotics: Science and systems, volume 2008. Citeseer, 2008.
Alex X. Lee, Anusha Nagabandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with latent variable model, 2020.
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. arXiv preprint urXiv:1509.02977 2015.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational reinforcement learning for complex observations, 2020a.
Xiao Ma, Peter Karkus, David Hsu, and Wee Sun Lee. Particle filter recurrent neural networks. In Proceedings of the AAAI Conference Artificial Intelligence, volume 34, pp. 5101-5108, 2020b.

---

Under review as a conference paper at ICLR 2022
Bogdan Mazoure, Thang Doan, Audrey Durand, Joelle Pineau, and R Devon Hjelm. Leveraging exploration in off- policy algorithms via normalizing flows. In Conference on Robot Learning, pp- 430-444. PMLR, 2020.
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap. Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning, 2016.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction. arXiv preprint arXiv:2007.14533 2020.
Masashi Okada and Tadahiro Taniguchi. Dreaming: Model-based reinforcement learning by latent imagination without reconstruction, 2021.
Masashi Okada, Norio Kosaka, and Tadahiro Taniguchi. Planet of the bayesians: Reconsidering and improving deep planning network by incorporating bayesian inference, 2020.
George Papamakarios, Theo Pavlakou, and lain Murray. Masked autoregressive flow for density estimation, 2018.
Danilo Jimenez Rezende and Shakir Mohamed. Variational inference with normalizing flows. In ICML, 2015.
Tim Salimans, Diederik Kingma, and Max Welling. Markov chain monte carlo and variational inference: Bridging the gap. In International Conference on Machine Learning, pp. 1218-1226. PMLR, 2015.
Guy Shani, Joelle Pineau, and Robert Kaplow, A survey of point-based pomdp solvers. Autonomous Agents and Multi-Agent Systems, 27(1):1-51, 2013.
David Silver and Joel Veness. Monte-carlo planning in large pomdps. Neural Information Processing Systems, 2010.
Richard D Smallwood and Edward J Sondik. The optimal control of partially observable markov processes over a finite horizon. Operations research, 21(5):1011-108; 1973.
Edward Jay Sondik. The optimal control of partially observable Markov processes. Stanford Uni- versity, 1971.
Richard Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Esteban G Tabak and Cristina V Turner. A family of nonparametric density estimation algorithms. Communications on Pure and Applied Mathematics, 66(2):145-164 2013.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq. Timothy Lillicrap, and Martin Ried- miller. Deepmind control suite, 2018
Takeshi Teshima, Isao Ishikawa, Koichi Tojo, Kenta Oono, Masahiro Ikeda, and Masashi Sugiyama. Coupling- based invertible neural networks are universal diffeomorphism approximators, 2020.
Naftali Tishby, Fernando Pereira, and William Bialek. The information bottleneck method. arXiv preprint physics/0004055, 2000.
Ahmed Touati, Harsh Satjja, Joshua Romoff, Joelle Pineau, and Pascal Vincent. Randomized value functions via multiplicative normalizing flows. In Uncertainty in Artificial Intelligence, pp. 422- 432. PMLR, 2020.
Jur Van Den Berg, Sachin Patil, and Ron Alterovitz. Motion planning under uncertainty using iterative local optimization in belief space. The International Journal of Robotics Research, 31 111):1263-1278, 2012.
Rianne van den Berg, Leonard Hasenclever, Jakub M. Tomezak, and Max Welling. Sylvester nor- malizing flows for variational inference, 2019.
12

---

Under review as a conference paper at ICLR 2022
Patrick Nadeem Ward, Ariella Smofsky, and Avishek Joey Bose. Improving exploration in soft- actor-critic with normalizing flows policies. arXiv preprint 1xxv:1996.02771120199
Guangxiang Zhu, Minghao Zhang, Honglak Lee, and Chongjie Zhang. Bridging imagination and reality for model-based deep reinforcement learning, 2020.
K.J Aström. Optimal control of markov processes with incomplete state information. Journal of Mathematical Analysis and Applications, 10(1):174-205, 1965. ISSN 0022-247X. doi:
URL https: /wwww scienced rect.
A APPENDIX
A.I DETAILS OF AFFINE COUPLING LAYER FOR NORMALIZING FLOW
In this section, we will introduce the details about the affine coupling layer (Dinh et 2017).
In the forward function, we split the input x E IK D into two parts according to the dimension: & I= Then, we let the first part X1:k stay identical, so that the first l dimensions in the output y E IR"" is Y1.k = Xi:k- After that, we use the identical part as the inputs to determine the transform parameters. In our case, define two neural network s, R - D- which stand for scale and translation functions. They receive X1. as inputs and output the affine parameters. As in (Dinh et al. 2017), the second part can be derived by:
Yk+1.D = X:D @ exp(s(x1:e)) + t(x1,)
(15)
Finally, the output, is the concatenation of the two parts: y [y Vk+1:D].
The affine coupling layer is an expressive transformation with easily-computed forward and reverse passes. The Jacobian of affine coupling layer is triangular matrix, and its log determinant can also be efficiently computed.
A.2 HYPER PARAMETERS AND IMPLEMENTATION DETAILS
Network Architecture We use the convolutional and deconvolutional networks that are similar to Dreamer(Hatner et al.) 2019a), GRU (Cho et al., 2014) with 200 units in the dynamics model, and implement all other functions as two fully connected layers of size 200 with ReLU activations. Base distributions in latent space are 30-dimensional diagonal Gaussians with predicted mean and standard deviation. As for the parameters network, we use a residual network composed of one fully connected layer, one residual block, and one fully connected layer. The residual network receives X, and e as input. The input is first concatenated with the context and passed into the network. The residual block passes the input through two fully connected layers and returns the sum of the input and the output. Finally the last layer outputs the parameters and we use 5 layers of affine coupling flows with LU layer between them.
In our case, we use samples from the belief distribution as the inputs to the actor and value function as an approximation
the actor and value function with belief distribution as input. Calculating Vbb) needs to integrate through both the observation model and state transition model. Our approx- imation makes an assumption like in Qmdp, to avoid integrating through the observation model.
We use a GRU as the recurrent neural network to summary to temporal information. We assume an initial state Su to be a zero vector. After taking action at, we concatenate o, with the previous state and pass it through small MLP to get y/ f(t, at), and use it as the input to the GRU: he+1. 1 GRU (he: yt). We pass žt41 through an MLP to get the base prior belief distribution Po (mean and variance) and then we sample from Po and pass it through sequence of Normalizing Flow to get a sample from PK- For the posterior distribution, we first use a CNN as encoder to encode the observation or into the feature Tt, and then concatenate 24+1 and 2; and pass them through an MLP to get the base posterior belief distribution go and a sequence of Normalizing Flow. Similarly, we finally get a sample 81+1 from (k-
13

---

Under review as a conference paper at ICLR 2022
Training Details We basically adopt the same data buffer updating strategy as in Dreamer Hafner et al. (2019)). First, we use a small amount of seed episodes (S = in DMC experiments) with random actions to collect data. After that, train the model for C/ update steps (C = 100 in DMC experiment) and conduct one additional episode to collect data with small Gaussian exploration noise added to the action. Algorithm 1 shows one update step in C' update steps. After C update steps, we conduct one additional episode to collect data (this is not shown in Algorithm 1). When the agent interacts with the environment, whole trajectory ((0,, rt)?mm) , and add it to data buffer B.
record the observations, actions, and rewards of the
Hyperparameters For DMControl tasks, we pre-process images by reducing the bit depth bits and draw batches of 50 sequences of length 50 to train the FORBES model, value model, and action model models using Adam (Kingma & Ba 2014) with learning rates ag 5 x 10- a1 8 x 10-5 02 8 x 10 respectively and scale down gradient norms that exceed 100. We clip the KL regularizers in Josdel below 3.0 free nats as Dreamer and PlaNet. The imagination horizon is H 15 and the same trajectories used update both action and value models. We compute the TD-A targets with = 0.99 and = 0.95. As for multiple imagined trajectories, we choose N 4 across all environments.
For digit writing experiments in Section 4.1 we decrease the GRU hidden size to be 20, let the base distributions be a 2-dimensional diagonal Gaussian and only use 3 layers of affine coupling flows. For the image processing, we simply divide the raw pixels by 255 and subtract 0.5 to make the inputs lie in 0.5,5 0.5].
A.3 EXTENDED INFORMATION OF BASELINES
For model-free baselines, compare with D4PG (Barth-Maron et al. 2018 distributed extension of DDPG, and A3C (Mnih et al., 2016), the distributed actor-critic approach. D4PG is an improved variant of DDPG (Lillicrap et al. (0155) that uses distributed collection, distributional Q-learning multi-step returns, and prioritized replay. We include the scores for D4PG with pixel inputs and A3C (Mnih al. 2016) with -estor-wise state inputs from DMCcontrol. For model-based baselines, use PlaNet (Hafner et al., 2019b) and Dreamer (Hafner et al., 2019a), two saate-of-the-ane model- based RL. PlaNet (Hafner et al. (2019b) selects actions via online planning without an action model and drastically improves over D4PG and A3C in data efficiency. Dreamer (Hafner et al., 2019a) further improve the data efficiency by generating imaginary rollouts in the latent space.
14

---

Under review as a conference paper at ICLR 2022
A.4
AN ABLATION STUDY ON THE NUMBER OF IMAGINED TRAJECTORIES
Cartpole Swingup
Cheetah Run
Finger Spin
1000
Episode Return
800
000
630
6:0
400
400
200
203
100
0
100
200
300
400
500
150
200
300
400
500
100
200
300
400
500
Hopper Stand
Walker Run
Walker Walk
700
10:0
6.0
Episode Return
8.00
500
600
400
6.00
400
300
400
200
100
100
200
350
400
500
100
500
100
200
300
45o
500
Environment Steps
Environment Steps (K)
Environment Steps
N=1
N=2
N=4
Figure 8: An ablation study on the effect of different N on DMC environments.
To show the effect of N, we adjust the number of imagined trajectories some DMC environments. We choose N = 2, 4 and run 500K environment steps. We run N = 1, with 3 different seeds, and N with different seeds (we use the main DMC experiment results, where N here). The result shows that, in Finger Spin, the performance gain caused by multiple imagined trajectories obvious. In finger spin, there are two objects and their interactions may result complex loco- motion patterns. When the environmental locomotion pattern itself is complex and flexible enough to incorporate diverse possibilities, then using FORBES allows the agent to make diverse predic- tions and using the multiple imagined trajectories technique will further exploit the advantages of FORBES. However, not all environments can show the advantages of multiple imaginations. In other environments, where there's 's only one agent and its behavior is relatively unimodal, larger N does not effectively improve the performance, and different N results in similar performances.
15

---

Under review as a conference paper at ICLR 2022
A.5
EXTENDED RESULTS ON DMC
Cartpole Swingup
Cheetah Run
Finger Spin
600
800
1000
Eplsode Retum
800
700
600
700
810
6:00
400
809
200
200
100
0.0
9.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.5
Hopper Stand
Walker Run
Walker Walk
Episode Return
000
600
6000
403
400
200
200
0.0
0.2
04
08
0.8
0.8
1.0
0.5
0.2
0.4
0.6
0.8
1.5
Environment Steps (M)
Environment Steps (M)
Environment Steps
FORBES
D4PG (1e9 steps)
PlaNet (1e6 steps)
Dreamer
A3C (1e9 steps, proprio)
Figure 9: The training curve on DMC environment for IM environment steps.
We run our algorithm for 1M environment steps and show the curve in Figure We choose 1M environment steps because most of the curves have converged in most of the environments. FORBES achieves higher scores than Dreamer in most of the tasks.
16

---

Under review as a conference paper at ICLR 2022
A.6 AN ABLATION STUDY ON THE MODEL PARAMETERS
Cheetah Run
Finger Spin
Episode Return
600-
600
500
400
200-
100
200
300
400
500
100 Environment Steps (K)
200
300
400
500
Environment Steps (K)
Dreamer
Dreamer+11 1L
Dreamer+22 2L
Figure 10: An ablation study on the effect of adding parameters to Dreamer on two DMC environ- ments.
In this section, we show that having flexible belief state distribution the key to improving perfor- mance, rather than introducing more parameters. Having more parameters do not necessarily mean better performance. Increasing parameters may also make it difficult to converge and negatively affect the sample efficiency.
We add an ablation study that adds more parameters to Dreamer to test the effectiveness having more parameters. We add 1, 2 hidden layer(s) all the MLP in RSSM, and the result shown in Figure TO The results show that simply adding parameters cannot improve the performance.
17

---

Under review as a conference paper at ICLR 2022
A.7
COMPARISON OF ELBO ON FORBES AND RSSM ON DMC
Cheetah Run
Walker Walk
-20
-20
-8
-25
50
100
150
200
250
50
100
150
200
250
Dreamer
FORBES
Figure 11: The ELBO of FORBES and RSSM.
We provide the ELBO in DMC environments and FORBES in Figure and FORBES has higher ELBO.
18

---

Under review as a conference paper at ICLR 2022
A.8 EVIDENCE LOWER BOUND DERIVATIONS
The
variational bound for latent dynamics models ( 81:7 (I:T)
Iee ((51(8t-1, dt- (s(e and variational posterior q ($1. V.T,"I.T) II. q (8t. oge, acr)
follows from importance weighting and Jensen's 's inequality as shown,
log (01:r (1rr) = log (ep
[ st) t=1
T
log Pyleurlour,oer)
p (ot p() 1 00-1 // logs, aes)
2 Eyceurlors.arr7
Σ log (or 81) logp (sc 8t-1, a111) log (8t act
(16)
19

---

Under review as a conference paper at ICLR 2022
A.9 PROOFS OF THEOREM
Theorem T:ehe approximation error of the lower bound is
T
logp(on:T; (uTlarr) JModel = Eax (owr)
DRI(a(SETT Dullll (c,
where p(st | Te) is the true posterior.
Proof:
(g(s Te, oollpp I st-1, 0(-1 01;T
= =f (((0 Te, Te) log
q(s Ti, T1, 0z)
d8g
p (s: 8t-1, a4 0t)
= faloe oi)
(8t Te, 0t)
ds;
pla 1s 1.91 1sp(
p(o
-f f /9(0 Te log Fe, 0,)ss + / g(se Tr, 01) logp(o: a1:r)ds;
- -fols q(ss 01) log(p(s; st a-11 ss)
=logp(o: (air)) / g(s) Ti, Tx, log Tx, ot)sd / g( Tes oi) loglu(s> 11-1, I st)]dse
= =logp(o: (1I7)++ + / a(s) nh log (ss Tr, ot - / g(0t
St-1,0e-1ee
-f q(s) o() log p(o: ss)ds
= =logp(og (I.T)) + / g(s: T()
log log a(s Te: 0t)
ds: f q(st Te, log st)ds; p(s; st-1, a-11)
= p(o (a.T)) + DKL (q (st Te, ot) lp (st S1-1, - E Egtorat [log (o: st)]
(17)
For sequence from time 1 to T, we have
DKL (st I oz) ||p (st $1-1,
t
log. (1.T) (T Σ (logp(or(ss) DRI(e(se/Te or)
(18)
Then we can derive the Theorem with equation 18
logp(on.,,Ie FI.T a1:T)
=B BB
DEL (a (8x Te, oc) |lp (011))1)000)))
+ Eylensrtr..rr)
Σ (logp(or(ss) logr(r((ss) - ( oagps:/)).-..llddiiiioo/dddaaad:
-Bug(eurrlroo
Σ DEL (q (st Te, (s: st-1,
-Mtddel
-Bqx(sir(1r.or))
Σ DKL. (q (st or) (s. Te, o1))
JMtdel
(19)
where p((e) 0t-1, p(s: Te, oz) given the sampled 8t-1 from (1))).
20

---

Under review as a conference paper at ICLR 2022
A.10
MORE RESULTS ON DIGIT WRITING EXPERIMENTS
In this section, we show more results of the predictions on the digit writing experiment in Figure 12
Ground Truth
Model Inputs
RSSM predictions: 030
FORBES predictions: 0g0
6
S
[
6
6
0o
01
05
00
015
Oo
01
05
10
055
7
6
6
6
6
6
6
6
6
6
6
6
6
6
6
6
4
4
6
6
01
05
015
7
7
9
7 3
9
9
7
9
7 9
9 9
7
7
7
9
9
9
015
6
0
0
0
6
0
0 0
(
0 0 5 5
0
0
10
15
5
5
5 9
5
G
5
7 9
5 9
9
9
5
5
5
5
0o
01
05
10
015
4
0
4 4
5 4
4
4
4 4
4
4
4
L
4
4
4
4
4
5
4
0g
01
05
O0
015
Figure 12: Additional prediction results on sequential MNIST of two models.
21

---

Under review as a conference paper at ICLR 2022
A.11 RECONSTRUCTIONS OF THE VISUAL CONTROL TASKS
(a) Cartpole Swing Up
T
T
7
π
(b) Cheetah Run
(c) Finger Spin
(d) Hopper Stand
(c) Walker Run
Walker Walk
Figure 13: The reconstruction results on of FORBES six environments from DeepMind Control Taisa et 2018).
In this section, we show the reconstructions of the visual control tasks during the evaluating phase.
22

---

Under review as a conference paper at ICLR 2022
For each environment, we use 10 frames. The left one is the original picture for each frame. and the right one is the reconstruction picture. The following results Figure 13 show that FORBES can make high-quality reconstructions. The corresponding videos can be found in the supplementary material.
23


Reviews:
[{"review_id": "x0knH3CtBOA", "reviewer": "Reviewer_8rin", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "* Strengths\nThe main contribution of this work is very straightforward: using a richer parametric distribution to learn environment dynamics leads to better sampled trajectories, and an algorithm that can optimize the policy of an agent more efficiently. I think the simplicity of the approach is the strength of the algorithm, and the method is tested and compared on a well known control tasks suite. Results show a consistent improvement in sample effiency of the RL framework proposed over closely related pieces of work and other standard RL algorithms (e.g. D4PG and A3C).\n\n* Weaknesses\n1) While most of this work is well written and the method is clearly described, I don't think the proposed contributions is particularly novel. In particular, the idea of using more flexible distributions for POMDPs have been explored (as mentioned by the authors) in the form of particle filters. While these methods are referenced (e.g. Ma et. al. 2020b, Igl et al. 2018), they are not compared or seriously discussed beyond an unsubstantiated claim that they suffer from insufficient sample efficiency. The authors also fail to cite an even more closely related work that directly addresses (and in my eyes, go beyond what the authors explore in this paper), see Gregor et al. Gregor et al. not only also suggests using normalizing flows, but find that using a convolutional DRAW (VAE model) outperforms flows for learning a model of complex environment dynamics.\n\n2) The second contribution proposed, namely the POMDP RL Framework is an effective way to leverage a model of the environment, as it allows to use imaginary trajectories to train TD(lambda) actor and critics. However, I fail to see exactly how this is different from the referenced Dreamer (Hafner et al. 2019a) work. The authors should make it clear how it differs from Dreamer beyond the use of normalizing flows. If there are no differences, then it should be clearly stated that they make use of an existing method. If there are differences, then I'd like to see further experiments to understand how these differences affect the algorithms performance.\n\n3) The FORBEs algorithm lacks detail on how the data buffer B is created, i.e. how is data collected during training? See PlaNet and Dreamer works for examples of algorithm description wrt. how experience buffers are updated as the agents train. Also, how are tau_t (which includes a_<t and o_<t) obtained if they are apparently not in the batch B?\n\n4) What number N of sampled trajectories are used during training? Furthermore, it would be interesting to see the performance of the method for different number N, to get a sense of how useful it is to increase the number of trajectories and whether there are diminishing returns.\n\nSome other minor issues:\n- 3 seeds seems like a small number and some experiments still have large error bars (e.g. Cheetah Run and Walker Run). I suggest increasing the number of runs from 5 to 10 if possible.\n- When comparing with PlaNet, it would be more informative to also show PlaNets training curve instead of converged performance after an arbitrary number of steps (1e6 steps). It may well be that PlaNet is just as sample efficient but somehow plateaued/converged a bit slower towards the end of training.\n\n**Post-rebuttal update**\n\nI appreciate the willingness of the authors to discuss and address our concerns, and their efforts to improve their work. The paper is now clearer in its contributions, limitations and additional experimental results and details will make it more useful to the community.\nI've revised the score to reflect all these improvements, but I still think this work missed out on an opportunity to truly tackle the challenges of POMDP: handling environment stochasticity and state uncertainty. During our discussion, the authors have agreed that the method makes some assumptions that limit learning policies and values that can reason about state uncertainty (which is now better captured thanks to the normalizing flow) -- it seems that this would be the main reason for using the proposed flow-based model.\nInstead, the environments in the experiments have essentially no stochastic temporal dynamics. I feel the updated paper still largely sidesteps discussing the above limitation.\nAnother minor point, is that I find it a bit puzzling that Dreamer+multiple imagined trajectories seemingly fairs quite a bit worse than Dreamer (Fig. 7) -- this is a bit counterintuitive (I would have thought that at worst, performance would be the same) and I would've liked further explanation.\n\n* References:\n\nGregor et al. Shaping Belief States with Generative Environment Models for RL. NeurIPS 2019", "summary_of_the_paper": "This paper proposes learning belief states for POMDPs using flexible posteriors and prior distributions for the state-space model. The authors state that related work often restrict the distributions to conditional diagonal Gaussians which causes the model to underfit the data. They propose a flow-based model to solve the well known issues of using inflexible distributions (i.e. inability of the aggregate posterior to match the prior) which can cause the model to generate implausible samples. In the context of POMDPs, the authors propose using the learned model of environment dynamics to train an actor and critic by sampling trajectories, resulting in a highly sample efficient RL algorithm. Results show the benefits of using the normalizing flow compared to baselines that only use a Gaussian distribution. ", "summary_of_the_review": "It is well known that VAEs with diagonal Gaussian distributions for the prior and posterior do not fit many datasets well. There are a large number of deep generative models proposed in the literature that are much more flexible and suited to different domains. The authors propose using one of these, normalizing flows, and apply it to learning a model of the environment dynamics in POMDPs settings, and in turn use it to train agents in a sample efficient way. While I think this is a good idea, and experiments indeed show the benefits of using such a model, the contribution in itself is not very novel, as prior work has suggested using richer models (see referenced works on particile filters, or the work of Gregor et al. which also refers to the idea of using flows). I think this paper would be more interesting if it compared a wide range of flexible models, and included a more extensive set of datasets. Alternatively, the contributions would also be strengthened if it provided novel ways to use the learned models when training agents in POMDP tasks. As it stands, I do not recommend accepting the paper due to its limited novelty.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "JwAr2g9If0I", "reviewer": "Reviewer_DDeH", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "Accurate belief tracking is important but often hard for continous-state POMDP.  This paper proposes a normalizing flow based belief approximation, and learning the belief approximation by maximizing an ELBO. An actor and a critic can be jointly learned together with the normalizing flow belief approximation. The experiments show that the normalizing flow based belief achieves higher ELBO than a baseline RSSM on the MNIST sequence dataset, and the actor-critic algorithm achieves good performance on several problems from DeepMind Control Suite.\n\nThe paper is generally easy to follow, and the idea of improving belief tracking by exploiting the power of normalizing flows to represente complex distribution is interesting. \n\nHowever, some important aspects of the idea is not clear to me, making the correctness of the idea unclear. In particular, comments and clarifications on the following are helpful:\n- In Eq. (7) and Theorem 1, expectation should be taken wrt $q(s_{1:T} | \\tau_{T}, o_{T})$. Also state clearly how $q(s_{1:T} | \\tau_{T}, o_{T})$ factorizes.\n- What's the motivation behind the architecture in Fig. 3? \n- How exactly are the observations and actions encoded into $q_{\\psi}(s_{t} | \\tau_{t}, o_{t})$? \n- Where do the transition dynamics $p_{\\psi}(s_{t} | s_{t-1}, a_{t-1})$ and the observation model $p_{\\psi}(o_{t} | s_{t})$ appear in the belief model? If they are used for encoding the actions and observations, then these are not designed to approximate the true transition dynamics and the true observation model, and using them as transition dynamics and observation model in Eq. (12) and Eq. (13) seem unjustified.\n- Eq. (9) actually deviates from Eq. (7), because while Eq. (7) uses a Markov transition model $p(s_{t} | s_{t-1}, a_{t-1})$, Eq. (9) uses a non-Markov transition model $p(s_{t} | \\tau_{t})$.\n- In the RL algorithm, the actor and critic are both functions of the state, not the belief. How is an action for a belief calculated then?\n- While there is a description of the learning objectives for the actor and the critic, the overall RL algorithm is not described. For example, how is training experience collected? Which experiences are stored in the experience buffer in Alg. 1? How often are the belief model, the actor and the critic updated?\n\nI also have some questions and concerns over the experiments \n- For the digit writing task, what's the state? The details of the networks used seem to be missing too.\n- RSSM is from the PlaNet paper, not from the paper of Okada et al. 2020.\n- The proposed belief model and RSSM are compared only on the digit writing task. Can they be compared on problems from DeepMind Control Suite?\n- For DeepMind Control Suite, only results for 6 out of the 20 problems are reported. For some problems, the results for Dreamer do not yet match those reported in the Dreamer paper yet. This may be because the algorithms are only run for 500k steps, while the results in the Dreamer paper are obtained for 5 million steps. It would be helpful to run the new algorithm longer to see whether it can eventually match Dreamer's performance as reported in the orginal paper.\n\n\nMinor comments:\n- pp. 3: usally -> usually\n- In the description of the POMDP model, $s_{1}$ denotes the initial state, and it doesn't produce an observation $o_{1}$, but $o_{1}$ appears in multiple notations. \n\n**Post-rebuttal**\n\nI appreciate the authors' effort in addressing the questions and concerns, and I have increased my score in view of the the improved clarity and more thorough experiments.\n\nA main weakness of the work is that it mainly replaces the latent model in Dreamer with normalizing flows, as noted by other reviewers too. In view of this, the paper could be improved in several aspects.\n* Current version is somewhat unclear and misleading about its novelty, and should provide a proper discussion on what it adds on top of PlaNet and Dreamer. \n* In addition, since the main novelty is in the use of a different latent model, I find the experiments not as thorough as it should be - it should compare with Dreamer on the same tasks, but results are only presented on a subset of them. \n* The authors have addressed most questions on clarification of the work well as far as I can see, but there are a few remaining concerns (e.g., those in other reviewers' post-rebuttal updates), and these need to be discussed/addressed.\n\n", "summary_of_the_paper": "This paper proposes using normalizing flow based belief approximation for continous-state POMDPs and learning the belief approximation via variational inference. In addition, the paper proposes learning an actor-critic reinformcement learning algorithm based on the proposed belief representation.\n", "summary_of_the_review": "Interesting idea, but there are some concerns over the motivation, explanation, and experiments.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "IRDHje6i9ln", "reviewer": "Reviewer_CAct", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "The authors propose and successfully use a more powerful form of distributions for learning belief states in POMDPs. To the best of my knowledge this is the first study on normalising flows in this context (the authors do cite other work using normalising flows in RL which do not apply in this context). The sequential MNIST results are convincing, but the results on DM Control Suite are not so strong - having other domains where the benefit of normalising flows are more pronounced would make this more attractive than simply sticking to simpler isotropic Gaussians. FORBES also introduces more learned parameters, but I doubt that this would contribute towards a change in performance.\n\nThe major weakness is the significance/novelty of this paper. Is it fair to say that this is Dreamer + normalising flows latent space + multiple imagined trajectories? Knowing the literature and from the experiments it would seem so, but the paper is written in a way that makes it unclear how exactly FORBES differs from prior work. On a side note, it is interesting that Deamer + multiple imagined trajectories does not seem to work that well - possibly due to the unimodal latent space. Further to this, while the DreamerV2 paper is cited with its argument for multimodality, DreamerV2 and its discrete latent space (which is also much more flexible than an isotropic Gaussian) are not discussed - which would seem an omission given the supposed benefit of normalising flows over isotropic Gaussians.\n\nOn a minor note, there are several spelling mistakes throughout, e.g., \"beilef\", \"poorfs\".\n\n**Post-rebuttal update:** I agree with the authors that FORBES is more general than an alteration to Dreamer, and there are theoretical contributions that are of relevance to the community. However, given my awarereness of the literature and the concerns of the other reviewers, I would hesitate to move my recommendation much higher.\n\n**Post-rebuttal update 2:** The authors have put considerable work into addressing the concerns of the other reviewers, and it is now a more solid paper. Although I would appreciate a more thorough discussion of related work (perhaps in supplementary material), e.g. of why certain baselines were not included, I think it is worth accepting given the updates.", "summary_of_the_paper": "Belief states are a common solution to learning in POMDPs, and in recent years people have applied deep recurrent generative models to approximating belief states in complex POMDPs. In this paper, the authors go beyond the isotropic Gaussian assumption used in most deep learning works and use normalising flows to represent more flexible and multimodal distributions. The authors demonstrate that their approach is better at modelling multimodal belief states and matches/outperforms Dreamer, which is a natural baseline for their method.", "summary_of_the_review": "The authors presented a relatively simple improvement to solving POMDPs by replacing isotropic Gaussians with normalising flows. While this does indeed work, the technical and empirical contributions are reasonably significant, and so I would recommend this for an accept.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "6Puer7jl_kX", "reviewer": "Reviewer_9het", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "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": "", "main_review": "Positives:\n- The paper demonstrates that normalizing flows can improve control performance and sample efficiency in MBRL over approaches working with a Gaussian assumption.\n- While VAEs and other generative models have been applied to belief inference, this appears to be the first work applying NFs. It is not groundbreaking novelty, but seems like a solid evaluation of the idea.\n- One of the advantages of POMDP planning is the ability to reason about future uncertainty. The sequence MNIST result where the proposed method is able to predict either 3's or 7's is quite exciting, since it shows the method has some capability for such reasoning. This aspect would be interesting to see even more details about.\n\nNegatives:\n- There are several missing baselines and comparisons in the empirical results, and some lack of clarity about parameter choices (details below).\n- Throughout, the paper seems to conflate in writing and symbol use the belief state and the hidden state. This is especially prevalent in Sect. 3.2., but should be revised overall. For conceptual clarity, the paper should in my opinion be careful to separate 1) hidden states or samples of hidden states (typically denoted s in POMDP literature) 2) belief states, or distributions over the hidden state, and their approximations (typically denoted b in POMDP literature).\n\n\nDetailed comments:\n\n\nThere are a couple of omissions in the evaluation, which should be justified or a comparison added:\n- The paper does not compare to PlaNet-Bayes (Okada et al., 2020) in the control tasks. This is the previous state-of-the-art in the area. Since the paper does compare to PlaNet-Bayes on the MNIST setting, I am wondering if it would not be possible and reasonable to compare on the control tasks as well.\n- Igl et al. (2018) \"Deep Variational Reinforcement Learning for POMDP\" is also very closely related to the proposed method. A comparison here would also be welcome, or some arguments why a comparison is not needed.\n- In some works addressing partially observable RL, belief states are inferred by using a recurrent neural network directly on the action-observation history. It would be nice to see a simple RNN baseline in the control tasks.\n- A missing ablation is how the number of imagined trajectories N affects the results. This should ideally be done in a domain where Gaussian or other simplifying assumptions on belief inference do not work, otherwise even with a single trajectory results could be OK. I also could not find information on how many trajectories were used in the experiments.\n- As mentioned in the strengths, it would be great to have some further details on the method's capability for reasoning about future uncertainty. For example, is the proposed method better than RSSM always, or especially for long prediction horizons?\n\n\nOther comments & minor:\n- Theorem 2 seems to not seem to add much to the paper, and does not merit a full reproduction. A citation could be sufficient.\n- Figure 3 needs some improvements. It is not really clear what the colors here are indicating, or what happens at the intersections of the arrows.\n- In 2.1., it would be helpful to define the POMDP objective considered. Is it maximization of expected sum of discounted rewards?\n- \"Dream to control\" is duplicated in bibliography\n\n**Post-rebuttal comments:**\nI thank the authors for answering my questions and concerns regarding the paper. The author rebuttal does a good job addressing my concerns on presentation and ablations. I am raising my score to reflect this. However, I also agree with many of the concerns raised by other reviewers regarding the significance of the work.\n\n", "summary_of_the_paper": "This work addresses model-based reinforcement learning (MBRL) in partially observable visual control tasks. The principal novelty of the work is the use of normalizing flows (NFs) in the belief inference model. The paper derives the evidence lower bound (ELBO) for the proposed inference model, and presents a related RL framework. Empirically, the proposed model's inference capabilities are tested in the sequential MNIST domain, and its control performance is evaluated in six continuous control problems. The results indicate improved belief fidelity, and improved sample complexity for learning to solve the control tasks.", "summary_of_the_review": "This paper proposes to use normalizing flows as belief representations in a POMDP, addressing drawbacks of VAEs or diagonal Gaussian approximations in representing multi-modal beliefs. While not highly original, the work is technically sound, and provides the necessary empirical evidence to merit acceptance. Nevertheless, I encourage the authors to address the concerns raised by other reviewers as well, and to provide comparisons to related works the authors also agree are highly relevant.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Meta-review: Summary: this is a difficult paper to meta-review, since it contains some insightful ideas and interesting experiments, while it also unfortunately contains omissions, confusions, and places where clarity is lacking (see below). One consistent theme is that the paper is too dismissive of prior work; the exposition is not as clear as it should be about what aspects of FORBES are present in previous papers, it uses too broad a brush to describe prior methods (resulting in too-general statements about what these methods can't do), and it skips important chunks of the extensive literature on POMDP belief representation and tracking. As a result, the paper doesn’t do a good job concisely and accurately stating its contribution; there is still reasonable concern about how significant this contribution is. On the other hand, the experimental results for FORBES are interesting; the new method seems to represent a better combination of techniques than at least many existing works, at least to the resolution of the experiments’ statistical power. So the end question is whether interesting experimental results and a new combination of techniques are enough to outweigh the problems outlined above. In the end we believe that the correct outcome is rejection; but we have every expectation that a future version of the paper will resolve the difficulties outlined here and will appear in a future conference.

A brief note about the discussion: the original scores for this paper were lower. While some reviewers raised their score later in the discussion, a thorough reading of the discussion and the revised paper indicates that a substantial fraction of the issues leading to the lower scores still remain.

More details:

There is a lot of prior work on tracking belief states, which should be cited more thoroughly. The paper's intro makes it sound like diagonal Gaussians were the only previous alternative. At least, the intro should cite older work on MCMC methods like particle filters (e.g., Thrun’s book Probabilistic Robotics, or Arnaud Doucet’s work), and prior deep-net papers that attempt non-Gaussian representations, even if these don’t perform as well as hoped (see below for examples). It is also important to compare to RKHS representations of beliefs, such as Nishiyama, Boularias, Gretton, Fukumizu 2012; these handle multimodality, and can behave similarly to deep nets if they use the neural tangent kernel. Accurately comparing to prior work is one of the most important functions of a paper, so it doesn’t make sense to be unfairly critical of prior work or to skip it.

The paper is also unclear about the effects of Gaussian distributions at different places in a variational approximation. Because of this lack of clarity, the criticisms it levels at previous variational methods seem to be true only of some of them. 

In particular, the introduction should distinguish between two uses of Gaussian approximations: first for the belief itself, and second for the distribution of observations given a belief. Some prior works make only one of these approximations. For example, a non-Gaussian distribution used as a belief state can predict multi-modal future behaviors, even if we approximate observations under a given belief as Gaussian.

The introduction should also distinguish between two common places that a Gaussian could enter into a variational approximation: at the input or at the output of a network. A Gaussian latent at the input of a variational network (even if it has diagonal covariance) can result in a highly non-Gaussian output distribution, while Gaussian noise added at the end will (if it is the only noise) lead to a Gaussian output. Again, some of the statements in the intro apply only to the latter use of a Gaussian, while some prior work focuses on the former use.

There is an important conceptual confusion in the paper about what it means to have a multimodal belief state: the paper presents the true belief as an inherent property of an environment, while in fact it is a property of an environment *model*. So, there can be two different equally accurate models of the same environment which differ in the belief representation; a simple example would be to use either a continuous state whose components are joint angles, or a discrete state obtained by finely discretizing this continuous one. In the first case the belief would be a distribution over the continuous space, while in the second it would be a categorical distribution (a point in a simplex). A consequence of such a difference is that beliefs can be multimodal in one representation and not another.

The importance of this confusion is that, since we are asking our network to learn a belief state, the learning process could potentially favor representations that lead to unimodal beliefs — so it’s not clear theoretically that forcing a unimodal belief representation is necessarily a disadvantage. The paper presents the situation as if the disadvantage is forced by theory, while instead the argument should be based on experiments: e.g., one could try to show that unimodal representations, even if given a higher latent dimension to work with, aren’t empirically able to capture the same information.

Some interesting prior deep-net POMDP papers that might need better discussion:
* Han, Doya, and Tani ICLR 2020 (which isn’t cited here) puts the Gaussian latent variable as an input to the network for predicting beliefs (eq 2), resulting in a possibly highly non-Gaussian output representing the belief.
* Tschiatschek et al, 2018 (also not cited) uses a Gaussian *mixture* as the variational distribution to approximate beliefs, again allowing multimodality.
* Igl et al. 2018 (which is cited only late in the current paper, and basically dismissed) uses a deep version of particle filters to allow non-Gaussian distributions for both beliefs and observations.
* More work that is potentially relevant but not adequately compared (even if briefly cited): Gregor et al. (2019), DreamerV2, Ha & Schmidhuber’s World Models. Each of these makes at least some choices to try to handle at least some kinds of multimodality, so a clear explanation of differences that avoids the confusions mentioned above would be very helpful.
* In general, the results of the search “variational encoder POMDP” seem to include a number of papers not cited in the current paper; another useful search is “normalizing flow POMDP"

Finally, in the experiments section, the paper needs to correctly report the reliability of its conclusions. In some places (e.g., Fig. 5) there’s no mention of reliability or repeatability of conclusions; the paper just says that its evidence “support[s] the claim that FORBES can better capture the complex belief states”. In other places (e.g., Fig. 6, 7), the paper displays uncertainty representations based on only a few replications of an experiment (e.g., 3 seeds for Fig. 6, or 5 seeds when a reviewer requested extra experiments). The corresponding uncertainty estimates almost certainly are strongly biased too low (too certain); e.g., three runs would have less than a 50% chance of even seeing failure modes that happen with probability as high as 20% (0.8^3 = 0.512 > 0.5), meaning that the estimated standard deviation could be almost arbitrarily badly biased downward. To be clear, experiments with few replicates can still be highly useful and informative, and it’s true that some experiments are too expensive to run many times; but in such cases the paper should add appropriate caveats to its conclusions. For example, instead of reporting the sample standard deviation based on a normal model, the paper could report a confidence interval based on a more robust model or test, such as a Wilcoxon test. (To illustrate the difference, confidence intervals at typical significant levels like p=0.05 would be vacuous (infinitely wide) under Wilcoxon with 3 seeds, but much-weaker p-values would still yield non-vacuous intervals.)

A few smaller questions:

The authors added a nice ablation study to compare to Dreamer; this is great to see. It would be good to discuss the connection to earlier methods such as PlaNet and Dreamer at places where the current method is similar or different (e.g., different from Dreamer in the belief state representation in sec 2.2, but similar in the RL framework in section 3.2). These comparisons would aid in the reader’s understanding of what is new in FORBES.

An unusual feature of FORBES is that the variational approximation to the belief at time t+1 is not a function of the belief at time t. Instead the belief inference network q_{\psi,\theta} takes as input the entire past trajectory, uses convolution and recurrence to reduce the variable-length input to fixed dimension, and passes this fixed-dimension representation through a normalizing flow mapping. It would be interesting to discuss the reason for this design decision. In particular, it seems like it would inhibit tracking — i.e., it could be hard to propagate information from one belief distribution to the immediate next one, particularly if there are a few unlikely observations scattered through a trajectory.

A minor point for clarity: in Fig 1 it's unclear what distributions the white and gray triangles refer to. They don't seem to correspond to a natural belief state: instead maybe they incorporate three simultaneous observations from the same starting belief? Correctly intersecting beliefs is an important issue though, so at a high level the point that the figure is trying to make fits well.

Another point for clarity: “there always exists a diffeomorphism that can turn one well-behaved distribution into another”: this is true for some definition of ”well-behaved”, but it’s misleading to say it this way. E.g., it is not true if the distributions in question can have atoms, or differ in dimension or topology; these exceptions are unfortunately important cases that do come up in practice.
Decision: Reject