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

• 2019-10-28 23:28:17
• Xudong Sun, Bernd Bischl
• 0

Abstract

Probabilistic Graphical Models and Variational Inference play an importantrole in recent advances in Deep Reinforcement Learning. As a self-inclusivetutorial survey, this article illustrates basic concepts of reinforcementlearning with Probabilistic Graphical Models and offers derivation of somebasic formula as a recap. Reviews and comparisons on recent advances in deepreinforcement learning are made from various aspects. We offer ProbabilisticGraphical Models, detailed explanation and derivation to several use cases ofGraphical Model and Variational Inference, which serves as a complementarymaterial on top of the original contributions.

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 Models and Variational Inference play an important role in recent advances in Deep Reinforcement Learning. As a self-inclusive tutorial survey, this article illustrates basic concepts of reinforcement learning with Probabilistic Graphical Models and offers derivation of some basic formula as a recap. Reviews and comparisons on recent advances in deep reinforcement learning are made from various aspects. We offer Probabilistic Graphical Models, detailed explanation and derivation to several use cases of Graphical Model and Variational Inference, which serves as a complementary material on top of the original contributions.

Probablistic Graphical Models; Variational Inference; Deep Reinforcement Learning

I Introduction

Reinforcement Learning, powered by Deep Neural Networks, has gaining increasing attention recently due to its great success in complicated tasks like games [19] and robot locomotion [23], as well as optimization tasks like Automatic Machine Learning [28]. Despite the existing survey [2] in Deep Reinforcement Learning, this paper, however, focuses on the application of Probabilistic Graphical Model and Variational Inference, especially Amortized Variational Inference [14]in Deep Reinforcement Learning. Specifically, we make the following contributions:

• We provide Probabilistic Graphical Models [3] for many basic concepts of Reinforcement Learning, which is rarely covered in literature. We also provide Graphcial Models to some recent works on Deep Reinforcement Learning [11, 7] which does not exist in the original contribution.

• We cover a taxonomy of Graphical Model and Variational Inference [4] used in Deep Reinforcement Learning and give detailed derivation for many of the critical equations, which is either not given in the original contributions [30, 15, 11], or in a slightly different way. This makes the paper in a relative standalone position to serve as a self-inclusive tutorial from beginner to advanced.

I-A Organization of the paper

In section I-B, we first introduce the fundamentals of Probabilistic Graphical Models and Variational Inference, then we review the basics about reinforcement learning by connecting probabilistic graphical models (PGM) in section II-A, as well as the basics and a incomplete overview about deep reinforcement learning, accompanied with a comparison of different methods in section II-E. 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] through d-separation [3]. 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 use $A^{act}$ explicitly to represent action. We use $(B\mathrel{\text{\scalebox{1.07}{\perp\mskip-10.0mu \perp}}}C)\mid A$ to represent $B$ is conditionally independent from $C$, given $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, 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 like variational resampling [27]. VI is also used in approximating the posterior on the weights distribution of neural networks for Thompson Sampling to tackle the exploration-exploitationtrade off in bandit problems [5], as well as approximating on the activations distribution like Variational AutoEncoder using Amortized VI [14].

II Reinforcement Learning and Deep Reinforcement Learning

II-A Basics about Reinforcement Learning with graphical model

II-A1 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 $S_{t}$, By executing an action $a_{t}$ according to a policy [30] $\pi(a|S_{t})$, the agent jumps to another state $S_{t+1}$, while receiving a reward $R_{t}$. Let discount factor $\gamma$ 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 [30], as well as reducing variance in Monte Carlo setting [15]. The goal is to maximize the accumulated rewards, $G=\sum_{t=0}^{T}\gamma^{t}R_{t}$ 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\rightarrow\infty$ correspond to continuous non-episodic reinforcement learning. We use another convention of $t\in\{0,\cdots,\infty\}$ 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 $\{(S_{t},A^{act}_{t},R_{t})\}$ in an episode, which is coined trajectory $\tau$ [33, 6]. Figure 2 illustrates part of a trajectory in one rollout. The state space $\mathcal{S}$ and action space $\mathcal{A}$, 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 $\mathcal{R}$.

II-A2 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 $\pi(a|s)=p(a|s)$, within a particular environment characterized by state transition probability $p(s_{t+1}|s_{t},a)$ and reward distribution function $p(r_{t}|s_{t},a_{t})$, a learning agent could observe different trajectories with different unrolling realizations. This is usually modeled as a Markov Decision Process [30], 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)

 $\displaystyle p(\tau)=\prod_{t=0}^{T}p(s_{t+1}|s_{t},a_{t})p(r_{t}|s_{t},a_{t}% )\pi(a_{t}|s_{t})$ (1)

Upon observation of a state $s_{t}$ in Figure 3, the action at the time step in question is conditionally independent with the state and action history $\mathcal{E}_{t}=\{S_{0},A^{act}_{0},\cdots,S_{t-1}\}$, which could be denoted as $(A^{act}_{t}\mathrel{\text{\scalebox{1.07}{\perp\mskip-10.0mu \perp}}}% \mathcal{E}_{t})\mid S_{t}$.

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 $O_{t}$ through a non invertible function of the latent state $S_{t}$ and the previous action $a_{t-1}$, as indicated the Figure by $p(o_{t}|s_{t},a_{t-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 $O_{t}$ is no longer Markov, but depends on the whole history. However, by introducing a probability distribution $b(S)$ over the hidden state $S$, with $\sum_{\mathcal{S}}b(S)=1$, which is termed belief state [12], where state $S$ takes value in range $\mathcal{S}$.

II-B Value Function, Bellman Equation, Policy Iteration

Define state value function of state $s\in\mathcal{S}$ in Equation (2), where the corresponding Bellman Equation is derived in Equation (3).

 $\displaystyle V^{\pi}(s)$ $\displaystyle=$ $\displaystyle E_{\pi,\varepsilon}[\sum_{i=0}^{\infty}\gamma^{i}R_{t+i}(S_{t+i}% ,A^{act}_{t+i})\mid\forall S_{t}=s]$ (2) $\displaystyle=$ $\displaystyle E_{\pi,\varepsilon}[R_{t}(S_{t},A^{act}_{t})+\gamma\sum_{i=1}^{% \infty}\gamma^{(i-1)}R_{t+i}(S_{t+i},A^{act}_{t+i})]$ $\displaystyle=$ $\displaystyle E_{\pi,\varepsilon}[R_{t}(S_{t},A^{act}_{t})+\gamma\sum_{i^{{}^{% \prime}}=0}^{\infty}\gamma^{i^{{}^{\prime}}}R_{t+1+i^{{}^{\prime}}}(S_{t+i^{{}% ^{\prime}}+1},A^{act}_{t+i^{{}^{\prime}}+1})]$ $\displaystyle=$ $\displaystyle E_{\pi,\varepsilon}[R_{t}(S_{t},A^{act}_{t})+\gamma V^{\pi}(S_{t% +1})]$ (3)

where $S_{t+i}\sim p(s_{t+i+1}|s_{t+i},a_{t+i})$ takes value from $\mathcal{S}$, $A^{act}_{t+i}\sim\pi(a|S_{t+i+1})$ taking value from $\mathcal{A}$, and we have used the $\pi$ and $\varepsilon$ 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 [30] is defined in Equation (4),

 $\displaystyle Q^{\pi}(s,a)\,(\forall S_{t}=s,A^{act}_{t}=a)$ $\displaystyle=$ $\displaystyle E_{\pi,\varepsilon}[R_{t}(S_{t}=s,A^{act}_{t}=a)+\sum_{i=1}^{% \infty}\gamma^{i}R_{t+i}(S_{t+i},A^{act}_{t+i})]$ (4) $\displaystyle=$ $\displaystyle E_{\pi,\varepsilon}[R_{t}(S_{t}=s,A^{act}_{t}=a)+\gamma V^{\pi}(% S_{t+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)=\sum_{a}\pi(a|s)Q(s,a)$ (6)

Define optimal policy [30] to be

 $\displaystyle\pi^{*}$ $\displaystyle=\underset{\pi}{arg\,max}\,V^{\pi}(s),\forall s\in S$ $\displaystyle=\underset{\pi}{arg\,max}\,E_{\pi}[R_{t}+\gamma V^{\pi}(S_{t+1})]$ (7)

Taking the optimal policy $\pi^{*}$ into the Bellman Equation in Equation (3), we have

 $\displaystyle V^{\pi^{*}}(s)=E_{\pi^{*},\varepsilon}\left[R_{t}(s,A^{act}_{t})% +\gamma V^{\pi^{*}}(S_{t+1})\right]$ (8)

Taking the optimal policy $\pi^{*}$ into Equation (4), we have

 $\displaystyle Q^{\pi^{*}}(s,a)=E_{\pi^{*},\varepsilon}[R_{t}(s,a)+\sum_{i=1}^{% \infty}\gamma^{i}R_{t+i}(S_{t+i},A^{act}_{t+i})]$ (9)

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

 $\displaystyle V^{\pi^{*}}(s)=\underset{a}{max}\,Q^{\pi^{*}}(s,a)$ (10)

and

 $\displaystyle Q^{\pi^{*}}(s,a)=E_{\varepsilon,\pi^{*}}\left[R_{t}(s,a)+\gamma% \underset{\bar{a}}{max}\,Q^{\pi^{*}}(S_{t+1},\bar{a})\right]$ (11)

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

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 [30], while time difference learning [30] is more ubiquitously used, like a bridge connecting the two extremities. Time difference learning is based on the Bellman update error $\delta_{t}=Q(s_{t},a_{t})-\left(R_{t}(s,a)+\gamma\underset{a}{max}\,Q(s_{t+1},% a)\right)$.

II-C Policy Gradient and Actor Critic

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

Suppose all episodes start from an auxiliary initial state $s_{0}$, which with probability $h(s)$, jumps to different state $s\in\mathcal{S}$ without reward. $h(s)$ characterizes the initial state distribution which only depends on the environment. Let $\eta(s)$ represent the expected number of steps spent on state $s$, which can be calculated by summing up the $\gamma$ discounted probability $P^{\pi}(s_{0}\rightarrow s,k+1)$ of entering state $s$ with $k+1$ steps from auxiliary state $s_{0}$, as stated in Equation (12), which can be thought of as the expectation of the R.V. $\gamma^{k}$ conditional on state $s$.

 $\displaystyle\eta(s)$ $\displaystyle=\sum_{k=0}\gamma^{k}P^{\pi}(s_{0}\rightarrow s,k+1)$ (12) $\displaystyle=h(s)+\sum_{\bar{s},a}\gamma\eta(\bar{s})\pi_{\theta}(a|\bar{s})P% (s|\bar{s},a)$ (13)

In Equation (13), the quantity is calculated by either directly starting from state $s$, which correspond to $k=0$ in Equation (12), or entering state $s$ from state $\bar{s}$ with one step, corresponding to $k+1\geq 2$ in Equation (12).

For an arbitrary state $s\in\mathcal{S}$, using $s^{{}^{\prime}}$ and $s^{{}^{\prime\prime}}$ to represent subsequent states as dummy index, we have

The terms in square brackets in Equation (17) are simply Equation (16) with $a$ and $s^{{}^{\prime}}$ replaced by $a^{{}^{\prime}}$ and $s^{{}^{\prime\prime}}$. Since $\nabla_{\theta}V^{\pi(\theta)}(s^{\infty})=0$, Equation (17) could be written as Equation (18),

 $\displaystyle\nabla_{\theta}V^{\pi(\theta)}(s)$ $\displaystyle=$ $\displaystyle\sum_{a}\nabla_{\theta}\pi_{\theta}(a|s)Q^{\pi(\theta)}(s,a)+$ $\displaystyle\sum_{k=1}\sum_{s_{k}}\sum_{a_{k}}\gamma^{k}P^{\pi}(s\rightarrow s% _{k},k)\nabla_{\theta}\pi_{\theta}(a_{k}|s_{k})Q^{\pi(\theta)}(s_{k},a_{k})$ (18)

, where $s_{k}$ represent the state of $k$ steps after $s$ and $P^{\pi}(s\rightarrow s_{k},k)$ already includes integration of intermediate state $s_{k-1},\ldots s_{1}$ before reaching state $s_{k}$.

Let objective function with respect to policy be defined to be the value function starting from auxiliary state $s_{0}$ as in Equation (19).

 $\displaystyle J(\pi_{\theta})=V^{\pi}(s_{0})=E_{\pi,\varepsilon}\sum_{t=0}^{% \infty}\gamma^{t}R_{t}(S_{0}=s)$ (19)

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

 $\displaystyle\nabla_{\theta}J(\pi_{\theta})$ $\displaystyle=$ $\displaystyle\nabla_{\theta}V^{\pi}(s_{0})$ $\displaystyle=$ $\displaystyle\sum_{k=0}\sum_{s_{k}}\sum_{a_{k}}\gamma^{k}P^{\pi}(s_{0}% \rightarrow s_{k},k)\nabla_{\theta}\pi_{\theta}(a_{k}|s_{k})Q^{\pi(\theta)}(s_% {k},a_{k})$ $\displaystyle=$ $\displaystyle\sum_{s}\sum_{a}\eta(s)\nabla_{\theta}\pi_{\theta}(a|s)Q^{\pi(% \theta)}(s,a)$ (20) $\displaystyle=$ $\displaystyle\frac{\sum_{s}\eta(s)}{\sum_{s}\eta(s)}\sum_{s}\sum_{a}\eta(s)% \nabla_{\theta}\pi_{\theta}(a|s)Q^{\pi(\theta)}(s,a)$ (21) $\displaystyle=$ $\displaystyle\sum_{\bar{s}}\eta(\bar{s})\sum_{s}\sum_{a}\mu(s)\nabla_{\theta}% \pi_{\theta}(a|s)Q^{\pi(\theta)}(s,a)$ (22) $\displaystyle=$ $\displaystyle\sum_{\bar{s}}\eta(\bar{s})\sum_{s}\sum_{a}\mu(s)\frac{\pi_{% \theta}(a|s)}{\pi_{\theta}(a|s)}\nabla_{\theta}\pi_{\theta}(a|s)Q^{\pi(\theta)% }(s,a)$ (23) $\displaystyle\propto$ $\displaystyle\,E_{\pi}\left[\frac{\nabla_{\theta}\pi_{\theta}(A|S)}{\pi_{% \theta}(A|S)}\hat{Q}^{\pi(\theta)}(S,A)\right]$ (24)

, where $\mu(s)$ is the relative occupancy of state $s$. The integration of $\mu(s)$ with respect to $s$ and $\pi_{\theta}(a|s)$ in the nominator with respect to $a$ in Equation (23) is replaced with expectation with respect to interaction with the environment in Equation (24), and $Q^{\pi(\theta)}(s,a)$ is replaced by an estimator $\hat{Q}^{\pi(\theta)}(S,A)$, which is usually $G_{t}=\sum_{i=0}^{\infty}\gamma^{i}R_{t+i}$. Functional approximation $f_{w}(s,a)$ to $Q$ should be compatible with policy parameterization in the sense that $\frac{\partial f_{w}(s,a)}{\partial w}=\frac{\partial\pi(s,a)}{\partial\theta}% \frac{1}{\pi(s,a)}$ [31].

The policy gradient could be augmented to include zero gradient baseline $b(s)$, with respect to objective function $J(\pi_{\theta})$ in Equation (23), as a function of state $s$, which does not include parameters for policy $\theta$, since $\sum_{a}\nabla_{\theta}\pi_{\theta}(a|s)=0$. To reduce variance of the gradient, the baseline is usually chosen to be the state value function estimator $\hat{V}_{w}(s)$ to smooth out the variation of $Q(s,a)$ at each state, while $\hat{V}_{w}(s)$ is updated in a Monte Carlo way by comparing with $\hat{Q}^{\pi_{\theta}}(S,A)=G_{t}$.

The actor-critic algorithm [30] decomposes $G_{t}-V_{w}(s_{t})$ to be $R_{t}+\gamma V_{w}(s_{t+1})-V_{w}(s_{t})$, so bootstrap is used instead of Monte Carlo.

II-D Policy Optimization and Trust Region Method

Another category of policy optimization methods is built on the identity from [13] in Equation (25), where $\eta_{\pi^{new}}(s)$ means the state visitation frequency under policy $\pi^{new}$ and advantage $A^{\pi^{old}}(a_{t},s_{t})=Q^{\pi^{old}}(a_{t},s_{t})-V^{\pi^{old}}(s_{t})$.

Based on Policy Advantage [13] $A_{\pi^{old}}(\pi^{new})=\sum_{s}\eta_{\pi^{old}}(s)\sum_{a}\pi^{new}(a|s)A^{% \pi^{old}}(a_{t},s_{t})$, a local approximation $L_{\pi^{old}}(\pi^{new})$ could be defined in Equation (26), which could form a surrogate function $M$ in Equation (27) that minorizes $J^{\pi^{new}}$ at $\pi^{old}$, where $D_{KL}^{max}(\pi^{old},\pi^{new})=\underset{s}{max}D_{KL}(\pi^{old}(a|s),\pi^{% new}(a|s))$ is the maximum KL divergence, so MM algorithm could be used to improve the policy, leading to the trust region method [23].

 $\displaystyle L_{\pi^{old}}(\pi^{new})$ $\displaystyle=$ $\displaystyle J(\pi^{old})+\sum_{s}\eta_{\pi^{old}}(s)\sum_{a}\pi^{new}(a|s)A^% {\pi^{old}}(a_{t},s_{t})$ (26)
 $\displaystyle M=L_{\pi^{old}}(\pi^{new})-\frac{4\underset{a,s}{max}|A_{\pi^{% old}}(s,a)|\gamma}{1-\gamma^{2}}D_{KL}^{max}(\pi^{old},\pi^{new})$ (27)

II-E 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 $R_{t}+\gamma\underset{a}{max}\,Q(s_{t+1},a)$ from the replay memory with state $s_{t}$ 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 [32], however, solves the problem by having two Q network and update the parameters in a alternating way.

II-E1 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-E2 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-E3 Goal based Reinforcement Learning

In robot manipulation tasks, the goal could be represented with state in some cases [33]. 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, similar to Recommendation System. Work of this direction include [1, 33], for example.

II-E4 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 [33] gives priority to those rarely occurred trajectory in sampling, which has been shown to improve over HER [33].

II-E5 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. 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.

II-E6 Surrogate optimization

Like surrogate model used in Bayesian Optimization [26], lower bound surrogate is also used in Reinforcement Learning, including Evidence Lower Bound (ELBO) based methods introduced below and Trust Region Policy Optimization (TRPO) [23].

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 $w_{k,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.

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

 $\displaystyle p(a|s)=\frac{1}{Z}e^{-\sum_{h}F(s,a,h)/T}=\frac{1}{Z}e^{\sum_{h}% Q(s,a)/T}$ (28)

, 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 (28) is learned as a critic [30], such that its associated policy is defined, MCMC sampling [3] could be used to sample actions, as an actor [30]. With the sampled actions, time difference learning method like SARSAR [30], 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(O_{t}=1|s_{t},a_{t})=\exp(r(s_{t},a_{t}))$, 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 (29):

 $\displaystyle p(\tau)$ $\displaystyle=p(s_{1})\prod_{t}[p(s_{t+1}|s_{t},a_{t})p(O_{t}=1|s_{t},a_{t})]$ $\displaystyle=\left[p(s_{1})\prod_{t}p(s_{t+1}|s_{t},a_{t})\right]exp\left(% \sum_{t}r(s_{t},a_{t})\right)$ (29)

.

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(a_{t}|s_{t},O_{t:T}=1)$, which reads as, conditional on optimality from current time step until end of episode, and the current current state to be $s_{t}$, the distribution of action $a_{t}$, and this posterior corresponds to the optimal policy. Observing the d-separation from Figure 7, $O_{1:{t-1}}$ is conditionally independent of $a_{t}$ given $s_{t}$, $(O_{1:{t-1}}\mathrel{\text{\scalebox{1.07}{\perp\mskip-10.0mu \perp}}}A^{act% }_{t})\mid{S_{t}}$, so $p(a_{t}|s_{t},O_{1:t-1}=,O_{t:T}=1)=p(a_{t}|s_{t},O_{t:T})$

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 $(a_{t},s_{t})$ as the latent variable counterpart of a HMM, with emission probability $p(O_{t}=1|s_{t},a_{t})=exp(r(s_{t},a_{t}))$, while the transition probability, is from the variable tuple $(a_{t},s_{t})$ to a subcomponent $s_{t+1}$ of the ”latent” variable tuple $(a_{t+1},s_{t+1})$.

Similar to the forward-backward message passing algorithm [3] in Hidden Markov Models [3], the posterior $p(a_{t}|s_{t},O_{t:T}=1)$ could also be calculated by passing messages. We offer a detailed derivation of the decomposition of the posterior $p(a_{t}|s_{t},O_{t:T}=1)$ in Equation (30), which is not available in [15].

In Equation (30), we define message $\beta(a_{t},s_{t})=p(O_{t:T}=1|a_{t},s_{t})p(a_{t}|s_{t})$ and message $\beta(s_{t})=\int_{a_{t}^{{}^{\prime}}}\beta(a_{t}^{{}^{\prime}},s_{t})d\{a_{t% }^{{}^{\prime}}\}$. If we consider $p(a_{t}|s_{t})$ as a prior with a trivial form [15], the only policy related term becomes $p(O_{t:T}=1|a_{t},s_{t})$.

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(S_{t}|O_{1:T})$ and $p(S_{t},S_{t+1}|O_{1:T})$, which is the marginal of the complete posterior $p(S_{1:T}|O_{1:T})$. The posterior marginal could be computed by the forward message $\alpha(S_{t})=p(O_{1:t},S_{t})$ and the backward message $\beta(S_{t})=p(O_{t:T}|S_{t})$, 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 $\beta(a_{t},s_{t})$ here is not a probability distribution as in HMM, instead, is just a probability. In Figure 7, the backward message $\beta(a_{t},s_{t})$ 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 (31).

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

IV-C Connection between Message Passing and Bellman equation

If we define Q function in Equation (32) and V function in Equation (33)

 $Q(s_{t},a_{t})=\log(\beta(a_{t},s_{t}))$ (32)

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

 $\pi(a_{t}|s_{t})=p(a_{t}|s_{t},O_{t:T}=1)=exp(Q(s_{t},a_{t})-V(s_{t}))$ (34)

.

Taking the logrithm of Equation (31), we get Equation (35)

 $\displaystyle\log(\beta(s_{t},a_{t}))$ $\displaystyle=\log\int\beta(s_{t+1})p(s_{t+1}|s_{t},a_{t})p(O_{t}=1|s_{t},a_{t% })ds_{t+1}$ $\displaystyle=\log\int exp[r(s_{t},a_{t})+V(s_{t+1})]p(s_{t+1}|s_{t},a_{t})ds_% {t+1}$ $\displaystyle=r(s_{t},a_{t})+\log\int exp(V(s_{t+1}))p(s_{t+1}|s_{t},a_{t})ds_% {t+1}$ (35)

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

 $Q(s_{t},a_{t})=r(s_{t},a_{t})+\log E_{s_{t+1}\sim p(s_{t+1}|s_{t},a_{t})}[\exp% (V(s_{t+1}))]$ (36)

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(a_{t}|s_{t})$, which could be represented by a neural network, to compose the proposal variational likelihood of the trajectory as in Equation (37):

 $\displaystyle q(\tau)$ $\displaystyle=p(s_{1})\prod_{t}[p(s_{t+1}|s_{t},a_{t})q(a_{t}|s_{t})]$ (37)

, where the initial state distribution $p(s_{1})$ 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 (38), which correspond to an interesting objective function of reward plus entropy return, as in Equation (39).

IV-E Examples

In [9], the state action value function is defined in Equation (40).

and a soft version of Bellman update similar to Q Learning [30] is carried out, which lead to policy improvement with respect to the corresponding functional objective in Equation (41).

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

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 on transition model

In Variational Information Maximizing Exploration (VIME) [11], where dynamic model $p_{\theta}(s_{t+1}|s_{t},a_{t})$ for the agent’s interaction with the environment is modeled using Bayesian Neural Network [5]. The R.V. for $\theta$ is denoted by $\Theta$, and is treated in a Bayesian way by modeling the weight uncertainty. We represent this model with the graphical model in Figure 8, which is not given in [11]. The belief about the environment is modeled as entropy of the neural network weights posterior distribution $H(\Theta|\xi_{t})$ based on trajectory observations $\xi_{t}=\{s_{1:t},a_{1:t-1}\}$. The method encourages taking exploratory actions by alleviating the average information gain of the agent’s belief about the environment after observing a new state $s_{t+1}$, which is $E_{p(s_{t+1}|\xi_{t},a_{t})}D_{KL}(p(\theta|\xi_{t+1})||p(\theta|\xi_{t}))$, and this is equivalent to the entropy minus conditional entropy $H(\Theta|\xi_{t},a_{t})-H(\Theta|\xi_{t},a_{t},s_{t+1})=H(\Theta|\xi_{t},a_{t}% )-H(\Theta|\xi_{t+1})$.

We derive in Equation (4344) that the entropy difference is actually the average information gain, which is equal to the mutual information $I(\Theta,S_{t+1}|\xi_{t},a_{t})$ between environmental parameter $\Theta$ and the new state $S_{t+1}$. Such a derivation is not given in [11].

Based on Equation (43), 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 $a_{t}$ and state $s_{t+1}$ pair on top of trajectory history $\xi_{t}$, the posterior on the distribution of the environmental parameter $\theta$, $p(\theta|\xi_{t})$, could be updated to be $p(\theta|\xi_{t+1})$ in a Bayesian way as derived in Equation (45), which is first proposed in [29].

In Equation (45), the denominator can be written as Equation (46), so that the dynamics of the environment modeled by neural network weights $\theta$, $p(s_{t+1}|\theta,a_{t},\xi_{t})$, could be used.

The last step of Equation (46) makes use of $p(\theta|\xi_{t},a_{t})=p(\theta|\xi_{t})$.

Since the integral in Equation (46) is not tractable, variational treatment over the neural network weights posterior distribution $p(\theta|\xi_{t})$ is used, characterized by variational parameter $\phi$, as shown in the dotted line in Figure 8. The variational posterior about the model parameter $\theta$, updated at each step, could than be used to calculate the intrinsic reward in Equation (43).

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 and discrete. The author assume a factorized form of observation and latent space joint probability, which we explicitly state in Equation (47).

 $\displaystyle p(O,S)=\pi_{\theta_{0}}(s_{0})\prod_{t=0}^{T}p_{\theta^{R}}(o_{t% }|s_{t})\prod_{t=1}^{T}p_{\theta^{T}}(s_{t}|s_{t-1},a_{t-1})$ (47)

Additionally, we characterize Equation (47) with the probabilistic graphical model in Figure 9 which does not exist in [7]. Compared to Figure 7, here the latent state $S$ is in discrete space instead of high dimension, and the observation is a high dimensional image instead of binary variable to indicate optimal action. By assuming a factorized form of the variational posterior in Equation (48),

 $\displaystyle q(S_{0:T}|O_{0:T})=\prod_{t=0}^{T}q_{\phi}(S_{t}|O_{t-k:t})$ (48)

The author assume the episode length to be $T$, and default frame prior observation to be blank frames. The Evidence Lower Bound (ELBO) of the observed trajectory of Equation (47) could be easily represented by a Varitional AutoEncoder [8] like architecture, where the encoder $q_{\phi}$, 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. Here prioritized sweeping [30] 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_{\theta^{T}}(s_{t}|s_{t-1},a_{t-1})$ in Equation (47). 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.

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. We offer some detailed derivations and Probabilistic Graphical Models to those methods using variational inference, which 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-E3, §II-E4.
• [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: 1st item, §I-B, §II-A2, §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: 2nd 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: §II-A1.
• [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: 1st item, §V-B.
• [8] C. Doersch (2016) Tutorial on variational autoencoders. arXiv preprint arXiv:1606.05908. Cited by: §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: 1st item, 2nd 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: §II-A2.
• [13] S. Kakade and J. Langford (2002) Approximately optimal approximate reinforcement learning. In ICML, Vol. 2, pp. 267–274. Cited by: §II-D, §II-D.
• [14] D. P. Kingma and M. Welling (2013) Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114. Cited by: §I-B, §I.
• [15] S. Levine (2018) Reinforcement learning and control as probabilistic inference: tutorial and review. arXiv preprint arXiv:1805.00909. Cited by: 2nd item, §II-A1, §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-E2.
• [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-E1.
• [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-E2, §II-E4, §II-E.
• [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-E3.
• [22] T. Schaul, J. Quan, I. Antonoglou, and D. Silver (2015) Prioritized experience replay. arXiv preprint arXiv:1511.05952. Cited by: §II-E4.
• [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: §I, §II-D, §II-E1, §II-E6, §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-E1.
• [25] D. Silver, G. Lever, N. Heess, T. Degris, D. Wierstra, and M. Riedmiller (2014) Deterministic policy gradient algorithms. Cited by: §II-E2.
• [26] X. Sun, A. Bommert, F. Pfisterer, J. Rahnenführer, M. Lang, and B. Bischl (2019) High dimensional restrictive federated model selection with multi-objective bayesian optimization over shifted distributions. arXiv preprint arXiv:1902.08999. Cited by: §II-E6.
• [27] X. Sun, A. Gossmann, Y. Wang, and B. Bischl (2019) Variational resampling based assessment of deep neural networks under distribution shift. Cited by: §I-B.
• [28] X. Sun, J. Lin, and B. Bischl (2019) ReinBo: machine learning pipeline search and configuration with bayesian optimization embedded reinforcement learning. Cited by: §I.
• [29] 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.
• [30] R. S. Sutton, A. G. Barto, et al. (1998) Introduction to reinforcement learning. Vol. 2, MIT press Cambridge. Cited by: 2nd item, §II-A1, §II-A2, §II-B, §II-B, §II-B, §II-C, §II-C, §III, §IV-E, §V-B.
• [31] R. S. Sutton, D. A. McAllester, S. P. Singh, and Y. Mansour (2000) Policy gradient methods for reinforcement learning with function approximation. In Advances in neural information processing systems, pp. 1057–1063. Cited by: §II-C.
• [32] 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-E.
• [33] R. Zhao, X. Sun, and V. Tresp (2019) Maximum entropy-regularized multi-goal reinforcement learning. arXiv preprint arXiv:1905.08786. Cited by: §II-A1, §II-E3, §II-E4.