A Closer Look at Gradient Estimators with Reinforcement Learning as Inference
Abstract
The concept of reinforcement learning as inference (RLAI) has led to the creation of a variety of popular algorithms in deep reinforcement learning.
Unfortunately, most research in this area relies on wider algorithmic innovations not necessarily relevant to such frameworks.
Additionally, many seemingly unimportant modiﬁcations made to these algorithms, actually produce inconsistencies with the original inference problem posed by RLAI.
Taking a divergence minimization perspective, this work considers some of the practical merits and theoretical issues created by the choice of loss function minimized in the policy update for off-policy reinforcement learning.
Our results show that while the choice of divergence rarely has a major affect on the sample efﬁciency of the algorithm, it can have important practical repercussions on ease of implementation, computational efﬁciency, and restrictions to the distribution over actions.
Introduction
The recent success of deep reinforcement learning (RL) across a wide range of applications is largely attributable to advances in actor-critic algorithms [25, 19, 22, 37, 49, 40, 50, 32].
While versatile, many of these algorithms are difﬁcult to implement and require careful hyper-parameter tuning [15].
Additionally, the most performant implementations of these algorithms rely on heuristics to stabilize learning that are not explicitly referenced within the algorithm deﬁnition or even informed by the relevant theory [15, 54].
In some cases these implementation details become so crucial to performance that reproducing the original results becomes almost impossible [30, 14, 26].
In light of these issues, some researchers refocused on algorithms informed by probabilistic methods, most notably a popular class of algorithms derived within the ‘RL as Inference’ (RLAI) framework [43, 35, 34, 59, 44, 17, 2, 51, 8, 47].
Under this framework, policy learning in a Markov decision process can be framed as an approximate posterior inference task, which has several key beneﬁts over alternative approaches.
First, policies trained using this approach can produce agents that transfer skills between tasks or even entire environments [16].
Second, RLAI couches the task of model learning (i.e. learning the dynamics of the environment) within the larger literature of generative modeling and can take advantage of existing tools and techniques [10, 2, 28, 8], allowing us to reason in the even wider class of differentiable simulation [39, 18, 3].
Lastly, Kullback–Leibler (KL) regularization introduces the notion of “default policies”, which induce different behaviors in the learned policy that are useful in solving certain environments and impriving exploration [20, 24, 23].
Practically speaking, RLAI provides insights into how to construct various policy losses which interact more intelligently with modern computation graph software [45, 1, 6].
In framing RL as an approximate inference task, we force this consideration in deciding what divergence we wish to minimize, and how gradients of that expectation will be estimated.
This is doubly important in settings where the the simulator itself can be fully or partially differentiable [6, 5].
While we will not consider how gradient estimates of the forward and backwards messages can be better estimated in general, we will look at how to provide learning signal to policies in a general setting (e.g. for anonymous policy distributions given only log-probability of samples).
We argue that closer consideration of this class of gradient estimators will allow for more composable reinforcement learning software, which in turn would allow for more stable policy learning over all model or distribution classes.
This paper focuses on the following: (1) gradient estimators of the forward and reverse KL divergences with respect to the policy parameters, (2) the practical implementations of popular algorithms based on these estimators, and (3) an empirical comparison of the two resulting policy losses.
The paper is organized as follows, Section 2 reviews RL and RLAI notation and background, sections 3 andderive policy gradient estimators based on divergence minimization, and section 5 considers a simple experiment to show how the performance of each objective differs under modiﬁed hyper-parameters.
We conclude with a discussion on the utility of each policy loss, and consider future directions.
Background
This section details three important ideas required to understand the differences between gradient estimators and target objectives in RLAI.
The ﬁrst is the probabilistic structure of a Markov Decision process, which will allow us to construct a corresponding inference task and model dependencies within various gradient estimators.
The second is the is how actor-critic algorithms amortize many of the previous costly sampling operations required in classical inference based methods and on-policy algorithms.
Lastly, we need a basic understanding of the RLAI inference task, and how it can be re-framed as divergence minimization.
2.1 Markov Decision Processes
We deﬁne a Markov decision process (MDP), M Φ ( S , A , R , T 0 , T , Π Φ ) , as a random process which produces a sequence of tuples τ t := { a t , s t , s t +1 , r t } , for a particular set of states s t ∈ S , actions a t ∈ A , initial state distribution p ( s 0 ) ∈ T 0 , transition distribution p ( s t +1 | s t , a t ) ∈ T , reward function r t : S ×
A → R , and policy π φ ∈ Π φ : S → A parameterized by φ .
The generative model for this ﬁnite horizon process is deﬁned:
[Equation]
Here we also denote the marginal distribution with respect to a state s , under the trajectory distribution given in equation 1, as q π φ ( s ) .
In this setting, RL seeks to recover the policy which maximizes the expected cumulative reward over a trajectory, φ ∗ = arg max φ ∈
{+Φ E q πφ ( τ ) [ (cid:80) Tt =
 r t ( s t , a t )] .+}
More generally, we can extend this notion to inﬁnite horizon problems through the use of a discount factor which signiﬁes the non-zero probability of termination after every time-step.
In this setting it is common to arrive at a policy through the solution to a ﬁxed point equation [4, 53], which can be rewritten as the following optimization problem with respect to the policy parameters φ :
[Equation]
Here, the distribution over states, d π φ ( s ) , is referred to as the state occupancy .
Methods which simultaneously learn the policy π φ and critic Q π φ are generally called actor-critic algorithms [40, 25, 50, 48], and represent the most popular class of reinforcement learning algorithms currently available.
2.2 Soft Actor Critic
One popular variant of this framework is known as soft actor-critic (SAC).
This algorithm is similar to most off-policy actor critic algorithms, in that it learns to estimate the expected reward ahead given the state and action (the Q-function), using the bellman error:
[Equation]
Here, D deﬁnes the set of examples gathered from the environment under a behavioral policy, or a sequence of behavioral policies.
The primary distinction between SAC and other AC algorithms [40, 19] however, is that r t is replaced with a surrogate ˆ r t =
r t − H t , where H t indicates the policy entropy at state s t .
Additionally in this setting, the bar over Q indicates a target network, generally deﬁned by the use of parameter averaging [41, 57], a model ensemble, or both [27].
Using this Q function estimate, we can iteratively update our policy to match the so called “softmax-optimal" policy [24, 25, 35].
This approach deﬁnes its policy objective based on a divergence between the current policy and a normalized target Q function:
[Equation]
Here the gradient is evaluated using the reparametrization trick [33] in cases where the distribution over actions is reparameterizable (e.g. normal, gamma ect.), and with REINFORCE [58] otherwise.
A description of SAC is included in Algorithm 1.
To ensure a fair comparison of the two divergences described within this paper, we use SAC as the base algorithm, and modify only the policy update with the corresponding gradient estimator under consideration (colored blue in Algorithm 1).
Initialize Q-functions, Q θ 1 ( s, a ) , Q θ 2 ( s, a ) and policy weights π φ ( a t
| s t )
Initialize target networks Q ¯ θ 1 ( s, a ) , Q ¯ θ 2 ( s, a )
Initialize replay buffer D for each iteration do for each environment step do
Sample action form policy π store transition into replay buffer end for for each step do
Update Q-function using J Q ( θ i ) for ∈ { 2 } θ i ← θ i − λ Q ˆ ∇ θ i J Q
( θ i )
Update policy weights using J π ( φ ) φ ← φ − λ π ˆ ∇ φ J π ( φ )
Adjust temperature using J ( α ← α − λ ˆ ∇ α J α )
Update target network weights using Polyak averaging ¯ θ
i ←
τθ i + 1 − i end for end for
2.3 Reinforcement Learning as Inference
RL can be posed as a posterior inference problem over the set of trajectories conditioned on how well we would like the agent to do.
This breaks down into a set of latent random variables (the trajectories), and a set of observed random variables describing how well the agent is performing in the classical sense (e.g. the reward).
We will refer to the observed random variables as “optimality" variables, denoted O t .
These optimality variables indicate if a given state action pair was optimal or not, and are assumed to be Bernoulli distributed.
We say that a state action was optimal by indicating O t
= 1 , while we say that a state action was sub-optimal by indicating O t = 0 .
In RLAI, O t is distributed as, p ( O t = 0 | a t , s t ) = 1 − 1 Z 0 exp r ( a t , s t ) , and p ( O t = 1 | [Equation]
[Equation]
We note that Z 0 is dependent on the range of values the reward can take on.
In general it is assumed that the reward is strictly negative, making a simple choice for Z 0 = 1 .
With optimality deﬁned, we can consider the probability that a given trajectory, a sequence of state action pairs sampled from a policy interacting with the environment, is itself optimal.
For ﬁnite horizon problems, this joint distribution over both latent and observed random variables is deﬁned,
[Equation]
where π 0 ( a t | s t ) denotes the ‘default’ policy or prior over actions given states.
For clarity, probabilistic graphical models with different reward and state assumptions are given in Figure 1.
However in control we don’t just want a metric on how likely a given trajectory is to be optimal, we want to produce a policy which itself induces trajectories that are optimal: π ∗ ( a t | s t ) =
π ( a t | s t , o t , ..., o T − 1 ) .
Unfortunately, arriving at an exact π ∗ ( a t | s t ) is intractable in all but the simplest cases.
Whats more, both our graphical model and our resulting joint distribution are deﬁned in terms of entire trajectories.
This means that not only must we perform a posterior inference task (e.g. determine what the distribution over states and actions given optimality), but we must also ﬁnd a way recover the marginal distribution over actions given the state and optimality for any state which we might see through environment interaction.
We address the ﬁrst issue (intractability of the posterior inference problem) by using variational inference to produce an approximate posterior efﬁciently.
We will then show through manipulation of the expected gradient and structure of the graphical model, we can deﬁne an objective who’s optimum characterizes the same policy as is deﬁned by π ( a t | s t , o t , ..., o T − 1 ) .
sasro..... sasro..... sasro.....
Policy Gradient Estimators Via The Reverse KL Divergence
This section contains derivations relevant to the reverse KL applied to RLAI.
More speciﬁcally, this section gives information on both the un-marginalized and marginalized objective and its associated gradient estimator.
This section will rely on classic tools found in variational inference, and makes frequent use of the REINFORCE trick [58], and the score function trick [52].
3.0.1 The Un-marginalized Reverse KL Objective
For the following we assume that τ t :=
( a t , s t , s t +1 , r t ) , and ( s 0: T , a 0: T − 1 , r 0: T − 1 ) :
= τ .
We begin by looking at the objective derived by integrating over all time-steps simultaneously and differentiating.
To derive this objective we start with the functional form of the reverse KL divergence,
[Equation]
Next we can look at the log term and note, that it can be broken up into each time step, and additionally that the dynamics between the the joint and approximate distribution cancel out.
We can then simplify the integral to exclude terms which are constant with respect to the policy parameters.
This allows us to convert a difﬁcult to compute objective which included the optimal policy, into an objective we can estimate through Monte-Carlo sampling and interaction with the environment:
[Equation]
3.0.2 A Marginalized Reverse KL Gradient Estimator
Unfortunately, the naive gradient estimator using REINFORCE on this loss will have variance which grows with the time-horizon T. However by considering its gradient estimator we can infer an objective whose minimum is the same as the one above up to a constant factor.
In a similar fashion to the policy gradient theorem [53] we consider the expected gradient of our loss and simplify:
[Equation]
where ˆ ∇ φ = ∇ φ log π φ ( s t , a t ) .
We start by passing the gradient into the expectation under the assumption that the integrand is bounded over the domain of interest,
[Equation]
Note that s t and s (cid:48)
t correspond to the same random variable (as do a t and a (cid:48) t for all t ).
The primes are introduced only to index the double sum below correctly.
Here we use the score function trick to transform the integral of a gradient into the expectation over a gradient.
We then combine terms and simplify, noting that the gradient of the dynamics is zero.
We then convolve these sums and further simplify using d separation of the probabilistic graphical model.
[Equation]
Where the last line follows from the Markov property of the graphical model which gives t (cid:48) <
t → r ( s t (cid:48) , a t (cid:48) ) − log π φ ( a t (cid:48) | s t (cid:48) )
+ log π 0 ( a t (cid:48) | s t (cid:48) ) is independent of ∇ φ log π φ ( a t
| s t ) .
This means that we can decompose the expectation such that for all t (cid:48) <
t ,
[Equation]
Therefore the inner expectation is a constant for all t (cid:48) less then t , and the score function estimator is zero in expectation at these terms in the series.
This sum of (entropy regularized) rewards ahead is often referred to as the Q function or the advantage (if we include a baseline or value function which integrates over all actions).
This term can be deﬁned more explicitly as,
[Equation]
Using this estimate of the expected rewards ahead, we can re-scale the log probability of certain actions and give them more weight in the gradient estimator.
Crucially, the variance of this gradient estimator is strictly lower, while maintaining the correct target distribution of interest.
The full marginalized objective and its estimator are given as the following:
[Equation]
where a mt ∼ π φ ( a t | s nt ) , and s nt ∼ q φ ( s t ) .
Reweighted Policy Gradient Estimators Via the Forward KL
This section details derivations for how to produce an approximately optimal policy under the forward
KL, and like the previous section we will produce a lower variance gradient estimator which targets the original distribution of interest.
Here we take advantage of tools from self normalized importance weighting to produce a biased but consistent estimator, and brieﬂy discuss practical issues.
4.0.1 Un-marginalized Objective Derivation for the Forward KL
This objective, represents the expected KL divergence between a target distribution, and an approximate distribution under samples from the target distribution given by KL ( p ( τ | O )
||
q φ ( τ )) .
Again, our latent variables are given as the associated state-action pairs, while the observed variables represents the optimality distribution.
In order to evaluate the expression above, we apply a standard importance weighting scheme, beginning with the removal constant terms within the expectation.
[Equation]
4.0.2 Un-marginalized gradient of the Forward KL Objective
Crucially, because the original problem is over a distribution not dependent on φ
[Equation]
we can directly pass the gradient into the expectation without the REINFORCE trick
[Equation]
or reparameterization:
[Equation]
This actually gives us a simple algorithm that can evaluate the gradient, and thereby minimize the expected KL between our two distributions following samples generated from our policy interacting with the simulator.
Based upon this approach, we can apply self normalized importance weighting to avoid explicitly computing the constant Z, at the cost of a biased but consistent estimator:
[Equation]
Where for optimal trajectories ( O = O t =0: T − 1 = 1 )
[Equation]
the un-normalized weights are deﬁned as,
[Equation]
=
[Equation]
We also include reward normalization constant Z 0 for generality.
In order to avoid computation of Z explicitly, we take advantage of the consistent but biased self normalized importance weighting estimator.
This means weights w i are replaced with ˆ w i using τ i ∼ q φ ( τ ) :
[Equation]
In this case we now have a biased estimator of the gradient, but also one that is generally low variance when compared to its un-marginalized RKL counterpart [34],
[Equation]
4.0.3 Marginalization Proof for forward KL
Like the reverse KL, we can further reduce variance of our gradient estimator.
In order to show that the objective can be marginalized with respect to actions given states and optimality, we start with the original deﬁnition of the objectives gradient. Using d separation of the probabilistic graphical model to ignore the optimality variables prior to the state conditioned on we have that:
[Equation]
Next looking at just one element of this sum at index t
[Equation]
=
[Equation]
, o T | s t ) dτ = (cid:90) s t ... (cid:90)
s T (cid:90) a t ... (cid:90) a T −p ( s t | O 0 ...
O T − 1 ) π 0 ( a t | s t ) p ( o t , ... [Equation] o T | s t ) ∇ φ log π φ ( a t | s t ) p ( s t +1 ...
s T , o t , ..., o T | s t , a t ) dτ
= (cid:90) s t ... (cid:90)
s T (cid:90) a t ... (cid:90) a T −p (
[Equation]
| O 0 ...
O T −
[Equation]
π 0 ( a t | s t ) p
( o
[Equation]
[Equation]
o
[Equation]
| s t ) ∇ φ log π
[Equation]
( a t | s t ) p ( s t +1 ...
s T
[Equation]
o t
[Equation]
[Equation]
o
[Equation]
| s t
[Equation]
a t ) dτ =
[Equation]
(cid:34)
[Equation]
(cid:20)
[Equation]
[Equation]
| a
[Equation]
) (cid:21)(cid:35)(cid:35)
Notably, we can deﬁne a speciﬁc expectation for convenience, which can be either be estimated through sampling or a learned function approximator as is done in practice [2, 47]:
[Equation]
We then note the following about the normalization constant sitting outside the action expectation:
[Equation]
This deﬁnition allows us to deﬁne our marginalized expectation more simply, and identify the crucial issue with direct estimation of the marginalized forward KL: [Equation] We can deal with the outer expectation in terms of p ( s t | O 0: T ) as well, though it induces a less satisfactory estimator whose variance again grows with time-horizon: [Equation]
Such an estimator would function by sampling N independent trajectories from q φ
[Equation]
and then for the speciﬁed time, re-weight each function f at its respective state by its associated regularized cumulative return. With this consideration in mind, we can deﬁne importance weights over states and actions for each time-step which produce a consistent estimate of the forward KL gradient: [Equation] with τ
[Equation]
where
[Equation] and
[Equation]
with x ti,k = exp
[Equation]
For our experiments, as done in practice in [2, 47], we re-weight the loss each at each state uniformly (i.e. ¯ y ti = 1 M ), to avoid high variance of the gradient estimator, ensure proper coverage of the state-space for better generalization, and to maintain sample efﬁciency within the larger RL algorithm.
Additionally we target the classic Q function instead of the soft Q function ¯ Q , as we found that using this lower bound on ¯ Q tended to produce more stable learning within the policy as in prior work [2].
Experiments
[Figure]
To display the relative merits of both objectives, we include experiments where a native implementation of soft actor critic was taken, and modiﬁed to include the forward KL objective (FKL).
We begin with an illustration of an instance where you can ﬁnd improvement by replacing the more common RKL for the FKL objective.
In Figure 2, we plot results of training on the Humanoid benchmark from MuJoCo using RKL and FKL.
Here, the FKL improves sample efﬁciency over the RKL, demonstrating approximately 60% increase in sample efﬁciency.
Humanoid is the most challenging of the ﬁve tasks considered from the MuJoCo suite.
Additional advantages of the FKL policy objective include its ease of implementation and its distribution agnostic property. This property means that in discrete control, or for hierarchical policies with discrete branching, the RKL objective would require further modiﬁcation, while the FKL functions as is. One drawback that we note, however, is that the FKL estimator scales with computation, i.e. its bias and thus its performance is dependent on the number of N re-samples from the policy.
Finally, for continuous control, we additionally show that the FKL can also improve the algorithm’s robustness to changes in hyperparameters.
All experiments are run with at 10 replicates and plotted with one-half standard deviation and smoothed using a 0.9 weighted EMA.
As seen in Figure 3, the forward KL can maintain policy performance, even when memory is severely limited.
Related Work
There is a growing literature on various methods that incorporate the RLAI framework discussed within this paper [2, 24, 35, 59, 60, 31, 43, 55].
In many cases these works use RLAI to give theoretical basis for the empirical effectiveness of well established implementation tricks such as entropy bonuses [24, 35, 59].
Very few works however, have considered the comparison between the forward and reverse KL as well as the relative merits of each in the context of reinforcement learning objectives [43, 38].
Even fewer produce scalable and sample efﬁcient implementations of both objectives, making a proper comparison difﬁcult.
We note for completeness that work has also [Figure] been done in formalizing methods relating to the original policy gradient algorithm [58, 50, 48] using tools from the optimization literature [21, 56], though these analysis do not directly relate the the RLAI framework.
Work by [7] more closely relates to this work in their analysis of improvement under the FKL, however they do not focus on practical gradient estimation.
Similarly there has been work analyzing critic and gradient estimation [9] in average return MDPs, as well as empirical studies for the more commonly utilized discounted MDPs discussed above [29]
[Figure]
Algorithmically, the closest class of algorithms to our work in the classical literature comes from policy search methods [43, 42, 36, 11, 46, 12].
Of the subset of these that have been made scalable, closely related algorithms to our work come from [2, 51, 8].
Unlike [2], we do not require an alternating update and directly learn the policy via the forward KL.
Additionally while we do not provide monotonic improvement guarantees which require additional assumptions and a trust-region update, we do provide a full proof of marginalization of the forward KL as it applies to RLAI in a ﬁnite horizon setting.
A similar proof for a marginalization can be found in [47], however this assumes access to the steady state distribution of the Markov chain as well as a model of the environment.
We do note that the objective for the target policy given in [51] does come very close to our own, however it requires additional parameterization of a value function, as well as a more complicated alternating update for the temperature coefﬁcient α , and a different inference formulation which culminates in an advantage based objective.
Most importantly, while these works represent end-to-end algorithms, our estimators can be applied a wide variety of existing frameworks.
Discussion
This paper describes in detail the differences in policy objective which occur both mathematically, and practically under a divergence minimization perspective of RLAI.
Crucially, we show that there are very real advantages to considering other objectives besides the classic policy gradient methods which are most common.
While these alternatives like the FKL are generally not going to drastically change performance, their ease of implementation, robustness to hyperparameters, and their distribution agnostic nature make them a very desirable alternative to classic methods.
In the future we would like to test how these objectives function on a wider variety of benchmarks and policy distributions.
Additionally, its possible that better non-uniform sampling methods might be derived following p ( s t | O 0: T ) .
Though estimating an expectation under such a distribution might be difﬁcult, and potentially require investigation into the particle smoothing literature [13].