Tutorial and Survey on Probabilistic Graphical Model and Variational Inference in Deep Reinforcement Learning

  • 2019-10-04 20:10:41
  • Xudong Sun, Bernd Bischl
  • 0

Abstract

Probabilistic Graphical Modeling and Variational Inference play an importantrole in recent advances in Deep Reinforcement Learning. Aiming at aself-consistent tutorial survey, this article illustrates basic concepts ofreinforcement learning with Probabilistic Graphical Models, as well asderivation of some basic formula as a recap. Reviews and comparisons on recentadvances in deep reinforcement learning with different research directions aremade from various aspects. We offer Probabilistic Graphical Models, detailedexplanation and derivation to several use cases of Variational Inference, whichserve as a complementary material on top of the original contributions.

 

Quick Read (beta)

Tutorial and Survey on Probabilistic Graphical Model and Variational Inference in Deep Reinforcement Learning

Xudong Sun Department of Statistics
Ludwig Maximillian University of Munich
Munich, Germany
Email: [email protected]
   Bernd Bischl Department of Statistics
Ludwig Maximillian University of Munich
Munich, Germany
Email: [email protected]
Abstract

Probabilistic Graphical Modeling and Variational Inference play an important role in recent advances in Deep Reinforcement Learning. Aiming at a self-consistent tutorial survey, this article illustrates basic concepts of reinforcement learning with Probabilistic Graphical Models, as well as derivation of some basic formula as a recap. Reviews and comparisons on recent advances in deep reinforcement learning with different research directions are made from various aspects. We offer Probabilistic Graphical Models, detailed explanation and derivation to several use cases of Variational Inference, which serve as a complementary material on top of the original contributions.

Probablistic Graphical Models; Variational Inference; Deep Reinforcement Learning

I Introduction

Deep Reinforcement Learning has gaining increasing attention recently due to its great success in complicated tasks [19], and it has developed in a rapid way. For a brief overview, see [2]. Despite the existing survey, this paper, however, focuses on Probabilistic Graphical Model and Variational Inference, especially Amortized Variational Inference [8] and their applications in Deep Reinforcement Learning.

Specifically,

  • We start from the basics of Reinforcement Learning with probabilistic graphical model [14] explanations and extend the discussion to complicated models using variational inference [4] in order to have a comprehensive yet brief summary of the topic.

  • We provide Probabilistic Graphical Models [14] for many basic concepts of Reinforcement Learning, as well as recent works on Deep Reinforcement Learning. To our best knowledge, such a comprehensive inclusion of Probabilistic Graphical Models in (Deep) Reinforcement Learning does not yet exist in literature.

  • We introduce a taxonomy of different Graphical Model and Variational Inference used in Deep Reinforcement Learning, which is also the first time to our best knowledge.

  • We give detailed derivation for some of the critical results, which is not explicitly stated in the original contributions like [28, 15, 11] or in a slightly different way. This makes the paper in a relative standalone position to be self understandable, which is another contribution.

I-A Organization of the paper

Since the paper serves as both a tutorial and survey, we keep the detailed derivation in the main text, instead of moving to appendix. In section I-B, we first introduce the fundamentals for Graphical Models and Variational Inference, then we review the basics about reinforcement learning by connecting probabilistic graphical models (PGM) in section I-C, as well as the basics and a incomplete overview on the advances about deep reinforcement learning, accompanied with a comparison of different methods in section II. In section III, we discuss how undirected graph could be used in modeling both the value function and the policy, which works well on high dimensional discrete state and action spaces. In section IV, we introduce the directed acyclic graph framework on how to treat the policy as posterior on actions, while adding many proofs that does not exist in the original contributions. In section V, we introduce works on how to use variational inference to approximate the environment model, while adding graphical models and proofs which does not exist in the original contributions.

I-B Prerequisite on Probabilistic Graphical Models and Variational Inference, Terminologies and Conventions

Directed Acyclic Graphs (DAG) [3] as a PGM offers an instinctive way of defining factorized join distributions of Random Variables (RV) by assuming the conditional independence [3] across the RV though d-separation [3]. In this paper, we use capital letter to denote a RV, while using the lower case letter to represent the realization of corresponding RV. To avoid symbol collision of using A to represent advantage in many RL literature, we sometimes use a to represent the RV A, as well as using Aact explicitly. For simplicity, we use p(c) to represent pC(C=c), the probability of RV C take value c, as well as use p(C) to represent PC(C). We use (B\scalebox1.07C)A to represent that RV B is conditionally independent from RV C, when conditional on observation of RV A, which is equivalent to write p(B|A,C)=p(B|A) or p(BC|A)=P(B|A)P(C|A).

Variational Inference (VI) approximates intractable posterior distribution of neural network, specified in a probabilistic graphical model usually, with a variational proposal posterior distribution, by optimizing the Evidence Lower Bound (ELBO) [4], which assigns the values of latent unobservables at the same time. Variational Inference is widely used in Deep Learning Community [26], including approximating the posterior on the weights [5] distribution of neural networks, as well as approximating on the activations distribution [8]. VI on the activations of neural networks has been used on Variational AutoEncoder [13], while VI on the weights of the Neural Network has lead to Bayesian Neural Networks [5]. Weight Uncertainty in Neural Networks [5] has been used for tackling the exploration-exploitation trade off in bandit problems, using Thompson sampling, which has also been shown to lead to systematic exploration by weights with higher variation [5].

I-C Basics about Reinforcement Learning with graphical model

\includegraphics

[scale=0.6]tikz_rl_def.pdf

Fig. 1: Concept of Reinforcement Learning

I-C1 RL Concepts, Terminology and Convention

As shown in Figure 1, Reinforcement Learning (RL) involves optimizing the behavior of an agent via interaction with the environment. At time t, the agent lives on state St, By executing an action at according to a policy [28] π(a|St), the agent jumps to another state St+1, while receiving a reward Rt. Let discount factor γ decides how much the immediate reward is favored compared to longer term return, with which one could also allow tractability in infinite horizon reinforcement learning [28], as well as reducing variance in Monte Carlo setting [15]. The goal is to maximize the accumulated rewards, G=t=0TγtRt which is usually termed return in RL literature.

For simplicity, we interchangeably use two conventions whenever convenient: Suppose an episode last from t=0:T, with T correspond to continuous non-episodic reinforcement learning. We use another convention of t{0,,} by assuming when episode ends, the agent stays at a self absorbing state with a null action, while receiving null reward.

By unrolling Figure 1, we get a sequence of state, action and reward tuples {(St,Atact,Rt)} in an episode, which is coined trajectory τ [30, 6]. Figure 2 illustrates part of a trajectory in one rollout. The state space 𝒮 and action space 𝒜, which can be either discrete or continuous and multi-dimensional, are each represented with one continuous dimension in Figure 2 and plotted in an orthogonal way with different colors, while we use the thickness of the plate to represent the reward space .

\includegraphics

[scale=0.9]tikz_drl_illustrator.pdf

Fig. 2: Illustration of State, Action and Reward Trajectory

I-C2 DAGs for (Partially Observed ) Markov Decision Process

Reinforcement Learning is a stochastic decision process, which usually comes with three folds of uncertainty. That is, under a particular stochastic policy characterized by π(a|s)=p(a|s), within a particular environment characterized by state transition probability p(st+1|st,a) and reward distribution function p(rt|st,at), a learning agent could observe different trajectories with different unrolling realizations. This is usually modeled as a Markov Decision Process [28], with its graphical model shown in Figure 3, where we could define a joint probability distribution over the trajectory of state , action and reward RVs. In Figure 3, we use dashed arrows connecting state and action to represent the policy, upon fixed policy, we have the trajectory likelihood in Equation (1)

p(τ)=t=0Tp(st+1|st,at)p(rt|st,at)π(at|st) (1)

Upon observation of a state st in Figure 3, the action at the time step in question is conditionally independent with the state and action history t={S0,A0act,,St-1}, which could be denoted as (Atact\scalebox1.07t)St.

\includegraphics

[scale=0.9]tikz_pgm_mdp.pdf

Fig. 3: Directed Acyclic Graph For Markov Decision Process

A more realistic model, however, is the Partially Observable Markov Decision process [12], with its Directed Acyclic Graph [3] representation shown in Figure 4, where the agent could only observe the state partially of getting Ot through a non invertible function of the latent state St and the previous action at-1, as indicated the Figure by p(ot|st,at-1), while the distributions on other edges are omitted since they are the same as in Figure 3. Under the graph specification of Figure 4, the observable Ot is no longer Markov, but depends on the whole history. However, by introducing a probability distribution b(S) over the hidden state S, with 𝒮b(S)=1, which is termed belief state [12], where state S takes value in range 𝒮.

\includegraphics

[scale=0.8]tikz_pgm_pomdp.pdf

Fig. 4: Probabilistic Graphical Model for POMDP

I-D Value Function, Bellman Equation, Policy Iteration

Define state value function of state s𝒮 in Equation (2), where the corresponding Bellman Equation is derived in Equation (3).

Vπ(s)
= Eπ,ε[i=0γiRt+i(St+i,At+iact)St=s] (2)
= Eπ,ε[Rt(St,Atact)+γi=1γ(i-1)Rt+i(St+i,At+iact)]
= Eπ,ε[Rt(St,Atact)+γi=0γiRt+1+i(St+i+1,At+i+1act)]
= Eπ,ε[Rt(St,Atact)+γVπ(St+1)] (3)

, where St+ip(st+i+1|st+i,at+i) takes value from 𝒮, At+iactπ(a|St+i+1) taking value from 𝒜, and we have used the π and ε in the subscript of the expectation E operation to represent the probability distribution of the policy and the environment (including transition probability and reward probability) respectively. State action value function [28] is defined in Equation (4),

Qπ(s,a)(St=s,Atact=a)
= Eπ,ε[Rt(St=s,Atact=a)+i=1γiRt+i(St+i,At+iact)] (4)
= Eπ,ε[Rt(St=s,Atact=a)+γVπ(St+1)] (5)

, where in Equation (5), its relationship to the state value function is stated.

Combining Equation (3) and Equation (4), we have

V(s)=aπ(a|s)Q(s,a) (6)

Define optimal policy [28] to be

π* =argmax𝜋Vπ(s),sS
=argmax𝜋Eπ[Rt+γVπ(St+1)] (7)

Taking the optimal policy π* into the Bellman Equation in Equation (3), we have

Vπ*(s)=Eπ*,ε[Rt(s,Atact)+γVπ*(St+1)] (8)

Taking the optimal policy π* into Equation (4), we have

Qπ*(s,a)=Eπ*,ε[Rt(s,a)+i=1γiRt+i(St+i,At+iact)] (9)

Based on Equation (9) and Equation (8), we get

Vπ*(s)=max𝑎Qπ*(s,a) (10)

and

Qπ*(s,a)=Eε,π*[Rt(s,a)+γmaxa¯Qπ*(St+1,a¯)] (11)

For learning the optimal policy and value function, General Policy Iteration [28] can be conducted, as shown in Figure 5, where a contracting process [28] is drawn. Starting from initial policy π0, the corresponding value function Vπ0 could be estimated, which could result in improved policy π1 by greedy maximization over actions. The contracting process is supposed to converge to the optimal policy π*.

As theoretically fundamentals of learning algorithms, Dynamic programming and Monte Carlo learning serve as two extremeties of complete knowledge of environment and complete model free [28], while time difference learning [28] is more ubiquitously used, like a bridge connecting the two extremities. Time difference learning is based on the Bellman update error in Equation (12).

δt=Q(st,at)-(rt+γmax𝑎Q(st+1,a)) (12)
\includegraphics

[scale=0.8]tikz_gpi.pdf

Fig. 5: General Policy Iteration

I-E Policy Gradient and Actor Critic

Reinforcement Learning could be viewed as a functional optimization process. We could define an objective function over a policy πθ(a|s), as a functional, characterized by parameter θ, which could correspond to the neural network weights, for example.

Suppose all episodes start from an auxiliary initial state s0, which with probability h(s), jumps to different state s𝒮 without reward. h(s) characterizes the initial state distribution which only depends on the environment. Let η(s) represent the expected number of steps spent on state s, which can be calculated by summing up the γ discounted probability Pπ(s0s,k+1) of entering state s with k+1 steps from auxiliary state s0, as stated in Equation (13), which can be thought of as the expectation of the R.V. γk conditional on state s.

η(s) =k=0γkPπ(s0s,k+1) (13)
=h(s)+s¯,aγη(s¯)πθ(a|s¯)P(s|s¯,a) (14)

In Equation (14), the quantity is calculated by either directly starting from state s, which correspond to k=0 in Equation (13), or entering state s from state s¯ with one step, corresponding to k+12 in Equation (13).

For an arbitrary state s𝒮, using s and s′′ to represent subsequent states as dummy index, we have

θVπ(θ)(s)
= θ[aQπ(θ)(s,a)πθ(a|s)] (15)
= a[θQπ(θ)(s,a)πθ(a|s)+θπθ(a|s)Qπ(θ)(s,a)]
= aθ[s,RP(s,R|s,a)(R+γVπ(θ)(s))]πθ(a|s)
+aθπθ(a|s)Qπ(θ)(s,a) (16)
= asγP(s|s,a)θVπ(θ)(s)πθ(a|s)+
aθπθ(a|s)Qπ(θ)(s,a) (17)
= aθπθ(a|s)Qπ(θ)(s,a)+asγP(s|s,a)πθ(a|s)
[as′′γP(s′′|s,a)θVπ(θ)(s′′)πθ(a|s)+
aθπθ(a|s)Qπ(θ)(s,a)] (18)

The terms in square brackets in Equation (18) are simply Equation (17) with a and s replaced by a and s′′. Since θVπ(θ)(s)=0, Equation (18) could be written as Equation (19),

θVπ(θ)(s)
= aθπθ(a|s)Qπ(θ)(s,a)+
k=1skakγkPπ(ssk,k)θπθ(ak|sk)Qπ(θ)(sk,ak) (19)

, where sk represent the state of k steps after s and Pπ(ssk,k) already includes integration of intermediate state sk-1,s1 before reaching state sk.

Let objective function with respect to policy be defined to be the value function starting from auxiliary state s0 as in Equation (20).

J(πθ)=Vπ(s0)=Eπ,εt=0γtRt(S0=s) (20)

The optimal policy could be obtained by gradient accent optimization, leading to the policy gradient algorithm [28], as in Equation (25).

θJ(πθ)
= θVπ(s0)
= k=0skakγkPπ(s0sk,k)θπθ(ak|sk)Qπ(θ)(sk,ak)
= saη(s)θπθ(a|s)Qπ(θ)(s,a) (21)
= sη(s)sη(s)saη(s)θπθ(a|s)Qπ(θ)(s,a) (22)
= s¯η(s¯)saμ(s)θπθ(a|s)Qπ(θ)(s,a) (23)
= s¯η(s¯)saμ(s)πθ(a|s)πθ(a|s)θπθ(a|s)Qπ(θ)(s,a) (24)
Eπ[θπθ(A|S)πθ(A|S)Q^π(θ)(S,A)] (25)

, where μ(s) is the relative occupancy of state s. The integration of μ(s) with respect to s and πθ(a|s) in the nominator with respect to a in Equation (24) is replaced with expectation with respect to interaction with the environment in Equation (25), and Qπ(θ)(s,a) is replaced by an estimator Q^π(θ)(S,A), which is usually Gt=i=0γiRt+i.

The policy gradient could be augmented to include zero gradient baseline b(s), with respect to objective function J(πθ) in Equation (24), as a function of state s, which does not include parameters for policy θ, since aθπθ(a|s)=0. To reduce variance of the gradient, the baseline is usually chosen to be the state value function estimator V^w(s) to smooth out the variation of Q(s,a) at each state, while V^w(s) is updated in a Monte Carlo way by comparing with Q^πθ(S,A)=Gt.

The actor-critic algorithm [28] decomposes Gt-Vw(st) to be Rt+γVw(st+1)-Vw(st), so bootstrap is used instead of Monte Carlo.

II Recent advances in Deep Reinforcement Learning

II-A Basics of Deep Reinforcement Learning

Deep Q learning [19] makes a breakthrough in using neural network as functional approximator on complicated tasks. It solves the experience correlation problem by using a reply memory and the instability of the target problem with a frozen target network. Specifically, the reinforcement learning is transformed in a supervised learning task by fitting on the target Rt+γmax𝑎Q(st+1,a) from the replay memory with state st as input. However, the target can get drifted easily which leads to unstable learning. In [19], a target network is used to provide a stable target for the updating network to be learned before getting updated occasionally. Double Deep Q learning [29], however, solves the problem by having two Q network and update the parameters in a alternating way.

II-B Taxonomy

While it is difficult to cover all aspects of recent advances in deep reinforcement learning. We pick some interesting research directions and list some contributions in these directions below.

II-B1 On Policy methods

A3C [18] stands out in the asynchronous methods in deep learning [18] which can be run in parallel on a single multi-core CPU. Trust Region Policy Optimization [23] and Proximal Policy Optimization [24] assimilates the natural policy gradient, which use a local approximation to the expected return. The local approximation could serve as a lower bound for the expected return, which can be optimized safely subject to the KL divergence constraint between two subsequent policies, while in practice, the constraint is relaxed to be a regularization.

II-B2 Off Policy methods

Except for Deep Q Learning [19] mentioned above, DDPG [16] extends Deterministic Policy Gradient (DPG) [25] with deep neural network functional approximator, which is an actor-critic algorithm and works well in continuous action spaces.

II-B3 Goal based Reinforcement Learning

In robot manipulation tasks, the goal could be represented with state in some cases [30]. Universal Value Function Approximator (UVFA) [21] incorporate the goal into the deep neural network, which let the neural network functional approximator also generalize to goal changes in tasks. Work of this direction include [1, 30], for example.

II-B4 Exploration with sparse reward

In complicated real environment, an agent has to explore for a long trajectory before it can get any reward as feedback. Due to lack to enough rewards, traditional Reinforcement Learning methods performs poorly, which lead to a lot of contributions in the sufficient exploration methods. The methods using graphical model and variational method we introduce later each use different mechanisms to explore the environments.

II-B5 Replay Memory Manipulation based Method

Replay memory is a critical component in Deep Reinforcement Learning, which solves the problem of correlated transition in one episode. Beyond the uniform sampling of replay memory in Deep Q Network [19], Prioritized Experience Replay [22] improves the performance by giving priority to those transitions with bigger TD error, while Hindsight Experience Replay (HER) [1] manipulate the replay memory with changing goals to transition so as to change reward to promote exploration. Maximum entropy regularized multi goal reinforcement learning [30] gives priority to those rarely occurred trajectory in sampling, which has been shown to improve over HER [30].

II-C Comparison

In the following sections, we give detailed explanation on how graphical model and variational inference could be used to model and optimize the reinforcement learning process with each category a different section. Together with the methods mentioned above, we make a comparison of them in Table I, where ”S” means state and ”A” means action, where ”c” means continuous, ”d” means discrete. ”standalone” means whether the algorithm needs to be combined with another algorithm to work or is a standalone method. ”var” means which probability the variational inference is approximating, ”p” means whether the method is on policy or off policy. ”na” means not applicable.

TABLE I: Comparison of deep reinforcement learning methods: ”S” means state and ”A” means action, where ”c” means continuous, ”d” means discrete. ”standalone” means whether the algorithm needs to be combined with another algorithm to work or is a standalone method. ”var” means which probability the variational inference is approximating, ”p” means whether the method is on policy or off policy. ”na” means not applicable
Algorithm S A standalone var p
Deep Q c d y na off
A3C c c/d y na on
TRPO/PPO c d y na on
DDPG c c y na off
Boltzmann d d y na on
VIME c c n pθ(st+1|st,at) na
VAST c d n p(st|ot-k) na
SoftQ c c/d y p(at|st) on

III Policy and value function with undirected graphs

We first discuss the application of undirected graphs in deep reinforcement learning, where we use deep belief network here. Rather than modeling conditional distribution, as in directed acyclic graphs, undirected graphs model joint distribution of variables in question and focus on cliques [3] with free energy associated with it, which could be used to model the value function in reinforcement learning. Restricted Boltzman Machine has nice property of tractable factorized posterior distribution over the latent variable conditioned on observables, instead of having to do gibbs sampling in general Boltzman Machine.

In [20], the authors use Restricted Bolzman Machine to deal with MDPs of large state and action spaces, by modeling the state-action value function with the negative free energy of the graph, where free energy of the graph could be easily calculated through the product of expert [20]. Specifically, the visible states of the Restricted Bolzmann Machine [20] consists of both state s and action a binary variables, as shown in Figure 6, where the hidden nodes consist of L binary variables, while state variable s are dark colored to represent it can be observed and action are light colored to represent it need to be sampled. Together with the auxilliary hidden variables, the undirected graph defines a joint probability distribution over state and action pairs, which defines a stochastic policy network that could sample actions out for on policy learning. Since it is pretty easy to calculate the derivative of the free energy F(s,a,h) with respect to the coefficient wk,j of the network, one could use temporal difference learning to update the coefficients in the network. Thanks to properties of Boltzmann Machine, the conditional distribution of action over state p(a|s) is still Boltzmann distributed, governed by the free energy, by adjusting the temperature, one could also change between different exploration strength.

\includegraphics

[scale=0.9]tikz_pgm_boltzmanmachine.pdf

Fig. 6: Restricted Boltzmann Machine Actor Critic

The conditional distribution of actions under state could serve as the policy, which is

p(a|s)=1Ze-hF(s,a,h)/T=1ZehQ(s,a)/T (26)

, where Z is the partition function [3] and we use the negative free energy to approximate the state action value function. Upon the state value function Q(s,a) in Equation (26) is learned as a critic [28], such that its associated policy is defined, MCMC sampling [3] could be used to sample actions, as an actor [28]. With the sampled actions, time difference learning method like SARSAR [28], could be carried out to update the state value function estimation. Such an on-policy process has been shown to be empirically effective in the large state actions spaces [20].

IV Variational Inference on Policies

IV-A policy as ”optimal” posterior

The Boltzmann Machine defined Product of Expert Model in [20] works well for large state and action spaces, but are limited to discrete specifically binary state and action variables. For continuous state and action spaces, in [9], the author proposed deep energy based models with Directed Acyclic Graphs (DAG) [3], which we re-organize in a different form in Figure 7 with annotations added. The difference with respect to Figure 3 is that, in Figure 7, the reward is not explicit expressed in the directed graphical model. Instead, an auxilliary binary Observable O is used to define whether the corresponding action at the current step is optimal or not. The conditional probability of the action being optimal is p(Ot=1|st,at)=exp(r(st,at)), which connects conditional optimality with the amount of award received by encouraging the agent to take highly rewarded actions in an exponential manner. Note that the reward here must be negative to ensure the validity of probability, which does not hurt generality since reward range can be translated [15].

The Graphical Model in Figure 7 in total defines the trajectory likelihood or the evidence in Equation (27):

p(τ) =p(s1)t[p(st+1|st,at)p(Ot=1|st,at)]
=[p(s1)tp(st+1|st,at)]exp(tr(st,at)) (27)

.

By doing so, the author is forcing a form of functional expression on top of the conditional independence structure of the graph by assigning a likelihood. In this way, calculating the optimal policy of actions distributions becomes an inference problem of calculating the posterior p(at|st,Ot:T=1), which reads as, conditional on optimality from current time step until end of episode, and the current current state to be st, the distribution of action at, and this posterior corresponds to the optimal policy. Observing the d-separation from Figure 7, O1:t-1 is conditionally independent of at given st, (O1:t-1\scalebox1.07Atact)St, so p(at|st,O1:t-1=,Ot:T=1)=p(at|st,Ot:T)

\includegraphics

[scale=0.9]tikz_pgm_soft_q.pdf

Fig. 7: Optimal Policy as posterior on actions: p(at|st,Ot:T=1)

IV-B Message passing for exact inference on the posterior

In this section, we give detailed derivation on doing exact inference on the policy posterior which is not given in [15]. Although the results are not used due to unexpected behavior, there is theoretical insights that is worth being noted.

The graph in Figure 7 is similar to Hidden Markov Models (HMM) [3], if we could treat the tuple of variable (at,st) as the latent variable counterpart of a HMM, with emission probability p(Ot=1|st,at)=exp(r(st,at)), while the transition probability, is from the variable tuple (at,st) to a subcomponent st+1 of the ”latent” variable tuple (at+1,st+1).

Similar to the forward-backward message passing algorithm [3] in Hidden Markov Models [3], the posterior p(at|st,Ot:T=1) could also be calculated by passing messages. We offer a detailed derivation of the decomposition of the posterior p(at|st,Ot:T=1) in Equation (28), which is not available in [15].

p(at|st,Ot:T=1)
= p(at,st,Ot:T=1)p(st,Ot:T=1)
= p(Ot:T=1|at,st)p(at,st)p(st,Ot:T=1)
= p(Ot:T=1|at,st)p(at|st)p(st)atp(st,at,Ot:T=1)d{at}
= p(Ot:T=1|at,st)p(at|st)p(st)atp(Ot:T=1|at,st)p(at|st)p(st)d{at}
= p(Ot:T=1|at,st)p(at|st)atp(Ot:T=1|at,st)p(at|st)d{at}
= β(at,st)atβ(at,st)d{at}
= β(at,st)β(st) (28)

In Equation (28), we define message β(at,st)=p(Ot:T=1|at,st)p(at|st) and message β(st)=atβ(at,st)d{at}. If we consider p(at|st) as a prior with a trivial form [15], the only policy related term becomes p(Ot:T=1|at,st).

In Hidden Markov Models (HMM) [3], if we use O to represent the visible observed state and S to represent the hidden latent state, T for the series length, then it is essential to calculate the posterior p(St|O1:T) and p(St,St+1|O1:T), which is the marginal of the complete posterior p(S1:T|O1:T). The posterior marginal could be computed by the forward message α(St)=p(O1:t,St) and the backward message β(St)=p(Ot:T|St), which is the probability distribution of observables from current time step until the end of the sequence, conditional on the current latent state.

In contrast, here, only the backward messages are relevant. Additionally, the backward message β(at,st) here is not a probability distribution as in HMM, instead, is just a probability. In Figure 7, the backward message β(at,st) could be decomposed recursively. Since in [15] the author only give the conclusion without derivation, we give a detailed derivaion of this recursion in Equation (29).

β(st,at)
= p(Ot=1,Ot+1:T=1|st,at)
= p(Ot=1,Ot+1:T=1,st,at,st+1,at+1)d{st+1,at+1}p(st,at)
= p(Ot+1:T=1,st+1,at+1,Ot=1|st,at)d{st+1,at+1}
= p(Ot+1:T=1,st+1,at+1|st,at)p(Ot=1|st,at)
d{st+1,at+1} ((Ot+1:T,St+1,At+1\scalebox1.07Ot)St,At)
= p(Ot+1:T=1,st+1,at+1)p(st+1,at+1)p(st+1,st,at)p(st,at)
p(Ot=1|st,at)d{st+1,at+1}
= p(Ot+1:T=1|st+1,at+1)p(st+1|st,at)p(Ot=1|st,at)
d{st+1,at+1}
= β(st+1)p(st+1|st,at)p(Ot=1|st,at)dst+1 (29)

The recursion in Equation (29) start from the last time point T of an episode.

IV-C Connection between Message Passing and Bellman equation

If we define

Q(st,at)=log(β(at,st)) (30)

and

V(st) =logβ(st)
=logβ(st,at)𝑑at
=logexp(Q(st,at))𝑑atmaxatQ(st,at) (31)

then the corresponding policy could be written as Equation (32).

π(at|st)=p(at|st,Ot:T=1)=exp(Q(st,at)-V(st)) (32)

.

Taking the logrithm of Equation (29), we get Equation (33)

log(β(st,at))
=logβ(st+1)p(st+1|st,at)p(Ot=1|st,at)dst+1
=logexp[r(st,at)+V(st+1)]p(st+1|st,at)𝑑st+1
=r(st,at)+logexp(V(st+1))p(st+1|st,at)𝑑st+1 (33)

which reduces to the risk seeking backup in Equation (34) as mentioned in [15]:

Q(st,at)=r(st,at)+logEst+1p(st+1|st,at)[exp(V(st+1))] (34)

The mathematical insight here is that if we define the messages passed on the Directed Acyclic Graph in Figure 7, then message passing correspond to a peculiar version Bellman Equation like backup, which lead to an unwanted risk seeking behavior [15].

IV-D Variational approximation to ”optimal” policy

Since the exact inference lead to unexpected behavior, approximate inference could be used. The optimization of the policy could be considered as a variational inference problem, and we use the variational policy of the action posterior distribution q(at|st), which could be represented by a neural network, to compose the proposal variational likelihood of the trajectory as in Equation (35):

q(τ) =p(s1)t[p(st+1|st,at)q(at|st)] (35)

, where the initial state distribution p(s1) and the environmental dynamics of state transmission is kept intact. Using the proposal trajectory as a pivot, we could derive the Evidence Lower Bound (ELBO) of the optimal trajectory as in Equation (36), which correspond to an interesting objective function of reward plus entropy return, as in Equation (37).

log(p(O1:T))
= logp(O1:T=1,s1:T,a1:T)q(s1:T,a1:T)q(s1:T,a1:T)ds1:Tda1:T
= logEq(s1:T,a1:T)p(O1:T=1,s1:T,a1:T)q(s1:T,a1:T)
Eq(s1:T,a1:T)[logp(O1:T=1,s1:T,a1:T)-logq(s1:T,a1:T)] (36)
= -DKL(q(τ)|p(τ)) (take q(at|st)=π(at|st))
= Eq(s1:T,a1:T)[t=1:T[r(st,at)-logq(at|st)]]
= t=1:TEst,at[r(st,at)+H(π(at|st))] (37)

IV-E Connection between policy gradient and Q learning

A representative method belonging to the above mentioned framework is Soft Q [9], where the state action value function is defined to be

Qsoftπ(s,a)=r0+Erπ,s0=s,a0=a[t=1γt(rt+H(π(.|st)))] (38)

Soft Q carries an soft version of Bellman update similar to Q Learning [28], which lead to policy improvement with respect to the corresponding functional objective in Equation (39).

J(π)
= tE(st,at)ρπl=tinfγl-tE(sl,al)[r(sl,al)+
αH(π(.|sl))|st,at]]
= tE(st,at)ρπ[Qsoftπ(st,at)+αH(π(.|st))] (39)

Setting policy as Equation (32) lead to policy improvement. We offer a detailed proof for a key formula in Equation (40), which is stated in Equation (19) of [9] without proof. In Equation (40), we use π(|s) to implicitly represent π(a|s) to avoid symbol aliasing whenever necessary.

H(π(|s))+Eaπ[Qsoftπ(s,a)]
= -aπ(a|s)[logπ(a|s)-Qsoftπ(s,a)]𝑑a
= -aπ(a|s)[logπ(a|s)-log[exp(Qsoftπ(s,a))]]𝑑a
= -aπ(a|s)[logπ(a|s)-log[exp(Qsoftπ(s,a))exp(Qsoftπ(s,a))𝑑a
exp(Qsoftπ(s,a))da]]da
= -aπ(a|s)[logπ(a|s)-log[π~(a|s)]-
logexp(Qsoftπ(s,a))]da
= -DKL(π(|s)||π~(|s))+logexp(Qsoftπ(s,a))da (40)

For the rest of the proof, we invite the reader to read the appendix of [9]. Algorithms of the this kind of maximum entropy family also include Soft Actor Critic [10].

V Variational Inference on the Environment

Another direction of using Variational Inference in Reinforcement Learning is to learn an environmental model, either on the dynamics or the latent state space posterior, instead of approximating the maximum entropy policy posterior in [15], explained in Section IV.

V-A Variational inference dynamics model parameter posterior

The weight distribution property of neural network has been exploited in deep reinforcement learning, in Variational Information Maximizing Exploration (VIME) [11], where dynamic model pθ(st+1|st,at) for the agent’s interaction with the environment is build with parameter θ, using Bayesian Neural Network [5], where the R.V. for θ is denoted by Θ, and is treated in a Bayesian way by modeling the weight uncertainty and the belief about the environment is modeled as entropy of the neural network weights posterior distribution H(Θ|ξt) based on trajectory observations ξt={s1:t,a1:t-1}. The method encourages taking explorative action of the environment by alleviating the information gain of the agent’s belief about the environment after observing a new state st+1, which is H(Θ|ξt,at)-H(Θ|ξt,at,St+1), and is equivalent to H(Θ|ξt)-H(Θ|ξt+1) since action does not affect belief.

\includegraphics

[scale=0.7]tikz_pgm_mdp_vime.pdf

Fig. 8: Probabilistic Graphical Model For VIME

We now derive in Equation (42) that entropy difference between beliefs over trajectory ξt and ξt+1 is actually the mutual information between environmental parameter θ and the new state St+1, as well as the expected KL divergence between the two beliefs, integrated over the dynamics p(st+1|st,at) in Equation (41). Such a derivation is not given in [11].

H(Θ|ξt,at)-H(Θ|ξt+1)
= -Θp(θ|ξt)logp(θ|ξt)𝑑θ+Θp(θ|ξt+1)logp(θ|ξt+1)𝑑θ
= -Θ𝒮p(st+1|ξt,at)p(θ|ξt,at,st+1)logp(θ|ξt)𝑑θ𝑑st+1+
Θ𝒮p(st+1|ξt,at)p(θ|ξt+1)logp(θ|ξt+1)𝑑θ𝑑st+1
= Ep(st+1|ξt,at)KL(p(θ|ξt+1)||p(θ|ξt)) (41)
= -Θ𝒮p(st+1,θ|ξt,at)log[p(θ|ξt+1)p(θ|ξt)p(st+1|ξt,at)p(st+1|ξt,at)]
= -Θ𝒮p(st+1,θ|ξt,at)log[p(st+1,θ|ξt,at)p(θ|ξt)p(st+1|ξt,at)]
= -I(Θ,St+1|ξt,at) (42)

Based on Equation (42), an intrinsic reward can be augmented from the environmental reward function, thus the method could be incorporated with any existing reinforcement learning algorithms for exploration, TRPO [23], for example. Upon additional observation of action at and state st+1 pair on top of trajectory history ξt, the posterior on the distribution of the environmental parameter θ, p(θ|ξt), could be updated to be p(θ|ξt+1) in a Bayesian way as derived in Equation (43), which is first proposed in [27].

p(θ|ξt+1)
= p(θ|ξt,at,st+1)
= p(θ,ξt,at,st+1)p(ξt,at,st+1)
= p(st+1|θ,ξt,at)p(θ,ξt,at)p(ξt,at,st+1)
= p(st+1|θ,ξt,at)p(θ,ξt,at)p(at,ξt)p(st+1|at,ξt)
= p(st+1|θ,ξt,at)p(θ|ξt,at)p(st+1|at,ξt)
= p(st+1|θ,ξt,at)p(θ|ξt)p(st+1|at,ξt) (43)

In Equation (43), the denominator can be written as Equation (44), so that the dynamics of the environment modeled by neural network weights θ, p(st+1|θ,at,ξt), could be used.

p(st+1|at,ξt)
=Θp(st+1,θ|at,ξt)𝑑θ
=Θp(st+1,θ,at,ξt)p(at,ξt)𝑑θ
=Θp(st+1|θ,at,ξt)p(θ,at,ξt)p(at,ξt)𝑑θ
=Θp(st+1|θ,at,ξt)p(θ|ξt)𝑑θ (44)

The last step of Equation (44) makes use of the fact that current action does not the environment model.

Since the integral in Equation (44) is not tractable, variational treatment over the neural network weights posterior distribution p(θ|ξt) is used, characterized by variational parameter ϕ, as shown in the dotted line in Figure 8. The variational posterior about the model parameter θ, updated at each step, could than be used to calculate the intrinsic reward in Equation (41).

V-B Variational Inference on hidden state posterior

In Variational State Tabulation (VaST) [7], the author assume the high dimensional observed state to be represented by Observable O, while the transition happens at the latent state space represented by S, which is finite discrete. The author assume a factorized form of observation and latent space joint probability, which we explicitly state in Equation (45).

p(O,S)=πθ0(s0)t=0TpθR(ot|st)t=1TpθT(st|st-1,at-1) (45)

Additionally, we characterize Equation (45) with the probabilistic graphical model in Figure 9 which is not in [7], where the difference compared to Figure 7 is that here the latent state S is in discrete space while the observation is a high dimensional image. By assuming a factorized form of the variational posterior in Equation (46),

q(S0:T|O0:T)=t=0Tqϕ(St|Ot-k:t) (46)

, where the author assume the episode length to be T, and default frame prior observation to blank frames, the Evidence Lower Bound (ELBO) of the observed trajectory of Equation (45) could be easily represented by a Varitional AutoEncoder [8] like architecture, where the encoder qϕ, together with the reparametrization trick [8], maps the observed state O into parameters for the Con-crete distribution [17], so backprobagation could be used on deterministic variables to update the weight of the network based on the ELBO, which is decomposed into different parts of the reconstruction losses of the variational autoencoder like architecture. Like VIME [11], VaSt could be combined with other reinforcement learning algorithms, where prioritized sweeping [28] is carried out on the Heviside activation of the encoder output directly, by counting the transition frequency, instead of waiting for the slowly learned environmental transition model pθT(st|st-1,at-1) in Equation (45). A potential problem of doing so is aliasing between latent state s and observed state o. To alleviate this problem, in [7], the author actively relabel the transition history in the replay memory once found the observable has been assigned a different latent discrete state.

\includegraphics

[scale=0.6]tikz_pgm_vast.pdf

Fig. 9: Graphical Model for Variation State Tabulation

VI Conclusion

As a tutorial survey, this paper introduces the application of Probabilistic Graphical Model and Variational Inference in Deep Reinforcement Learning. We reformulates some key concepts in Reinforcement Learning with Probabilistic Graphical Models, summarizes recent advances of Deep Reinforcement Learning and compares some representative methods from different aspects. Furthermore, we offer some detailed derivations and Probabilistic Graphical Models to those methods using variational inference, when such detailed derivations and Probabilistic Graphical Models are not included in the original contribution.

References

  • [1] M. Andrychowicz, F. Wolski, A. Ray, J. Schneider, R. Fong, P. Welinder, B. McGrew, J. Tobin, O. P. Abbeel, and W. Zaremba (2017) Hindsight experience replay. In Advances in Neural Information Processing Systems, pp. 5048–5058. Cited by: §II-B3, §II-B5.
  • [2] K. Arulkumaran, M. P. Deisenroth, M. Brundage, and A. A. Bharath (2017) A brief survey of deep reinforcement learning. arXiv preprint arXiv:1708.05866. Cited by: §I.
  • [3] C. M. Bishop (2006) Pattern recognition and machine learning. springer. Cited by: §I-B, §I-C2, §III, §III, §IV-A, §IV-B, §IV-B, §IV-B.
  • [4] D. M. Blei, A. Kucukelbir, and J. D. McAuliffe (2017) Variational inference: a review for statisticians. Journal of the American Statistical Association 112 (518), pp. 859–877. Cited by: 1st item, §I-B.
  • [5] C. Blundell, J. Cornebise, K. Kavukcuoglu, and D. Wierstra (2015) Weight uncertainty in neural networks. arXiv preprint arXiv:1505.05424. Cited by: §I-B, §V-A.
  • [6] J. D. Co-Reyes, Y. Liu, A. Gupta, B. Eysenbach, P. Abbeel, and S. Levine (2018) Self-consistent trajectory autoencoder: hierarchical reinforcement learning with trajectory embeddings. arXiv preprint arXiv:1806.02813. Cited by: §I-C1.
  • [7] D. Corneil, W. Gerstner, and J. Brea (2018) Efficient model-based deep reinforcement learning with variational state tabulation. arXiv preprint arXiv:1802.04325. Cited by: §V-B.
  • [8] C. Doersch (2016) Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908. Cited by: §I-B, §I, §V-B.
  • [9] T. Haarnoja, H. Tang, P. Abbeel, and S. Levine (2017) Reinforcement learning with deep energy-based policies. In Proceedings of the 34th International Conference on Machine Learning-Volume 70, pp. 1352–1361. Cited by: §IV-A, §IV-E, §IV-E, §IV-E.
  • [10] T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine (2018) Soft actor-critic: off-policy maximum entropy deep reinforcement learning with a stochastic actor. arXiv preprint arXiv:1801.01290. Cited by: §IV-E.
  • [11] R. Houthooft, X. Chen, Y. Duan, J. Schulman, F. De Turck, and P. Abbeel (2016) Vime: variational information maximizing exploration. In Advances in Neural Information Processing Systems, pp. 1109–1117. Cited by: 4th item, §V-A, §V-A, §V-B.
  • [12] L. P. Kaelbling, M. L. Littman, and A. R. Cassandra (1998) Planning and acting in partially observable stochastic domains. Artificial i ntelligence 101 (1-2), pp. 99–134. Cited by: §I-C2.
  • [13] D. P. Kingma and M. Welling (2013) Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114. Cited by: §I-B.
  • [14] D. Koller and N. Friedman (2009) Probabilistic graphical models: principles and techniques. MIT press. Cited by: 1st item, 2nd item.
  • [15] S. Levine (2018) Reinforcement learning and control as probabilistic inference: tutorial and review. arXiv preprint arXiv:1805.00909. Cited by: 4th item, §I-C1, §IV-A, §IV-B, §IV-B, §IV-B, §IV-B, §IV-C, §IV-C, §V.
  • [16] T. P. Lillicrap, J. J. Hunt, A. Pritzel, N. Heess, T. Erez, Y. Tassa, D. Silver, and D. Wierstra (2015) Continuous control with deep reinforcement learning. arXiv preprint arXiv:1509.02971. Cited by: §II-B2.
  • [17] C. J. Maddison, A. Mnih, and Y. W. Teh (2016) The concrete distribution: a continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712. Cited by: §V-B.
  • [18] V. Mnih, A. P. Badia, M. Mirza, A. Graves, T. Lillicrap, T. Harley, D. Silver, and K. Kavukcuoglu (2016) Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928–1937. Cited by: §II-B1.
  • [19] V. Mnih, K. Kavukcuoglu, D. Silver, A. A. Rusu, J. Veness, M. G. Bellemare, A. Graves, M. Riedmiller, A. K. Fidjeland, G. Ostrovski, et al. (2015) Human-level control through deep reinforcement learning. Nature 518 (7540), pp. 529. Cited by: §I, §II-A, §II-B2, §II-B5.
  • [20] B. Sallans and G. E. Hinton (2004) Reinforcement learning with factored states and actions. Journal of Machine Learning Research 5 (Aug), pp. 1063–1088. Cited by: §III, §III, §IV-A.
  • [21] T. Schaul, D. Horgan, K. Gregor, and D. Silver (2015) Universal value function approximators. In International Conference on Machine Learning, pp. 1312–1320. Cited by: §II-B3.
  • [22] T. Schaul, J. Quan, I. Antonoglou, and D. Silver (2015) Prioritized experience replay. arXiv preprint arXiv:1511.05952. Cited by: §II-B5.
  • [23] J. Schulman, S. Levine, P. Abbeel, M. Jordan, and P. Moritz (2015) Trust region policy optimization. In International conference on machine learning, pp. 1889–1897. Cited by: §II-B1, §V-A.
  • [24] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov (2017) Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. Cited by: §II-B1.
  • [25] D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, and M. Riedmiller (2014) Deterministic policy gradient algorithms. Cited by: §II-B2.
  • [26] X. Sun, Y. Wang, A. Gossmann, and B. Bischl (2019) Resampling-based assessment of robustness to distribution shift for deep neural networks. arXiv preprint arXiv:1906.02972. Cited by: §I-B.
  • [27] Y. Sun, F. Gomez, and J. Schmidhuber (2011) Planning to be surprised: optimal bayesian exploration in dynamic environments. In International Conference on Artificial General Intelligence, pp. 41–51. Cited by: §V-A.
  • [28] R. S. Sutton, A. G. Barto, et al. (1998) Introduction to reinforcement learning. Vol. 2, MIT press Cambridge. Cited by: 4th item, §I-C1, §I-C2, §I-D, §I-D, §I-D, §I-E, §I-E, §III, §IV-E, §V-B.
  • [29] H. Van Hasselt, A. Guez, and D. Silver (2016) Deep reinforcement learning with double q-learning. In Thirtieth AAAI conference on artificial intelligence, Cited by: §II-A.
  • [30] R. Zhao, X. Sun, and V. Tresp (2019) Maximum entropy-regularized multi-goal reinforcement learning. arXiv preprint arXiv:1905.08786. Cited by: §I-C1, §II-B3, §II-B5.