Online Expectation Maximization for Reinforcement Learning in POMDPs

1 Introduction

The policy of a partially observable Markov decision process (POMDP) is a mapping from belief-states to actions [Kaelbling et al., 1998]. For a given POMDP, the optimal policy can be obtained by solving a Markov decision process in belief-states. In reinforcement learning (RL), however, the agent must learn a policy via experiencing with the POMDP, without assuming knowledge of the process's true underlying model. In this case, the policy can be obtained by using either of the following approaches: (i) a model-based approach [Poupart and Vlassis, 2008; Doshi-Velez et al., 2009; Doshi-Velez, 2010], in which the agent first learns the underlying model using its experience with the POMDP and then solves the learned model to obtain an approximate policy; (ii) a model-free approach [Li et al., 2009; Cai et al., 2009; Liu et al., 2011], in which the agent learns the policy directly from the experience, skipping the model-learning step.

The two approaches each have their advantages and disadvantages. Since the underlying model is a complete description of the POMDP, a learned model can be used not only to find the optimal policy but for other purposes also; thus, a model-based approach is versatile. However, the agent needs a comprehensive set of experiences with the POMDP to make the learned model complete, implying the modellearning step is not efficient in time. In addition, the agent needs a second step to solve the learned model to the get the policy. Although modern POMDP solvers can find approximate policies efficiently, the efficiency can be plagued when the solvers must be invoked repeatedly, as in sample-based methods [Doshi-Velez et al., 2009; Doshi-Velez, 2010].

A model-free approach, by contrast, focuses exclusively on the policy. Though one can find the optimal policy given the model, inverting the learned policy into the model is generally not possible. The fact that the policy is not as rich in information as the model implies: (i) learning a policy may need only a subset of the experience required in learning a complete underlying model of the POMDP; (ii) the policy can only be used for a subset of the purposes which the model can be used for. For these reasons, a model-free approach is typically more efficient in time, but less versatile, than its modelbased counterpart; moreover, the avoidance of a second step of model-to-policy conversion enhances the time efficiency.

The choice between a model-based approach and a modelfree approach is, therefore, a tradeoff between versatility and time efficiency. When the policy is the primary interest, the right choice would be an approach that is just versatile enough for finding the optimal policy, but not more. This idea has been pursued in model-based methods. For example, the work in [Wierstra and Wiering, 2004] approximates a POMDP as an utile distinction hidden Markov model which creates memory to preserve perceptual and utility distinction only when necessary. The recent work in [Grady et al., 2013] solves a less complex model, with reduced state and action spaces, to find approximate policies for the original POMDP. Both these methods solve a compressed model of the original POMDP to find approximate policies, enhancing time efficiency while reducing the model's versatility. A model-free approach can be considered as an extreme case, in which the POMDP is compressed into the policy itself, i.e., the most compressed "model" that preserves the policy.

Most model-free methods for RL in a POMDP are based on representing the policy as a finite state controller (FSC). Early work learns FSCs using stochastic gradient descent [Meuleau et al., 1999; Aberdeen and Baxter, 2002]. Recently, it has been shown in [Li et al., 2009] that expectation maximization (EM), a popular algorithm in statistics [Dempster et al., 1977], provides an effective tool for learning a family of policies that include FSCs as a special case. The policies considered in [Li et al., 2009], collectively referred to as regionalized policy representation (RPR), treat belief-state as a latent random variable and integrate it out to yield a marginalized policy that is expressed as a probability distribution of the current action conditional on the history of past actions and observations. The RPR is a general form of FSC, where each internal memory unit (called a decision state in the RPR and a machine node in the FSC) is associated with a distribution of actions, instead of a single action; as a result, the transition from unit z to unit z′ depends jointly on the action at z and the observation at z′ , instead of on the observation only. The action-dependency reflects the inherent uncertainty of the action choice in each belief region, which, as discussed in depth in [Sondik, 1978], cannot be resolved unless the policy is finite transient, in which case the RPR specializes to a FSC.

The optimal RPR can be learned by nested expectation maximization [Li et al., 2009], where the outer-loop EM evaluates the policy in the E-step and improves the policy in the M-step, and the inner-loop EM realizes the outer-loop Mstep. It is noteworthy that the nested EMmaximizes a value function, which cannot be interpreted as a probability of the observed variables as in the case of a likelihood function.

The nested EM is better interpreted as a consequence of the policy improvement theorem of a POMDP [Blackwell, 1965]. As detailed in [Li et al., 2009, pages 1140-1143], the outer-loop E-step re-computes the immediate rewards such that a discounted sum of the new rewards, averaged over the episodes, represents the value of the current policy; in addition, the inner-loop EM updates the policy by maximizing the weighted log-likelihood function of a cost-sensitive hidden Markov model, with the weights constituted by the recomputed new rewards. The weighted log-likelihood function corresponds to the logarithm of the action-value function in [Blackwell, 1965], parameterized by the RPR and averaged over all belief points, and the average value is maximized by the inner-loop EM to yield an improved policy.

The RPR has been introduced in [Li et al., 2009] as a tool for multi-task RL across multiple POMDPs. Subsequent work has focused on using the RPR for reinforcement learning in a single POMDP [Cai et al., 2009; Liu et al., 2011]. All previous work assumes batch-mode learning, updating the policy based on the entire set of experiences collected so far. Consequently, as new experiences arrive sequentially, the computational cost and memory demand increase fast.

In this paper, we aim to reduce the time and memory complexity of learning an RPR by using online nested expectation maximization. The major change from the batch-mode EM is that the policy is updated each time based on a partial policy evaluation, with only the latest episode of agent-environment interaction used in the outer-loop E-step and inner-loop Estep. Therefore each episode's contribution to the sufficient statistics is computed only when the episode is new, instead of computed repeatedly when each subsequent episode comes in; this makes the E-steps computationally more efficient. Furthermore, as each episode is used only when it is new, the algorithm need only memorize the latest episode, discarding all older ones, leading to tremendous memory savings.

2 Regionalized Policy Representation

Definition 1. [Li et al., 2009] A regionalized policy representation is a tuple (A,O,Z,W, μ, π), where A, O, and Z are respectively a finite set of actions, observations, and decision states; W is a set of Markov transition matrices, with W zao z′ denoting the probability of transiting from decision state z to z′ when action a in z results in observation o in z′ ; μ is the initial distribution of decision states, with μz the probability of initially being in z; π is a set of stochastic policies, with πza the probability of taking action a in z.

For simplicity, Z is denoted as {1, 2, · · · , |Z|}, where |Z| is the cardinality; similar notations are used for A and O. The set of RPR parameters is denoted as Θ = {π, μ,W}. A consecutively indexed variable is abbreviated as the variable with its index range, which is indicated in the subscript of the variable; for example, a0:T = (a0, a1, . . . , aT ), W zao 1:|Z| = (W zao 1 ,W zao 2 , · · · , W zao |Z| ), etc.

Given ht = {a0:t-1, o1:t}, the history of actions and observations up to t, the RPR chooses action at according to

where the second equality arises because ot has no influence on the actions before t, and p(a0:t|o1:t, Θ) is a marginal of

by integrating out latent decision states z0:t. It follows from (1) that p(a0:t|o1:t, Θ) = ∏tτ=0 p(aτ |hτ , Θ).

The RPR parameters are learned from the agent experiences by using an empirical value function defined below. Assuming the agent-POMDP interaction is episodic [Sutton and Barto, 1998], the experiences are represented as a set of episodes. An episode of length Tk is denoted by (ak0 rk 0 ok1 ak1 rk 1 · · · okTk akTk rkTk ), where r is a nonnegative immediate reward, k indexes the episodes, and the subscripts index discrete time steps.

Definition 2. [Li et al., 2009] Let D(K) = {(ak0 rk 0 ok1 ak1 rk 1 · · · okTk akTk rkTk )}Kk=1 be a set of episodes resulting from the interaction between the POMDP and an agent who chooses actions according to Π, an arbitrary stochastic policy with action-selecting distributions pΠ(a|h) > 0, ∀ action a, ∀ history h. The empirical value function is defined as

where hkt = (ak0:t-1 , ok1:t ), 0 < γ < 1 is the discount as defined in the POMDP.

It has been shown in [Li et al., 2009] that limK→∞ ˆ V (D(K) ;Θ) is the expected sum of discounted rewards by following the RPR parameterized by Θ for an infinite number of steps. Therefore, the RPR resulting from maximization of the empirical value function is an approximation of the optimal policy, assuming the number of decision states, i.e., |Z|, is large enough to accommodate the optimal policy. The optimal policy of a POMDP can be represented by a RPR because the RPR subsumes as a special case the finite state controller (FSC) [Li et al., 2009], which is known to approximate the optimal policy of any POMDP

to a arbitrary precision [Sondik, 1978]. The optimal |Z| can be inferred by the method in [Liu et al., 2011].

The batch-mode nested EM algorithm in [Li et al., 2009] is re-stated in Theorem 3.

Theorem 3. Let {Θn}n≥0 be a sequence produced by the iterative update Θn+1 = arg max ˆΘ∈Ξ LB( ˆΘ|Θn), where Θ0 is an arbitrary initialization and

then {Θn}n≥0 monotonically increases (3), until convergence to a maxima.

The quantity ν k t in (4) is the re-computed reward at time t in episode k; it represents the expected reward of following policy Θn, and is computed in the outer-loop E-step. The quantities in (5)-(6) are standard soft counts resulting from the Baum-Welch formula for learning an HMM [Rabiner, 1989], and they are computed in the inner-loop E-step. Therefore, (4) is computed less frequently than (5)-(6); that is, once (4) is computed, it is not computed again until the inner-loop EM converges to the desired precision. In practice, one may compute (4) after running the inner-loop EM only a few iterations, instead of waiting for its convergence; we find this works well in our experiments.

3 Online Nested ExpectationMaximization

To prepare for the online EM algorithm, we rewrite the local objective in Theorem 3 equivalently as

where, for any i, j ∈ Z, a ∈ A, and o ∈ O,

ωiao j = 1K∑Kk=1∑Tk t=0γ tν k t ∑t-1 τ=1 ξ k t,τ(i, j)δ(akτ , a)δ(okτ+1 , o),(10) and δ(a, b) is one if a = b and zero otherwise. The problem max ˆΘ∈Ξ LB( ˆΘ|Θn) has a unique solution given by

for i, j ∈ Z, a ∈ A, and o ∈ O. Define

One may readily verify that f(ˆΘ; υ, ρ, ω) has a unique stationary point that is equal to ˆΘ as given in (11).

Let θ = [μ, π,w]T and s = [υ, ρ,ω]T , with the definitions μ = [μi]i=1:|Z|, π = [πia ]a=1:|A|,i=1:|Z|, w = [W iao j ]j=1:|Z|,i=1:|Z|,a=1:|A|,o=1:|O|, and similarly υ = [υi]i=1:|Z|, ρ = [ρia ]a=1:|A|,i=1:|Z|, ω = [ωiao j ]j=1:|Z|,i=1:|Z|,a=1:|A|,o=1:|O|, where the elements in a vectorization are arranged by following a left-to-right order in sorting the indices. For example, [πia ]a=1:|A|,i=1:|Z| =

Define C = diag(Cμ,Cπ,CW), where Cμ is a |Z| × |Z| matrix of ones, Cπ = diag{(Ciπ )i=1:|Z|} with Ciπ a |A| × |A| matrix of ones, and CW = diag{(Ciao W )i=1:|Z|,a=1:|A|,o=1:|O|} with Ciao W a |Z| × |Z| matrix of ones. By construction, C is a block-diagonal matrix with each block a matrix of all ones, and each block is associated with the corresponding sub-vector in θ whose elements sum up to one.

Let ln(θ) denote a vector resulting from element-wise application of the natural logarithm to θ. Define

Then one can rewrite (12) as

and the solution in (11) is a solution to ∇ˆ θ f ( ˆ θ; s) = 0 or equivalently ∇ˆ θψT ( ˆ θ)s = 0, where ∇θ denotes the gradient with respect to θ.

Using the vector notation, one can restate the algorithm in Theorem 3 as:

Outer- and/or inner- loop E-steps: Update the sufficient statistic s = s(D(K) , θn-1) using (8)-(10) and (4)-(6).

We have written s = s(D(K) , θn-1) to emphasize that the sufficient statistic is a function of the previous parameter vector θn-1 and episodes D(K) . Note that the constraint ˆΘ ∈ Ξ in Theorem 3 has been accounted for in (12) and (14).

To see the possibility of an online version of the EM algorithm, we note from (8)-(10) that the sufficient statistic is a sum over the episodes, and therefore one can write

where s(εk, θn-1), a vector collecting only the terms with index k in (8)-(10), represents the contribution from εk, and D(n) = {ε1, ε2, · · · , εn} are episodes coming in sequentially.

The online EM algorithm employs a noisy (indicated by the hat above s) version of the true sufficient statistic in (15),

where βk = αk ∏nj=k+1 (1 - αj), and {0 < αk < 1}nk=1 , referred to as learning rates, discount the contribution from each episode, with earlier episodes giving greater discounts because they are computed using earlier versions of θ and are thus less accurate. One can equivalently write (16) as

with the recursive definition

The online nested EM algorithm processes the episodes on the fly, and removes an episode from the memory once its contribution to the sufficient statistic has been recorded using (18). The contribution from the k-th episode, s(εk, θk-1), is computed using (8)-(10) and (4)-(6), considering only the terms with index k. Recalling that (4) constitutes the outerloop E-step and (5)-(6) the inner-loop E-step, one need only compute (4) every δn iterations of the inner-loop EM, where δn may be determined by the convergence of the inner-loop EM. When n is small, however, the sufficient statistics is very noisy, and in this case it is preferable to run only a few iterations for both the outer-loop EM and the inner-loop EM, to prevent premature convergence. The RPR parameter vector is updated in the inner-loop M-step, as the solution to

The solution, given in closed-form in (11), is essentially a set of normalized sub-vectors of the sufficient statistic.

The complete online algorithm is given in Algorithm 3. The algorithm performs on-policy learning [Sutton and Barto, 1998], where the agent learns the RPR while also using it to choose actions in collecting experiences. In order to achieve a balance between exploration and exploitation, the behavior policy used in the algorithm, Πn, is a mixture of two experts [Jordan and Jacobs, 1994], i.e.,

where the first expert, p(a|h, Θn), is the RPR policy as given by (1), the second is a random policy, and the gating network p(y|h) is learned, along with the EM algorithm, using the dual-policy approach described in [Cai et al., 2009]. The behavior policy used here generalizes the ϵ-greedy policy [Sutton and Barto, 1998] in such a sense that the probability of choosing random actions is not a constant ϵ; rather, it depends on the history of past actions and observations.

3.1 Time and memory complexity

The batch-mode EM algorithm has a time complexity of O(|A||O||Z|2 ∑nm=1 ∑mk=1 T 2 k ) when there is a nonzero reward at every step in an episode (dense reward), or O(|A||O||Z|2 ∑nm=1 ∑mk=1 Tk ) when nonzero reward is received only at the terminal step in each episode.

In comparison, the proposed online EM algorithm has a time complexity of O(|A||O||Z|2 ∑nk=1 T 2 k ) in the densely-rewarded case, and O(|A||O||Z|2 ∑nk=1 Tk ) in the terminally-rewarded case.

When all episodes have the same number of steps, say T, the complexity of batch-mode EM becomes

Algorithm 1 The online nested EM algorithm for RPR

Initialize n = 0, α0 ≪ 1, ˆs0 = 0, Π0 as a random policy, while Pexplore = E[p(y = 1|h)] < ϵ do

1. collect the n-th episode εn following Πn-1
2. (outer-loop E-step) re-compute rewards with (4).

while inner-loop EM not convergent do

1. (inner-loop E-step) compute (5)-(6).
2. update the sufficient statistic using (18).
3. (inner-loop M-step) update the RPR with (11).

end while

1. (exploitation-exploration balance) update Πn (the gating network) using the method in [Cai et al., 2009].
2. adjust learning rate αn
3. increase pointer n := n + 1.

end while

O(|A||O||Z|2T 2n2 ) in the case of dense rewards and O(|A||O||Z|2Tn2 ) in the case of terminal rewards, while the complexity of online EM is O(|A||O||Z|2T 2n) and O(|A||O||Z|2Tn), in the two respective cases.

Therefore, the accumulative learning time, as a function of the number of episodes n, is approximately on the order of O(n2 ) for batch-mode EM, while it is approximately on the order of O(n) for online EM.

The batch algorithm stores all episodes, and thus it has a memory complexity of O(n). The online algorithm stores only the latest episode; its memory complexity is O(1).

3.2 Convergence Analysis

The convergence of the online algorithm depends on the asymptotic behavior of the noisy sufficient statistic in (18). If ˆ sn approaches to the accurate (noise-free) sufficient statistic s(D(∞) , θn) as n is sufficiently large, the online algorithm converges to the batch-mode solution after n iterations; thus, the online algorithms also converges to the optimal RPR, considering that the batch-mode solution gives the optimal RPR when the number of episodes is sufficiently large (see the discussion below Definition 2).

Theorem 4 shows that the noise in ˆ sn indeed vanishes as n →∞, under certain conditions on the learning rates {αn}. Thus convergence of the online algorithm is guaranteed.

Theorem 4 can be proven using the techniques in [Delyon et al., 1999] and [Capp´e and Moulines, 2009]. The details are omitted here to conserve space.

4 Experiments

We investigate the empirical performance of the online algorithm on the five benchmark problems described in Table 1. The POMDP models for these problems are available at http://www.cs.brown.edu/research/ai/pomdp/examples. These models are assumed unknown to the agent; they are used as black boxes to simulate the episodes of agent-environment interaction. For all problems, an episode starts from a random position and ends when the goal is achieved.

4.1 Online RPR versus Batch-mode RPR

We compare the performance of the online RPR against that of the batch-mode RPR. The performance is measured by the discounted accumulative reward averaged over 5 independent runs; in each run, we use 101 (or 251) test episodes of at most 101 (or 251) steps for the first three (or other) problems. The computational efficiency is compared based on the accumulative CPU time in the learning phase. The number of states is set to |Z| = 20 for the first three small problems, |Z| = 40 for Hallway and Tag. The comparison results are plotted in Figure 1. The performance for the online-learned policy as a function of |Z| (not provided here due to the page limit) shows that the policy is robust over a wide range of |Z|. Automatic inference of |Z| can be performed using the infiniteRPR framework in [Liu et al., 2011].

As can be seen from the leftmost three columns of Figure 1, the online EM algorithm shows some advantages in performance, especially during the early stage of learning when episodes are collected by exploration. A comparison of CPU time shows that online EM is significantly more efficient than the batch EM, and the time agrees with the complexity analysis in Section 3.1. Moreover it is noted that, in the Shuttle and Cheesemaze problems, batch EM is prone to local optima, whereas online EM is a stochastic algorithm which seems capable of jumping out of local optima to obtain better solutions.

The rightmost two columns in Figure 1 demonstrate that the online algorithm makes the RPR scale to larger problems, and achieves a better tradeoff between the number of episodes (n) and the policy quality. When n is small, batch EM is better than online EM, which can be explained by the fact that the batch method updates the sufficient statistics using all available data. However, given the same amount of learning time, the online RPR achieves better policy performance; for example, by the time that online EM has processed 50, 000 episodes and achieved a test reward of 0.3 for Hallway and -8 for Tag, batch EM has processed less than 1000 episodes, achieving only 0.1 and -14 for these two problems, respectively. These comparisons show the RPR can scale up to larger POMDP problems with online learning.

4.2 Comparison with Other Methods

We compare the RPR policy learned with the online algorithm to three existing algorithms, which all assumes the POMDP models are unknown. The first one is McCallum's U-tree [McCallum, 1995], including its modified version [Zheng and Cho, 2011], which, like the RPR, is a policy-based approach and does not attempt to learn the POMDP models. The second is iPOMDP [Doshi-Velez et al., 2009], including its variant Policy Prior (PP) [Doshi-Velez, 2010], which learn both POMDP models and control policies. The last competing algorithm is mixture learning [Vlassis and Toussaint, 2009], referred to MDP-EM, which is an online EM algorithm treating the observation in a POMDP as the state in an MDP. We do not repeat the experiments for the three competing methods, but rather cite the results from the literature. Since the cited results are based on different testing procedures, we report the comparison to each method separately, using the same test procedure as in the literature for each respective method.

We compare to the U-tree results as reported in [Zheng and Cho, 2011]. To make the results comparable, we strictly follow the same procedures to set up our experiment for the RPR. Specifically, the experiment for each problem consists of 21 independent trials. In each trial, we consider three settings of |Z| for the RPR policy, i.e., |Z| ∈ {10, 20, 30}; for each setting, the RPR policy is learned using the online EM algorithm, based on episodes collected by following the behavior policy for 75000 steps, and then the learned RPR policy is followed 25000 steps to test the its quality. The quality in each trial is measured by the average reward earned over the 25000 testing steps. The results from the 21 trials are used to calculate the mean and standard deviation for the RPR in each setting of |Z|, which are reported in Table 2, along with the results of U-trees cited from [Zheng and Cho, 2011].

The online RPR performs comparably as U-trees on Shuttle and better on 4 × 4 Grid. The RPR performs worse than U-trees on Cheese Maze for the given budget of 75000 steps of learning. However, we notice that the RPR's performance can catch up after additional learning steps. Since the decision states in the RPR correspond to the nodes in U-trees, it is interesting to compare |Z| used by the RPR and the nodes generated by U-trees, which are shown as the parenthesized numbers in Table 2. As can be seen, the RPR generally needs less decision states to achieve superior performance, while U-trees generate more nodes than necessary.

To compare to iPOMDP [Doshi-Velez et al., 2009] and policy prior [Doshi-Velez, 2010], we follow their experimental setting and report the cumulative rewards as a function of the learning step (3500 steps for Hallway, 5000 steps for Tag). The results are summarized in Table 3, which show that online RPR achieves significantly larger accumulative rewards than iPOMDP, and performs comparably to the policy prior methods which utilize expert information. These results demonstrate the potential benefits of learning a policy directly from the episodes, as opposed to solving a (possibly inaccurate) POMDP estimated from the episodes.

The MDP-EM method learns a reactive policy (it does not memorize past observations). To make the comparison fair, we consider the RPR policy conditioning on different lengths of history. The policy evaluation procedure is the same as described in Section 4.1, with the MDP-EM policy provided by the authors of [Vlassis and Toussaint, 2009]. The results are plotted in Figure 2, where |h| indicates the length of history, e.g. |h| = 1 means only the current observation is used, which makes the RPR a reactive policy as the MDP-EM. As shown by Figure 2, the online RPR performs generally better than the MDP-EM, and the performance increases as |h| increases from one to three, and then becomes saturated when longer histories are used. The improvement achieved by the RPR with |h| > 1 is expected, considering the MDPEM uses only the current observation while the RPR with |h| > 1 also uses past actions and observations. However, even with |h| > 1, the RPR can also achieve significant improvements; the improvements can be attributed to the reward re-computation in (4) which, as discussed in the introduction and shown in (7), produce a weighted log-likelihood function whose maximization gives an improved policy.

5 Conclusions

We have presented an online nested EM algorithm for learning the RPR, a parametric policy for RL in POMDPs. The online algorithm uses only the latest learning episode to update the policy's value in the outer-loop E-step and update the sufficient statistic in the inner-loop E-step, and computes the improved policy in closed-form given the sufficient statistic. The algorithm reduces both the time and memory complexity an order of magnitude, in comparison with the corresponding batch-mode algorithm. Under standard conditions on the learning rates, the online algorithm provably converges to the optimal batch-mode policy when the number of learning episodes becomes sufficiently large.

Experimental results on five benchmark POMDP problems show that: (i) the online RPR produces policies as good as the batch-mode RPR by using only a small fraction of the time; (ii) the stochastic nature of the online algorithm can help it to jump out of local optima; (iii) model-free methods have the potential to learn better policies than model-based methods; (vi) the RPR can produce high-quality policies by using fewer internal memory units than the U-tree; (v) the RPR has a performance increasing with the history length, and is a competitive reactive policy when the length is one. Future work includes online inference of the number of decision states and multi-agent online learning.

Acknowledgments This research was supported by the US Office of Naval Research (ONR), under the Active Transfer Learning (ATL) and Science of Autonomy programs.