F LOW - BASED R ECURRENT B ELIEF S TATE L EARNING FOR POMDP S
Anonymous authors
Paper under double-blind review
A BSTRACT
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 F l O w-based R ecurrent BE lief S tate model (FORBES), which incorporates normalizing ﬂows 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 efﬁciency.
I NTRODUCTION
Partially Observable Markov Decision Process (POMDP) ( ˚Astr¨om, 1965) provides a principled and generic framework to model real world sequential decision making processes.
Unlike Markov Decision Process (MDP), the observations of a POMDP are generally non-Markovian.
Therefore, to make optimal decisions, the agent needs to consider all historical information, which is usually intractable.
One effective solution is to obtain the belief state.
The belief state is deﬁned 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 ﬁnite discrete space with a known model.
In many real world problems, however, the underlying model remains unknown, and the state space is large and even continuous.
With the recent advances of deep learning technologies, a branch of works have been proposed to learn the belief states of POMDPs with unknown model and continuous state space (Krishnan et al., 2015; Gregor et al., 2019; Lee et al., 2020; Hafner et al., 2019b;a; 2021).
These works solve the belief state learning problem by sequentially maximizing the observation probability at each timestep using the variational inference and achieve the state-of-the-art performance on many visual-motor control tasks (Hafner et al., 2019a; Zhu et al., 2020; Okada et al., 2020; Ma et al., 2020a).
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 Gaussians.
This approximation imposes strong restrictions and is problematic.
As shown in Figure 1, the blue area denotes the unobservable state space of the POMDP.
Given the past information τ , the agent maintains a prior distribution of the state s , denoted as p ( s | τ ) (the distribution in white).
Each colored distribution corresponds to the belief state after receiving a different new observation o , denoted as the posterior distribution q ( s | τ, o ) .
Consider an example of the true beliefs as shown in Figure 1(b), with their Gaussian approximations shown in Figure 1(a).
The approximation error of Gaussian distributions will easily result in problems of intersecting belief which leads to a mixed-up state (e.g., the white triangle), and empty belief, which leads to a meaningless state (e.g., the grey triangle).
This also explains the poor reconstruction problems in interactive environments observed by Okada & Taniguchi (2021).
Furthermore, as mentioned in Hafner et al.
(2021), the Gaussian approximation of belief states also makes it difﬁcult to predict multi-modal future behaviours.
Therefore, it is preferable to relax the Gaussian assumptions and use a more ﬂexible family of distributions to learn accurate belief states as shown in Figure 1(b).
In this paper, we propose a new method called F l O w-based R ecurrent BE lief S tate model (FORBES) that is able to learn general continuous belief states for POMDPs.
FORBES incorporates Normalizing Flows (Tabak & Turner, 2013; Rezende & Mohamed, 2015;
Dinh et al., 2017) into the variational inference step to construct ﬂexible belief states.
In experiments, we show that
FORBES allows the agent to maintain ﬂexible 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 on challenging visual-motor control tasks (DeepMind
Control Suite, Tassa et al.
The results show the efﬁcacy of FORBES in terms of improving both performance and sample efﬁciency.
Our contributions can be summarized as follows:
• We propose FORBES, the ﬁrst ﬂow-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 to learn ﬂexible belief states that enable multi-modal predictions as well as high quality reconstructions and help improve both performance and sample efﬁciency for challenging visual-
motor control tasks.
P RELIMINARIES 2. P ARTIALLY O BSERVABLE M ARKOV D ECISION P ROCESS
Formally, a Partially Observable Markov Decision Process (POMDP) is a 7-tuple ( S , A , T, R, Ω , O, γ ) , where S is a set of states, A is a set of actions, T is a set of conditional transition probabilities between states, R is the reward function, Ω is a set of observations, O is a set of conditional observation probabilities, and γ is the discount factor.
At each timestep t − 1 , the state of the environment is s t − 1 ∈ S .
The agent takes an action a t − 1 ∈ A , which causes the environment to transit to state s t with probability T ( s t | s t − 1 , a t − 1 ) .
The agent then receives an observation o t ∈ Ω which depends on the new state of the environment s t with probability
O ( o t | s t ) .
Finally, the agent receives a reward r t − 1 equal to R ( s t − 1 ) .
{+The agent’s goal is to maximize the the expected sum of discounted rewards E [ (cid:80) ∞ t =
+}
γ t r t ] . Such a POMDP model can also be described as a probabilistic graphical model (PGM) as shown in Figure 2.
After having taken action a t − 1 and observing o t , an agent needs to update its belief state, which is deﬁned as the probability distribution of the environment state conditioned on all historical information:
[Equation]
where τ t = { o 1 , a 1 , . . .
, o t − 1 , a t − 1 } .
2. N ORMALIZING F LOW
Instead of using the Gaussian family to approximate the prior and posterior belief distributions, we believe it is more desirable to use a family of distributions that is highly ﬂexible and preferably ﬂexible enough to describe all possible true belief states.
Therefore, we use Normalizing Flows (Tabak & Turner, 2013; Rezende & Mohamed, 2015) to parameterize those distributions.
Rather than directly parameterizing statistics of the distribution itself, Normalizing Flows model the transformations, or the “ﬂow” progress, needed to derive such a distribution.
More speciﬁcally, it describes a sequence of invertible mappings that gradually transform a relatively simple probability density to a more ﬂexible and complex one.
Let f θ : R D
→ R D to be an invertible and differentiable mapping in state space parameterized by θ .
Given a random variable x ∈ R D with probability distribution p ( x ) , we can derive the probability of the transformed random variable z = f θ ( x ) by applying the change of variable formula:
[Equation]
To construct a highly ﬂexible family of distributions, we can propagate the random variable at beginning z 0 through a sequence of K mappings and get z K = f θ K ◦ f θ K
[Equation] 1 ◦ · · · ◦ f θ [Equation] z 0 ) with the probability
[Equation]
Given a relatively simple distribution of z 0 , say, Gaussian distribution, by iteratively applying the transformations, the ﬂow is capable of representing a highly complex distribution with the probability that remains tractable.
The parameters θ 1 , . .
, θ K determine the transformations of the ﬂow.
An effective transformation that is widely accepted is afﬁne coupling layer (Dinh et al., 2017;
Kingma & Dhariwal, 2018; Kingma et al., 2017).
Given the input x ∈ R D , let s and t stand for scale and translation functions which are usually parameterized by neural networks, where s, t : R k
→ R D − k , k < D .
The output, y , can be viewed as a concatenation of its ﬁrst k dimensions y 1: k and the remaining part y k
+1:
D :
[Equation]
where (cid:12) denotes the element-wise product (see details about afﬁne coupling layer in Appendix A.1).
F LOW - BASED R ECURRENT B ELIEF S TATE L EARNING 3. F LOW - BASED R ECURRENT B ELIEF S TATE MODEL
We propose the F l O w-based R ecurrent BE lief S tate model (FORBES) which learns general continuous belief states via normalizing ﬂows under the variational inference framework.
Speciﬁcally, the FORBES model consists of components needed to construct the PGM of POMDP as shown in
Figure 2:
[Equation]
In addition, we have a belief inference model q ( s t | τ t , o t ) to approximate the true posterior distribution p ( s t | τ t , o t ) , where τ t = { o 1 , a 1 , . . .
, o t − 1 , a t − 1 } is the past information.
The above components 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 et al., 2016):
[Equation]
Detailed derivations can be found in Appendix.
A.9.
In practice, the state transition model, observation model, reward model, and belief inference model can be represented by stochastic deep neural networks parameterized by ψ :
[Equation]
where their outputs usually follow simple distributions such as diagonal Gaussians.
The parameterized belief inference model q ψ ( s t | τ t , o t ) acts as an encoder that encodes the historical information using a combination of convolutional neural networks and recurrent neural networks.
In FORBES we provide special treatments for the belief inference model and the state transition model to represent more complex and ﬂexible posterior and prior distributions.
As shown in Figure 3(a), the input images o 1: t and actions a 1: t − 1 are encoded with q ψ ( s t | τ t , o t ) (the blue and the red path).
Then our ﬁnal inferred belief is obtained by propagating q ψ ( s t | τ t , o t ) through a set of normalizing ﬂow mappings denoted f θ K ◦ · · · ◦ f θ 1 to get a representative posterior distribution q ψ,θ ( s t | τ t , o t ) .
For convenience, we denote q 0 = q ψ and q K = q ψ,θ .
On the other hand, o 1: t −and a 1: t − 2 are encoded with q ψ ( s t − 1 | τ t − 1 , o t − 1 ) (the blue path), then the state transition model is used to obtain the prior guess of the state p ψ ( s t | τ t ) =
E q ψ ( s t − 1 | τ t − 1 ,o t − 1 ) [ p ψ ( s t | s t − 1 , a t − 1 )]
(the green path).
Then our ﬁnal prior is obtained by propagating p ψ ( s t | τ t ) through another set of normalizing ﬂow mappings denoted f ω K ◦· · ·◦ f ω 1 to get a representative prior distribution p ψ,ω
( s t | τ t ) .
For convenience, we denote p
= p ψ and p K = p ψ,ω .
Then as shown in Figure 3(b), we can sample the initial state s t (the yellow and purple triangles) from the belief states q K ( s t | τ t , o t ) .
For each sampled initial state, we can use the state transition model to predict the future states
ˆ s t
+ h given the future actions a t : t + h − 1 , and then use the observation model to reconstruct the observations ˆ o t
+ h , where h is the prediction horizon.
[Figure]
With the above settings, we can substitute the density probability inside the KL-divergence term in
Equation 7 with Normalizing Flow:
[Equation]
[Equation]
where p K ( s t | s t − 1 , a t − 1 )
=
p K ( s t | τ t ) given the sampled s t − 1 from q K ( s 1: t | τ t , o t ) .
s t,k is the state variable s t transformed by k layers of normalizing ﬂows, and s t, 0 = s t .
To further demonstrate the properties of FORBES, we provide the following theorems.
Theorem 1 The approximation error of the log-likelihood when maximizing the J Model (the derived
ELBO) deﬁned in Equation 7 is:
[Equation]
t , o t )) (cid:35) (10) where p ( s t | τ t , o t ) denotes the true belief states.
Detailed proofs can be found in Appendix.
A.9.
Theorem 1 suggests that, when the learning algorithm maximizes the J Model (the derived ELBO), then the D KL terms in the right-hand side are minimized, which indicate the KL-divergence between the learned belief states q ( s t | τ t , o t ) and the true belief states p ( s t | τ t , o t ) .
Clearly, if p ( s t | τ t , o t ) is a complex distribution and q ( s t | τ t , o t ) is chosen from a restricted distribution class such as diagonal Gaussian, then when the algorithm maximizes the J Model (the derived ELBO), there will still be a potentially large KL-divergence between the learned and the true belief states.
Therefore, naturally there raises the problem that is normalizing ﬂow a universal distributional approximator that is capable of accurately representing arbitrarily complex belief states, so the KL-divergence terms in the right-hand side of Equation (10) can be minimized to approach zero?
The answer is yes for a wide range of normalizing ﬂows.
To be speciﬁc, Teshima et al.
(2020) provides theoretical results for the family of the ﬂow used in
FORBES.
In fact, there always exists a diffeomorphism that can turn one well-behaved distribution to another.
Besides the aforementioned afﬁne coupling ﬂow, many works show the distributional universality of other ﬂows (Kong & Chaudhuri, 2020; Huang et al., 2018).
Ideally, the universal approximation property of the ﬂow model q K ( s t | τ t , o t ) allows us to approximate the true posterior p ( s t | τ t , o t ) 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 a higher ELBO than previous works.
We verify this statement in section 4.1.
3. POMDP RL FRAMEWORK BASED ON FORBES
To show the advantage of the belief states inferred by the FORBES model compared to the existing belief inference method in visual-motor control tasks, we design a ﬂow-based belief reinforcement learning algorithm for learning the optimal policy in 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 learned belief states and thus beneﬁt from the accurate representations learned by the FORBES model.
The critic v ξ ( s τ ) aims to predict the discounted sum of future rewards that the actor can achieve given an initial state s t , known as the state value E (cid:16) (cid:80) t + ∞ τ =
t γ τ − t r τ (cid:17) , where ξ denote the parameters of the critic network and H 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 state.
In order to trade-off the bias and the variance of the state value estimation, we use the more general TD( λ ) target (Sutton & Barto, 2018), which is a weighted average of n-step returns for different horizons and is deﬁned as follows:
[Equation]
To better utilize the ﬂexibility belief states from FORBES, we run the sampling method multiple times to capture the diverse predictions.
Speciﬁcally, we sample N states from the belief state given by FORBES and then rollout trajectories of future states and rewards using the state transition model and the reward model.
Finally, we train the critic to regress the TD( λ ) target return using a mean squared error loss:
[Equation]
where sg ( · ) is the stop gradient operation.
The actor a τ ∼ q φ ( a τ | s τ ) aims to output actions that maximize the prediction of long-term future rewards made by the critic and is trained directly by backpropagating the value gradients through the sequence of sampled states and actions, i.e., maximize:
[Equation]
We jointly optimize the model loss J Model with respect to the model parameters ψ , θ and ω , the critic loss J Critic with respect to the critic parameters ξ and the actor J Actor loss with respect to the actor parameters φ using the Adam optimizer with different learning rates:
[Equation]
α 2 J Model ( ψ, θ, ω ) (14) where α 0 , α 1 , α 2 are coefﬁcients for different components, and we summarize the whole framework of optimizing in Algorithm 1.
Input: buffer B , imagination horizon H , interacting step T , batch size B , batch length L , number of trajetories N .
Draw B data sequences { ( o t , a t , r t ) } k + L t = k from B
Infer belief state q K ( s t | τ t , o t ) .
for i = 1 , . . .
, N do
Rollout imaginary trajectories { ( s i,τ , a i,τ ) } t
+ H τ = t with belief transition model.
end for
For each s i,τ , predict rewards p ψ ( r i,τ | s i,τ ) and values v φ ( s i,τ ) (cid:46)
Calculate imaginary returns
Update θ, ω, ξ, φ, ψ using Equation (7), (9), (12), (13) and (14) (cid:46)
Optimize the model, critic and actor
E XPERIMENTS
Our experiments evaluate FORBES on two image-based tasks.
We ﬁrst demonstrate the belief learning capacity on a digit writing task in Section 4.1, and show that FORBES captures beliefs that allow for multi-modal yet precise long-term predictions as well as higher ELBO.
For large-scale experiments, we test the proposed POMDP RL framework based on FORBES in Section 4.2.
The results of multiple challenging visual-motor control tasks from DeepMind Control Suite (Tassa et al., 2018) show that FORBES outperforms baselines in terms of performance and sample efﬁciency.
In Section 4.3, we further provide ablation studies of the multiple imagined trajectories technique used in our method.
4. D IGIT W RITING T ASKS
In this experiment, we validate the capacity of FORBES by modelling the partially observable sequence with visual inputs.
We adopt the MNIST Sequence Dataset (D. De Jong, 2016) that consists of sequences of handwriting MNIST digit stokes.
This problem can be viewed as a special case of POMDP, whose action space is Ø and rewards remain 0 .
Such a problem setting separates the belief learning and policy optimizing problem and allows us to concentrate on the former one in this section.
We convert the digit stroke to a sequence of images of size 28 × 28 to simulate the writing process.
At time step t , the agent can observe o t that has already written t pixels, and we train the agent maximizing J Model in Equation 7 except for the reward reconstruction term.
[Figure]
As shown in Figure 4, we randomly select three digits as 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-the-art 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 { o 1 , o 2 , · · · , o 15 } to the model, and the columns in grey exhibit a part of the inputs.
Then we sample several states from the inferred belief state and rollout via the learned state transition model (Equation (6)) for 15 steps and show the reconstruction results of the predictions.
As shown in the blue and green columns on the right of Figure 4, 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.
However, FORBES can make precise yet diverse predictions.
Each prediction is clear and distinct from other digits.
Given the beginning of the digit 7, FORBES successfully predicts both 7 and 3 since they have a similar beginning.
The results can be partially explained via the mixed-up belief and the empty belief as shown in Figure 1, which support the claim that FORBES can better capture the complex belief states.
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 a tighter ELBO, which veriﬁes the theoretical results in 3.1.
Details of the network can be found in Appendix A.2.
V ISUAL - MOTOR CONTROL TASKS
We experimentally evaluate the performance of FORBES on Reinforcement Learning on a variety of visual-motor control tasks from the DeepMind Control Suite (Tassa et al., 2018), illustrated in
Figure 6.
Across all the tasks, the observations are 64 × 64 × 3 images.
These environments provide different challenges.
The Cartpole-Swingup task requires a long planning horizon and memorizing the state of the cart when it is out of view; Finger-Spinning includes contact dynamics between the ﬁnger and the object; Cheetah-Run exhibits high-dimensional state and action spaces; the WalkerWalk and Walker-Run are challenging because the robot has to learn to ﬁrst stand up and then walk;
Hopper Stand is based on a single-legged robot, which is sensitive to the reaction force on the ground and thus needs more accurate control.
As for baselines, we include the scores for A3C Mnih et al.
(2016) with state inputs (1e9 steps), D4PG Barth-Maron et al.
(2018) (1e9 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 efﬁciency on challenging visual-motor control tasks.
As illustrated in Figure 6, FORBES achieves [Figure] higher scores than Dreamer (Hafner et al., 2019a) in most of the tasks and achieves better performance 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
[Figure]
while FORBES can better capture the general belief states.
Furthermore, multiple imagined trajectories can better utilize the diversity in the rollout.
Therefore
[Figure]
the inner coherency within the model components allows the agent a better performance.
We further discuss the role of multiple imagined trajectories in the next section
[Figure]
4. A BLATION S TUDY
[Figure]
In order to verify that the outperformance of FORBES is not simply due to increasing the number of imagined trajectories [Figure] 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 7, 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 in FORBES.
The Gaussian assumptions lead to the lack of trajectory diversity, so that increasing the number of imagined trajectories will not effectively help.
R ELATED W ORK
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 approximate the value function using point-based methods (Kurniawati et al., 2008;
Shani et al., 2013) or using Monte-Carlo sampling in the belief space (Silver & Veness, 2010; Kurniawati & 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., 2012; Indelman et al., 2015) or using random sampling in the belief space (Agha-Mohammadi et al., 2014; Hollinger &
Sukhatme, 2014).
However, the classical POMDP methods mentioned above are all based on 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 highdimensional 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 et al., 2019b;a).
Among which, the recurrent state-space model (RSSM) based methods (Hafner et al., 2019b;a) provide a 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 on 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, including mode collapse, posterior collapse and object vanishing in reconstruction(Bowman et al., 2016;
Salimans et al., 2015; Okada & Taniguchi, 2020).
A few works propose particle ﬁlter based methods that use samples to approximate the belief states (Ma et al., 2020b; Igl et al., 2018).
However, they suffer from insufﬁcient sample efﬁciency 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 computational time cost of the log determinant is O
( D 3 ) .
Thus most previous works choose to make the computation more tractable.
Rezende & Mohamed (2015); van den Berg et al.
(2019) propose to use restricted functional form of f .
Another choice is to force the Jacobian of f 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.
(2015; 2017); Kingma & Dhariwal (2018) propose the coupling method to make the Jacobian triangular and ensure the forward and inverse can be computed with a single pass.
The applications of
NF include image generation (Ho et al., 2019; Kingma & Dhariwal, 2018), video generation (Kumar et al., 2019) and reinforcement learning (Mazoure et al., 2020; Ward et al., 2019; Touati et al., 2020).
C ONCLUSION
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
F l O w-based R ecurrent BE lief S tate model (FORBES) that can learn general continuous belief states by incorporating normalizing ﬂows into the variational inference framework and then effectively utilize 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 ﬁnal performance of our extended POMDP RL framework 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 reconstructions.
Experiments on challenging visual control tasks demonstrate that the ﬂow-based belief state inferred by FORBES improves the sample efﬁciency and ﬁnal performance of downstream model-based reinforcement learning algorithms.
General belief inference plays a vital role in solving the POMDP, and our method paves a way towards it.
In the future, we will explore further approaches to improve the accuracy of belief states inference, such as combining contrastive learning and using advanced network architectures such as transformers to build normalizing ﬂows.
REPRODUCIBILITY
We implement FORBES based on the Pytorch version of Dreamer, provided by “ https://github.com/yusukeurakami/dreamer-pytorch ” and our code will be released after the author notiﬁcation in “ https://anonymous.4open.science/r/
Implementation-of-FORBES-E45B ”.
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.
A A PPENDIX A. D ETAILS OF AFFINE COUPLING LAYER FOR NORMALIZING FLOW
In this section, we will introduce the details about the afﬁne coupling layer (Dinh et al., 2017).
In the forward function, we split the input x ∈ R D into two parts according to the dimension: x = [ x 1: k , x k +1: D ] .
Then, we let the ﬁrst part x 1: k stay identical, so that the ﬁrst k dimensions in the output y ∈ R D is y 1: k
= x 1: k .
After that, we use the identical part as the inputs to determine the transform parameters.
In our case, we deﬁne two neural network s, t : R k
→ R D − k , which stand for scale and translation functions.
They receive x 1: k as inputs and output the afﬁne parameters.
As in (Dinh et al., 2017), the second part can be derived by:
[Equation]
Finally, the output, y is the concatenation of the two parts: y =
[ y 1: k , y k
+1: D ] .
The afﬁne coupling layer is an expressive transformation with easily-computed forward and reverse passes.
The Jacobian of afﬁne coupling layer is a triangular matrix, and its log determinant can also be efﬁciently computed.
A. H YPER P ARAMETERS AND IMPLEMENTATION DETAILS
We use the convolutional and deconvolutional networks that are similar to Dreamer(Hafner et al., 2019a), a 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 a and c as input.
The input is ﬁrst 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 afﬁne coupling ﬂows with a LU layer between them.
We use a GRU as the recurrent neural network to summary to temporal information.
We assume an initial state s 0 to be a zero vector. After taking action a t , we concatenate a t with the previous state s t and pass it through a small MLP to get y t = f ( s t , a t ) , and use it as the input to the GRU: h t +1 , z t +1 = GRU ( h t , y t ) .
We {+pass z t +1 through an MLP to get the base prior belief distribution p 
 (mean and variance) and then we sample from p 0 and pass it through a sequence+} of Normalizing Flow to get a sample from p K .
For the posterior distribution, we ﬁrst use a CNN as encoder to encode the observation o t into the feature x t , and then concatenate z t +1 and x t and pass them through an MLP to get the base posterior belief distribution q 0 and a sequence of Normalizing Flow. Similarly, we ﬁnally get a sample s t +1 from q K .
We basically adopt the same data buffer updating strategy as in Dreamer Hafner et al. (2019a).
First, we use a small amount of S seed episodes ( S = 5 in DMC experiments) with random actions to collect data.
After that, we 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, we record the observations, actions, and rewards of the whole trajectory ( ( o t , a t , r t )
Tt =1 ) and add it to data buffer B .
For DMControl tasks, we pre-process images by reducing the bit depth to 5 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 α 0 = 5 × 10 − 4 , α 1 = 8 × 10 − 5 , α 2 = 8 × 10 − 5 , respectively and scale down gradient norms that exceed 100 .
We clip the KL regularizers in J Model below 3 . 0 free nats as in Dreamer and PlaNet.
The imagination horizon is H = 15 and the same trajectories are used to update both action and value models.
We compute the TD- λ targets with γ = 0 .
99 and λ = 0 .
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 afﬁne coupling ﬂows.
For the image processing, we simply divide the raw pixels by 255 and subtract 0 .
5 to make the inputs lie in [ − 0 .
A. E XTENDED INFORMATION OF B ASELINES
For model-free baselines, we compare with D4PG (Barth-Maron et al., 2018), a 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., 2015) 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 et al., 2016) with vector-wise state inputs from DMCcontrol.
For model-based baselines, we use PlaNet (Hafner et al., 2019b) and Dreamer (Hafner et al., 2019a), two state-of-the-art modelbased RL.
PlaNet (Hafner et al., 2019b) selects actions via online planning without an action model and drastically improves over D4PG and A3C in data efﬁciency.
Dreamer (Hafner et al., 2019a) further improve the data efﬁciency by generating imaginary rollouts in the latent space.
A. A N A BLATION S TUDY ON THE N UMBER OF I MAGINED T RAJECTORIES
[Figure]
Finger Spin
N=
N=
N=
To show the effect of N , we adjust the number of imagined trajectories on the Finger Spin environment.
We choose Finger Spin since the contact between the ﬁnger and the object incorporates many subtle patterns.
So we believe Finger Spin is representative.
We choose N = 1 , 2 , 4 and run 500K environment steps.
We run N = 1 , 2 with 3 different seeds, and N = 4 with 5 different seeds (we use the main DMC experiment results, where N = 4 here).
The result shows that, in
Finger Spin, the performance gain caused by multiple imagined trajectories is obvious.
When the environmental locomotion pattern itself is complex and ﬂexible enough to incorporate diverse possibilities, then using FORBES allows the agent to make diverse predictions and using the multiple imagined trajectories technique will further exploit the advantages of FORBES.
A larger N means more computation time and more memory cost, so we choose N = 4 as a trade-off.
A. E XTENDED R ESULTS ON DMC
[Figure]
We run our algorithm for 1M environment steps and show the curve in Figure 9.
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.
A.
A N A
BLATION S TUDY ON THE M ODEL P ARAMETERS
[Figure]
In this section, we show that having a ﬂexible belief state distribution is the key to improving performance, rather than introducing more parameters.
Having more parameters do not necessarily mean better performance.
Increasing parameters may also make it difﬁcult to converge and negatively affect the sample efﬁciency.
We add an ablation study that adds more parameters to Dreamer to test the effectiveness of having more parameters.
We add 1 , 2 hidden layer(s) to all the MLP in RSSM, and the result is shown in
Figure 10.
The results show that simply adding parameters cannot improve the performance.
A. C OMPARISON OF ELBO ON FORBES AND RSSM ON DMC
[Figure]
Cheetah Run
[Figure]
Walker Walk
[Figure]
We provide the ELBO in DMC environments and FORBES in Figure 11, and FORBES has higher
ELBO.
A. E VIDENCE L OWER B OUND D ERIVATIONS
The variational bound for latent dynamics models p ( o 1: T , s 1: T | a 1: T ) =
(cid:81) t p ( s t | s t − 1 , a t − 1 ) p
( o t | s t ) and a variational posterior q ( s 1: T | o 1: T , a 1: T ) = (cid:81) t q ( s t | o ≤ t , a <t ) follows from importance weighting and Jensen’s inequality as shown, log p ( o 1: T | a 1: T )
=
log E p ( s 1: T | a 1: T ) (cid:34) T (cid:89) t = p ( o t | s t ) (cid:35) =
log E
q
( s 1: T | o 1: T ,a 1: T ) (cid:34) T (cid:89) t = p ( o t | s t ) p ( s t | s t − 1 , a t − 1 ) /q
( s t | o ≤ t , a <t ) (cid:35) ≥ E q ( s 1: T | o 1: T ,a 1: T ) (cid:34) T (cid:88) t = log p ( o t | s t ) + log p ( s t | s t − 1 , a t − 1 ) − log q ( s t | o ≤ t , a <t ) (cid:35) (16)
A. P ROOFS OF T HEOREM
Theorem 1: The approximation error of the lower bound is
[Equation]
where p ( s t | τ t , o t ) is the true posterior [Equation]
Proof:
[Equation]
For a sequence from time 1 to T
[Equation]
we have
[Equation]
Then we can derive the Theorem 1 with equation 18:
[Equation]
(cid:34)(cid:88) t D KL ( q ( s t | τ t , o t ) (cid:107) p ( s t | s t − 1 , a t − 1 , o t )) (cid:35)
+ E q ( s 1: T | τ T ,o T ) (cid:34) T (cid:88) t = (log p ( o t | s t ) + log p ( r t | s t ) − D KL ( q ( s t | τ t , o t ) (cid:107) p ( s t | s t − 1 , a t − 1 )))
(cid:35) =
E q K ( s 1: T | τ T ,o T ) (cid:34)(cid:88) t D KL ( q ( s t | τ t , o t ) (cid:107) p ( s t | s t − 1 , a t − 1 , o t )) (cid:35) + J Model =
E q K ( s 1: T | τ T ,o T ) (cid:34)(cid:88) t
D KL ( q ( s t | τ t , o t ) (cid:107) p
( s t | τ t , o t ))
(cid:35) + J Model (19) where p ( s t | s t − 1 , a t − 1 , o t ) =
p ( s t | τ t , o t ) given the sampled s t − 1 from q ( s 1: t | τ t , o t ) .
A. M ORE R ESULTS ON D IGIT W RITING E XPERIMENTS
In this section, we show more results of the predictions on the digit writing experiment in Figure 12.
[Figure]
A. R ECONSTRUCTIONS OF THE VISUAL CONTROL TASKS
[Figure]
In this section, we show the reconstructions of the visual control tasks during the evaluating phase.