V ARIANCE R EDUCTION FOR P OLICY G RADIENT M ETHODS WITH A CTION -D EPENDENT B ASELINES
Anonymous authors
Paper under double-blind review
A BSTRACT
Policy gradient methods have enjoyed success in deep reinforcement learning but suffer from high variance of gradient estimates.
The high variance problem is particularly exasperated in problems with long horizons or high-dimensional action spaces.
To mitigate this issue, we derive an action-dependent baseline for variance reduction which fully exploits the structural form of the stochastic policy itself, and does not make any additional assumptions about the MDP.
We demonstrate and quantify the beneﬁt of the action-dependent baseline through both theoretical analysis as well as numerical results.
Our experimental results indicate that actiondependent baselines allow for faster learning on standard reinforcement learning benchmarks as well as on high dimensional manipulation and multi-agent communication tasks.
I NTRODUCTION
Deep reinforcement learning has achieved impressive results in recent years in domains such as video games from raw visual inputs (Mnih et al., 2015), board games (Silver et al., 2016), simulated control tasks (Schulman et al., 2016; Lillicrap et al., 2016; Rajeswaran & V. Kumar, 2017), and robotics (Levine et al., 2016).
An important class of methods behind many of these success stories are policy gradient methods (Williams, 1992; Sutton et al., 2000; Kakade, 2002; Schulman et al., 2015; Mnih et al., 2016), which directly optimize parameters of a stochastic policy through local gradient information obtained by interacting with the environment using the current policy.
Policy gradient methods operate by increasing the log probability of actions proportional to the future rewards inﬂuenced by these actions.
On average, actions which perform better will acquire higher probability, and the policy’s expected performance improves.
A critical challenge of policy gradient methods is the high variance of the gradient estimator.
This high variance is caused in part due to difﬁculty in credit assignment to the actions which affected the future rewards.
Such issues are further exacerbated in long horizon problems, where assigning credits properly becomes even more challenging.
To reduce variance, a “baseline” is often employed, which allows us to increase or decrease the log probability of actions based on whether they perform better or worse than the average performance when starting from the same state.
This is particularly useful in long horizon problems, since the baseline helps with temporal credit assignment by removing the inﬂuence of future actions from the total reward.
A better baseline, which predicts the average performance more accurately, will lead to lower variance of the gradient estimator.
The key insight of this paper is that when the individual actions produced by the policy can be decomposed into multiple factors, we can incorporate this additional information into the baseline to further reduce variance.
In particular, when these factors are conditionally independent given the current state, we can compute a separate baseline for each factor, whose value can depend on all quantities of interest except that factor.
This serves to further help credit assignment by removing the inﬂuence of other factors on the rewards, thereby reducing variance.
In other words, information about the other factors can provide a better evaluation of how well a speciﬁc factor performs.
Such factorized policies are very common, with some examples listed below.
• In continuous control and robotics tasks, multivariate Gaussian policies with a diagonal covariance matrix are often used.
In such cases, each action coordinate can be considered a factor.
Similarly, factorized categorical policies are used in game domains like board games and Atari.
• In multi-agent and distributed systems, each agent deploys its own policy, and thus the actions of each agent can be considered a factor of the union of all actions (by all agents).
This is particularly useful in the recent emerging paradigm of centralized learning and decentralized execution (Foerster et al., 2017; Lowe et al., 2017).
In contrast to the previous example, where factorized policies are a common design choice, in these problems they are dictated by the problem setting.
We demonstrate that action-dependent baselines consistently improve the performance compared to baselines that use only state information.
The relative performance gain is task-speciﬁc, but in certain tasks, we observe signiﬁcant speed-up in the learning process.
We evaluate our proposed method on standard benchmark continuous control tasks, as well as on a high dimensional door opening task with a ﬁve-ﬁngered hand, and on a blind peg insertion POMDP task.
We believe that our method will facilitate further applications of reinforcement learning methods in domains with extremely high-dimensional actions, including multi-agent systems.
Videos and additional results of the paper are available at https://sites.google.com/view/ad-baselines .
R ELATED WORKS
Three main classes of methods for reinforcement learning include value-based methods (Watkins &
Dayan, 1992), policy-based methods (Williams, 1992; Kakade, 2002; Schulman et al., 2015), and actor-critic methods (Konda & Tsitsiklis, 2000; Peters & Schaal, 2008; Mnih et al., 2016).
Valuebased and actor-critic methods usually compute a gradient of the objective through the use of critics, which are often biased, unless strict compatibility conditions are met (Sutton et al., 2000; Konda &
Tsitsiklis, 2000).
Such conditions are rarely satisﬁed in practice due to the use of stochastic gradient methods and powerful function approximators.
In comparison, policy gradient methods are able to compute an unbiased gradient, but suffer from high variance.
These methods are therefore usually less sample efﬁcient, but can be more stable than critic-based methods (Duan et al., 2016).
A large body of work has investigated variance reduction techniques for policy gradient methods.
One effective method to reduce variance without introducing bias is through using a baseline, which has been widely studied (Sutton & Barto, 1998; Weaver & Tao, 2001;
Greensmith et al., 2004; Schulman et al., 2016).
this factorization has not been studied in detail.
A recently proposed algorithm,
Q-Prop (Gu et al., 2017), makes use of an action-dependent control variate, a technique commonly used in Monte Carlo methods and recently adopted for RL.
Since Q-Prop utilizes off-policy data, it has the potential to be more sample efﬁcient than pure on-policy methods.
However, Q-prop is signiﬁcantly more computationally expensive, since it needs to perform a large number of gradient updates on the critic using the off-policy data, thus not suitable with fast simulators.
In contrast, our formulation of action-dependent baselines has little computational overhead, and improves the sample efﬁciency compared to on-policy methods with state-only baseline.
The idea of using additional information in the baseline or critic has also been studied in other contexts.
Methods such as Guided Policy Search (Levine & Koltun, 2013; Mordatch et al., 2015) and variants train policies that act on high dimensional observations like images, but use a more low dimensional encoding of the problem like joint positions during the training process.
Recent efforts in multi-agent systems (Foerster et al., 2017; Lowe et al., 2017) also use additional information in the centralized training phase to speed-up learning.
However, using the structure in the policy parameterization itself to enhance the learning speed, as we do in this work, has not been explored.
P RELIMINARIES
In this section, we establish the notations used throughout this paper, as well as basic results for policy gradient methods, and variance reduction via baselines.
3. N OTATION
This paper assumes a discrete-time Markov decision process (MDP), deﬁned by ( S , A , P , r, ρ 0 , γ ) ,
in which S ⊆ R n is an n -dimensional state space,
A ⊆
R m an m -dimensional action space,
P : S ×
A × S → R + a transition probability function, r : S ×
A → R a bounded reward function, ρ 0 : S → R + an initial state distribution, and γ ∈ (0 , 1] a discount factor.
The presented models are based on the optimization of a stochastic policy π θ : S ×
A → R + parameterized by θ .
Let η ( π θ ) denote its expected return: η ( π θ ) =
E τ [ (cid:80) ∞ t =
γ t r ( s t , a t )] , where τ = ( s 0 , a 0 , . . . ) denotes the whole trajectory, s 0 ∼ ρ 0 ( s 0 ) , a t ∼ π θ ( a t
| s t ) , and s t +1 ∼ P ( s t +1 | s t , a t ) for all t .
Our goal is to ﬁnd the optimal policy arg max θ η ( π θ ) .
For a partially observable Markov decision process (POMDP), two more components are required, namely Ω , a set of observations, and O : S × Ω
→ R ≥ 0 , the observation probability distribution.
In the fully observable case, Ω ≡ S .
Though the analysis in this article is written for policies over states, the same analysis can be done for policies over observations.
3. T HE S CORE F UNCTION (SF) E STIMATOR
An important technique used in the derivation of the policy gradient is known as the score function (SF) estimator (Williams, 1992), which also comes up in the justiﬁcation of baselines.
Suppose that we want to estimate ∇ θ E x
[ f ( x )] where x ∼ p θ ( x ) , and the family of distributions { p θ ( x ) : θ ∈ Θ } has common support.
Further suppose that log p θ ( x ) is continuous in θ .
In this case we have
[Equation]
[Equation]
3. P OLICY G RADIENT
The Policy Gradient Theorem (Sutton et al., 2000) states that
[Equation]
For convenience, deﬁne ρ π ( s ) =
(1 − γ ) (cid:80) ∞
t =0 γ t p ( s t = s ) as the normalized state visitation frequency, and ˆ Q ( s t , a t ) = (cid:80) ∞
t (cid:48)
=
t γ t (cid:48) − t r t (cid:48) .
We can rewrite the above equation as
[Equation]
It is further shown that we can reduce the variance of this gradient estimator without introducing bias by subtracting off a quantity dependent on s t from ˆ Q ( s t , a t ) (Williams, 1992; Greensmith et al., 2004).
[Equation] [Equation]
This is valid because, applying the SF estimator in the opposite direction, we have
[Equation]
A CTION - DEPENDENT BASELINES FOR FACTORIZED POLICIES
In practice there can be rich internal structure in the policy parameterization.
For example, for continuous control tasks, a very common parameterization is to make π θ ( a t
| s t ) a multivariate Gaussian with diagonal variance, in which case each dimension a it of the action a t is conditionally independent of other dimensions, given the current state s t .
Another example is when the policy outputs a tuple of discrete actions with factorized categorical distributions.
In the following subsections, we show that such structure can be exploited to further reduce the variance of the gradient estimator without introducing bias by changing the form of the baseline.
Then, we derive the optimal action-dependent baseline for a class of problems and analyze the suboptimality of non-optimal baselines in terms of variance reduction.
We then propose several practical baselines for implementation purposes.
Even if this conditional independence does not hold (say for Gaussians with general covariance structure), as long as we can decompose the action into multiple factors, our analysis still holds, despite yielding a different baseline.
Finally, we give an exposition on how action-dependent baselines can be combined with the Generalized Advantage Estimator (GAE) (Schulman et al., 2016) to smoothly interpolate the bias-variance trade-off curve.
4. B ASELINES FOR C ONDITIONALLY I NDEPENDENT A CTIONS
First, we start with the conditionally independent case.
Assuming an m -dimensional action space, we have π θ ( a t | s t ) =
(cid:81) mi =1 π θ ( a it | s t ) .
Hence ∇
[Equation]
In this case, we can set b
i , the baseline for the i th factor, to depend on all other actions in addition to the state.
Let a − i t denote all dimensions other than i in a t and denote the i th baseline by b i ( s t , a − i t ) .
Due to conditional independence, we have
[Equation]
Hence we can use the following gradient estimator
[Equation]
4. O PTIMAL ACTION - DEPENDENT BASELINE
In this section, we derive the optimal action-dependent baseline and show that it is better than the state-only baseline.
We seek the optimal baseline to minimize the variance of the policy gradient estimate.
First, we write out the variance of the policy gradient under any actiondependent baseline.
Let us deﬁne ∇ η
i
[Equation]
=
E
[Equation]
s
[Equation]
i t )
[Equation]
and z i := ∇ θ log π θ ( a it | s t ) .
For simplicity of exposition
[Equation]
we make the following assumption:
[Equation]
which translates to meaning that different subsets of parameters strongly inﬂuence different action dimensions or factors
This is true in case of distributed systems by construction, and also true in a single agent system if different action coordinates are strongly inﬂuenced by different policy network channels.
Under this assumption, we have:
[Equation]
The overall variance is minimized when each component variance is minimized.
We now derive the optimal baselines b ∗ i ( s t , a − i t ) which minimize each respective component.
[Equation]
Having written down the expression for variance under any action-dependent baseline
[Equation]
we seek the optimal baseline that would minimize this variance
[Equation] [Equation]
The optimal action-dependent baseline is:
[Equation]
Since the optimal action-dependent baseline is different for different action coordinates
[Equation]
it is outside the family of state-dependent baselines barring pathological cases
[Equation]
4. V ARIANCE REDUCTION IMPROVEMENT
We now turn to quantifying the reduction in variance of the policy gradient estimate under the optimal baseline derived above.
Let Var ∗ ( ∇ θ η ( π θ )) denote the variance resulting from the optimal action-dependent baseline
[Equation]
and let Var ( ∇ θ η ( π θ )) denote the variance resulting from another baseline b ( s t , a t ) , which may be suboptimal or action-independent.
We use the notations:
[Equation] [Equation]
Finally
[Equation]
deﬁne the variance improvement I :=
Var ( ∇ θ η ( π θ ))
− Var ∗ ( ∇ θ η i ( π θ )) .
Using these deﬁnitions, the variance can be re-written as:
[Equation]
Furthermore
[Equation]
the variance of the gradient with the optimal baseline can be written as
[Equation]
The difference in variance can be calculated as:
[Equation]
4. S UBOPTIMALITY OF THE OPTIMAL STATE - DEPENDENT BASELINE
How much do we reduce variance over a traditional baseline that only depends on state?
Using
Equation (25), we show the following improvement
[Equation]
This suggests that the variance difference to be a weighted sum of the deviation of the per-component score-weighted marginalized Q (denoted Y i ) from the weighted average of all the component scoredweighted marginalized Q values.
This suggests that the difference is particularly large when the Q function is highly sensitive to the actions, especially along those directions that inﬂuence the gradient the most.
Our empirical results in Section 5 additionally demonstrate the beneﬁt of action-dependent over state-only baselines.
4. M ARGINALIZATION OF THE GLOBAL ACTION - VALUE FUNCTION
Using the previous theory, we now consider various baselines that could be used in practice, and associated computational cost.
Marginalized Q baseline Even though the optimal state-only baseline is known, it is rarely used in practice (Duan et al., 2016).
Rather, for both computational and conceptual beneﬁt, the choice of b ( s t ) =
E a t [ Q ( s t , a t )]
= V ( s t ) is often used.
Similarly, we propose to use b
i ( s t , a − i t ) =
E a it
[ Q π θ ( s t , a t )] which is the action-dependent analogue.
In particular, when log probability of each policy factor is loosely correlated with the action-value function, then the proposed baseline is close to the optimal baseline.
[Equation]
when E a it (cid:104) z
Ti z i ˆ Q ( s t , a t ) (cid:105)
≈ E a it (cid:2)
z Ti z i (cid:3) E a it (cid:104)
ˆ Q ( s t , a t ) (cid:105) .
This has the added beneﬁt of only needing to learn one function approximator, for estimating
Q ( s t , a t ) , and implicitly using it to obtain the baselines for each action coordinate.
Monte Carlo marginalized Q baseline After learning Q π θ ( s t , a t ) we can obtain the baselines through Monte Carlo estimates:
[Equation]
where a j ∼ π θ ( a it | s t ) are samples of the action coordinate i .
Mean marginalized Q baseline Though we reduced the computational burden from learning m functions to one function, the use of Monte Carlo samples can still be computationally expensive.
In particular, when using deep neural networks to approximate the Q-function, forward propagation through the network can be even more computationally expensive than stepping through a fast simulator (e.g. MuJoCo).
In such settings, we further propose the following more computationally practical baseline:
[Equation]
where ¯ a i = E π θ (cid:2) a it (cid:3) is the average action for coordinate i .
4. B ASELINES FOR G ENERAL A CTIONS
In the preceding derivations, we have assumed policy actions are conditionally-independent across dimensions.
In the more general case, we only assume that there are m factors a 1 t through a mt which altogether forms the action a t .
Conditioned on s t , the different factors form a certain directed acyclic graphical model (including the fully dependent case).
Without loss of generality, we assume that the following factorization holds:
[Equation]
where f ( i ) denotes the indices of the parents of the i th factor.
Let D ( i ) denote the indices of descendants of i in the graphical model (including i itself).
In this case, we can set the i th baseline to be b i ( s t , a [ m ] \ D ( i ) t ) , where [ m ] = { 1 , 2 , . . .
, m } .
In other words, the i th baseline can depend on all other factors which the i th factor does not inﬂuence.
The overall gradient estimator is given by
[Equation]
In the most general case without any conditional independence assumptions, we have f ( i ) =
, i − 1 } , and D ( i ) =
{ i, i + 1 , . . .
, m } .
The above equation reduces to
[Equation]
The above analysis for optimal baselines and variance suboptimality transfers also to the case of general actions.
Computing action-dependent baselines for general actions The marginalization presented in
Section 4.5 does not apply for the general action setting.
Instead, m individual baselines can be trained according to the factorization, and each of them can be ﬁtted from data collected from the previous iteration.
In the general case, this means ﬁtting m functions
b i ( s t , a 1 t , . . .
, a it ) , for i ∈ { 1 , . . .
, m } .
4. C OMPATIBILITY WITH GAE
Temporal Difference (TD) learning methods such as GAE Schulman et al.
(2016) allow us to smoothly interpolate between high-bias, low-variance estimates; and low-bias, high-variance estimates of the policy gradient.
These methods are based on the idea of being able to predict future returns, thereby bootstrapping the learning procedure.
In particular, when using the value function as baseline, we have A ( s t , a t ) =
E [ r t + γV
( s t +1 )
− V ( s t )] =
[ r t + γb ( s t +1 ) − b ( s t )]] if b ( s ) is an unbiased estimator for V ( s ) .
GAE proposed in Schulman et al.
(2016) use an exponential averaging of such temporal difference terms over a trajectory to signiﬁcantly reduce the variance of the advantage at the cost of a small bias (it allows us to pick where we want to be on the bias-variance curve).
Similarly, if we use b i ( s t , a − i t ) as an unbiased estimator for E a i Q ( s, a ) , we have:
[Equation]
Thus, the temporal difference error with the action dependent baselines is an unbiased estimator for the advantage function as well.
This allows us to use the GAE procedure to further reduce variance at the cost of a bias.
E XPERIMENTS AND R ESULTS
Continuous Control Benchmarks Firstly, we present the results of the proposed actiondependent baselines on popular benchmark tasks.
These tasks have been widely studied in the deep reinforcement learning community (Duan et al., 2016; Gu et al., 2017; Lillicrap et al., 2016;
Rajeswaran et al., 2017).
The studied tasks include the hopper, half-cheetah, and ant locomotion tasks simulated in MuJoCo (Todorov et al., 2012).
In addition to these tasks, we also consider a door opening task with a high dimensional multi-ﬁngered hand, to study the effectiveness of the proposed approach in high dimensional tasks.
Figure 1 presents the learning curves on these tasks.
We compare the action-dependent baseline with a baseline that uses only information about the states, which is the most common approach in literature.
We observe that the action-dependent baselines perform consistently better.
A popular baseline parameterization choice is a linear function on a small number of non-linear features of the state Duan et al.
(2016), especially for policy gradient methods.
In this work, to enable a fair comparison, we use a Random Fourier Feature representation for the baseline (Rajeswaran et al., 2017).
The features are constructed as: y ( x ) = sin( 1 ν Px + φ ) where P is a matrix with each element independently drawn from the standard normal distribution, φ is a random phase shift in [ − π, π ) and, and ν is a bandwidth parameter.
These features approximate the RKHS features under an RBF kernel.
Using these features, the baseline is parameterized as b = w T y ( x ) where x are the appropriate inputs to the baseline, and w are trainable parameters.
P and φ are not trained in this parameterization.
Such a representation was chosen for two reasons: (a) we wish to have the same number of trainable parameters for all the baseline architectures, and not have more parameters in the action-dependent case (which has a larger number of inputs to the baseline); (b) since the ﬁnal representation is linear, it is possible to accurately estimate the optimal parameters with a Newton step, thereby alleviating the results from confounding optimization issues.
For the experiments, we used 250 random Fourier features.
For policy optimization, we use a variant of the natural policy gradient and TRPO methods as described in Rajeswaran et al.
[Figure]
Choice of Action-Dependent Baseline Form Next, we study the inﬂuence of computing the baseline by using empirical averages sampled from the Q-function versus using the mean-action of the action-coordinate for computing the baseline (both described in 4.5).
In our experiments, as shown in Figure 2 we ﬁnd that the two variants perform comparably, with the latter performing slightly better towards the end of the learning process.
This suggests that though sampling from the Q-function might provide a better estimate of the conditional expectation in theory, function approximation from ﬁnite samples injects errors that may degrade the quality of estimates.
In particular, sub-sampling from the Q-function is likely to produce better results if the learned Q-function is accurate for a large fraction of the action space
[Figure]
but getting such high quality approximations might be hard in practice
[Figure]
Compatibility with GAE Temporal Difference (TD) based methods including GAE (Schulman et al., 2016) allow for a smooth interpolation between high-bias, low-variance estimates; and lowbias, high-variance estimates of the policy gradient.
As shown in Section 4, the action-dependent baselines are consistent with TD procedures with their temporal differences being estimates of the advantage function.
Our results summarized in Figure 3 suggests that slightly biasing the gradient to reduce variance produces the best results, while high-bias estimates perform poorly.
Prior work with baselines that utilize global information (Foerster et al., 2017) employ the high-bias variant.
The results here suggest that there is potential to further improve upon those results by carefully studying the bias-variance trade-off.
[Figure]
High-Dimensional Action Spaces Intuitively, the beneﬁt of the action-dependent baseline can be greater for higher dimensional problems.
We show this effect on a simple synthetic example.
The example is a one-step MDP comprising of a single state, S = { 0 } and an m -dimensional action space, A = R m .
The reward is given as the negative (cid:96) 2 loss of the action vector, r ( s,
[Equation] [Equation]
a
[Equation]
The optimal action is thus to select the zero vector a = 0 .
The results for a demonstrative example are shown in Figure 4, which shows that the action-dependent baseline successfully improves convergent more for higher dimensional problems than lower dimensional problems.
Due to the lack of state information, the linear baseline reduces to whitening the discounted returns.
The actiondependent baseline, on the other hand, allows the learning algorithm to assess the advantage of each individual action dimension by utilizing information from all other action dimensions.
[Figure]
Partial Observability Finally, we also consider the extension of the core idea of using global information, by studying a POMDP task and a multi-agent task.
We use the blind peg-insertion task which is widely studied in robot learning literature Montgomery & Levine (2016).
The task requires the robot to insert the peg into the hole (slot), but the robot is blind to the location of the hole.
Thus, we expect a searching behavior to emerge from the robot, where it learns that the hole is present on the table and performs appropriate sweeping motions till it is able to ﬁnd the hole.
In this case, we consider a baseline that knows the location of the hole.
We observe that a baseline with this additional information enables faster learning.
For the multi-agent setting, we analyze a two-agent particle environment task in which the goal is for each agent to reach their goal, where their goal is known by the other agent and they have a continuous communication channel.
Figure 5 shows that including the inclusion of information from other agents into the baseline improves the training performance, indicating that variance reduction may be key for multi-agent reinforcement learning.
[Figure]
[Figure] [Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
[Figure] [Figure]
C ONCLUSION
An action-dependent baseline enables using additional signals beyond the state to achieve bias-free variance reduction.
In this work, we consider both conditionally independent action spaces and general action spaces, and derive an optimal action-dependent baseline for a wide class of problems. We prodive analysis of the variance reduction improvement over non-optimal baselines, including the traditional optimal baseline that only depends on state.
We additionally propose several practical action-dependent baselines which perform well on a variety of continuous control tasks and are demonstrated to give greater improvement for synthetic high-dimensional action problems.
The use of additional signals beyond the local state generalizes to other problem settings, for instance in POMDP and multi-agent tasks.
In future work, we propose to investigate related methods in such settings on large scale problems.