Title: SHRINKAGE-BASED BIAS-VARIANCE TRADE-OFF FOR DEEP REINFORCEMENT LEARNING

Abstract: Deep reinforcement learning has achieved remarkable successes in solving various challenging artificial intelligence tasks. A variety of different algorithms have been introduced and improved towards human-level performance. Although technical advances have been developed for each individual algorithms, there has been strong evidence showing that further substantial improvements can be achieved by properly combining multiple approaches with difference biases and variances. In this work, we propose to use the James-Stein (JS) shrinkage estimator to combine on-policy policy gradient estimators which have low bias but high variance, with low-variance high-bias gradient estimates such as those constructed based on model-based methods or temporally smoothed averaging of historical gradients. Empirical results show that our simple shrinkage approach is very effective in practice and substantially improve the sample efficiency of the state-of-the-art on-policy methods on various continuous control tasks.

1. INTRODUCTION:

Deep reinforcement learning (RL) has achieved remarkable successes recently, as witnessed by the human/super-human level performance achieved in various challenging artificial intelligence tasks; notable examples include AlphaGo (Silver et al., 2017; 2016) , Atari games (Mnih et al., 2013) , robotic learning (Levine et al., 2016; Andrychowicz et al., 2017) , among many others. A remarkable nature of RL is that multiple types of algorithms have been developed, based on different principles, highlighting different features and trade-offs. This includes model-based approaches which derive optimal policies from estimated dynamic models of the unknown environment (Deisenroth & Rasmussen, 2011; Nagabandi et al., 2017) , model-free off-policy methods such as Q-learning (Mnih et al., 2013) , as well as model-free on policy methods such as Proximal Policy Optimization (PPO) (Schulman et al., 2017) and Trust Region Policy Optimization (TRPO) (Schulman et al., 2015) . These methods exploit different properties of the Markov decision process, and have different pros and cons. Model-based and off-policy methods are known to be more sample efficient when implemented successfully, while on-policy methods are less sample efficient, but are simpler and easier to train, especially for problems with continuous or high dimensional action spaces. Although there have been emerging advances in improving each individual algorithms, more substantial improvements are likely achieved by adaptively integrating multiple strategies into an overall intelligence system. This is evident from human brain, which has been discovered to contain multiple distinct or even competing pathways for learning from reward (Daw et al., 2005; Rangel et al., 2008) , including both model-free strategies consisting of error-and-trial learning to repeat rewarded actions, and model-based learning that learns models of the environment, and use it to evaluate candidate actions by mental simulations of their consequences. We would expect integrating different methods in RL will improve the performance, especially when integrated into an adaptive fashion. Algorithmically, a key challenge for an adaptive combination is to evaluate the reliability of each individual strategy, without knowing the ground truth answer while performing the task.

In this work, we focus on the problem of adaptively combining policy gradient estimators from various sources. We leverage the fact that different strategies yield different gradient estimators with different bias-variance trade-off, which can be exploited by statistical shrinkage estimators to obtain optimal combinations that have better mean square error than any individual estimators. In particular, policy gradient estimators provide (nearly) unbiased estimators of the true gradient of the expected reward, but often suffers from high variance and is sample inefficient. Model-based and on-policy methods, on the other hand, provide more biased estimates but can have much smaller variance by taking the advantage of the large amount of imaginary data efficiently simulated from the model, or off-policy data stored in the replay buffer.

In particular, we propose to use James-Stein (JS) shrinkage estimator (Stein, 1956; James & Stein, 1961; Efron & Morris, 1975) to combine the unbiased on-policy gradient estimator with biased but low variance estimators constructed based on models or off-policy data. The basic idea of JS estimator is to access the quality of the biased estimator by inspecting its consistency with unbiased estimators, and use it as a clue for constructing an optimally weighted combination to trade-off bias and variance. JS estimator admits strong theoretical guarantees for standard Gaussian cases, and more importantly, also provides a simple practical heuristic for more complex practical non-Gaussian cases. It can potentially provide an ideal tool for solving the very challenge of bias-variance trade-off in various components of deep reinforcement learning, but seems to be still largely under-explored in the literature. In this work, we explore the practical application of JS estimators in deep RL, illustrated by two examples: Shrinking Towards Model-based Gradient We develop a hybrid algorithm which adaptively integrates on-policy policy gradient with Dyna-style model-based learning agent (Sutton, 1991) . We train a transition model using the existing rollout data in parallel to the on-policy gradient descent; imaginary rollouts are simulated from the model, providing a biased, lower variance gradient estimator. This model-based gradient estimator compared against the unbiased, high variance on-policy gradient, and a weighted combination is constructed using JS estimator based on the consistency of the two estimators. Our method improves the sample efficiency by exploiting the data more efficiently using model learning while avoiding significant bias introduced by the model using adaptive JS estimator.

Shrinking Towards Temporally Smoothed Gradient It has been observed heuristically that proper bias is helpful for training (Sutskever et al., 2013; Nesterov, 2013) . We present a simple method that explicitly constructs an optimal bias-variance trade-off by shrinking the on-policy gradient estimator towards a smoothed gradient estimator obtained by a weighted sum over previous iterations. The smoothed gradient provides a reasonable guess of the true gradient of the current estimation and forms a biased, low variance estimator. By shrinking towards the smoothed gradient, we significantly decrease the variance of the on-policy gradient, at the cost of introducing a small bias. This effectively creates a statistical momentum, which shares similarity with the traditional momentum methods derived from optimization perspectives, but with an adaptive coefficient determined by the JS estimator. It can also be combined with traditional momentum to yield gains on both statistical and optimization sides. This simple algorithm effectively leverages the historic (off-policy) data via the temporally smoothed gradient, yielding higher sample efficiency than standard policy gradient.

We perform empirical studies to test the performance of both above examples, showing our adaptive combination strategies can yield better results than any fixed-weight combination, significantly improving the sample efficiency on standard on-policy methods. This illustrates the practical power of JS estimators, which we expect to be broadly applicable to more settings in deep RL in future works.

2. BACKGROUND:

We first introduce the basics of reinforcement learning and model-free policy gradient, and set up the notation that we will use in the paper. We then introduce a model-based policy gradient method.

2.1. REINFORCEMENT LEARNING:

Reinforcement learning considers the problem of finding an optimal policy for an agent which interacts with an uncertain environment and collects reward per action. The goal of the agent is to maximize the long-term cumulative reward. Formally, this problem can be formulated as a Markov decision process over the environment states s ∈ S and agent actions a ∈ A, under an unknown environmental dynamic defined by a transition probability T (s |s, a) and a reward signal r(s, a) immediately following the action a performed at state s. The agent's action a is selected by a conditional probability distribution π(a|s) called policy. In policy optimization, we consider a set of candidate policies π θ (a|s) parameterized by θ and obtain the optimal policy by maximizing the expected cumulative reward or return

J(θ) = E s∼ρπ,a∼π(a|s) [r(s, a)] ,

where ρ π (s) = ∞ t=1 γ t-1 Pr(s t = s) is the normalized discounted state visitation distribution with discount factor γ ∈ [0, 1). To simplify the notation, we denote E s∼ρπ,a∼π(a|s) [•] by simply E π [•] in the rest of paper.

2.2. ON-POLICY POLICY GRADIENT:

According to the policy gradient theorem (Sutton & Barto, 1998) , the gradient of J(θ) equals

EQUATION

where

Q π (s, a) = E π ∞ t=1 γ t-1 r(s t , a t )|s 1 = s, a 1 =

a denotes the expected return under policy π starting from state s and action a.

Different policy gradient methods are based on different stochastic estimations of the expected gradient in Eq (1). Perhaps the most straightforward way is to simulate the environment with the current policy π to obtain a trajectory {(s t , a t , r t )} n t=1 and estimate ∇ θ J(θ) using the Monte Carlo estimation:

EQUATION

where Qπ (s t , a t ) is an empirical estimate of Q π (s t , a t ), e.g., Qπ (s t , a t ) = j≥t γ j-t r j . This is known as an on-policy gradient estimator because it only uses the data from the current policy.

Proximal Policy Optimization (PPO) (Schulman et al., 2017; Heess et al., 2017) is one of the state-theart model-free policy gradient methods for policy optimization. It uses a proximal Kullback-Leibler (KL) divergence penalty to regularize and stabilize the policy gradient update. Given an existing policy π old , PPO obtains a new policy by maximizing the following surrogate loss function

J ppo (θ) = E πold π θ (a|s) π old (a|s) Q π (s, a) -λKL [π old (•|s) || π θ (•|s)] ,

where the first term is an approximation of the expected reward, and the second term enforces the the updated policy to be close to the previous policy under KL divergence. The gradient of J ppo (θ) can be rewritten as

∇ θ J ppo (θ) = E πold w π (s, a)∇ θ log π(a|s)Q π λ (s, a)

where w π (s, a) := π θ (a|s)/π old (a|s) is the density ratio of the two polices, and Q π λ (s, a) := Q π (s, a) + λw π (s, a) -1 where the second term comes from the KL penalty.

The advantage of using on-policy estimation is that it provides (nearly) unbiased estimation of the gradient of expected reward. A well-known challenge, however, is that it has large variance, given that it is only possible to use a small number of data from the current policy at each iteration because the cost of simulation from the real environment is high. In addition, it does not take advantage of the off-policy data collected previously. Using off-policy data allows us to obtain low variance but potentially highly biased estimates, which can be combined with on-policy gradient with James-Stein estimator as we propose in this work. In the sequel, we introduce model-based policy gradient as an important way for leveraging off-policy data.

2.3. MODEL-BASED POLICY GRADIENT:

The model-based methods derive optimal polices by learning a transition model of the underlying dynamic T (s |s, a) of the Markov Decision Process. Assume the states s is continuous, it is common to assume s = s+f φ (s, a)+ξ where f φ is a parametric function (e.g., neural network) with learnable parameter φ, and ξ is a zero-mean noise. The model parameter φ is often trained by minimizing the L 2 one-step prediction loss,

EQUATION

where D := (s t , a t , s t+1 ) n t=1 is a set of transition pairs collected from previous rollouts. See Deisenroth & Rasmussen (2011) ; Fu et al. (2016) ; Nagabandi et al. (2017) ; Thanard Kurutach (2018); Depeweg et al. (2017) for variety of the recent variants of model-based learning methods.

With a learned model, the problem of estimating the optimal policy becomes a planning problem, which can be solved either using various classical planning and control methods, typically based on variants of dynamic programming (see e.g., Sutton & Barto, 1998) , or sample-based methods that generate fictitious data by simulating the model forward and obtain optimal policies by applying standard RL algorithm on the fictitious data.

In this work, we will mainly consider a simple sample-based policy gradient method whose main idea is rooted from Dyna (Sutton, 1991) . The idea is simply generating a fictitious data (s t , ãt , rt , st+1 ) by simulating the model, and plug it into the policy gradient formula (2),

EQUATION

)

By simulating a large number of fictitious data from the model, which is much less expensive than simulating from the real environment, we can make the variance of ∇model θ J(θ) significantly smaller than the variance of the on-policy gradient in (2).

The idea of using ∇model θ J(θ) can be found in (Gu et al., 2016; Thanard Kurutach, 2018) . The problem, however, is that the transition model can deviate significantly from the true model, introducing a potentially large bias that makes the training problem unstable. Thanard Kurutach (2018) found that by training an ensemble of models can help stabilize the training. As we discuss in the sequel, one of our key idea is to calculate on-policy and model-based gradient simultaneously, and use the unbiased information of the on-policy gradient to evaluate the biasness of model-based gradient, and weigh its importance accordingly to ensure the stability while exploiting the model-based information.

3. GRADIENT ESTIMATION WITH SHRINKAGE:

We introduce our main method which uses shrinkage estimator to trade-off the bias and variance to adaptively combine multiple gradient estimators. We start with reviewing James-Stein Shrinkage estimator in Section 3.1, and then discuss its two applications in reinforcement learning, including shrinkage estimator towards model-based policy gradient in Section 3.2, and temporally smoothed gradients in Section 3.3.

3.1. JAMES-STEIN SHRINKAGE ESTIMATOR:

James-Stein estimator (Stein, 1956; James & Stein, 1961 ) is an estimator of the mean of a multivariate normal distribution that achieves a smaller mean square error (MSE) than the maximum likelihood estimator (MLE) in dimensions more than three. This estimator shrinks the standard maximum likelihood estimator towards a specified target by an adaptive shrinkage factor, which introduces bias whilst obtains lower variance and MSE.

Consider the problem of estimating the mean of a p-dimensional multivariate normal distribution N (θ, Q), where θ ∈ R d is the mean and Q ∈ R p×p is the covariance matrix. Given a sample x drawn from N (θ, Q), a natural estimator of the maximum likelihood estimation (MLE) which predicts θ using θ ≈ x. This estimator is unbiased, and gives the minimum variance among all the possible unbiased estimators.

MLE was used to be believed to be admissible, that is, no other estimators can always achieve lower MSE than MLE. However, this was proved to be false by James and Stein, who constructed an estimator, now known as James-Stein estimator, that always dominants MLE when p > 2:

EQUATION

where θ ∈ R p is any fixed constant relative to x, and α is a combination coefficient defined as where p = trace(Q)/λ max (Q) is considered as an effective dimension of covariance matrix Q, which equals p if Q is an identity matrix. The intuition is that := (x x x -θ) Q -1 (x x x -θ) the bias of θ in respect to the variation of x. If θ is close to the true mean θ, α is small and θ is highly weighted; if θ is far away from the true mean, α is large and the unbiased estimator x is highly weighted. Overall, the adaptive strategy ensures the bias and variance is always traded off optimally.

EQUATION

More generally, we can relax c = p -2 in equation ( 5) to 0 < c < 2(p -2), which still dominates the MLE when p > 2. Baranchik (1964) also suggests that JS estimator can be further improved by limiting the combination coefficient to be positive. This gives a general positive James-Stein estimator: (JSE + ) as

EQUATION

Although the standard theoretical results of JS estimators are established for normal distribution, which may not hold for settings in deep RL, however, JS still provides a simple yet highly useful practical strategy for adaptive combination. In practice, when the covariance matrix Q is not known, it can be replaced by an empirical estimation Q, which is estimated over multiple trajectories. In fact, in our subsequent experiments, we approximate Q with its diagonal matrix to avoid matrix inverse. This introduces error in estimating α + , but can be practically compensated choosing a proper parameter c using trial tests. This work focuses on demonstrating a novel empirical application of JS estimators for deep RL, and defers further theoretical work to future works.

3.2. SHRINKING TOWARDS MODEL-BASED GRADIENT ESTIMATOR:

We describe our method that adaptively combines model-based and on-policy strategies using shrinkage estimator. For notation, we denote by g c := ∇θ J(θ) the on-policy gradient in (2) obtained from the "real" trajectories, and g mb := ∇model θ J(θ) the model-based gradient obtained from "imaginary" trajectories rolled out from learned dynamics. With these two gradient estimators, we can obtain a new estimator as

EQUATION

)

The only difference between these two estimators g mb and g c is the data used for gradient calculation.

Here α + serves as a measurement of the quality of dynamic model training. A good dynamic gives imaginary trajectories that are similar to real trajectories, making g mb are close to g c and hence a small α + according to (6).

One subtle difficulty is that when simulating long horizon trajectories, the error in the model is amplified, causing significantly bad model-based estimator (yielding α ≈ 1 in ( 8)). To address this problem, we simulate only short horizon imaginary trajectories initialized with different states randomly drawn from true samples. This trick allows us to obtain higher quality model-based estimators. We integrate our method with Proximal Policy Optimization (PPO) (Schulman et al., 2017; Heess et al., 2017) , which we refer as PPO-MBS, and is summarized in algorithm 1. 

3.3. SHRINKING TOWARDS TEMPORALLY SMOOTHED GRADIENT:

In addition to the model-based method, we present another simpler method that constructs shrinkage estimator based on a smoothed gradient obtained by averaging previously gradients. This simple strategy is significantly simpler and computationally efficient since no additional model learning or other expensive calculation is needed. It turns out this strategy works surprisingly well.

Denote by g t the on-policy gradient at iteration t, and g c := g T the gradient at the current iteration (T ). We calculate a smoothed gradient from the previous iterations by a weighted average:

EQUATION

Assume the parameters do not change much over the iterations, g w provides a reasonable guess of the current gradient. Therefore, we can shrink the current noisy gradient towards the g w :

EQUATION

We refer our method as PPO-STS and summarize it in algorithm 2. Interestingly, one can view the update in (10) as a type of momentum method, because it forces the parameters to update along the previous gradient directions. Compared with the standard momentum methods developed from the optimization perspective (referred as optimization momentum) (Sutskever et al., 2013; Nesterov, 2013; Kingma & Ba, 2014; Duchi et al., 2011) , our method is derived for the statistical purpose of obtaining a better estimator of gradient (and hence forms statistical momentum). A unique feature of the statistical momentum is that its magnitude changes adaptively with the consistency between the momentum and the current gradient estimator, reducing to zero if the gradient estimator is deterministic (zero variance). As we do in our experiments, our statistical momentum can be directly applied on the top of classical optimization momentum to combine the benefits of both.

4. EXPERIMENTS:

We evaluate PPO-MBS and PPO-STS on continuous control tasks (Brockman et al., 2016) using the MuJoCo physics simulator (Todorov et al., 2012) . We report the variance, bias, and MSE of gradient estimation with and without PPO-STS, PPO-MBS in Section 4.1, and then we show the results of extensive experiments in Section 4.2. We choose Proximal Policy Gradient (PPO) as the baseline, which is one of the state-of-the-art policy gradient methods on various continuous control tasks. Implementation details 1 and hyperparameters values are provided in Appendix A and Table 1 .

Extensive experiments show that PPO-STS and PPO-MBS can reduce the MSE and the variance of the gradient estimation while we can obtain significantly better performance comparing standard on-policy methods.

4.1. EVALUATION OF PPO-STS AND PPO-MBS:

We start with evaluating the effect of PPO-STS and PPO-MBS on gradient estimation, namely, we demonstrate the bias variance decomposition of gradient estimations, and also compare the average reward of adaptive combination with different fixed combinations. We choose the Walker2d as the evaluation environment, and calculate the MSE, variance, and bias of gradient estimation in PPO-MBS (Algo 1), PPO-STS (Algo 2), and PPO during the training process.

Figure 1 shows the results on Walker2d. Similarly, Figure 2 shows the results of PPO-MBS on Walker2d-v1. The empirical results show that our methods obtain significantly lower MSE, and variance than the typical model-free policy gradient method. It also shows that adaptive combination (adaptively to learn value of α) is essential to get a better estimation of gradient, and can lead to higher reward in policy optimization. The value of James-Stein coefficient α can also indicate the difference between temporal gradient and current gradient, or model-free policy gradient and model-based policy gradient estimates.

4.2. COMPARISON PPO-STS, PPO-MBS WITH STATE-OF-THE-ART METHODS:

Finally, we evaluate PPO-STS, PPO-MBS on a more extensive list of tasks shown in Figure 3 . It shows that PPO-STS and PPO-MBS have superior performance over baseline method PPO. Empirically, we find that PPO-STS has better performance on high dimensional tasks, such as Walker2d while PPO-MBS performs better on low dimensional environments such as Reacher, Swimmer or Hopper. Ww further investigate PPO-STS on more higher dimensional tasks such as Humanoid and Ant, the results of which are in Appendix B.

5. RELATED WORK:

As a classical method, James-Stein (JS) estimator has been extensively studied and used in statistics and machine learning. See e.g., Gruber (2017) ; Young & Smith (2005) for overviews. However, we are not aware of previous works on applying JS estimators for gradient estimation in RL.

The idea of combining multiple algorithms to achieve better results have been widely considered in reinforcement learning. For example, Wiering & Van Hasselt (2008) by combining the policies derived from the value functions of the different RL algorithms. Hessel et al. (2017) combines several extensions of DQN to improve Atari games. Closely related to our work is Gu et al. (2017) , which uses a linear combination of on-policy and off-policy gradients, but with a fixed, user-specified combination coefficient. O'Donoghue et al. (2017) unifies policy gradient and Q-Learning and empirically observes better data efficiency and stability.

Another large body of works have explored the model-based/ model-free integration. The classic Dyna framework (Sutton, 1991) has primarily been used with small and discrete systems, although its extensions for high dimensional continuous control tasks exist. (Gu et al., 2016; Feinberg et al., 2018) . There is also a line of research using imaginary trajectories for training policy (Kalweit & Boedecker, 2017; Thanard Kurutach, 2018) . Heess et al. (2015) uses models to improve the accuracy of model-free value function backups. Model-based training has also been used to produce a good initialization for the model-free training (Farshidian et al., 2014; Nagabandi et al., 2017) , most recently, Clavera et al. (2018) proposes to meta learn a policy to adapt to an ensemble of dynamic models.

Our work provides a unified adaptive framework for various gradient estimators, with the purpose of utilizing advantages of each component, which results in lower variance and MSE over onpolicy model free optimization. Two illustrated examples in our experiments show that our adaptive combination strategy achieves better results than on-policy methods or fixed-weighted combinations.

6. CONCLUSION:

We propose to use James-Stein shrinkage estimator to combine the unbiased on-policy gradient estimator with biased but low variance estimators constructed based on models-based methods or historic temporally smoothed gradients. Empirical experiments show that our adaptive combination strategy can reduce variance and MSE, yielding significant improvement over standard on-policy methods. Future work includes investigating more powerful shrinkage estimators and alternative biased estimators for combination in policy optimization.

:



Figures and Tables:

Figure 1: Comparison of MSE, Variance, and Bias of PPO and PPO-MBS gradient estimation on Walker2d. Bias Square

Figure 2: Comparison of MSE, Variance, and Bias of PPO and PPO-STS gradient estimation on Walker2d.

Figure 3: Comparison of PPO-STS, PPO-MBS, and PPO on various continuous control task.

Algorithm 1 PPO-MBS: PPO with model-based gradient shrinkage Initialize policy π θ (a|s), model-based dynamic f φ (s, a), and replay buffer D ← ∅. repeat Collect real trajectories using π θ and add them to replay buffer D. for K iterations do Update φ with data from replay buffer D using (3). end for Sample imaginary trajectories {τi} n i=1 with max time-step Tp using π θ and f φ . for M iterations do Calculate model free PPO gradient gc w.r.t θ with real samples.

Table content: <table><tr><td colspan="3">Calculate model-based PPO gradient g mb w.r.t θ with imaginary trajectories {τi} n i=1 .</td></tr><tr><td>Get JSE ĝJS +</td><td>using (8) and update θ with ĝJS +</td><td>.</td></tr><tr><td>end for</td><td/><td/></tr><tr><td>until Convergence</td><td/><td/></tr></table>

Algorithm 2 PPO-STS: PPO with temporally smoothed gradient shrinkage Initialize policy π θ (a|s) and historic gradient set β ← ∅. repeat Collect real trajectories using π θ . for M iterations do Calculate model free PPO gradient gc w.r.t θ.

Table content: <table><tr><td colspan="2">Calculate weighted average historic gradient gw.</td><td/></tr><tr><td>Get JSE ĝJS +</td><td>using (10) and update θ with ĝJS +</td><td>.</td></tr><tr><td colspan="2">Add gw to the gradient set β.</td><td/></tr><tr><td>end for</td><td/><td/></tr><tr><td>until Convergence</td><td/><td/></tr></table>

Peer review (iclr2019):
 The paper suggest a shrinkage-based estimator (James-Stein estimator) to compute policy gradients in reinforcement learning to reduce the variance by trading some bias. Two versions are suggested: The on-policy gradients is shrinked either towards (i) model based gradient, or towards (ii) a delayed average of previous on-policy gradients. Empirically, both methods have better performance than the baseline. The paper is clearly written and well motivated. Some details are lacking that would be of interest to the reader and to make the results reproducible. For example how is .hat Q estimated? The trick that is referred to in the end of page about only simulating short horizon trajectories deserves more detail. I would suggest providing more details, in the text and/or in the two algorithms. The authors claim that JS estimator for gradient estimation in RL has not been used before. I am also not aware of any other work, but have also not been looking after that line of work. The paper seems to be a good contribution to the ever increasing literature of how to improve deep RL. Minors: .hat .theta on RHS in eq (7) should be .bar .theta ? Otherwise, what is .hat .theta? section 4.2 Ww -> We ======= After revision ========= I still think this is a very interesting, novel and relevant idea that desires attention. However, on the same time, I agree with the points raised by the other two reviewers which are all well-motivated and relevant concerns. Therefore, I join the view that the paper is not yet ready for publication but I do encourage the authors to improve their work and resubmit to another venue.