I TERATIVE A MORTIZED P OLICY O PTIMIZATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Policy networks are a central feature of deep reinforcement learning (RL) algorithms for continuous control, enabling the estimation and sampling of high-value actions.
From the variational inference perspective on RL, policy networks, when employed with entropy or KL regularization, are a form of amortized optimization, optimizing network parameters rather than the policy distributions directly.
However, this direct amortized mapping can empirically yield suboptimal policy estimates.
Given this perspective, we consider the more ﬂexible class of iterative amortized optimizers.
We demonstrate that the resulting technique, iterative amortized policy optimization, yields performance improvements over conventional direct amortization methods on benchmark continuous control tasks.
I NTRODUCTION
Reinforcement learning (RL) algorithms involve policy evaluation and policy optimization (Sutton & Barto, 2018).
Given a policy, one can estimate the value for each state or state-action pair following that policy, and given a value estimate, one can improve the policy to maximize the value.
This latter procedure, policy optimization, can be challenging in continuous control due to instability and poor asymptotic performance.
In deep RL, where policies over continuous actions are often parameterized by deep networks, such issues are typically tackled using regularization from previous policies (Schulman et al., 2015; 2017) or by maximizing policy entropy (Mnih et al., 2016;
Fox et al., 2016).
These techniques can be interpreted as variational inference (Levine, 2018), using optimization to infer a policy that yields high expected return while satisfying prior policy constraints.
This smooths the optimization landscape, improving stability and performance (Ahmed et al., 2019).
However, one subtlety arises: when used with entropy or KL regularization, policy networks perform amortized optimization (Gershman & Goodman, 2014).
That is, rather than optimizing the action distribution, e.g. mean and variance, many deep RL algorithms, such as soft actor-critic (SAC) (Haarnoja et al., 2018b;c), instead optimize a network to output these parameters, learning to optimize the policy.
Typically, this is implemented as a direct mapping from states to action distribution parameters.
While direct amortization schemes have improved the efﬁciency of variational inference as encoder networks (Kingma & Welling, 2014; Rezende et al., 2014; Mnih & Gregor, 2014), they are also suboptimal (Cremer et al., 2018;
Kim et al., 2018; Marino et al., 2018b).
This suboptimality is referred to as the amortization gap (Cremer et al., 2018), translating into a gap in the RL objective.
Inspired by techniques and improvements from variational inference, we investigate iterative amortized policy optimization.
Iterative amortization (Marino et al., 2018b) uses gradients or errors to iteratively update the parameters of a distribution.
Unlike direct amortization, which receives gradients only after outputting the distribution, iterative amortization uses these gradients online , thereby learning to perform iterative optimization.
In generative modeling settings, iterative amortization tends to empirically outperform direct amortization (Marino et al., 2018b;a), with the added beneﬁt of ﬁnding multiple modes of the optimization landscape (Greff et al., 2019).
Using MuJoCo environments (Todorov et al., 2012) from OpenAI gym (Brockman et al., 2016), we demonstrate performance improvements of iterative amortized policy optimization over direct amortization in model-free and model-based settings.
We analyze various aspects of policy optimization, including iterative policy reﬁnement, adaptive computation, and zero-shot optimizer transfer.
Identifying policy networks as a form of amortization clariﬁes suboptimal aspects of direct approaches to policy optimization.
Iterative amortization, by harnessing gradient-based feedback during policy optimization, offers an effective and principled improvement.
B ACKGROUND 2. P RELIMINARIES
We consider Markov decision processes (MDPs), where s t ∈ S and a t ∈
A are the state and action at time t , resulting in reward r t = r
( s t , a t ) .
Environment state transitions are given by s t
+1 ∼ p env ( s t +1 | s t , a t ) , and the agent is deﬁned by a parametric distribution, p θ ( a t | s t ) , with parameters θ .
The discounted sum of rewards is denoted as
[Equation]
t γ t r t , where γ ∈ (0 , 1] is the discount factor, and τ = ( s 1 , a 1 , . . . ) is a trajectory.
The distribution over trajectories is:
[Equation]
where the initial state is drawn from the distribution ρ ( s 1 ) .
The standard RL objective consists of maximizing the expected discounted return, E p ( τ ) [ R ( τ )] .
For convenience of presentation, we use the undiscounted setting ( γ = 1 ), though the formulation can be applied with any valid γ .
2. KL-R EGULARIZED R EINFORCEMENT L EARNING
Various works have formulated RL, planning, and control problems in terms of probabilistic inference (Dayan & Hinton, 1997; Attias, 2003; Toussaint & Storkey, 2006; Todorov, 2008; Botvinick &
Toussaint, 2012; Levine, 2018).
These approaches consider the agent-environment interaction as a graphical model, then convert reward maximization into maximum marginal likelihood estimation, learning and inferring a policy that results in maximal reward.
This conversion is accomplished by introducing one or more binary observed variables (Cooper, 1988), denoted as O , with
[Equation]
, where α is a temperature hyper-parameter.
These new variables are often referred to as “optimality” variables (Levine, 2018).
We would like to infer latent variables, τ , and learn parameters, θ , that yield the maximum log-likelihood of optimality, i.e. log p ( O = 1) .
Evaluating this likelihood requires marginalizing the joint distribution, p ( O = 1)
=
(cid:82) p ( τ, O = 1) dτ .
This involves averaging over all trajectories, which is intractable in high-dimensional spaces.
Instead, we can use variational inference to lower bound this objective, introducing a structured approximate posterior distribution:
[Equation]
This provides the following lower bound on the objective, log p ( O = 1) :
[Equation]
Equivalently, we can multiply by α , deﬁning the variational RL objective as:
[Equation]
This objective consists of the expected return (i.e., the standard RL objective) and a KL divergence between π ( τ |O ) and p ( τ ) .
In terms of states and actions, this objective is written as:
[Equation]
At a given timestep, t , one can optimize this objective by estimating the future terms in the summation using a “soft” action-value ( Q π ) network (Haarnoja et al., 2017) or model (Pich´e et al., 2019).
For instance, sampling s t ∼ p env , slightly abusing notation, we can write the objective at time t as:
[Equation]
Policy optimization in the KL-regularized setting corresponds to maximizing J w.r.t.
π .
We often consider parametric policies, in which π is deﬁned by distribution parameters, λ , e.g. Gaussian mean, µ , and variance, σ 2 .
In this case, policy optimization corresponds to maximizing:
[Equation]
Optionally, we can then also learn the policy prior parameters, θ (Abdolmaleki et al., 2018).
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
tanh( µ
[Table]
[Figure] [Figure]
[Table]
[Figure]
0 .1 .t
a nh ( µ 3 )
Direct Policy Network Iterative Policy Network
Optimal Estimate
J
Direct <latexit sha1_base64="NHTlNGKuf9jvAAkt9p+DKT273ZM=">AAAChnicdVHLTgIxFC3jC/EB6NJNIzFxRWZ8BJdEXbjERMAEJqRTLtDYTiftHQOZ8CVu9aP8GzvIQkRv0uTk3NfpuVEihUXf/yx4G5tb2zvF3dLe/sFhuVI96lidGg5trqU2zxGzIEUMbRQo4TkxwFQkoRu93OX57isYK3T8hLMEQsXGsRgJztBRg0q5jzBFo7J7YYDjfFCp+XV/
[Table]
+
[Table]
/7TJU2KBpjzrI9Ugu4tUwsHXYb0E/3j2t59snQZSPeu2++cyRKsOMJ/q/Y7mjSirDDNhqNoHC2MboJ4BUZkZUfZTpAlM81axWtkEqydxpHB1EGDgkneDZLWcgPsHEo+9bAGxW3qFofr6K5nZrTQjX810gV7scKBsnaucp/Zr2svx3ryb7Fpi8XL1InatMhrthxUtJKipr0KdCYazlDOPQDWCL8rZRX4M6PXaq3TcZy6frm+zdr4XHWDXXqR6dcwqD6u54EstZ9QqX/Qgm0WGMtbf1U98wf0ksSXBdgEJwfj+Nn44P3z0eGrlThb5CF5RPZITF6QQ/KWHJEJYeRX8CAYBY/Dz+HX8Fv4fZkaBqua+2TNwh [Figure] /AQJ
[Figure]
Y
[Figure]
<
[Figure]
2. KL-R EGULARIZED P OLICY N ETWORKS P ERFORM D IRECT A MORTIZATION
Policy-based approaches to RL typically do not directly optimize the action distribution parameters, e.g. through gradient-based optimization.
Instead, the action distribution parameters are output by a function approximator (deep network), f φ , which is trained using deterministic (Silver et al., 2014; Lillicrap et al., 2016) or stochastic gradients (Williams, 1992; Heess et al., 2015).
When combined with entropy or KL regularization, this policy network is a form of amortized optimization (Gershman & Goodman, 2014), learning to estimate policies.
Again, denoting the action distribution parameters, e.g. mean and variance, as λ , for a given state, s , we can express this direct mapping as
[Equation]
and we denote the corresponding policy as π φ ( a | s , O ; λ ) .
Thus, f φ attempts to learn to optimize
Eq. 8.
This setup is shown in Figure 1 (Right).
Without entropy or KL regularization, i.e. π φ ( a | s )
=
p θ ( a | s ) , we can instead interpret the network as directly integrating the LHS of Eq. 3, which is less efﬁcient and more challenging.
Adding regularization smooths the optimization landscape, resulting in more stable improvement and higher asymptotic performance (Ahmed et al., 2019).
Viewing policy networks as a form of amortized variational optimizer (inference model) (Eq. 9) allows us to see that they are similar to encoder networks in variational autoencoders (VAEs) (Kingma & Welling, 2014; Rezende et al., 2014).
This raises the following question: are policy networks providing fully-optimized policy objectives?
In VAEs, it is empirically observed that amortization results in suboptimal approximate posterior estimates, with the resulting gap in the variational bound referred to as the amortization gap (Cremer et al., 2018).
In the RL setting, this means that an amortized policy, π φ , results in worse performance than the optimal policy within the parametric policy class, which we denote as [Equation] π .
Thus, the amortization gap is the gap in following inequality:
[Equation]
Because J is a variational bound on the RL objective, i.e. expected return, a looser bound, due to amortization, prevents an agent from more completely optimizing this objective.
To visualize the RL amortization gap, in Figure 1 (Left), we display the optimization surface, J , for two dimensions of the policy mean at a particular state in the MuJoCo environment Hopper-v2 .
We see that the estimate of a direct amortized policy (diamond) is suboptimal, far from the optimal estimate (star).
Additional 2D plots are shown in Figure B.3. To improve upon this suboptimal scheme, in Section 3, we turn to a technique developed in generative modeling, iterative amortization (Marino et al., 2018b), which retains the efﬁciency beneﬁts of amortization while employing a more ﬂexible iterative estimation procedure.
Initialize φ for each environment step do λ ← f φ ( s t ) a t ∼ π φ ( a t | s t , O ; λ ) s t
+1 ∼ p env ( s t +1 | s t , a t ) end for for each training step do φ ← φ + η ∇ φ J end for
Initialize φ for each environment step do
Initialize λ for each policy optimization iteration do λ ← f φ ( s t , λ , ∇ λ J ) end for a t ∼ π φ ( a t | s t , O ; λ ) s t
+1 ∼ p env ( s t +1 | s t , a t ) end for for each training step do φ ← φ + η ∇ φ J end for 2. R ELATED W ORK Previous works have investigated methods for improving policy optimization. QT-Opt (Kalashnikov et al., 2018) uses the cross-entropy method (CEM) (Rubinstein & Kroese, 2013), an iterative derivative-free optimizer, to optimize a Q -value estimator for robotic grasping.
CEM and related methods are also used in model-based RL for performing model-predictive control (Nagabandi et al., 2018; Chua et al., 2018; Pich ´ e et al., 2019; Hafner et al., 2019).
Gradient-based policy optimization, in contrast, is less common (Henaff et al., 2017; Srinivas et al., 2018; Bharadhwaj et al., 2020).
Most policy-based methods use direct amortization, either using a feedforward (Haarnoja et al., 2018b) or recurrent (Guez et al., 2019) network.
Similar approaches have also been applied to model-based value estimates (Byravan et al., 2020; Clavera et al., 2020), as well as combining direct amortization with model predictive control (Lee et al., 2019) and planning (Rivi ` ere et al., 2020).
A separate line of work has explored improving the policy distribution, using normalizing ﬂows (Haarnoja et al., 2018a; Tang & Agrawal, 2018) and latent variables (Tirumala et al., 2019).
In principle, iterative amortization can perform policy optimization in each of these settings.
I TERATIVE A MORTIZED P OLICY O PTIMIZATION 3. F ORMULATION
Iterative amortized optimizers (Marino et al., 2018b) utilize some form of error or gradient to update the approximate posterior distribution parameters.
While various forms exist, we consider gradient-encoding models (Andrychowicz et al., 2016) due to their generality.
Compared with direct amortization in Eq.
9, we use iterative amortized optimizers of the general form
[Equation]
also shown in Figure 1 (Right), where f φ is a deep network and λ are the action distribution parameters.
For example, if π = N ( a ; µ , diag( σ 2 )) , then λ ≡ [ µ , σ ] .
Technically, s is redundant, as the state dependence is already captured in J , but this can empirically improve performance (Marino et al., 2018b).
In practice, the update is carried out using a “highway” gating operation (Hochreiter & Schmidhuber, 1997; Srivastava et al., 2015).
Denoting ω φ ∈ [0 , 1] as the gate and δ φ as the update, both of which are output by f φ , the gating operation is expressed as
[Equation]
(11) where (cid:12) denotes element-wise multiplication.
This update is typically run for a ﬁxed number of steps, and, as with a direct policy, the iterative optimizer is trained using stochastic gradient estimates of ∇ φ J , obtained through the path-wise derivative estimator (Kingma & Welling, 2014;
Rezende et al., 2014; Heess et al., 2015).
Because the gradients ∇ λ
J must be estimated online, i.e. during policy optimization, this scheme requires some way of estimating J online, e.g. through a parameterized Q -value network (Mnih et al., 2013) or a differentiable model (Heess et al., 2015).
− 1 . 0 − 0 .0 .0 .1 .
tanh( µ 2 ) − 1 .− 0 .0 .0 .
.t
a nh ( µ 6 )
J − 1 Act.
Dim.0 .2 .5 .
.12 .π
φ ( a | s ,
O ) − 1 Act.
Dim.
Opt. {+Run
+}
Iterative amortized optimizers are more ﬂexible than their direct counterparts, incorporating feedback from the objective during policy optimization (Algorithm 2), rather than only after optimization (Algorithm 1).
Increased ﬂexibility improves the accuracy of optimization, thereby tightening the variational bound (Marino et al., 2018b;a).
We see this ﬂexibility in Figure 1 (Left), where an iterative amortized policy network, despite being trained with a different value estimator, is capable of iteratively optimizing the policy estimate (blue dots), quickly arriving near the optimal estimate.
Direct amortization is restricted to a single estimate.
In contrast, iterative amortized optimizers, by using stochastic gradients, can traverse the optimization landscape.
Thus, even for a uni-modal distribution class, over multiple optimization runs, an iterative amortized optimizer can effectively sample from a multi-modal distribution (Greff et al., 2019).
We illustrate this capability across two action dimensions in Figure 2 for a state in the Ant-v2 MuJoCo environment.
Over multiple policy optimization runs, iterative amortization ﬁnds multiple modes, sampling from two high-value regions of the action space.
This provides increased ﬂexibility in action exploration.
3.2. M ITIGATING V ALUE O VERESTIMATION
Model-free approaches generally estimate Q π using function approximation and temporal difference learning.
However, this comes with the pitfall of value overestimation, i.e. positive bias in the estimate, (cid:98)
Q [Equation] (Thrun & Schwartz, 1993).
This issue is tied to uncertainty in the value estimate, though it is distinct from optimism under uncertainty.
If the policy can exploit regions of high uncertainty, the resulting target values will introduce positive bias into the estimate.
More ﬂexible policy optimizers may exacerbate the problem, exploiting this uncertainty to a greater degree.
Further, a rapidly changing policy increases the difﬁculty of value estimation (Rajeswaran et al., 2020).
Various techniques have been proposed for mitigating value overestimation in deep RL.
The most prominent technique, double deep Q -network (Van Hasselt et al., 2016) maintains two Q -value estimates (Van Hasselt, 2010), attempting to decouple policy optimization from value estimation.
Fujimoto et al.
(2018) apply and improve upon this technique for actor-critic settings, estimating the target Q -value as the minimum of two Q -networks, Q ψ 1 and Q ψ 2 :
[Equation]
where ψ (cid:48)
i denotes the “target” network parameters.
As noted by Fujimoto et al.
(2018), this not only counteracts value overestimation, but also penalizes high-variance value estimates, because the minimum decreases with the variance of the estimate.
Ciosek et al.
(2019) noted that, for a bootstrapped ensemble of two Q -networks, the minimum operation can be interpreted as estimating
[Equation]
[Equation]
[Figure] [Figure]
[Equation]
[Figure]
[Equation]
[Figure] (cid:48) i ( s , a ) , standard deviation σ Q ( s , a ) ≡ ((cid:80) i =1 , 2 ( Q ψ (cid:48)
i ( s , a ) − µ Q ( s , a )) 2 ) 1 / 2 , and β = 1 .
Thus, to further penalize high-variance value estimates, preventing value overestimation, we can increase β .
For large β , however, value estimates become overly pessimistic, negatively impacting training.
Thus, β reduces target value variance at the cost of increased bias.
Due to the ﬂexibility of iterative amortization, the default β = 1 results in increased value bias (Figure 3a) and a more rapidly changing policy (Figure 3b) as compared with direct amortization.
Further penalizing high-variance target values with β = 2 .
5 reduces value overestimation and improves policy stability.
For details, see Appendix A.2.
Recent techniques for mitigating overestimation have been proposed, such as adjusting the temperature, α (Fox, 2019).
In ofﬂine RL, this issue has been tackled through the action prior (Fujimoto et al., 2019; Kumar et al., 2019; Wu et al., 2019) or by altering Q -network training (Agarwal et al., 2019; Kumar et al., 2020).
While such techniques could be used here, increasing β provides a simple solution with no additional computational overhead.
E XPERIMENTS 4. S ETUP
To focus on policy optimization, we implement iterative amortized policy optimization using the soft actor-critic (SAC) setup described by Haarnoja et al.
(2018c).
This uses two Q -networks, uniform action prior, p θ ( a | s ) =
U ( − 1 , 1) , and a tuning scheme for the temperature, α .
In our experiments, “direct” refers to direct amortization employed in SAC, i.e. a direct policy network, and “iterative” refers to iterative amortization.
Both approaches use the same network architecture, adjusting only the number of inputs and outputs to accommodate gradients, current policy estimates, and gated updates (Sec. 3.1) [Figure] t i
[Figure] [Figure] t i
[Figure]
G a
[Figure]
4. A NALYSIS 4.2. V ISUALIZING P OLICY O PTIMIZATION
In Figure 1 (Left), we visualize the trajectory of iterative amortized policy optimization along two dimensions of the policy mean on a state from Hopper-v2 .
Through iterative optimization, the network arrives near the optimum.
Notably, this is performed with a value function trained using a different , direct policy, demonstrating generalization to other optimization landscapes.
Additional 2D plots comparing direct and iterative amortization are shown in Figure B.3 in the appendix.
In
Figure 4, we visualize iterative reﬁnement using a single action dimension from Ant-v2 across time steps.
The reﬁnements in Figure 4a give rise to the objective improvements in Figure 4b.
We compare with Adam (Kingma & Ba, 2014)
(gradient-based) and CEM (Rubinstein & Kroese, 2013) (gradient-free) in Figure 4c (see Appendix B.2), where iterative amortization is substantially more efﬁcient. This trend is consistent across environments, as shown in Figure B.2 in the appendix.
4.2. P ERFORMANCE C OMPARISON
We evaluate iterative amortized policy optimization on MuJoCo (Todorov et al., 2012) continuous control tasks from OpenAI gym
(Brockman et al., 2016).
In Figure 5, we compare the cumulative reward of direct and iterative amortized policy optimization across environments.
Each curve shows the mean and ± standard deviation of 4 random seeds.
In all cases, iterative amortized policy optimization matches or outperforms the baseline direct method, both in sample efﬁciency and ﬁnal performance.
Across environments, iterative amortization also yields more consistent performance.
4.2. D ECREASED A MORTIZATION G AP
To evaluate policy optimization, we estimate per-step amortization gaps using the experiments from
Figure 5, performing additional iterations of gradient ascent on J , w.r.t.
the policy parameters, λ ≡ [ µ , σ ] (see Appendix A.3).
We also evaluate the iterative agents trained with 5 iterations for an additional 5 amortized iterations.
Results are shown in Figure 6.
We emphasize that it is challenging to directly compare amortization gaps across optimization schemes, as these involve different value functions, and therefore different objectives.
Likewise, we estimate the amortization gap using the learned Q -value networks, which may be biased (Figure 3).
Nevertheless, we ﬁnd that iterative [Figure] [Figure] The estimated amortized policy optimization achieves, on average, lower amortization gaps than direct amortization across all environments.
Further amortized iterations at evaluation yield further estimated improvement, demonstrating generalization.
However, we note that the amortization gaps are relatively small compared to the estimated discounted objective.
Accordingly, when we evaluate the more fully optimized policies in the environment, we do not observe a noticeable increase in performance (see Appendix B.4).
This demonstrates that policy suboptimality is not a signiﬁcant concern for evaluation .
Rather, improved policy optimization is helpful for training .
This allows the agent to collect data where value estimates are highest and ultimately improve these value estimates. Indeed, when we train iterative amortization while varying the policy optimization iterations per step (Section 4.2.4), we observe that the estimated amortization gap again decreases with increasing iterations, but now with a corresponding increase in performance (Figure B.6). Thus, reducing the amortization gap only indirectly improves task performance by improving training, with the relationship depending on the Q -value estimator and other factors. 4.2. V ARYING I TERATIONS
Direct amortized policy optimization is restricted to a single forward pass through the network.
Iterative amortization, in contrast, is capable of improving during policy optimization with additional computation time.
To demonstrate this capability, we train iterative amortized policy optimization while varying the number of iterations in { 1 , 2 , 5 } .
In Figure 7, we see that increasing the number of amortized optimization iterations generally improves sample efﬁciency ( Walker2d-v2 ), asymptotic performance ( Ant-v2 ), or both ( Hopper-v2 & HalfCheetah-v2 ).
Thus, the quality of the policy optimizer can play a signiﬁcant role in determining performance through training.
4.2. I TERATIVE A MORTIZATION WITH M ODEL -B ASED V ALUE E STIMATES
While our analysis has centered on the model-free setting, iterative amortized policy optimization can also be applied to model-based value estimates.
As model-based RL remains an active research area (Janner et al., 2019), we provide a proof-of-concept in this setting, using a learned deterministic model on HalfCheetah-v2 (see Appendix A.5).
As shown in Figure 8a, iterative amortization outperforms direct amortization in this setting.
Iterative amortization reﬁnes planned trajectories, shown for a single state dimension in Figure 8b, yielding corresponding improvements (Figure 8c).
Further, because we are learning an iterative policy optimizer, we can zero-shot transfer a policy optimizer trained with a model-free value estimator to a model-based value estimator (Figure 8d).
This is not possible with a direct amortized optimizer, which does not use value estimates online during policy optimization.
Iterative amortization is capable of generalizing to new value estimates, instantly incorporating updated value estimates in policy optimization.
This demonstrates and highlights the opportunity for improving model-based planning through iterative amortization.
D ISCUSSION
We have introduced iterative amortized policy optimization, a ﬂexible and powerful policy optimization technique.
Using the MuJoCo continuous control suite, we have demonstrated improved performance over direct amortization with both model-based and model-free value estimates.
Iterative amortization provides a drop-in replacement and improvement over direct policy networks in deep RL.
Although iterative amortized policy optimization requires additional computation, this could be combined with some form of adaptive computation time (Graves, 2016;
Figurnov et al., 2018), gauging the required iterations.
Likewise, efﬁciency depends, in part, on the policy initialization, which could be improved by learning the action prior, p θ ( a | s ) (Abdolmaleki et al., 2018).
The power of iterative amortization is in using the value estimate during policy optimization to iteratively improve the policy online .
This is a form of negative feedback control (Astrom & Murray, 2008), using errors to guide policy optimization.
Beyond providing a more powerful optimizer, we are hopeful that iterative amortized policy optimization, by using online feedback, will enable a range of improved RL algorithms, capable of instantly adapting to different value estimates, as shown in
Figure 8d.