Flow-based Recurrent Belief State Learning for POMDPs

1 INTRODUCTION . Partially Observable Markov Decision Process ( POMDP ) ( Åström , 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 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 . 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 can not 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 difficult to predict multi-modal future behaviours . Therefore , it is preferable to relax the Gaussian assumptions and use a more flexible family of distributions to learn accurate belief states as shown in Figure 1 ( b ) . In this paper , we propose a new method called FlOw-based Recurrent BElief State 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 flexible belief 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 on challenging visual-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 to 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 , 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 st−1 ∈ S. The agent takes an action at−1 ∈ A , which causes the environment to transit to state st with probability T ( st | st−1 , at−1 ) . The agent then receives an observation ot ∈ Ω which depends on the new state of the environment st with probability O ( ot | st ) . Finally , the agent receives a reward rt−1 equal to R ( st−1 ) . The agent ’ s goal is to maximize the the expected sum of discounted rewards E [ ∑∞ t=0 γ trt ] . Such a POMDP model can also be described as a probabilistic graphical model ( PGM ) as shown in Figure 2 . After having taken action at−1 and observing ot , an agent needs to update its belief state , which is defined as the probability distribution of the environment state conditioned on all historical information : b ( st ) = p ( st | τt , ot ) ( 1 ) where τt = { o1 , a1 , . . . , ot−1 , at−1 } . 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 a family of distributions that is highly flexible and preferably flexible 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 “ flow ” progress , needed to derive such a distribution . More specifically , it describes a sequence of invertible mappings that gradually transform a relatively simple probability density to a more flexible and complex one . Let fθ : RD → RD to be an invertible and differentiable mapping in state space parameterized by θ . Given a random variable x ∈ RD 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 : p ( z ) = p ( x ) ∣∣∣∣det ∂f−1θ∂z ∣∣∣∣ ( 2 ) log p ( z ) = log p ( x ) − log ∣∣∣∣det ∂fθ∂z ∣∣∣∣ ( 3 ) To construct a highly flexible family of distributions , we can propagate the random variable at beginning z0 through a sequence of K mappings and get zK = fθK ◦ fθK−1 ◦ · · · ◦ fθ1 ( z0 ) with the probability log pK ( zK ) = log p ( z0 ) − K∑ k=1 log ∣∣∣∣det ∂fθk∂zk−1 ∣∣∣∣ ( 4 ) Given a relatively simple distribution of z0 , say , Gaussian distribution , by iteratively applying the transformations , the flow is capable of representing a highly complex distribution with the probability that remains tractable . The parameters θ1 , . . . , θK 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 x ∈ RD , let s and t stand for scale and translation functions which are usually parameterized by neural networks , where s , t : Rk → RD−k , k < D. The output , y , can be viewed as a concatenation of its first k dimensions y1 : k and the remaining part yk+1 : D : y1 : k = x1 : k , yk+1 : D = xk+1 : D exp ( s ( x1 : k ) ) + t ( x1 : k ) ( 5 ) where 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-based Recurrent BElief State model ( FORBES ) which learns general continuous 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 2 : State transition model : p ( st|st−1 , at−1 ) . Observation model : p ( ot|st ) Reward model : p ( rt|st ) ( 6 ) In addition , we have a belief inference model q ( st|τt , ot ) to approximate the true posterior distribution p ( st|τt , ot ) , where τt = { o1 , a1 , . . . , ot−1 , at−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 ) : log p ( o1 : T , r1 : T |a1 : T ) ≥ Eq ( s1 : T |o1 : T , a1 : T−1 ) [ T∑ t=1 ( ln p ( ot|st ) + ln p ( rt|st ) −DKL ( q ( st|st−1 , at−1 , ot ) ‖p ( st|st−1 , at−1 ) ) ) ] = T∑ t=1 [ Eq ( st|st−1 , at−1 , ot ) ( ln p ( ot|st ) + ln p ( rt|st ) ) − Eq ( st−1|st−2 , at−2 , ot−1 ) ( DKL ( q ( st|st−1 , at−1 , ot ) ‖p ( st|st−1 , at−1 ) ) ) ] . = JModel ( 7 ) here we use the factorization of q ( s1 : T |o1 : T , a1 : T−1 ) = ∏ t q ( st|st−1 , at−1 , ot ) . 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 ψ : pψ ( st|st−1 , at−1 ) , pψ ( ot|st ) , pψ ( rt|st ) , qψ ( st|τt , ot ) ( 8 ) where their outputs usually follow simple distributions such as diagonal Gaussians . The parameterized belief inference model qψ ( st|τt , ot ) acts as an encoder that encodes the historical information using a combination of convolutional neural networks and recurrent neural networks . Note that q below is obtained through the factorization of q ( s1 : T |o1 : T , a1 : T−1 ) = ∏ t q ( st|st−1 , at−1 , ot ) . 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 Figure 3 ( a ) , the input images o1 : t and actions a1 : t−1 are encoded with qψ ( st|τt , ot ) ( the blue and the red path ) . Then our final inferred belief is obtained by propagating qψ ( st|τt , ot ) through a set of normalizing flow mappings denoted fθK ◦ · · · ◦ fθ1 to get a representative posterior distribution qψ , θ ( st|τt , ot ) . For convenience , we denote q0 = qψ and qK = qψ , θ . On the other hand , o1 : t−1 and a1 : t−2 are encoded with qψ ( st−1|τt−1 , ot−1 ) ( the blue path ) , then the state transition model is used to obtain the prior guess of the state pψ ( st | τt ) = Eqψ ( st−1|τt−1 , ot−1 ) [ pψ ( st | st−1 , at−1 ) ] ( the green path ) . Then our final prior is obtained by propagating pψ ( st|τt ) through another set of normalizing flow mappings denoted fωK ◦· · ·◦fω1 to get a representative prior distribution pψ , ω ( st|τt ) . For convenience , we denote p0 = pψ and pK = pψ , ω . Then as shown in Figure 3 ( b ) , we can sample the initial state st ( the yellow and purple triangles ) from the belief states qK ( st | τt , ot ) . For each sampled initial state , we can use the state transition model to predict the future states ŝt+h given the future actions at : t+h−1 , and then use the observation model to reconstruct the observations ôt+h , where h is the prediction horizon . With the above settings , we can substitute the density probability inside the KL-divergence term in Equation 7 with Normalizing Flow : DKL ( qK ( st|τt , ot ) ‖pK ( st|st−1 , at−1 ) ) = EqK ( s1 : t|τt , ot ) [ log qK ( st|τt , ot ) − log pK ( st|τt ) ] = EqK ( s1 : t|τt , ot ) [ log q0 ( st|τt , ot ) − K∑ k=1 log ∣∣∣∣det ∂fθk∂st , k−1 ∣∣∣∣− log p0 ( st|τt ) + log ∣∣∣∣det ∂fωk∂st , k−1 ∣∣∣∣ ] ( 9 ) where pK ( st | st−1 , at−1 ) = pK ( st | τt ) given the sampled st−1 from qK ( s1 : t|τt , ot ) . st , k is the state variable st transformed by k layers of normalizing flows , and st,0 = st. To further demonstrate the properties of FORBES , we provide the following theorems . Theorem 1 The approximation error of the log-likelihood when maximizing the JModel ( the derived ELBO ) defined in Equation 7 is : log p ( o1 : T , r1 : T |a1 : T ) − JModel = EqK ( s1 : T |τT , oT ) [ T∑ t=1 DKL ( q ( st|τt , ot ) ‖p ( st | τt , ot ) ) ] ( 10 ) where p ( st | τt , ot ) denotes the true belief states . Detailed proofs can be found in Appendix.A.9 . Theorem 1 suggests that , when the learning algorithm 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 q ( st|τt , ot ) and the true belief states p ( st | τt , ot ) . Clearly , if p ( st | τt , ot ) is a complex distribution and q ( st|τt , ot ) is chosen from a restricted distribution class such as diagonal Gaussian , then when the algorithm maximizes the JModel ( 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 flow 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 flows . To be specific , Teshima et al . ( 2020 ) provides theoretical results for the family of the flow used in FORBES . In fact , there always exists a diffeomorphism that can turn one well-behaved distribution to another . Besides the aforementioned affine coupling flow , many works show the distributional universality of other flows ( Kong & Chaudhuri , 2020 ; Huang et al. , 2018 ) . Ideally , the universal approximation property of the flow model qK ( st | τt , ot ) allows us to approximate the true posterior p ( st | τt , ot ) 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 can not 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 .

Summary: 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.