A Policy Update Strategy in Model-free Policy Search: Expectation-Maximization

Expectation-Maximization Algorithm

Policy gradient methods require the user to specify the learning rate which can be problematic and often results in an unstable learning process or slow convergence. By formualting policy search as an inference problem with latent variables and using the EM algorithm to infer a new policy, this problem can be avoided since no learning rate is required.

The standard EM algorithm, which is well-known for determining the maximum likelihood solution of a probabilitistic latent variable model, takes the parameter update as a weighted maximum likelihood estimate which has a closed form solution for most of the used polices.

Let’s assume that :

yzpθ(y,z): observed random variable : unobserved random variable : parameterized joint distribution $$\begin{array}{rl}y& :\text{ observed random variable }\\ z& :\text{ unobserved random variable }\\ p_{\theta }(y,z)& :\text{ parameterized joint distribution }\end{array}$$

Given a data set Y=[y[1],…,y[N]]T$Y=[y^{[1]},\dots ,y^{[N]}{]}^T$ , we wanna approximate the parameter θ$\theta$ which means maximizing the log likelihood:

maxθlogpθ(Y,Z)$$\underset{\theta }{max}\log p_{\theta }(Y,Z)$$

Since Z$Z$ is latent variable we cannot solve the maximization problem directly. But by computing the expectation of Z$Z$ , we can maximize the log-marginal likelihood of Y$Y$ :

logpθ(Y)=∑i=1Nlogpθ(y[i])=∑i=1Nlog∫pθ(y[i],z)dz$$\log p_{\theta }(Y)=\sum _{i=1}^N\log p_{\theta }(y^{[i]})=\sum _{i=1}^N\log \int p_{\theta }(y^{[i]},z)dz$$

However, it is evident that we cannot obtain a closed-form solution for the parameter θ$\theta$ of our probability model pθ(y,z)$p_{\theta }(y,z)$ . Do you know what is closed-form solution?

Closed form solution :

An equation is said to be a closed-form solution if it solves a given problem in terms of functions and mathematical operations from a given generally-accepted set. For example, an infinite sum would generally not be considered closed-form. However, the choice of what to call closed-form and what not is rather arbitrary since a new “closed-form” function could simply be defined in terms of the infinite sum.

EM(Expectation-Maximization) is a powerful method to estimate parameterized latent variables. The basic idea behind is that if parameter θ$\theta$ is known then we can estimate the optimal latent variable Z$Z$ in view of Y$Y$ (E-Step); If latent variable Z$Z$ is known, we can estimateθ$\theta$ by maximum likelihood estimation. EM method can be seen as a kind of coordinate descent method to maximize the lower-bound of the log likelihood.

The iterative procedure for estimating the maximum log-likelihood consists of two main segments: Expectation Steps and Maximization Steps as mentioned above. Assume that we begin at the θ0${\theta }^0$ . Then we execute the following iterative steps:

• Based on θt${\theta }^t$ estimating the expectation of the latent variable Zt$Z^t$ .
• Based on Y$Y$ and Zt$Z^t$ estimating the parameter θt+1${\theta }^{t+1}$ by maximum likelihood estimation.

In general, we do not feel like the expectation of Z$Z$ but the distribution of Z$Z$ , i.e. pθt(Z|Y)$p_{{\theta }^t}(Z|Y)$ . To be specific, let’s introduce an auxiliary distribution q(Z)$q(Z)$ , which is variational, to decompose the marginal log-likelihood by using the identity pθ(Y)=pθ(Y,Z)/pθ(Z|Y)$p_{\theta }(Y)=p_{\theta }(Y,Z)/p_{\theta }(Z|Y)$ :

logpθ(Y)=====logpθ(Y)∫q(Z)dZ∫q(Z)logpθ(Y)dZ∫q(Z)logq(Z)pθ(Y,Z)q(Z)pθ(Z|Y)dZ∫q(Z)logpθ(Y,Z)q(Z)dZ+∫q(Z)logq(Z)pθ(Z|Y)dZLθ(q)+KL(q(Z)∥pθ(Z|Y))$$\begin{array}{rl}\log p_{\theta }(Y)=& \log p_{\theta }(Y)\int q(Z)dZ\\ =& \int q(Z)\log p_{\theta }(Y)dZ\\ =& \int q(Z)\log \frac{q(Z)p_{\theta }(Y,Z)}{q(Z)p_{\theta }(Z|Y)}dZ\\ =& \int q(Z)\log \frac{p_{\theta }(Y,Z)}{q(Z)}dZ+\int q(Z)\log \frac{q(Z)}{p_{\theta }(Z|Y)}dZ\\ =& {\mathcal{L}}_{\theta }(q)+KL(q(Z)\Vert p_{\theta }(Z|Y))\end{array}$$

Since the KL divergence is always larger or equal to zero, the term Lθ(q)${\mathcal{L}}_{\theta }(q)$ is a lower bound of the log-marginal likelihood.

E-Step

In E-step we update the variational distribution q(Z)$q(Z)$ by minimizing the KL divergence KL(q(Z)∥pθ(Z|Y))$KL(q(Z)\Vert p_{\theta }(Z|Y))$ , i.e. setting q(Z)=pθ(Z|Y)$q(Z)=p_{\theta }(Z|Y)$ . Note that the value of the log-likelihood logpθ(Y)$\log p_{\theta }(Y)$ has nothing to do with the variational distribution q(Z)$q(Z)$ . In summary, E-step :

Update q(Z)⇐ Minimize KL(q(Z)∥pθ(Z|Y))⇐ Set q(Z)=pθ(Z|Y)$$\text{Update }q(Z)\Leftarrow \text{ Minimize }KL(q(Z)\Vert p_{\theta }(Z|Y))\Leftarrow \text{ Set }q(Z)=p_{\theta }(Z|Y)$$

M-Step

In M-step we optimize the lower bound w.r.t. θ$\theta$ , i.e.

θnew=====argmaxθLθ(q)argmaxθ∫q(Z)logpθ(Y,Z)q(Z)dZargmaxθ∫q(Z)logpθ(Y,Z)dZ+H(q)argmaxθEq(Z)[logpθ(Y,Z)]argmaxθQθ(q)$$\begin{array}{rl}{\theta }_{new}=& \arg \underset{\theta }{max}{\mathcal{L}}_{\theta }(q)\\ =& \arg \underset{\theta }{max}\int q(Z)\log \frac{p_{\theta }(Y,Z)}{q(Z)}dZ\\ =& \arg \underset{\theta }{max}\int q(Z)\log p_{\theta }(Y,Z)dZ+H(q)\\ =& \arg \underset{\theta }{max}{\mathbb{E}}_{q(Z)}[\log p_{\theta }(Y,Z)]\\ =& \arg \underset{\theta }{max}{\mathcal{Q}}_{\theta }(q)\end{array}$$

where H(q)$H(q)$ denotes the entropy of q$q$ , Q$\mathcal{Q}$ is the expected complete data log-likelihood. The log now acts on the joint distribution directly. So M-step can be obtained in closed form. Moreover,

Qθ(q)=∫q(Z)logpθ(Y,Z)dZ=∑i=1N∫qi(z)logpθ(y[i],z)dZ$${\mathcal{Q}}_{\theta }(q)=\int q(Z)\log p_{\theta }(Y,Z)dZ=\sum _{i=1}^N\int q_i(z)\log p_{\theta }(y^{[i]},z)dZ$$

The M-step is based on a weighted maximum likelihood estimate of θ$\theta$ using the complete data points [y[i],z]$[y^{[i]},z]$ weighted by qi(z)$q_i(z)$ . In summary, M-step:

Update θ⇐ Maximize Qθ(q)$$\text{Update }\theta \Leftarrow \text{ Maximize }{\mathcal{Q}}_{\theta }(q)$$

Reformulate Policy Search as an Inference Problem

Let’s assume that :

Binary reward event RTrajectory τ: observed variable: unobserved variable$$\begin{array}{rl}\text{Binary reward event }R& :\text{ observed variable}\\ \text{Trajectory }\tau & :\text{ unobserved variable}\end{array}$$

Maximizing the reward implies maximizing the probability of the reward event, and, hence, our trjectory distribution pθ(τ)$p_{\theta }(\tau )$ needs to assign high probability to trajectories with high reward probability p(R=1|τ)$p(R=1|\tau )$ .

We would like to find a parameter vector θ$\theta$ that maximizes the probability of the reward event, i.e.

logpθ(R)=∫τp(R|τ)pθ(τ)dτ$$\log p_{\theta }(R)={\int }_{\tau }p(R|\tau )p_{\theta }(\tau )d\tau$$

As for the standard EM algorithm, a variational distribution q(τ)$q(\tau )$ is used to decompose the log-marginal likelihood into tow terms:

logpθ(R)=Lθ(q)+KL(q(τ)∥pθ(τ|R))$$\log p_{\theta }(R)={\mathcal{L}}_{\theta }(q)+KL(q(\tau )\Vert p_{\theta }(\tau |R))$$

where the reward-weighted trajectory distribution:

pθ(τ|R)=p(R|τ)pθ(τ)pθ(R)=p(R|τ)pθ(τ)∫p(R|τ)pθ(τ)dτ∝p(R|τ)pθ(τ)$$p_{\theta }(\tau |R)=\frac{p(R|\tau )p_{\theta }(\tau )}{p_{\theta }(R)}=\frac{p(R|\tau )p_{\theta }(\tau )}{\int p(R|\tau )p_{\theta }(\tau )d\tau }\propto p(R|\tau )p_{\theta }(\tau )$$

E-Step

Update q(τ)⇐ Minimize KL(q(τ)∥pθ(τ|R))⇐ Set q(τ)=pθ(τ|R)$$\text{Update }q(\tau )\Leftarrow \text{ Minimize }KL(q(\tau )\Vert p_{\theta }(\tau |R))\Leftarrow \text{ Set }q(\tau )=p_{\theta }(\tau |R)$$

M-Step

θnew=argmaxθLθ(q)=argmaxθ∫q(τ)logpθ(R,τ)q(τ)dτ=argmaxθ∫q(τ)logpθ(R,τ)dτ+H(q)=argmaxθ∫q(τ)log(p(R|τ)pθ(τ))dτQθ(q)=argmaxθ∫q(τ)logpθ(τ)dτ+f(q)=argminθ∫q(τ)(−logpθ(τ))dτ=argminθ[∫q(τ)logq(τ)pθ(τ)dτ+∫q(τ)log1q(τ)dτ]=argminθKL(q(τ)∥pθ(τ))$$\begin{array}{rl}{\theta }_{new}& =\arg \underset{\theta }{max}{\mathcal{L}}_{\theta }(q)\\ & =\arg \underset{\theta }{max}\int q(\tau )\log \frac{p_{\theta }(R,\tau )}{q(\tau )}d\tau \\ & =\arg \underset{\theta }{max}\int q(\tau )\log p_{\theta }(R,\tau )d\tau +H(q)\\ & =\arg \underset{\theta }{max}\underset{{\mathcal{Q}}_{\theta }(q)}{\underbrace{\int q(\tau )\log (p(R|\tau )p_{\theta }(\tau ))d\tau }}\\ & =\arg \underset{\theta }{max}\int q(\tau )\log p_{\theta }(\tau )d\tau +f(q)\\ & =\arg \underset{\theta }{min}\int q(\tau )(-\log p_{\theta }(\tau ))d\tau \\ & =\arg \underset{\theta }{min}\left[\int q(\tau )\log \frac{q(\tau )}{p_{\theta }(\tau )}d\tau +\int q(\tau )\log \frac{1}{q(\tau )}d\tau \right]\\ & =\arg \underset{\theta }{min}KL(q(\tau )\Vert p_{\theta }(\tau ))\end{array}$$

i.e.

Update θ⇐ Maximize Qθ(q)⇐ Minimize KL(q(τ)∥pθ(τ))$$\text{Update }\theta \Leftarrow \text{ Maximize }{\mathcal{Q}}_{\theta }(q)\Leftarrow \text{ Minimize }KL(q(\tau )\Vert p_{\theta }(\tau ))$$

EM-based Policy Search Algorithms

Monte-Carlo EM-based Policy Search

MC-EM-algorithm uses a sample-based approximation for the variational distribution q$q$ , i.e. in the E-step, MC-EM minimizes the KL divergence KL(q(Z)∥pθ(Z|Y))$KL(q(Z)\Vert p_{\theta }(Z|Y))$ by using samples Zj∼pθ(Z|Y)$Z_j\sim p_{\theta }(Z|Y)$ . Subsequently, these samples Zj$Z_j$ are used to estimate the expectation of the complete data log-liklihood:

Qθ(q)=∑j=1Klogpθ(Y,Zj)$${\mathcal{Q}}_{\theta }(q)=\sum _{j=1}^K\log p_{\theta }(Y,Z_j)$$

In terms of policy search, MC_EM methods use samples τ[i]${\tau }^{[i]}$ from the old trajectory distribution pθ′$p_{{\theta }^{\prime }}$ to represent the variational distribution q(τ)∝p(R|τ)pθ′(τ)$q(\tau )\propto p(R|\tau )p_{{\theta }^{\prime }}(\tau )$ over trajectories. As τ[i]${\tau }^{[i]}$ has already been sampled from pθ′(τ)$p_{{\theta }^{\prime }}(\tau )$ , q(τ[i])∝p(R|τ[i])$q({\tau }^{[i]})\propto p(R|{\tau }^{[i]})$ . Consequently, in the M-step, we maximize:

Qθ(θ′)=∑τ[i]∼pθ′(τ)p(R|τ[i])logpθ(τ[i])$${\mathcal{Q}}_{\theta }({\theta }^{\prime })=\sum _{{\tau }^{[i]}\sim p_{{\theta }^{\prime }}(\tau )}p(R|{\tau }^{[i]})\log p_{\theta }({\tau }^{[i]})$$

There are Episode-based EM-algorithms such as Reward-Weighted Regression(RWR) and Cost-Regularized Kernel Regression(CrKR), and Step-based EM-algorithms such as Episodic Reward-Weighted Regression(eRWR) and Policy Learning by Weighting Exploration with Returns(PoWER).

Variational Inference-based Methods

The MC-EM approach uses a weighted maximum likelihood estimate to obtain the new parameters θ$\theta$ of the policy. It averages over several modes of the reward function. Such a behavior might result in slow convergence to good policies as the average of several modes might be in an area with low reward.

The maximization used for the MC-EM approach is equivalent to minimizing:

KL(p(R|τ)pθ′(τ)∥pθ(τ))=∫p(R|τ[i])pθ′(τ[i])logp(R|τ)pθ′(τ[i])pθ(τ[i])$$KL(p(R|\tau )p_{{\theta }^{\prime }}(\tau )\Vert p_{\theta }(\tau ))=\int p(R|{\tau }^{[i]})p_{{\theta }^{\prime }}({\tau }^{[i]})\log \frac{p(R|\tau )p_{{\theta }^{\prime }}({\tau }^{[i]})}{p_{\theta }({\tau }^{[i]})}$$

w.r.t. parameter θ$\theta$ . This minimization is also called the Moment Projection of the reward-weighted trajectory distribution as it matches the moments of pθ(τ)$p_{\theta }(\tau )$ with the moments of p(R|τ)pθ′(τ)$p(R|\tau )p_{{\theta }^{\prime }}(\tau )$ .

Alternatively, we can use the Information projection argminθKL(pθ(τ)∥p(R|τ)pθ′(τ))$\arg \underset{\theta }{min}KL(p_{\theta }(\tau )\Vert p(R|\tau )p_{{\theta }^{\prime }}(\tau ))$ to update the policy. This projection forces the new trajectory distribution pθ(τ)$p_{\theta }(\tau )$ to be zero everywhere where the reward-weighted trajectory distribution is zero.

• Thanks J. Peters et al for their great work of A Survey on Policy Search for Robotics .
• 感謝周志華——《機器學習》清華大學出版社