Variational Model-based Policy Optimization

1 Introduction

Model-free reinforcement learning (RL) algorithms that learn a good policy without constructing an explicit model of the system's dynamics have shown promising results in complex simulated problems [Mnih et al., 2013; Mnih et al., 2015; Schulman et al., 2015; Haarnoja et al., 2018]. However, these methods are not sample efficient, and thus, not suitable for problems in which data collection is burdensome. Model-based RL algorithms address the data efficiency issue of the model-free methods by learning a model, and combining model-generated data with those collected from interaction with the real system [Sutton, 1990; Janner et al., 2019]. However, designing model-based RL algorithms is often challenging because the bias in model may affect the process of learning policies and result in worse asymptotic performance than the model-free counterparts. A potential solution to this challenge is to incorporate the policy/value optimization method in the process of learning the model. An ideal case here would be to have a universal objective function that is used to learn and improve model and policy jointly.

Casting RL as a probabilistic inference has a long history (e.g., [Todorov, 2008; Toussaint, 2009; Kappen et al., 2012; Rawlik et al., 2013]). This formulation has the advantage that allows powerful tools for approximate inference to be employed in RL. One such class of tools are variational techniques that have been successfully used in RL (e.g., [Neumann, 2011; Levine and Koltun, 2013; Abdolmaleki et al., 2018]). Another formulation of RL with strong connection to probabilistic inference is the formulation of policy search as an expectation maximization (EM) style algorithm (e.g., [Dayan and Hinton, 1997; Peters and Schaal, 2007; Peters et al., 2010; Chebotar et al., 2017; Abdolmaleki et al., 2018]). The main idea here is to write the expected return of a policy as a (pseudo)-likelihood function, and then conditioning on the success in maximizing the return, find the policy that most likely would have been taken. Another class of RL algorithms that are related to the inference formulation are entropy-regularized algorithms that add an entropy term to the reward and find the soft-max optimal policy (e.g., [Levine and Abbeel, 2014; Nachum et al., 2017; Haarnoja et al., 2018; Fellows et al., 2019]). For a comprehensive tutorial on RL as probabilistic inference, we refer readers to [Levine, 2018].

In this paper, we leverage the connection between RL and probabilistic inference, and formulate an objective function for jointly learning and improving model and polciy as a variational lower-bound of a log-likelihood. This allows us to use EM: iteratively fix a baseline policy and learn a variational distribution, consisting of a model and a policy (E-step), followed by improving the baseline policy given the learned variational distribution (M-step). We propose model-based and model-free policy iteration (PI) style algorithms for the E-step and show how the variational distribution that they learn can be used to optimize the M-step only from modelgenerated samples. It is important to note that both algorithms are model-based and only differ in using model-based and model-free algorithms for the E-step. We call our algorithm that uses model-based PI for the E-step, variational modelbased policy optimization (VMBPO). Our experiments on a number of continuous control tasks show that VMBPO is more sample-efficient and robust to hyper-parameter tuning than its model-free counterpart. Using the same control tasks, we also compare VMBPO with several state-of-the-art model-based and model-free RL algorithms, including model-based policy optimization (MBPO) [Janner et al., 2019] and maximum a posteriori policy optimization (MPO) [Abdolmaleki et al., 2018], and show its sample efficiency and performance.

2 Preliminaries

We study the RL problem in which the agent's interaction with the environment is modeled as a Markov decision process (MDP) M = 〈X , A, r, p, p0〉, where X and A are state and action spaces; r : X ×A → R is the reward function; p : X × A → ∆X is the transition kernel (∆X is the set of probability distributions over X); and p0 : X → ∆X is the initial state distribution. A stationary Markovian policy π : X → ∆A is a probabilistic mapping from states to actions. Each policy π is evaluated by its expected return, i.e., J(π) = E[∑T-1 t=0 r(xt, at) | p0, p, π], where T is the (random) time of hitting a terminal state. 1 We denote by X 0 the set of all terminal states. The agent's goal is to find a policy with maximum expected return, i.e, π∗ ∈ argmaxπ∈∆A J(π). We denote by ξ = (x0, a0, . . . , xT-1, aT-1, xT ), a system trajectory of length T, whose probability under a policy π is defined as pπ(ξ) = p0(x0) ∏T-1 t=0 π(at|xt)p(xt+1|xt, at). Finally, we define [T ] := {0, . . . , T - 1}.

3 Policy Optimization as Inference

Policy search in RL can be formulated as a probabilistic inference problem (e.g., [Todorov, 2008; Toussaint, 2009; Kappen et al., 2012; Levine, 2018]). The goal in the conventional RL formulation is to find a policy whose generated trajectories maximize the expected return. In contrast, in the inference formulation, we start with a prior over trajectories and then estimate the posterior conditioned on a desired outcome, such as reaching a goal state. In this formulation, the notion of a desired (optimal) outcome is introduced via independent binary random variables Ot, t ∈ [T], where Ot = 1 denotes that we acted optimally at time t. The likelihood of Ot, given the state xt and action at, is modeled as

where η > 0 is a temperature parameter. This allows us to define the log-likelihood of π being optimal as

where p(O0:T-1 = 1 | ξ) is the likelihood of trajectory ξ being optimal and is defined as

1 Similar to [Levine, 2018], our setting can be easily extended to infinite-horizon γ-discounted MDPs. This can be done by modifying the transition kernels, such that any action transitions the system to a terminal state with probability 1 - γ, and all standard transition probabilities are multiplied by γ.

As a result, finding an optimal policy in this setting would be equivalent to maximizing the log-likelihood in (1), i.e., π∗ soft ∈ argmaxπ log pπ(O0:T-1 = 1). A potential advantage of formulating RL as an inference problem is the possibility of using a wide range of approximate inference algorithms, including variational methods. In variational inference, we approximate a distribution p(·) with a potentially simpler (e.g., tractable factored) distribution q(·) in order to make the whole inference process more tractable. If we approximate pπ(ξ) with a variational distribution q(ξ), we will obtain the following variational lower-bound for the log-likelihood in (1):

(a) is from Jensen's inequality and J(q; π) is the evidence lower-bound (ELBO) of the log-likelihood function. A variety of algorithms have been proposed (e.g., [Peters and Schaal, 2007; Hachiya et al., 2009; Neumann, 2011; Levine and Koltun, 2013; Abdolmaleki et al., 2018; Fellows et al., 2019]), whose main idea is to approximate π∗ soft bymaximizing J (q; π) w.r.t. both q and π. This often results in an EM-style algorithm in which we first fix π and maximize J (·; π) for q (E-step), and then given the q obtained in the E-step, we maximize J (q; ·) for π (M-step).

4 Variational Model-based Policy Optimization

In this section, we describe the ELBO objective function used by our algorithms, study the properties of the resulted optimization problem, and propose algorithms to solve it. We propose to use the variational distribution q(ξ) = p0(x0) ∏T-1 t=0 qc(at|xt)qd(xt+1|xt, at) to approximate pπ(ξ). Note that q has the same initial state distribution as pπ (defined in Section 2), but has different control strategy (policy), qc, and dynamics, qd. Using this variational distribution, we may write the ELBO objective (3) as

To maximize J (q; π) w.r.t. q and π, we first fix π and compute the variational distribution (E-step):

and then optimize π given q∗ , i.e., argmaxπ J(q∗ ; π) (Mstep). Note that in (5), q∗ c and q∗ d are both functions of π, but we remove π from the notation to keep it lighter.

Remark 1. In our formulation (our choice of the variational distribution q), the M-step is independent of the true dynamics, p, and thus, can be implemented offline (using samples generated by the model qd). Moreover, as we will see in Section 5, we also use the model, qd, in the E-step. As discussed throughout the paper, using simulated samples (from qd) and reducing the need for real samples (from p) is an important feature of our proposed model-based formulation and algorithms.

Remark 2 (relationship withMPO). There are similarities between our variational formulation and the one used in the maximum a posteriori policy optimization (MPO) algorithm [Abdolmaleki et al., 2018]. However, MPO sets its variational dynamics, qd, to be the dynamics of the real system, p, which results in a model-free algorithm, while our approach is modelbased, since we learn qd and use it to generate samples in both E-step and M-step of our algorithms.

In the rest of this section, we study the E-step optimization (5) and propose algorithms to solve it.

4.1 Properties of the E-step Optimization

We start by defining two Bellman-like operators related to the E-step optimization (5). For any variational policy qc : X → ∆A and any value function V : X → R, such that V (x) = 0, ∀x ∈ X 0 , we define the qc-induced operator, Tqc , and the optimal operator, T , as

We define the optimal value function of the E-step, Vπ, as

For any value function V , we define its associated action-value function Q : X × A → R as

The following lemmas, whose proofs are reported in Appendices A.1 and A.2, show properties of operators Tqc and T , and their relation with the (E-step) optimal value function, Vπ.

Lemma 1. The qc-induced and optimal operators, defined by (6) and (7), can be rewritten as

Lemma 2. The qc-induced and optimal operators are monotonic and contractive. Moreover, the optimal value function Vπ is the unique fixed-point of T (T [Vπ](x) = Vπ(x), ∀x ∈ X).

From the definition of Q-function in (9) and Lemma 2, we prove (in Appendix A.3) the following proposition for the action-value function associated with the E-step optimal value function Vπ.

Proposition 1. The E-step optimal value function Vπ and its associated action-value function Qπ, defined by (9), are related as Vπ(x) = log Ea∼π(·|x) [ exp (Qπ(x, a))], ∀x ∈ X .

In the rest of this section, we show how to derive a closedform expression for the variational distribution q∗ = (q∗ c , q∗ d ). For any value function V , we define its corresponding variational dynamics, qV d , as the solution to the maximization problem in the definition of Tqc (see Eq. 6), i.e.,

and its corresponding variational policy qQ c (Q is the actionvalue function associated with V , defined by Eq. 9), as the solution to the maximization problem in the definition of T (see Eqs. 7 and 10), i.e.,

We now derive (proof in Appendix A.4) closed-form expressions for the variational distributions qV d and qQ c .

Lemma 3. The variational dynamics and policy corresponding to a value function V and its associated action-value function Q can be written in closed-form as

Equations 14 and 15 show that the variational dynamics, qV d , and policy, qQ c , can be seen as an exponential twisting of the dynamics p and policy π with weights V and Q, respectively. For the special case V = Vπ (the E-step optimal value function), these distributions can be written in closed-form as

where the denominator of q∗ c is obtained by applying Proposition 1 to replace Qπ with Vπ.

4.2 Policy and Value Iteration for the E-step

Using the results of Section 4.1, we now propose model-based and model-free dynamic programming (DP) style algorithms, i.e., policy iteration (PI) and value iteration (VI), for solving the E-step problem (5). The model-based algorithms compute the variational dynamics, qd, at each iteration, while the model-free counterparts compute qd only at the end (upon convergence). Having access to qd at each iteration has the advantage that we may generate samples from the model, qd, when we implement the sample-based version (RL version) of these DP algorithms in Section 5.

In the model-based PI algorithm, at each iteration k, given the current variational policy q(k) c , we

Policy Evaluation: Compute the q(k) c -induced value function Vq(k) c (the fixed-point of the operator Tq(k) c ) by iteratively applying Tq(k) c from (6), i.e., Vq(k) c (x) = limn→∞ T n q(k) c [V ](x), ∀x ∈ X, where the variational model qd in (6) is computed using (14) with V = V (n) . We then compute the corresponding Q-function Qq(k) c using (9).

Policy Improvement: Update the variational distribution q(k+1) c using (15) with Q = Qq(k) c . 2

Upon convergence, i.e., q(∞) c = q∗ c , we compute q∗ d from (14).

The model-free PI algorithm is exactly the same, except in its policy evaluation step, the q(k) c -induced operator Tq(k) c is applied using (10) (without the variational dynamics qd). In this case, the variational dynamics qd is computed only upon convergence, q∗ d , using (14).

Lemma 4. The model-based and model-free PI algorithms converge to their optimal values, q∗ c and q∗ d , defined by (5), i.e., q(∞) c = q∗ c and q(∞) d = q∗ d . (proof in Appendix ??)

We can similarly derive model-based and model-free (VI) algorithms for the E-step. These algorithms start from an arbitrary value function V and iteratively apply the optimal operator T from (6) and (7) (model-based) and (11) (model-free) until convergence, i.e., Vπ(x) = limn→∞T n [V ](x), ∀x ∈ X. Given Vπ, these algorithms first compute Qπ from Proposition 1, and then compute (q∗ c , q∗ d ) using (16). From the properties of the optimal operator T in Lemma 2, both model-based and model-free VI algorithms converge to q∗ c and q∗ d .

In the rest of the paper, we focus on the PI approach, in particular the model-based one, and only report the details of the VI algorithms in Appendix B. In the next section, we show how the PI algorithms can be implemented and combined with a routine for solving the M-step, when the true MDP model p is unknown (the RL setting) and the state and action spaces are large that require using function approximation.

5 Variational Model-based RL Algorithm

In this section, we propose a RL algorithm, called variational model-based policy optimization (VMBPO). It is an EM-style algorithm based on the variational formulation proposed in Section 4. The E-step of VMBPO is the sample-based implementation of the model-based PI algorithm, described in Section 4.2. We describe the E-step and M-step of VMBPO in details in Sections 5.1 and 5.2, and report its pseudo-code in Algorithm 1 in Appendix C. VMBPO uses 8 neural networks to represent: policy π, variational dynamics qd, variational policy qc, log-likelihood ratio ν = log(qd/p), value function V , action-value function Q, target value function V ′ , and target action-value function Q′ , with parameters θπ, θd, θc, θν, θv, θq, θ′v , and θ′q , respectively.

2 When the number of actions is large, the denominator of (15) cannot be computed efficiently. In this case, we replace (15) in the policy improvement step of our PI algorithms with q(k+1) c = arg minqc KL(qc||qQ c ), where Q = Qq(k) c . We also prove the convergence of our PI algorithms with this update in Appendix A.5.

5.1 The E-step of VMBPO

At the beginning of the E-step, we generate a number of samples (x, a, r, x′ ) from the current baseline policy π, i.e., a ∼ π(·|x) and r = r(x, a), and add them to the buffer D. The E-step consists of four updates: 1) computing the variational dynamics qd, 2) estimating the log-likelihood ratio log(qd/p), 3) computing the qc-induced value and action-value functions Vqc and Qqc (critic update), and finally 4) computing the new variational policy qc (actor update). We describe the details of each step below.

Step 1. (Computing qd) We find qd as a solution to the optimization problem (12) for V equal to the target value network V ′ . Since the qV d in (14) is the solution of (12), we compute qd by minimizing KL(qV ′ d ||qd), which results in the following forward KL loss (for all x ∈ X and a ∈ A):

where (a) is by removing the θ-independent terms from (17). We update θd by taking several steps in the direction of the gradient of a sample average of the loss function (18), i.e.,

where (x, a, r, x′ ) are randomly sampled from D. The intuition here is to focus on learning the dynamics model in the regions of the state-action space that have higher temporal difference (regions with higher anticipated future return).

Step 2. (Computing log(qd/p)) Using the duality of f-divergence [Nguyen et al., 2008] w.r.t. the reverse KLdivergence, the log-likelihood ratio log(qd(·|x, a; θd)/p(·|x, a)) is a solution to

for all x, x′ ∈ X and a ∈ A. Note that the optimizer of (20) is unique almost surely (at (x, a, x′ ) with P(x′ |x, a) > 0), because qd is absolutely continuous w.r.t. p (see the definition of qd in Eq. 14) and the objective function of (20) is strictly concave. The optimization problem (20) allows us to compute ν(·|·; θν) as an approximation to the log-likelihood ratio log(qd(·; θd)/p). We update θν by taking several steps in the direction of the gradient of a sample average of (20), i.e.,

where E is the set of samples for which x′ is drawn from the variational dynamics, i.e., x′ ∼ qd(·|x, a). Here we first sample (x, a, x′ ) randomly from D and use them in the second sum. Then, for all (x, a) that have been sampled, we generate x′ from qd and use the resulting samples in the first sum.

Step 3. (critic update) To compute Vqc and its actionvalue Qqc , we rewrite (6) with the maximizer qd from Step 1 and the log-likelihood ratio log(qd/p) from Step 2:

Since Tqc can be written as both (10) and (22), we compute the qc-induced Q-function by setting the RHS of these equations equal to each other, i.e., for all x ∈ X and a ∼ qc(·|x; θc),

Since the expectation in (22) is w.r.t. the variational dynamics (model) qd, we can estimate Qqc only with samples generated from the model. We do this by taking several steps in the direction of the gradient of a sample average of the square-loss obtained by setting both sides of (22) equal, i.e.,

Note that in (22), the actions are generated by qc. Thus, in (23), we first randomly sample x, then sample a from qc(·|x; θc), and finally draw x′ from qd(·|x, a; θd). If the reward function is known (chosen by the designer of the system), then it is used to generate the reward signals r = r(x, a) in (23), otherwise, a reward model has to be learned.

After estimating Qqc , we approximate Vqc , the fixed-point of Tqc , using Tqc definition in (10) as Tqc [V ](x) ≈ V (x) ≈ Ea∼qc(·|x) [Q(x, a; θq) - log qc(a|x;θc) π(a|x;θπ ) ]. This results in updating Vqc by taking several steps in the direction of the gradient of a sample average of the square-loss obtained by setting the two sides of the above equation equal, i.e.,

where x is randomly sampled and a ∼ qc(·|x; θc) (without sampling from the true environment).

Step 4. (actor update) We update the variational policy qc (policy improvement) by solving the optimization problem (13) for the Q estimated by the critic in Step 3. Since the qc that optimizes (13) can be written as (15), we update it by minimizing KL(qc||qQ c ). This results in the following reverse KL loss (for all x ∈ X):

If we reparameterize qc using a transformation a = f(x, ϵ; θc), where ϵ is a Gaussian noise, we can update θc by taking several steps in the direction of the gradient of a sample average of the above loss, i.e.,

We can also compute qc as the closed-form solution to (15), as described in [Abdolmaleki et al., 2018]. They refer to this as non-parametric representation of the variational distribution.

5.2 The M-step of VMBPO

As described in Section 4, the goal of the M-step is to improve the baseline policy π, given the variational model q∗ = (q∗ c , q∗ d ) learned in the E-step, by solving the following optimization problem:

A nice feature of (26) is that it can be solved using only the variational model q∗ , without the need for samples from the true environment p. However, it is easy to see that if the policy space considered in the M-step, Π, contains the policy space used for qc in the E-step, then we can trivially solve the Mstep by setting π = q∗ c . Although this is an option, it is more efficient in practice to solve a regularized version of (26). A practical way to regularize (26) is to make sure that the new baseline policy π remains close to the old one, which results in the following optimization problem:

This is equivalent to the weighted MAP formulation used in the M-step of MPO [Abdolmaleki et al., 2018]. In MPO, they define a prior over the parameter θ and add it as logP(θ) to the objective function of (26). Then, they set the prior P(θ) to a specific Gaussian and obtain an optimization problem similar to (27) (see Section 3.3 in [Abdolmaleki et al., 2018]). However, in the absence of a variational dynamics model (i.e., qd = p), they need real samples to solve their optimization problem, while our model-based approach uses simulated samples.

6 Experiments

To illustrate the effectiveness of VMBPO, we (i) compare it with several state-of-the-art RL methods, and (ii) evaluate sample efficiency of MBRL via ablation analysis.

Comparison with Baseline RL Algorithms. We compare VMBPO with five baselines, two popular model-free algorithms: MPO [Abdolmaleki et al., 2018] and SAC [Haarnoja et al., 2018], and three recent model-based algorithms: MBPO [Janner et al., 2019], PETS [Chua et al., 2019], and STEVE [Buckman et al., 2018]. We also compare VMBPO with its variant that uses a model-free PI algorithm to solve the E-step (see Section 4.2). We refer to this variant as VMBPOMFE and describe it in details in Appendix D. We evaluate all the algorithms on a classical control task: Pendulum, and five MuJoCo tasks: Hopper, Walker2D, HalfCheetah, Reacher, and Reacher7DoF. We use similar neural network architectures (for the dynamics model, value functions, and policies) for VMBPO and MBPO. The detailed description of the network architectures and hyper-parameters is reported in Appendix E. Since we use a parametric representation for qc in the E-step of VMBPO, as discussed in Section 5.2, we simply set π = q∗ c in its M-step. We set the number of training steps to 400, 000 for the difficult environments (Walker2D, HalfCheetah), to 150, 000 for the medium one (Hopper), and to 50, 000 for the simpler ones (Pendulum, Reacher, Reacher7DOF). We evaluate policy performance every 1, 000 training steps. Each measurement is an average return over 5 episodes, generated with a separate random seed.

To illustrate the relative performance of the algorithms, we report the average return of VMBPO, VMBPO-MFE, and the baselines, with their best hyper-parameters, in Table 1 and Figure 1 (see Appendix E.1). The results show that VMBPO performs better than the baselines in most of the tasks, and usually converges faster even when the final performances are similar. The data-efficiency of VMBPO is mainly the result of using synthetic data generated by the learned model, and its extra performance can be attributed to jointly learning model and policy using a universal objective function.

The results also show that VMBPO-MFE outperforms MPO in 4 out of the 6 domains. However, in some cases its learning curve degrades and results in poor final performance. This is partly due to the instability in critic learning caused by sample variance amplification in exponential-TD minimization (see Eq. 39 in Sec. D.1). A way to alleviate this issue is to add a temperature term τ to the exponential-TD update [Borkar, 2002], although tuning this hyper-parameter is often non-trivial. 3

To study the sensitivity of the algorithms w.r.t. the hyperparameters, we report their performance averaged over all hyper-parameter/random-seed configurations in Table 2 and Figure 2 (see Appendix E.1). These results show that VMBPO

3 The variance is further amplified with a large τ, but the critic learning is hampered by a small τ .

is much more robust to change in hyper-parameters than the other algorithms, with the best average performance over all the tasks.

Ablation Study. We study the dependence of the VMBPO performance on the number of samples generated from the dynamics model qd. Here we only experiment with two tasks, Hopper and HalfCheetah, and with fewer training steps 100, 000. At each step, we update the actor and critic using {128, 256, 512} synthetic samples. Figure 1 shows the learning performance averaged over all hyper-parameter/randomseed configurations and illustrates how synthetic data can help with policy learning. The results show that increasing the amount of synthetic data generally improves the policy convergence rate. In the early phases, when the model is inaccurate, sampling from it may slow down learning, while in the later phases, with an improved model, adding more synthetic data can lead to a more significant performance boost.

7 Conclusion

We formulated the problem of jointly learning and improving model and policy in RL as a variational lower-bound of a loglikelihood and proposed EM-style algorithms to solve it. Our algorithm, called variational model-based policy optimization (VMBPO), uses model-based policy iteration for solving the Estep. We compared our (E-step) model-based and model-free algorithms with each other, and with a number of state-of-theart model-based (e.g., MBPO) and model-free (e.g., MPO) RL algorithms, and showed its sample efficiency and performance.

We briefly discussed VMBPO algorithms in which the E-step is solved by value iteration methods. However, full implementation of these algorithms and studying their relationship with the existing methods requires more work that we leave for future. Another future directions are: 1) finding more efficient implementation for VMBPO, and 2) using VMBPO style algorithms in solving control problems from high-dimensional observations, by learning a low-dimensional latent space and a latent dynamics, and perform control there. This class of algorithms is referred to as learning controllable embedding [Levine et al., 2020].