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 and limited exploration.
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 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.
Likewise, direct amortization is typically restricted to a single estimate of the distribution, limiting the ability to sample diverse solutions.
In RL, this translates into a deﬁciency in exploration.
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]
[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+/AQJ1+Y4=</latexit>
Policy optimization in the
-regularized corresponds to maximizing J π .
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).
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 .
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
Algorithm
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
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.
However, note that the absolute difference in the objective due to direct amortization is relatively small compared with the objective itself.
That is, suboptimal estimates tend to have only a minor impact on evaluation performance, as we show in Appendix B.4.
Rather, policy suboptimality hinders data collection, sampling fewer actions with high value estimates.
Further, direct amortization is typically limited to a single static estimate of the policy, unable to directly adapt to the RL objective and therefore limiting exploration.
To improve upon this 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.
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), however, gradient-based optimization can also be combined with CEM (Amos & Yarats, 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;
Amos 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]
(iterative amortization) (10) 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
− 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
et ω φ
1] as of which are output by φ , 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).
3. C ONSIDERATIONS 3.2. A DDED F LEXIBILITY
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 typically restricted to a single estimate, inherently limiting exploration.
In contrast, iterative amortized optimizers, by using stochastic gradients and random initialization, can traverse the optimization landscape.
As with any iterative optimization scheme, this allows iterative amortization to obtain multiple valid estimates (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 π (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).
[Figure]
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]
with mean µ Q ( s , a ) ≡
[Equation] [Equation]
, standard deviation
[Equation]
µ 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).
Unless otherwise stated, we use 5 iterations per time step for iterative amortization, following Marino et al.
(2018b).
For details, refer to Appendix A and Haarnoja et al.
(2018b;c).
[Figure] [Figure]
t
i
[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 5 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]
Increasing the number of iterations generally results in improvements.
Curves show the mean and ± standard deviation over 4 random seeds.
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,
[Figure]
shown for a single state dimension in Figure 8b, yielding corresponding improvements (Figure 8c).
Further, because we are learning an iterative policy optimizer
[Figure]
we can zero-shot transfer a policy optimizer trained with a model-free value estimator to a model-based value estimator (Figure 8d)
[Figure]
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.
A E XPERIMENT D ETAILS A. 2D P LOTS
In Figures 1 and 2, we plot the estimated variational objective, J , as a function of two dimensions of the policy mean, µ .
To create these plots, we ﬁrst perform policy optimization (direct amortization in Figure 1 and iterative amortization in Figure 2), estimating the policy mean and variance.
This is performed using on-policy trajectories from evaluation episodes (for a direct agent in Figureand an iterative agent in Figure 2).
While holding all other dimensions of the policy constant, we then estimate the variational objective while varying two dimensions of the mean ( 1 & 3 in Figure 1 and 2 & 6 in Figure 2).
Iterative amortization is additionally performed while preventing any updates to the constant dimensions.
Even in this restricted setting, iterative amortization is capable of optimizing the policy.
Additional 2D plots comparing direct vs. iterative amortization on other environments are shown in Figure 14, where we see similar trends.
A. V ALUE B IAS E STIMATION
We estimate the bias in the Q -value estimator using a similar procedure as Fujimoto et al.
(2018), comparing the estimate of the Q -networks ( (cid:98)
Q π ) with a Monte Carlo estimate of the future objective in the actual environment, Q π , using a set of state-action pairs.
To enable comparison across setups, we collect 100 state-action pairs using a uniform random policy, then evaluate the estimator’s bias,
E s , a
[Equation]
, throughout training.
To obtain the Monte Carlo estimate of Q π , we use 100 action samples, which are propagated through all future time steps.
The result is discounted using the same discounting factor as used during training, γ = 0 . 99 , as well as the same Lagrange multiplier, α .
Figure 3 shows the mean and ± standard deviation across the 100 state-action pairs.
A. A MORTIZATION G AP E STIMATION
Calculating the amortization gap in the RL setting is challenging, as properly evaluating the variational objective, J , involves unrolling the environment.
During training, the objective is estimated using a set of Q -networks and/or a learned model.
However, ﬁnding the optimal policy distribution, (cid:98) π , under these learned value estimates may not accurately reﬂect the amortization gap, as the value estimator likely contains positive bias (Figure 3).
Because the value estimator is typically locally accurate near the current policy, we estimate the amortization gap by performing gradient ascent on
J w.r.t.
the policy distribution parameters, λ , initializing from the amortized estimate (from π φ ).
This is a form semi-amortized variational inference (Hjelm et al., 2016; Krishnan et al., 2018;
Kim et al., 2018).
We use the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 5 × 10 − 3 for 100 gradient steps, which we found consistently converged.
This results in the estimated optimized (cid:98) π .
We estimate the gap using 100 on-policy states, calculating J ( θ, (cid:98) π ) − J ( θ, π ) , i.e. the improve- ment in the objective after gradient-based optimization.
Figure 6 shows the resulting mean and ± standard deviation.
We also run iterative amortized policy optimization for an additional 5 iterations during this evaluation, empirically yielding an additional decrease in the estimated amortization gap.
A. H YPERPARAMETERS
Our setup follows that of soft actor-critic (SAC) (Haarnoja et al., 2018b;c), using a uniform action prior, i.e. entropy regularization, and two Q -networks (Fujimoto et al., 2018).
Off-policy training is performed using a replay buffer (Lin, 1992; Mnih et al., 2013).
Training hyperparameters are given in Table 6.
Temperature Following Haarnoja et al.
(2018c), we adjust the temperature, α , to maintain a speciﬁed entropy constraint, (cid:15) α = |A| , where |A| is the size of the action space, i.e. the dimensionality.
Policy We use the same network architecture (number of layers, units/layer, non-linearity) for both direct and iterative amortized policy optimizers (Table 2).
Each policy network results in Gaussian distribution parameters, and we apply a tanh transform to ensure a ∈ [ − 1 , 1] (Haarnoja et al., 2018b).
In the case of a Gaussian, the distribution parameters are λ = [ µ , σ ] .
The inputs and
[Table] [Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
a
[Table]
+KdeRffpV5tpNlAY+FdfwEBCc3T</latexit> r  J <latexit sha1_base64="5JfT0SmUyNEjLkgtHj+pAbvQ0cE=">AAACkHicdVHLSgMxFE3Hd31VXYmbYBFclRkVdOdrI64UrAqdYbiTpm0wL5KMWIbBr3Gr3+PfmKldWKsXAodzXyfnZpoz68LwsxbMzM7NLywu1ZdXVtfWGxub91blhtA2UVyZxwws5UzStmOO00dtKIiM04fs6bLKPzxTY5mSd26oaSKgL1mPEXCeShvbsYSMQ1rEesBKHAtwAwK8uC7TRjNshaPA0yAagyYax026UUvjriK5oNIRDtZ2olC7pADjGOG0rMe5pRrIE/Rpx0MJgtqkGP2hxHue6eKeMv5Jh0fsz44ChLVDkfnKSqP9navIv3Kd3PVOkoJJnTsqyfeiXs6xU7gyBHeZocTxoQdADPNaMRmAAeK8bROT7qKkqMRVYybWZ6Ks7+GfTCVDO/EyWQe8r/yGgfiHZmS6QVuae1dV1xvoTxL9PsA0uD9oRYetg9uj5tnF+DiLaAfton0UoWN0hq7QDWojgl7RG3pHH8FmcBKcBuffpUFt3LOFJiK4/gIdPMyV</latexit>
[Figure]
∇ φ J .
This generally requires some form of stochastic gradient estimation to differentiate through a t
[Equation]
outputs of each optimizer form are given in Table 1.
Again, δ and ω are respectively the update and gate of the iterative amortized optimizer (Eq. 11), each of which are deﬁned for both µ and σ .
Following Marino et al.
(2018b), we apply layer normalization (Ba et al., 2016) individually to each of the inputs to iterative amortized optimizers.
We initialize iterative amortization with µ = 0 and σ = 1 , however, these could be initialized from a learned action prior (Marino et al., 2018a).
[Table] [Table]
Q -value
We investigated two Q -value network architectures.
Architecture A (Table 3) is the same as that from Haarnoja et al.
(2018b).
Architecture B (Table 4) is a wider, deeper network with highway connectivity (Srivastava et al., 2015), layer normalization (Ba et al., 2016), and ELU nonlinearities (Clevert et al., 2015).
We initially compared each Q -value network architecture using each policy optimizer on each environment, as shown in Figure 10.
The results in Figure 5 were obtained using the better performing architecture in each case, given in Table 5.
As in Fujimoto et al.
(2018), we use an ensemble of 2 separate Q -networks in each experiment.
Value Pessimism ( β ) As discussed in Section 3.2.2, the increased ﬂexibility of iterative amortization allows it to potentially exploit inaccurate value estimates.
We increased the pessimism hyperpa-
[Figure]
l a t i v e R [Figure] w a r
[Figure]
rameter, β , to further penalize variance in the value estimate.
Experiments with direct amortization use the default β = 1 in all environments, as we did not ﬁnd that increasing β helped in this setup.
For iterative amortization, we use β = 1 . 5 on Hopper-v2 and β = 2 . 5 on all other environments.
This is only applied during training; while collecting data in the environment, we use β = 1 to not overly penalize exploration
[Figure]
[Table] [Table]
M ODEL -B ASED V ALUE E STIMATION
For model-based experiments, we use a single, deterministic model together with the ensemble of
Q -value networks (discussed above).
Model We use separate networks to estimate the state transition dynamics, p env ( s t +1 | s t , a t ) , and reward function, r ( s t , a t ) .
The network architecture is given in Table 7.
Each network outputs the mean of a Gaussian distribution; the standard deviation is a separate, learnable parameter.
The reward network directly outputs the mean estimate, whereas the state transition network outputs a residual estimate, ∆ s t , yielding an updated mean estimate through:
[Equation]
Model Training The state transition and reward networks are both trained using maximum loglikelihood training, using data examples from the replay buffer.
Training is performed at the same frequency as policy and Q -network training, using the same batch size ( 256 ) and network optimizer.
However, we perform 10 3 updates at the beginning of training, using the initial random steps, in order to start with a reasonable model estimate.
[Table]
[Table]
Value Estimation To estimate Q -values, we combine short model rollouts with the model-free estimates from the Q -networks.
Speciﬁcally, we unroll the model and policy, obtaining state, reward, and policy estimates at current and future time steps.
We then apply the Q -value networks to these future state-action estimates.
Future rewards and value estimates are combined using the Retrace estimator (Munos et al., 2016).
Denoting the estimate from the Q -network as (cid:98)
Q ψ ( s , a ) and the reward estimate as
[Equation]
r ( s , a ) , we calculate the Q -value estimate at the current time step as (cid:98)
[Equation]
13) where λ is an exponential weighting factor, h is the rollout horizon, and the expectation is evaluated under the model and policy.
In the variational RL setting, the state-value, V π ( s ) , is
[Equation]
In Eq.
, we approximate V π using the Q -network to approximate Q π in Eq.
14, yielding (cid:98)
[Equation]
Finally, to ensure consistency between the model and the Q -value networks, we use the model-based estimate from Eq.
13 to provide target values for the Q -networks, as in Janner et al.
Future Policy Estimates Evaluating the expectation in Eq.
13 requires estimates of π at future time steps.
This is straightforward with direct amortization, which employs a feedforward policy, however, with iterative amortization, this entails recursively applying an iterative optimization procedure.
Alternatively, we could use the prior, p θ ( a | s ) , at future time steps, but this does not apply in the max-entropy setting, where the prior is uniform.
For computational efﬁciency, we instead learn a separate direct (amortized) policy for model-based rollouts.
That is, with iterative amortization, we create a separate direct network using the same hyperparameters from Table 2.
This network distills iterative amortization into a direct amortized optimizer, through the KL divergence, D KL ( π it.
||
π dir. ) .
Rollout policy networks are common in model-based RL (Silver et al., 2016; Pich ´ e et al., 2019).
B A DDITIONAL R ESULTS B. I MPROVEMENT PER S TEP
In Figure 12, we plot the average improvement in the variational objective per step throughout training, with each curve showing a different random seed.
That is, each plot shows the average change in the variational objective after running 5 iterations of iterative amortized policy optimization.
With the exception of HalfCheetah-v2 , the improvement remains relatively constant throughout training and consistent across seeds.
B. C OMPARISON WITH I TERATIVE O PTIMIZERS
Iterative amortized policy optimization obtains the accuracy beneﬁts of iterative optimization while retaining the efﬁciency beneﬁts of amortization.
In Section 4, we compared the accuracy of iterative and direct amortization, seeing that iterative amortization yields reduced amortization gaps (Figure 6) and improved performance (Figure 5).
In this section, we compare iterative amortization with two popular iterative optimizers: Adam (Kingma & Ba, 2014), a gradient-based optimizer, and cross-entropy method (CEM) (Rubinstein & Kroese, 2013), a gradient-free optimizer.
[Table]
/bYDIZkjtiGeY/3Opf+TdmahfW6oULh3Nfh3uiRAqLvv9Z8paWV1bXyuuVjc2t7Z3dvf221anh0OJaavMYMQtSxNBCgRIeEwNMRRI60fNNUe+8gLFCxw84SSBUbBSLoeAMHfXUQ3hFo7I7PQCZ93erfs2fBl0EwQxUySya/
b1SvzfQPFUQI5fM2m7gJxhmzKDgEvJKL7WQMP7MRtB1MGYKbJhNZef02DEDOtTGZYx0yv6cyJiydqIi16kYju3vWkH+VeumOLwIMxEnKULMvw8NU0lR0+IHdCAMcJQTBxg3wmmlfMwM4+g+NbfpIQizQlyxZu58pPLKMf3JFDISVK/zfUyOtLswVv/Qgi8OJBZS91VnSV5xlgS/DVgE7XotOKvV7xvVq+uZOWVySI7ICQnIObkit6RJWoQTQ97IO/nwVrxTr+41vlu90mzmgMyFd/kFY4rH+Q==</latexit>
[Figure] [Figure]
To compare the accuracy and efﬁciency of the optimizers, we collect 100 states for each seed in each environment from the model-free experiments in Section 4.2.2.
For each optimizer, we optimize the variational objective, J , starting from the same initialization.
Tuning the step size, we found that 0 .
01 yielded the steepest improvement without diverging for both Adam and CEM.
Gradients are evaluated with 10 action samples.
For CEM, we sample 100 actions and ﬁt a Gaussian mean and variance to the top 10 samples.
This is comparable with QT-Opt
(Kalashnikov et al., 2018), which draws 64 samples and retains the top 6 samples.
− 1 .− 0 .0 .0 .1 .
tanh( µ 3 ) − 1 .− 0 .0 .0 .
.t
a nh ( µ 5 )
DirectPolicyNetworkIterativePolicyNetworkOptimalEstimate
J (a) HalfCheetah-v − 1 .− 0 .0 .0 .1 .
tanh( µ 4 ) − 1 .− 0 .0 .0 .
.t
a nh ( µ 6 )
DirectPolicyNetworkIterativePolicyNetworkOptimalEstimate
J (b) Walker2d-v − 1 .− 0 .0 .0 .1 .
tanh( µ 6 ) − 1 .− 0 .0 .0 .
.t
a nh ( µ 7 )
DirectPolicyNetworkIterativePolicyNetworkOptimalEstimate
J (c) Ant-v
The results, averaged across states and random seeds, are shown in Figure 13.
CEM (gradientfree) is less efﬁcient than Adam (gradient-based), which is unsurprising, especially considering that
Adam effectively approximates higher-order curvature through momentum terms.
However, Adam and CEM both require over an order of magnitude more iterations to reach comparable performance with iterative amortization.
While iterative amortized policy optimization does not always obtain asymptotically optimal estimates, we note that these networks were trained with only 5 iterations, yet continue to improve and remain stable far beyond this limit.
Finally, comparing wall clock time for each optimizer, iterative amortization is only roughly 1 .
25 × slower than CEM and 1 .
15 × slower than Adam, making iterative amortization still substantially more efﬁcient.
B. A DDITIONAL 2D O PTIMIZATION P LOTS
In Figure 1, we provided an example of the suboptimal optimization resulting from direct amortization on the Hopper-v2 environment.
We also demonstrated that iterative amortization is capable of automatically generalizing to this optimization surface, outperforming the direct optimizer.
To show that this is a general phenomenon, in Figure 14, we present examples of corresponding 2D plots for each of the other environments considered in this paper.
As before, we see that direct amortization is near-optimal, but, with the exception of HalfCheetah-v2 , does not match the optimal estimate.
In contrast, iterative amortization is able to ﬁnd the optimal estimate, again, generalizing the unseen optimization surfaces.
B. A DDITIONAL O PTIMIZATION & THE A MORTIZATION G AP
In Section 4, we compared the performance of direct and iterative amortization, as well as their estimated amortization gaps.
In this section, we provide additional results analyzing the relationship between policy optimization and the performance in the actual environment.
As we have emphasized previously (see Section 3.2.2), this relationship is complex, as optimizing an inaccurate Q -value estimate does not improve task performance.
The amortization gap quantiﬁes the suboptimality in the objective, J , of the policy estimate.
As described in Section A.3, we estimate the optimized policy by performing additional gradient-based optimization on the policy distribution parameters (mean and variance).
However, when we deploy this optimized policy for evaluation in the actual environment, as shown for direct amortization in
Figure 15, we do not observe a noticeable difference in performance.
Thus, while amortization may ﬁnd suboptimal policy estimates, we observe that the actual difference in the objective is either too small or inaccurate to affect performance at test time.
Likewise, in Section 4.2.2, we observed that using additional amortized iterations during evaluation further decreased the amortization gap for iterative amortization.
However, when we deploy this
[Figure]
[Figure] [Figure]
Increasing the iterations improves performance and decreases the estimated amortization gap.
more fully optimized policy in the environment, as shown in Figure 16, we do not generally observe a corresponding performance improvement.
In fact, on HalfCheetah-v2 and Walker2d-v2 , we observe a slight decrease in performance.
This further highlights the fact that additional policy optimization may exploit inaccurate Q -value estimates.
However, importantly, in Figures 15 and 16, the additional policy optimization is only performed for evaluation.
That is, the data collected with the more fully optimized policy is not used for training and therefore cannot be used to correct the inaccurate value estimates.
Thus, while more accurate policy optimization, as quantiﬁed by the amortization gap, may not substantially affect evaluation performance, it does play a signiﬁcant role in improving training.
This was shown in Section 4.2.4, where we observed that training with additional iterative amortized policy optimization iterations, i.e., a more ﬂexible policy optimizer, generally results in improved performance.
By using a more accurate (or exploitative) policy for data collection, the agent is able to better evaluate its Q -value estimates, which accrues over the course of training.
This trend is shown for HalfCheetah-v2 in Figure 17, where we observed the largest difference in performance across numbers of iterations.
We generally observe that increasing the number of iterations during training improves performance and decreases the amortization gap.
Interestingly, when performance dips for the agents trained with 2 iterations, there is a corresponding increase in the amortization gap.
[Figure]
[Figure] [Figure]
P o i c y
M e a n L 2 D
i s t a n ce
[Figure] [Figure]
P o i c y
M e a n L 2 D
i s t a n ce
[Figure] [Figure]
[Figure] [Figure]
P o i c y
M e a n L 2 D
i s t a n ce
[Figure] [Figure]
[Figure] [Figure]
P
[Figure]
a
[Figure]
i
[Figure]
B. M ULTIPLE P OLICY E STIMATES
As discussed in Section 3.2.1, iterative amortization has the added beneﬁt of potentially obtaining multiple policy distribution estimates, due to stochasticity in the optimization procedure (as well as initialization).
In contrast, unless latent variables or normalizing ﬂows are incorporated into the policy, direct amortization is limited to a single policy estimate.
To estimate the degree to which iterative amortization obtains multiple policy estimates during training, we perform two separate runs of policy optimization per state and evaluate the L2 distance between the means of these policy estimates (after applying the tanh ).
Note that in MuJoCo action spaces, which are bounded to [ − 1 , 1] , the maximum distance is
[Equation]
|A| , where |A| is the size of the action space.
We aver- age the policy mean distance over 100 states and all experiment seeds, with the results shown in
Figure 18.
In all environments, we see that the average distance initially increases during training, remaining relatively constant for Hopper-v2 and HalfCheetah-v2 and decreasing slightly for
Walker2d-v2 and Ant-v2 .
Note that the distance for direct amortization would be exactlythroughout.
This indicates that iterative amortization does indeed obtain multiple policy estimates, maintaining some portion of multi-estimate policies throughout training.