Iterative Amortized Policy Optimization
Authors: Joseph Marino
Alexandre Piché
Alessandro Davide Ialongo
Yisong Yue
Venue:  (2021)
Keywords: Reinforcement Learning
Policy Optimization
Amortization
Variational Inference

Abstract: 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 flexible 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.

Under review as a conference paper at ICLR 2021
ITERATIVE AMORTIZED POLICY OPTIMIZATION
Anonymous authors
Paper under double-blind review
ABSTRACT
Policy networks are a central feature of deep reinforcement learning (RL) algo- rithms for continuous control, enabling the estimation and sampling of high-value actions. From the variational inference perspective on RI., policy networks, when employed with entropy or KL regularization, are a form of amortized optimiza- tion, 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 flexible class of iterative amortized optimizers. We demonstrate that the resulting technique, iterative amortized policy optimization, yields performance improve- ments over direct amortization methods on benchmark continuous control tasks.
INTRODUCTION
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 fol- lowing that policy, and given value estimate, This latter procedure, policy optimization, can be challenging in continuous control due to instabil- ity 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 previ- ous policies (Schulman al. 2015 2017 or by maximizing policy entropy (Mnih al. 2016 Fox et al. (2166) These techniques can be interpreted as variational inference (Levine (018)) using op- timization 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).
can improve the policy to maximize the value.
However, one subtlety arises: when used with entropy or KI. regularization, policy networks per- form amortized optimization (Gershman & Goodman] 2014). That is, rather than optimizing the ac- tion distribution, c.g. mean and variance, many deep RL algorithms, such as soft actor-critic (SAC) (Haarnoj et al.) 20188 c), instead optimize network to output these parameters, learning to opti- mize the policy. Typically, this is implemented as a direct mapping from states to action distribution parameters. While direct amortization schemes have improved the efficiency of variational inference as encoder networks (Kingma & Welling 2014 Rezende et al. 2014 Mnih & Gregor| 2014), they
also suboptimal (Cremer al. 2018]] Kim et al. 2018 Marino et al.) (018b). This suboptimality is referred to as the amortization gap (Cremer et al. 2018). translating into gap in the RL objective. Likewise, direct amortization is typically restricted
single estimate of the distribution, limiting
the ability to sample diverse solutions. In RL, this translates into a deficiency in exploration.
Inspired by techniques and improvements from variational inference, we investigate iterative amor- tized policy optimization. Iterative amortization (Marino et al., 2018b) uses gradients or errors to iteratively update the parameters of distribution. Unlike direct amortization, which receives gradi- ents 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.) 201 a), with the added benefit of finding multiple modes of the optimization landscape (Greff et al. 2019).
Using MuJoCo environments (Toorov et al. 2012) from OpenAl gym (Brockman et al. 2016), we demonstrate performance improvements of iterative amortized policy optimization over direct amor- tization in model-free and model-based settings. We analyze various aspects of policy optimization, including iterative policy refinement, adaptive computation, and zero-shot optimizer transfer. Iden- tifying policy networks as form of amortization elarifies suboptimal aspects of direct approaches

---

Under review as a conference paper at ICLR 2021
to policy optimization. Iterative amortization, by harnessing gradient-base based feedback during policy optimization, offers an effective and principled improvement.
BACKGROUND
2.1 PRELIMINARIES
We consider Markov decision processes (MDPs), where E S and E A are the state and action at time t, resulting in reward Tt = r(s, as). Environment state transitions are given by Si+1 Pea (s, at). and the agent is defined by parametrie distribution, pa(as)s;)) with parameters 0. The discounted sum of rewards is denoted as R(r) Exrhr, 7'r,, where y E (0, 1] is the discount factor, and r ($1, a1,...) is trajectory. The distribution over trajectories is:
T
p(r) = p(s)) II
Penv (St+ a )po (ar)s))
(1)
where the initial state drawn from the distribution p(s)). The standard RL objective consists of maximizing the expected discounted return, Ep(r) [R(r)]. For convenience of presentation, we use undiscounted setting (7 1), though the formulation can be applied with any valid 7.
2.2 KL-REGULARIZED REINFORCEMEN LEARNING
Various works have formulated RL, planning, and control problems in terms of probabilistic infer-
ence (Dayan Hinton 1997 Attias 2003 Toussaint & Storkey 2006 Todorov 2008) Botvinick &
Toousaint 2012 Levine 2018). These approaches consider the agent-environment interaction
graphical model, then convert reward maximization into maximum marginal likelihood estimation, learning and inferring policy that results in maximal reward. This conversion accomplished by introducing one or more binary observed variables (Cooper 1988), denoted as O, with
p(O 1|T) exp (R(r)/a),
where a is temperature variables (Levine 2018). We would like to infer latent variables, and learn parameters, 6, that yield the maximum likelihodoo of optimality, i.e. log 1). Evaluating this likelihood re- quires marginalizing the joint distribution, p(O = p(T, O = 1)dr. This involves averaging over trajectories, which intractable high-dimensional spaces. Instead, we can use variational inference to lower bound this objective, introducing structured approximate posterior distribution:
These new variables are often referred to as "'pptimalit""
(T|0) II Pene (St+1 |st, az) (s, O)
(2)
This provides the following lower bound the objective, log 1):
log / M10 1/T)p(7)dT 2 [T]| log p( 1/7) + log p(()) dr
(3)
[T[
= _E (=) (l((T))
(4)
Equivalently, we can multiply by a, defining the variational RL objective as:
J 0) = E, [R(T)] aDKL (т(T (O)[(T)))((()
This objective consists of the expected return (i.c., the standard RL objective) and KL divergence between π (T(O) and p(r). In terms of states and actions, this objective is written as:
T rs a log
T (Ta St, O) Pe (ar(()))
I J(T,0) = Ess T Peert de
Σ
(6)
At a given timestep. f, one can optimize this objective by estimating the future terms in the summa- tion using "soft" aation-valuee (Qx) network (Haarnoja et al.) (2017) or model (Piché et al.) 2019). For instance, sampling Penv slightly abusing notation, we can write the objective at time / as:
J 0) E,, (st, a,)] aDx (7 O) (a,ls.))
2

---

Under review as a conference paper at ICLR 2021
301
B, (9.
Qr( ac)
ILG
taahpp)
0.0
5003
lox xolax (s.
Fi (sc, O pulat Sc)
po(a
Durct Poley Netwoek
Puñcy Netsnck
Peo 9c-1,
Optinaal Excinoda
Direct
Iteaatve
Lm
Figure I: Amortization. Left: Optimization over two dimension of the policy mean, P1 and pa for particular state. A direct amortized policy network outputs suboptimal estimate, yielding an amortization gap in performance. An iterative amortized policy network finds an improved estimate. Right: Diagrams of direct and iterative amortization. Larger circles denote distributions, and smaller red circles denote terms in the objective, I (Eq. (7). Dashed arrows denote amortization. Iterative amortization uses gradient feedback during optimization, while direct amortization does not.
Policy optimization in the KL-regularized setting corresponds to maximizing J w.r.t. We often consider parametric policies, in which æ is defined by distribution parameters, X, e.g. Gaussian mean, μu, and variance, o². In this case, policy optimization corresponds to maximizing:
x arg max [xx,)
(8)
Optionally, then also learn the policy prior parameters, 0 (Abdolmaleki al. 2018).
2.3 KL-REGULARIZER POLICY NETWORKS PERFORM DIRECT AMORTIZATION
Policy-based approaches to RL typically do not directly optimize the action distribution parameters, e.g. through gradient- optimization, Instead, the action distribution parameters are output by function approximator (deep network), fø. which is trained using deterministic (Silver et al. 2014 Lillicrap al. 2016) or stochastic gradients (Williams 1992 Heess et al. 2015). When combined with entropy KL regularization, this policy network is form of amorlized optimization (Ger- shman & Goodman 2014), learning estimate policies. Again, denoting the action distribution express this direct mapping as
parameters, c.g. mcan and variance, X, for given state, s,
A + fo(s),
(direct amortization)
(9)
and we denote the corresponding policy as (als,,, O; Thus, fo attempts to learn to optimize Eq. g This setup is shown in Figure (Right). Without entropy KI regularization, i.e. ng (a(s))
(a(a s), we can instead interpret the network as directly integrating the L.SS of Eq. 3..3. which less efficient 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 form of amortized variational optimizer (inference model) (Eq. al- lows 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 VAFs, it is empirically observed that amortization results in suboptimal approximate posterior estimates, with the resulting gap the variational bound referred to as the amortization gap (Cremer et al., 2018). 8). In the RL setting, this means that amor- tized policy, nR. results in worse performance than the optimal policy within the parametric policy class, which we denote as न. Thus, the amortization gap is the gap in following inequality:
[ 0) < ().
Because J is a variational bound on the RL objective, iee. expected return, a looser bound, due to amortization, prevents agent from more completely optimizing this objective.
To visualize the RL amortization gap, in Figure (Left), we display the optimization surface. J, for two dimensions of the policy mean particular state in the MuJoCo environment Hopper-v2.
3

---

Under review as a conference paper at ICLR 2021
Algorithm 2 Iterative Amortization
Algorithm Direct Amortization
Initialize ø
for each environment step do Initialize À
Initialize é for each environment step do
for each policy optimization iteration do
X
(st)
À + fo( X, VXx))
a, ~ (a St, 0;
end for
1 ~ Penv 1 Se, ar)
(at O:
end for for each training step do
S:41 ~ Peav (S4+1 as)
end for
Φ + + V T
for each training step do
end for
Φ Φ + 17 V end for
We see that the estimate of 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 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 Appendix B.4. Rather, policy suboptimality hinders data collection, sampling fewer actions with high value estimates. Further, direct amortization is typically limited to single static estimate of the policy, unable 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 al.] 2018b), which retains the efficiency benefits of amortization while employing more flexible iterative estimation procedure.
2.4 RELATED WORK
Previous works have investigated methods for improving policy optimization. QT-Opt (Kalash-
nikov et al. 20188) uses the cross-entropy method (CEM) (Rubinstein & Kroese 2013), an iterative derivative-free optimizer,
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é et al. 2019 Hafner et al. 2019). Gradient-based policy optimization,
in contrast, is less common (Henaff et al.] 2017/viinivas et al. 2018 Bhhraddwajj et al. 2020). how-
ever, 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.) (211)) network. Similar approaches have also been applied to model-based value estimates (Byravan et al. 2020) Clavera et al. 2020 [Amos et al., (020). as well as combin- ing direct amortization with model predictive control (Lee al. 2019) and planning (Rivière al., 2020). A separate line of work has explored improving the policy distribution, using normalizing flows (Haarnoja et al.) 22018a] Tang Agrawall 2018) and latent variables (Tirumala et al. 2019). In principle, iterative amortization can perform policy optimization in each of these settings.
ITERATIVE AMORTIZED POLICY OPTIMIZATION
3.1 FORMULATION
Iterative amortized optimizers MMrinn et al. 2018b) utilize some form of error gradient to up- date the approximate posterior distribution parameters. While various forms exist, we consider gradient-encoding models (Andrychowicz et al. 2016) due their generality. Compared with di- rect amortization in Eq.
9 we use iterative amortized optimizers of the general form
A + [o
(iterative amortization)
(10)
] (Right), where fo is deep network and X are the action distribution param- eters. For example, if π = A(a: μ., diag(o?)) then = [M, o). Technically, s is redundant, as the state dependence is already captured in J, but this can empirically improve performance (Marino)
also shown in Figure
4

---

Under review as a conference paper at ICLR 2021
1,0
126
0,5
10.0
Oyt Run
tash(pp)
7.5
0.9
406 5.0
-0.5
-1.0
-10
-0.5
0.0
0.5
1.0
Lanla(pa))
Act. Dim. 2
Act. Dim.
Figure 2: Estimating Multiple Policy Modes. Unlike direct amortization, which restricted to a single estimate, iterative amortization can effectively sample from multiple high-value action modes. This is shown for particular state in Ant-v2, showing multiple optimization runs across two action dimensions (Left). Each square denotes initialization. The optimizer finds both modes, with the densities plotted on the Right. This capability provides increased flexibility in action exploration.
et al. 80) In practice, the update is carried out using "highway" gating operation (Hochreiter & Schmidhuber 1997; Srivastava et al., 2015). Denoting wi E [0, 1] as the gate and dg as the update, both of which are output by far the gating operation expressed as
A A w. C A (1 wo) o 6g.
(11)
where @ denotes element-wing wise multiplication. This update is typically run for a fixed number of steps, and, as with direct policy, the iterative optimizer is trained using stochastic gradient
timates of V J.J. obtained through the path-wise derivative estimator (Kingma & Welling, [0144]
Rezende et [004] Heess al., 2015). Because the gradients VXJ must be estimated online, i.c. during policy optimization, this scheme requires some way of estimating J online, e.g. through a parameterized Q-value network (Mnih al. 2013) or differentiable model (Heess et al. 2015).
3.2
CONSIDERATIONS
3.2. ADDED FLEXIBILITY
Iterative amortized optimizers more flexible than their direct counterparts, incorporating feed- back from the objective during policy optimization (Algorithm{)) ( rather than only after optimiza- tion (Algorithm []. Increased flexibility improves the accuracy optimization, thereby tightening the variational bound (Marino et al. a0) We see this flexibility in Figure [ ([ where an iterative amortized policy network, despite being trained with a different value estimator, 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 vat-v2 MuJoCo environment. Over multiple policy optimization runs, iterative amortization finds multiple modes, sampling from two high-value regions of the action space. This provides increased flexibility in action exploration.
3.2.2 MITIGATING VALUE OVERESTIMATION
Model-free approaches generally estimate Qx using function approximation and temporal differ- ence learning. However, this comes with the pitfall of value overestimation, i.c. positive bias in the estimate, Qr (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 uncer- tainty, the resulting target values will introduce positive bias into the estimate. More flexible policy optimizers may exacerbate the problem. exploiting this uncertainty to a greater degree. Further, a rapidly changing policy increases the difficulty of value estimation (Rajeswaran et al. 2020).
5

---

Under review as a conference paper at ICLR 2021
Cirert,
Irerative -
Diet (fteo (nteollli (npa))
Berative = 2,5
109
100
<
50
-58
01.5
1.0
1.5
2.0
2.5
3.0
107
0.0
0.5
1.0
1.5
2.
2.5
3.0
Milloom Stegg
Million Shops
(a)
(b)
Figure 3: Mitigating Value Overestimation. Using the same value estimation setup (3 = 1 in Eq. [2]. shown on Ant-v2 iterative amortization results in (a) higher value overestimation bias (closer to zero is better) and (b) a more rapidly changing policy as compared with direct amortiza- tion. Increasing helps to mitigate these issues by further penalizing variance in the value estimate.
Various techniques have been proposed for mitigating value overestimation in deep RL. The most prominent technique. double deep Q-netwrrk (Van Hasselt et al. 2016) maintains two Q-value estimates VVan 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, Qu,, and Q
(s. a) = min (s,a)) 1=1.
where vig denotes the "target" network parameters. As noted by Fujimoto al. (2018). this only counteracts value overestimation, but also penalizes high-variance value estimates, because the minimum decreases with the variance of the estimate. Ciosek al. (2019) noted that, for bootstrapped ensemble of two Q-networks, the minimum operation can be interpreted as estimating
(s, po(s, a) (s)
(12)
with mean (s a) Ei=12 Quy a). standard deviation (s a) = E,
(s, a))3)1/2, and 1. Thus, further penalize high-variance value estimates, preventing value overestimation, we can increase 3. For large B, however, value estimates become overly pessimistic, negatively impacting training. Thus, 3 reduces target value variance at the cost of increased bias.
Due to the flexibility of iterative amortization, the default = results in increased value bias (Fig- ure Bab and more rapidly changing policy (Figure 3b) as compared with direct amortization. Fur- ther penalizing high-variance target values with 3 = 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. a (Fox 2019). In offline RL, this issue has been tackled through the action prior (Fujimoto al. 2019 Kumar et al., 2019 Wu et al. 2019) or by altering Q-network training (Agarwal et 2209 Kumar et [l.0) While such techniques could be used here, increasing 3 provides simple solution with no additional computational overhead.
EXPERIMENTS
4.1
SETUP
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, Pa (a(s) = (1-1,,11 1), and a tuning scheme for the temperature, a. our experiments, "direct" refers to direct amortization employed in SAC, i.e. 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 up- dates (Sec, 3.I.. Unless otherwise stated, we use 5 iterations per time step for iterative amortization, following Marino et al. (2018b). For details, refer Appendix A and Haarnoja al. (2018hlc).

---

Under review as a conference paper at ICLR 2021
Opt. Ber.
Oje Bar
Action Dum
015
0.0
<
310
It Amon.
-0.5
Adban
350
Envinonmen
10
Time Stopp
20
Enviroment
10 15
Time Step
21
25
100 Opt Benatim
200
(a) Policy
(b) Improvement
(c) Comparison
Figure 4: Policy Optimization. Visualization over time steps of (a) dimension of the policy distribution and (b) the improvement in the objective. AJ, across policy optimization iterations. (c)
Comparison of iterative amortization with Adam (Kingma & Ba 2014) (gradient-based) and CEM
(Rubinstein & Kroese
2013) (gradient-free). Iterative amortization is substantially more efficient.
Chunlatge Rosurd
Cumndalie Rasond
Comndation Toward
Cundatiee Roward
a00
.0000
L000
3000
600
5000
200
dined
tenolivv
Milion
Million Stope
Million Steps
Million
(a) Hopper-v2
(b) Faaeeahha
Walker2d-v2
(d) Ant-v2
Figure 5: Performance Comparison. Iterative amortized policy optimization performs comparably with better than direct amortized policies across range of MuJoCo environments. Performance curves show the mean and + standard deviation over 5 random seeds.
4.2 ANALYSIS
4.2.1
VISUALIZING POLICY OPTIMIZATION
In Figure [Left), we visualize the trajectory of iterative amortized policy optimization along two dimensions of the policy mean 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 visualize iterative refinement using a single action dimension from Ant-v2 across time steps. The refinements in Figure 4a give rise to the objective improvements in Figure 4b We compare with Adam (Kingma & Ba 2014) (gradient-
and CEM (Rubinstein & Kroese 2013)
(gradient-frre) freeen in Figure [4c (see Appendix B.2), where iterative amortization is substantially more efficient. This trend consistent across environments, shown in Figure B.2 in the appendix.
4.2.2 PERFORMANCE COMPARISON
We evaluate iterative amortized policy optimization on MuJoCo (Toorov et al. 2012) continuous control tasks from OpenAl 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 efficiency and final performance. Across environments, iterative amortization also yields more consistent performance.
4.2.3 DECREASED AMORTIZATION GAP
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, x= [|,, o (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 We emphasize that it challenging to directly compare amortization gaps across optimization schemes, these involve different value functions, and therefore different objectives. Likewise, we estimate the amortization gap using the learned Qavuuu networks, which may be biased (Figure 3). Nevertheless, we find that iterative

---

Under review as a conference paper at ICLR 2021
In
12
107
Anorlaati Cup
Amerliatbe Cap
Amoritation Cap
Amoricationn
00-1
direct (BAC)
lir1
ar-
iteratio Sitee
Bendlee 10 lax.
Miliio
Million Strps
Million Stepe
Million
(a) Hopper-v2
(b) Hall ECheetab-22
Walker2d-v2
(d) Ant-v2
Figure 6: Decreased Amortization Gap. Estimated amortization gaps per step for direct and itera- tive amortized policy optimization. Iterative amortization achieves comparable lower gaps across environments. Gaps are estimated using stochastic gradient-based optimization over 100 random states. Curves show the mean and + standard deviation over random seeds.
Ohumlate Roward
Chomlaaio Roward
5500
Ommelaiv Malie Powerd
Qumdatee Deward
3000
L00
5000
2 lar
ttrr.
Milloe
Million Stope
Millinn Steps
Million
(a) Hopper-v2
(b) Hall ECheetah-v2
(c) Walker2d-v2
(d) Ant-v2
Figure 7: Iterations During Training. Performance of iterative amortized policy optimization for varying numbers of iterations during training. 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 amorti- zation across all environments. Further amortized iterations at evaluation yield further estimated improvement, demonstrating generalization.
However, note that the amortization gaps relatively small compared the estimated dis- counted objective. Accordingly, when we evaluate the more fully optimized policies in the environ- ment, do not observe noticeable increase in performance (see Appendix B.4). This demon- strates that policy suboptimality is not a significant concern for evaluation. Rather, improved policy optimization is helpful for training. This allows the agent collect data where value estimates 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 esti- mated 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 the Q-value estimator and other factors.
4.2.4 VARYING ITERATIONS
Direct amortized policy optimization is restricted to a single forward pass through the network. It- erative 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, 5). In Figure amortized optimization iterations generally improves sample efficiency (Walked wwwld-errddddd Hal HaleCheet ah-v2). Thus, the quality of
see that increasing the number of asymp-
totic performance (Ant-v2) or both (Hopper-u the policy optimizer can play significant role in determining performance through training.
4.2.5 ITERATIVE AMORTIZATION WITH MODEL-BASED VALUE ESTIMATES
While 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), provide a proof-of-concepe in this setting, using learned deterministic model on Hal @Cheeehh-hht (see Appendix A.5). As shown in Figure Sal iterative amortization outperforms direct amortization in this setting. Iterative amortization refines planned trajectories,
8

---

Under review as a conference paper at ICLR 2021
Conuletivn Recoad
Ope. Ies.
Cuanlative Remored
12000
10000
Din.
E2000
0030
direct
-0.
370
2000
MB
MP
0.30 Million Stops
0.25
Model Dollont Step
Opt. Jerakk
0.00
Million
0.75
100
(a) Performance
(b) Model-Based Planning (c) Improvement
(d) Ceto-Shot Transfer
Figure 8: Optimizing Model-Basee Value Estimates. (a) Performance comparison of direct and
iterative amortization using model-based value estimates. (b) Planned trajectories over policy opti-
mization iterations. (c) The corresponding estimated objective increases over iterations. (d) Zero-
shot transfer of iterative amortization from model-free (MF) to model-based (MB) estimates.
shown for single state dimension in Figure 8b yielding corresponding improvements (Figure &c Further, because we are learning an iterative policy optimizer, we can zero-shot transfer policy optimizer trained with model-free value estimator to a model-based value estimator (Figure 3d 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 instantly incorporating updated value estimates in policy optimization. This demonstrates and high- lights the opportunity for improving model-based planning through iterative amortization.
value estimates,
DISCUSSION
We have introduced iterative amortized policy optimization, a flexible and powerful policy opti- mization technique. Using the MuJoCo continuous control suite, we have demonstrated improved performance over direct amortization with both model-based and model-free value estimates. It- erative amortization provides drop-in replacement and improvement 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 al., 2018), gauging the required iterations. Likewise, efficiency depends. in part, on the policy initializa- tion, which could be improved by learning the action prior, po (a|s) (Abdolmaleki al. 2018). The power of iterative amortization is in using the value estimate during policy optimization to iteratively improve the policy online. This is form of negative feedhack 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 range of improved RL algorithms, capable of instantly adapting to different value estimates, shown in Figure 8d
direct policy networks
REFERENCES
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Mar- Riedmiller. Maximum a posteriori policy optimisation. International Conference on Learn- ing Representations, 2018.
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. arXiv preprint 17X40:1907.04533, 2019.
Zafarali Ahmed, Nicolas Le Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy policy optimization. In International Conference Machine Learning, pp- 151-160, 2019.
Brandon Amos and Denis Yarats. The differentiable cross-entropy method. In International Con- ference on Machine Learning, 2020.
Brandon Amos. Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model- based stochastic value gradient for continuous reinforcement learning. arXiv preprint arXiv:2008. 12775, 2020.

---

Under review as a conference paper at ICLR 2021
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoftman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. Advances in Neural Information Processing Systems (NIPS), pp. 3981-3989, 2016.
Karl Johan Astrom and Richard M Murray. Feedback Systems: An Introduction for Scientists and Engineers. Princeton University Press, 2008.
Hagai Attias. Planning by probabilistic inference. In AISTATS. Citeseer, 2003.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint 1rXiv:1607.064500 2016.
Homanga Bharadhwaj, Kevin Xie, and Florian Shkurti, Model-predictive planning via cross-entropy and gradient-based optimization. In Learning for Dynamics and Control, pp. 277-286, 2020.
Matthew Botvinick and Mare Toussaint. Planning as inference. Trends in cognitive sciences, 16
(10):485-488) 488, 2012.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 0rXiv:1666.015400 2016.
Arunkumar Byravan, Jost Tobias Springenberg, Abbas Abdolmaleki, Roland Hafner, Michael Ne- unert, Thomas Lampe, Noah Siegel, Nicolas Heess, and Martin Riedmiller. Imagined value gra- dients: Model-based policy optimization with tranferable latent dynamics models. Conference Robot Learning, pp.
2020.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learn- ing in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, pp. 4754-4765, 2018
Kamil Ciosek, Quan Vuong, Robert Loftin, and Katja Hofmann. Better exploration with optimistic actor critic. In Advances in Neural Information Processing Systems, pp. 1787-1798, 2019.
Ignasi Clavera, Yao Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. In International Conference on Learning Representations, 2020.
Djork-Aré Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network
learning by exponential linear units (elus). arXiv preprint arXiv:1511.07281 2015.
Gregory F Cooper. A method for using belief networks as influence diagrams. In Fourth Workshop on Uncertainty in Artificial Intelligence, 1988.
Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality variational autoen- coders. In International Conference Machine Learning, pp. 1078-1086, 2018.
Peter Dayan and Geoffrey Hinton. Using iopectatioo-maxxiizztiion for reinforcement learning.
Neural Computation, 9():2711-278 1997.
Michael Figurnov, Artem Sobolev, and Dmitry Vetrov. Probabilistic adaptive computation time. Bulletin of the Polish Academy of Sciences. Technical Sciences, 66(6), 2018.
Roy Fox. Toward provably unbiased lemporal-difference value estimation. Optimization Founda- tions for Reinforcement Learning Workshop at NeurIPS, 2019.
Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. In Proceedings of the Thirty-Secons Conference on Uncertainty in Artificial Intelligence, pp. 202-211. AUAI Press, 2016.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, pp. 1587-1596. 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In International Conference Machine Learning, pp. 2052-2062. 2019.
10

---

Under review as a conference paper at ICLR 2021
Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Pro- ceedings of the Cognitive Science Sociery, volume 36, 2014.
Alex Graves. Adaptive computation time for recurrent neural networks. arXiv preprint vrXiv:1603.08983 2016.
Klaus Greff, Raphaël Lopez. Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-ohject representation learning with iterative variational inference. International Conference on Machine Learning, pp. 2424-2433, 2019.
Arthur Guez, Mchdi Mirza, Karol Gregor, Rishabh Kabra, Sebastien Racaniere, Theophane Weber, David Raposo. Adam Santoro, Laurent Orseau, Tom Eccles, et al. An investigation of model-free planning. In International Conference on Machine Learning, pp. 2464-2473 2019.
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. Proceedings of the 34th International Conference on Machine Learning-Volumee Volume 70, pp. 1352-1361. JMLR. org, 2017.
Tuomas Haarnoja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine. Latent space policies for hicrarchical reinforcement learning. In International Conference on Machine Learning, pp- 1851-1860, 2018a.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International Confer- ence on Machine Learning, pp. 1856-1865, 2018b.
Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Schoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and appli- cations. arXiv preprint arXiv:1812.05500
2018c.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pp. 2555-2565. PMLR, 2019.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learn- ing continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems, pp. 2944-2952, 2015.
Mikael Henaff, William F Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint 1755::1705.07117. 2017.
Devon Hjelm, Ruslan R Salakhutdinov, Kyunghyun Cho, Nebojsa Jojic, Vince Calhoun, and Juny- oung Chung. Iterative refinement of the approximate posterior for directed belief networks. In Advances in Neural Information Processing Systems (NIPS), pp. 4691-4699, 2016.
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8): 1735-1780, 1997.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. In Advances in Neural Information Processing Systems, pp. 12519- 12530, 2019.
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Iharz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation. arXiv preprint arXiv:1806.1003 10293, 2018.
Yoon Kim, Sam Wiseman, Andrew Miller, David Sontag, and Alexander Rush. Semi-amortized variational autoencoders. In Proceedings of the International Conference Machine Learning (ICML), 2018
Diederik Kingma and Max Welling. Stochastic gradient vb and the variational auto-encoder. In Proceedings of the International Conference Learning Representations, 2014.
11

---

Under review as a conference paper at ICLR 2021
Durk P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6900, 2014.
Rahul G Krishnan, Dawen Liang, and Matthew Hoffman. On the challenges of learning with infer- ence networks on sparse, high- high-dimensional data. In Proceedings of the International Conference Artificial Intelligence and Statistics (AISTATS), pp. 143-151, 2018.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy q-learning via bootstrapping error reduction. In Advances in Neural Information Processing Sys- tems, pp. 1 11784-11794, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. arXiv preprint arXiv:2006. 204,799900.07779, 2020.
Keuntack Lee. Kamil Saigol, and Evangelos Theodorou. Safe end-to-end imitation learning for
model predictive control. In International Conference on Robotics and Automation, 2019.
Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: 08059 2018
Timothy Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Inter- national Conference on Learning Representations, 2016.
Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3-4))293-321, 1992.
Joseph Marino, Milan Cvitkovic, and Yisong Yuc. A general method for amortizing variational filtering. In Advances in Neural Information Processing Systems, 2018a.
Joseph Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Conference on Machine Learning, pp. 3403-3412, 2018b.
Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning, pp. 1791-1799, 2014.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. In NIPS Deep Learn- ing Workshop, 2013.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pp. 1928-1937, 2016.
Rémi Munos, Tom Stepleton, Anna Harutyunyan, and Mare Bellemare. Safe and efficient off-policy reinforcement learning. In Advances in Neural Information Processing Systems, pp. 1054-1062, 2016.
Anusha Nagabandi, Gregory Kahn, Ronald S Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning, In 2018 IEEE Interna- lional Conference on Robotics and Automation (ICRA), pp. 7559-7566. IEEE, 2018.
Alexandre Piché, Valentin Thomas, Cyril Ibrahim, Yoshua Bengio, and Chris Pal. Probabilistic planning with sequential monte carlo methods. In International Conference Learning Repre- sentations, 2019.
Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. arXiv preprint MxXi:2004.07800 2020.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. Proceedings of the International Conference Machine Learning, pp. 1278-1286 2014.
12

---

Under review as a conference paper at ICLR 2021
Benjamin Rivière, Wolfgang Hönig, Yisong Yue, and Soon-Jo Chung. Glas: Global-to-local safe autonomy synthesis for multi-robot motion planning with end-to-end learning. IEEE Robotics and Automation Letters, 5(3):4249-442) 2020.
Reuven Y Rubinstein and Dirk Kroese. The cross-entropy method: unified approach combina- torial optimization, Monte-Carlo simulation and machine learning. Springer Science & Business Media, 2013.
John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International conference machine learning, pp. 1889-1897, 2015.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv: 2017.
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International Conference on Machine Learning, pp- 387-395, 2014.
David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schrittwieser, loannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484 489- 2016.
Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal plan- ning networks: Learning generalizable representations for visuomotor control. In International Conference on Machine Learning, pp. 4732-4741, 2018.
Rupesh K Srivastava, Klaus Greff, and Jürgen Schmidhuber. Training very deep networks. In Advances in neural information processing systems (NIPS), pp. 2377-2385, 2015.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Yunhao Tang and Shipra Agrawal. Boosting trust region policy optimization by normalizing flows policy. arXiv preprint arXiv: 1809.1832666 10326, 2018
Sebastian Thrun and Anton Schwartz. Issues in using function approximation for reinforcement learning. In Proceedings of the 1993 Connectionist Models Summer School Hillsdale, NJ. Lawrence Erlbaum, 1993.
Dhruva Tirumala, Hyeonwoo Noh, Alexandre Galashov, Leonard Hasenclever, Arun Ahuja, Greg Wayne, Razvan Pascanu, Yee Whye Teh, and Nicolas Heess. Exploiting hierarchy for learning and transfer in kl-regularized rl arXiv preprint arXiv: 1903.003890. 2019.
Emanuel Todorov. General duality between optimal control and estimation. In Decision and Control, 2008. CDC 2008. 47th IEEE Conference on, pp. 0298. IEEE, 2008.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In Intelligent Robots and Systems (IROS), 2012 IEEE/RSJ International Conference on, pp. 5026- 5033. IEEE, 2012.
Mare Toussaint and Amos Storkey, Probabilistic inference for solving discrete and continuous state markov decision processes. Proceedings of the 23rd international conference on Machine learning. pp. 945-952. ACM, 2006.
Hado Van Hasselt. Double q-learning. In Advances in neural information processing systems, pp- 2613-2621, 2010.
Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double q- learning. Thirtieth AAAI conference on artificial intelligence, 2016.
Ronald J Williams. Simple statistical pradient-follovin algorithms for connectionist reinforcement learning. In Reinforcement Learning, pp. 5-32. Springer, 1992.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint 131, 2019.
13

---

Under review as a conference paper at ICLR 2021
A EXPERIMENT DETAILS
A.1 2D PLOTS
In Figures and 2, we plot the estimated variational objective, J, as a function of two dimensions of the policy mean, μ.. To create these plots, we first perform policy optimization (direct amortization in Figure I and iterative amortization in Figure 2), estimating the policy mean and variance. This performed using on-policy trajectories from evaluation episodes (for direct agent in Figure and 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 I and 2 & 6 in Figure 2). Iterative amortization is additionally performed while preventing any updates the constant dimensions. Even in this restricted setting, iterative amortization capable of optimizing the policy. Additional 2D plots comparing direct vs. iterative amortization on other environments are shown in Figure [4] where we see similar trends.
A.2 VALUE BIAS ESTIMATION
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 (Qx) with a Monte Carlo estimate of the future objective in the actual environment, Qu 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,
Es,a ②x 2a-QQ. Qm throughout training. To obtain the Monte Carlo estimate of QT. 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, a. Figure shows the mean and + standard deviation across the 100 state-action pairs.
A.3 AMORTIZ.ATION GAP ESTIMATION
Calculating the amortization gap in the RL setting is challenging, as properly evaluating the varia- tional objective, J, involves unrolling the environment. During training, the objective is estimated using a set of Q-networks and/or a learned model. However, finding the optimal policy distribution, म. under these learned value estimates may not accurately reflect the amortization gap, 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 J w.r.t. the policy distribution parameters, X, initializing from the amortized estimate (from T&). This is form semi-amortized variational inference (Hielm et al. (0166) Krishnan al. 2018 Kim et al. 2018). We use the Adam optimizer (Kingma & Ba 2014) with a learning rate of 5 x 10 for 100 gradient steps, which we found consistently converged. This results in the estimated optimized न. We estimate the gap using 100 on-policy states, calculating J(0, f) I (0, T). i.c. the improve- ment in the objective after gradient-based optimization. Figure 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.4 HYPERPARAMETERS
Our setup follows that of soft actor-critic (SAC) (Haarnoja et al., 2018b(c) using a uniform action prior, i.c. 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
Temperature Following Haarnoaa et al. (2018c). we adjust the temperature, a, maintain spec- ified entropy constraint, 6a =I [|] where A is the size of the action space, i.c. the dimensionality.
Policy We use the same network architecture (number of layers, units/layer, non- -inearierr 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 E [-1,1] (Haarnoja et al.
2018b). In the case of Gaussian, the distribution parameters are X [1, The inputs and
14

---

Under review as a conference paper at ICLR 2021
Table 2: Policy Networks.
Table 1: Policy Inputs & Outputs.
Hyperparameter
Value
Inputs
Outputs
Number of Layers Number of Units / Layer Non-linearity
Direct Iterative
2566 ReLU
8, A,
VAJ
8, w
Direct
caspoon
Iterative
coapooa
V
V
OODD090 GOORROO
GO00099 C00A600
V.J
(a) Direct Amortization
(b) Iterative Amortization
Figure 9: Amortized Optimizers. Diagrams of (a) direct and (b) iterative amortized policy opti- mization. As in Figure 1, larger circles represent probability distributions, and smaller red circles represent terms in the objective. Red dotted arrows represent gradients. In addition the state, St, iterative amortization uses the current policy distribution estimate, A, and the policy optimization gradient, VxJ, to iteratively optimize J. Like direct amortization, the optimizer network parame- ters, Φ, are updated using VoJ. This generally requires some form of stochastic gradient estimation differentiate through a, [aalss O: A).
outputs of each optimizer form are given Table [ Again, ö and w are respectively the update and gate of the iterative amortized optimizer (Eq. 11), each of which are defined for both μa and a. 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 learned action prior (Marino et al. 2018)).
Table 3: Q-value Network Architecture A.
Table 4: Q-value Network Architecture
Hyperparameter
Value
Hyperparameter
Value
Number of Layers
2
Number of Layers
Number of Units/ / Layer
256
Number of Units / Laaer Non-linearity
512
Non-linearity
ReLU
ELU
Layer Normalization
False
Layer Normalization Connectivity
True
Connectivity
Sequential
Highway
Q-value We investigated two Q Q-value network architectures. Architecture A (Table is the same as that from Haarnoja et al.] (2018)). Architecture (Table 4 is a wider, deeper network with highway connectivity (Srivastava et al., (015), layer normalization (Ba et 2016), and ELU non- lincaritics (Clevert et al. 2015). We initially compared each Q-value network architecture using each policy optimizer each environment, obtained using the better performing architecture in each case, given Table [ As in Fujimoto ensemble of 2 separate Q-networks in each experiment.
shown in Figure I0 The results in Figure 5 were
et al. 2018),
Value Pessimism (3) As discussed in Section 3.2.2. the increased flexibility of iterative amortiza- tion allows it to potentially exploit inaccurate value estimates. We increased the pessimism hyperpa-
15

---

Under review as a conference paper at ICLR 2021
Table 5: Q-value Network Architecture by Environment.
Hopper-v2
Hal fCheetah-v
Wa Walker2d-v2
Ant-v2
Direct Iterative
B
A
A
B
B
Caunietio Benard
Curele Rorcod
Comnletion tiv Reroad
Caunlatrv Beeard
20000
L00
direct (SMC)
000
5000
nonaioo + A
Renbion
Stops
Stape
Million Stops
Stepe
(a) Hopper-v2
(h) fCheetah-v2
walker2d-v2
(d) Ant-v2
Figure 10: Value Architecture Comparison. Plots show performance for 2 3 seeds for each value architecture (A or B) for each policy optimization technique (direct or iterative). Note: results for iterative + B on Hopper-ve r-v2 were obtained with an overly pessimistic value estimate (3 2.5
rather than 3 1.5) and are consequently worse.
rameter, 3, further penalize variance in the value estimate. Experiments with direct amortization use the default 3 in all environments, as we did not find 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, use 3 1 to not overly penalize exploration.
Table 6: Training Hyperparameters.
Hyperparameter
Value
Discount Factor (y)
0.99
Q-network Update Rate (r)
5 10-3 33
Network Optimizer
Adam
Learning Rate
3 10 4
Batch Size
256
Initial Random Steps
5 10
Replay Buffer Size
106
A.5 MODEL-BASED VALUE ESTIMATION
For model-based experiments, we use single, deterministic model together with the ensemble of 2 Q-value networks (discussed above).
Model We use separate networks to estimate the state transition dynamics, Penv (s |s,,, az), and reward function, r(st, a.). The network architecture is given in Table 7 Each network outputs the mean Gaussian distribution; the standard deviation is a separate, learnable parameter. The reward network directly outputs the mean estimate, whereas the state transition network outputs residual estimate, As,, yielding an updated mean estimate through:
St A
Model Training The state transition and reward networks are both trained using maximum log- likelihood training, using data examples from the replay buffer. Training is performed at the same frequency policy and Q-network training, using the same batch size (256) and network optimizer. However, we perform 103 updates at the beginning training, using the initial random steps, in order to start with reasonable model estimate.
16

---

Under review as a conference paper at ICLR 2021
Table 7: Model Network Architectures.
Table 8: Model-Based Hyperparameters.
Hyperparameter
Value
Hyperparameter
Value
Number of Layers
Rollout Horizon, h
2
Number of Units Layer Non-linearity Layer Normalization
256
Retrace À
0.9
Leaky
ReLU
Pre-train Model Updates
10
True
Model-Basse Value Targets
True
Value Estimation To estimate Q-values, we combine short model rollouts with the model-free estimates from the Q-networks. Specifically, 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 (s, a) and the reward estimate as r(s, a), caleulate the Q-value estimate at the current time step
I++h
Σ -ex- (F(se, ar) + + (su44))) - i Qu(sy,ar)) av))
(n a.) = ( a.) +
(13)
where À is an exponential weighting factor, 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
(aals, O) O(a
V (s) Ex Q=
a
(14)
pa (als)
In Eq. 13 Finally, to ensure consistency between the model and the Q-value networks, we use the model-based estimate from Eq. 13
we approximate Vn using the Q-network to approximate Qx in Eq. 14 yielding Ve
provide target values for the Q-networks, as in Janner al. (2019).
Future Policy Estimates Evaluating the expectation in Eq. 13 requires estimates of future time steps. This is straightforward with direct amortization, which employs feedforward policy, however, with iterative amortization, this entails recursively applying iterative optimization pro- cedure. Alternatively, we could use the prior, (als).. at future time steps, but this does not apply in the max-entropy setting, where the prior uniform. For computational efficiency, separate direct (amortized) policy for model-based rollouts. That is, with iterative amortization, we create separate direct network using the same hyperparameters from Table [ This network distills iterative amortization into direct amortized optimizer, through the KL divergence, DRI. (Ti. Tair.) Rollout policy networks are common model-based RL (Silver et al. 2016 Piché et al. 2019).
instead learn
B ADDITIONAL RESULTS
B.1 IMPROVEMENT PER STEP
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 aver- age change in the variational objective after running 5 iterations of iterative amortized policy opti- mization. With the exception of Ha FCheetah-v22 the improvement remains relatively constant throughout training and consistent across seeds.
B.2 COMPARISON WITH ITERATIVE OPTIMIZERS
Iterative amortized policy optimization obtains the accuracy benefits of iterative optimization while retaining the efficiency benefits of amortization. In Section 4, 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, compare iterative amortization with two popular iterative optimizers: Adam (Kingma & Ba 2014), a gradient- based optimizer, and gradient-free optimizer.
cross-entropy method (CEM) (Rubinstein & Kroese 2013
17

---

Under review as a conference paper at ICLR 2021
rss ax)
0000000
000000a 0000000
coooaco 0000000 000000o
ooooaco
Value
Policy
Model
0000000
G000000
SL+11 [8,, B
Figure 11: Model-Based Value Estimation. Diagram of model-based value estimation (shown with direct amortization). For clarity, the diagram shown without the policy prior network, (a (s.). The model consists of a deterministic reward estimate, r(s,, az). (green diamond) and state esti- mate, 81+1 (8t, at. (orange diamond). The model unrolled over horizon, H, and the Q-value estimated using the Retrace estimator MMuoos et al. 2016), given in Eq. 13
AT Stop
A3 Soo
Aa cuep
Sadl
Millin Shape
Miliin Sape
Millia Stepes
Millis Stage
(a) Hopper-v2
(b) Hal fCheetah-v2
(c) Walker2d-v2
(d) Ant-v2
Figure 12: Per-Step Improvement. Each plot shows the per-step improvement in the estimated variational RL objective, J. throughout training resulting from iterative amortized policy optimiza- tion. Each curve denotes different random seed.
971
500
113
1, 272
Amoct
to 185
Adam
020
441
500
1:0
700
Opc. Itcacioa
Opt. henation
Ost. Iteratoon
Opt. Decarico
(a) Hopper-v2
(b) Hall FCheetah-v2
(c) Walkkrdd=v2
(d) Ant-v2
Figure 13: Comparison with Iterative Optimizers. Average estimated objective over policy opti- mization iterations, comparing with Adam (Kingma & Ba 2014) and CEM (Rubinstein & Kroese 2013]. These iterative optimizers require over an order of magnitude more iterations to reach com- parable performance with iterative amortization, making them impractical in many applications.
To compare the accuracy and efficiency of the optimizers, we collect 100 states for each seed in each environment from the model-free experiments 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, sample 100 actions and fit 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 samples.
18

---

Under review as a conference paper at ICLR 2021
10
Danst Fulky Network
Dae Pday Stvekk
Dave: Pulex Sturch
Besstive Pikky Setwek
Devalive Pilks Netwols
Jeriii Picky Newock
Oprinall
82
Optival
Faits
ted
(a) Hal FCheetah-v2
(b) Wa hallke22dvv2
(c) Ant-v2
Figure 14: 2D Optimization Plots. Each plot shows the optimization objective over two dimen- sions of the policy mean, μs. This optimization surface contains the value function trained using direct amortized policy. The black diamond, denoting the estimate of this direct policy, is generally near-optimal, but does not match the optimal estimate (red star). Iterative amortized optimizers capable of generalizing to these surfaces in each case, reaching optimal policy estimates.
The results, averaged across states and random seeds, are shown in Figure 13 CEM (gradient- free) is less efficient than Adam crradient-basedde which is unsurprising, especially considering that Adam effectively approximates higher-order curvature through momentum terms. However, Adam and CEM both require over order of magnitude more iterations 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 only roughly 1.25x slower than CEM and 1.15x slower than Adam, making iterative amortization still substantially more efficient.
B.3 ADDITIONAL 2D OPTIMIZATION PLOTS
In Figure I, we provided an example of the suboptimal optimization resulting from direct amortiza- tion the Hopper-va v2 environment. We also demonstrated that iterative amortization capable of automatically generalizing to this optimization surface, outperforming the direct optimizer. To show that this general phenomenon, Figure [4] 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 Ha 1 FCheet ah does not match the optimal estimate. In contrast, iterative amortization is able to find the optimal estimate, again, generalizing the unseen optimization surfaces.
B.4 ADDITIONAL OPTIMIZATION & THE AMORTIZATION GAP
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 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 quantifies 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 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 do not observe noticeable difference in performance. Thus, while amortization may find suboptimal policy estimates, we observe that the actual difference in the objective is either too small inaccurate to affect performance 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
19

---

Under review as a conference paper at ICLR 2021
Caumlattg Beoard
0000
Cunslelee Rocant
5000
Chmmbaaion Reward
Cambatize Reward
1000
Mv Ma papp
2000
600
5000
5000
2000
dinet
dinad
ganliai
Milion
Million Stepe
Million Stepe
Millio
(a) Hopper-v2
(b) HallHhCheee ECheetab-v2
(c) Walker2d-v2
(d) Ant-v2
Figure 15: Test-Time Gradient-Basse Optimization. Each plot compares the performance of direct amortization vs. direct amortization with 50 additional gradient-base policy optimization
iterations. Note that this additional optimization is only performed at test time.
Chumultte Roward
8000
Cummisiion Boward
5000
Commistion Reward
Chamulatvv Reward 2000
3000
8000
5000
200
Txestim 5 the
Millon Strps
Million Stope
Millin Stope
Million Stepe
(a) Hopper-v2
(b) Hall ECheetah-ee2
(c) Walker2d-v2
(d) Ant
Figure 16: Additional Amortized Test-Time Iterations. Each plot compares the performance of
iterative amortization (trained with 5 iterations) vs. the same agent with an additional 5 iterations at
evaluation. Performance remains similar or slightly worse.
Cumdatbe Rawwnr
Ammioatte Gap
.0000
ti-11
Hien
Itr.
Million Stope
Milbon Stops
(a)
Figure 17: Iterations During Training. (a) Performance and (b) estimated amortization gap for varying numbers of policy optimization iterations per step during training on HalfCheetahr ah-v2. Increasing the iterations improves performance and decreases the estimated amortization gap.
more fully optimized policy in the environment, as shown in Figure 16 do not generally observe a corresponding performance improvement. In fact, on Hal fCheetah-ve and Wa welleee2ddaa observe 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 correct the inaccurate value estimates. Thus, while more accurate policy optimization, as quantified by the amortization gap, may not substantially affect evaluation performance, it does play significant role 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 flexible policy optimizer, generally results in improved performance. By using 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 shown for Hal fallfhheetah-ae 2- in Figure 17 where we observed the largest difference in per- formance across numbers of iterations. We generally observe that increasing the number of itera- tions 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.
20

---

Under review as a conference paper at ICLR 2021
Poley Menn 12 Die
Bang Mouu to Di mas
Poicy Moon Diesnros
pase Moan 12 Distage
0.4
Million Stope
Million Steps
Million Strge
(a) Hopper-v2
(b) Hall FaLfChhetab-v2
(c) Walker2d-v2
(d) Ant-v2
Figure 18: Distance Between Policy Means. Each plot shows the L2 distance between the estimated policy means from two separate policy optimization runs given state. Results are averaged over 100 on-policy states at each point in training and over experiment seeds.
B.5 MULTIPLE POLICY ESTIMATES
As discussed in Section 3.2.1, iterative amortization has the added benefit 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 flows are incorporated into the policy, direct amortization is limited to 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 1.2 distance between the means of these pol- icy estimates (after applying the tanh). Note that MuloCo action spaces, which are bounded to [-1,1], 1], the maximum distance is 2 /AA, where [A] is the size the action space. We aver- age the policy mean distance over 100 states and all experiment seeds, with the results shown in Figure 18 all environments, we see that the average distance initially increases during training, remaining relatively constant for Hopper-v2 and
and decreasing slightly for Wa 2erkee2d-v2 and Ant-v2. Note that the distance for direct amortization would be exactly throughout. This indicates that iterative amortization does indeed obtain multiple policy estimates, maintaining some portion of multi-estimate policies throughout training.
21


Reviews:
[{"review_id": "_23hHUPn0XZ", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Contributions**: The authors propose to use iterative amortization for policy optimization to help reduce suboptimality in policy optimization. They find that like in variational inference for generative modeling, iterative amortization is able cover multiple modes of distributions and lower the amortization gap, and show slightly improved performance on mujoco gym benchmarks. Additionally, they find that the iteratively amortized policy is able to better exploit the learned Q function and lead to more overestimation bias, and address this by tuning a parameter to make the Q-function backups more pessimistic to compensate.\n\n**Questions and Concerns**: It is not entirely clear to me where the better performance observed from using iterative amortization derives from in the benchmark results. For example, in Hopper and HalfCheetah, we see that the iterative amortization slightly outperforms the direct baseline, yet the the difference in amortization gaps seems negligible throughout training. On the other hand, in the Ant environment, the amortization gap of the direct method is much higher, yet the learning curves of both methods are practically identical. One potential issue is that the objective of the policy update depends on the *learned* Q-function instead of the true returns from the environment, and so a better inference procedure may not even necessarily be a benefit. For this reason, we often consider schemes like trust regions for policy updates that explicitly limit policy optimization.\n\nSome additional questions:\nTo what extent is the amortization gap alleviated by direct amortization but larger, more expressive networks? Does doing so provide similar improvements in performance?\n\nIn Figure 3b, it appears that with the iterative schemes, the new policy consistently differs the old one (what exactly is pi_old in this case?) by a significant amount. Should we be concerned here that the learned policy (and presumably also the Q-functions) are not really coming close to converging? This further suggests to me strictly trying to have the policy better approximate this noisy oscillating Q may not actually be beneficial behavior.\n\n**Summary**: The authors demonstrate that iterative amortization enables some improvement to policy optimization in a maximum entropy RL algorithm. However, it is not entirely well understood to me how precisely iterative amortization leads to such benefits. Regarding novelty, the presented method is a straightforward application of a technique from variational inference to maxent RL, so is perhaps a bit incremental in that sense. Overall, I like the method since it provides a very simple way to improve learning by plugging in iterative amortization, but think it needs a to provide a better understanding of why it helps.", "extra_scores": {}, "extra_text": {}}, {"review_id": "VAmQQXF_evW", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary\n\nThe paper proposes to replace the actor/policy network with an iterative version to encode the action distribution parameters, which is inspired by prior work on iterative amortized optimization. This scheme generates the action distribution parameters for each state at the end of an inner loop which takes the objective gradient wrt to the parameters also as input (as opposed to a regular actor network which does one forward pass). This seems similar to prior work on learning to learn/meta learning but the presentation can be improved (more detailed comments below).  Experimental evidence is given for benefits over SAC in mujoco environments.\n\nOverall, the general idea seems interesting and useful, but the novelty is unclear given the prior work; the clarity of the presentation can be improved which makes an accurate assessment of the significance harder. \n\n### Detailed Comments\n\n* In the introduction, authors motivate the idea of amortization by making a distinction between optimizing the distribution versus a network whose parameters output the distribution. Can the authors provide an example of an algorithm that does the former for clarity (i.e. an example of a non-amortized policy -- what would that look like hypothetically?) \n\n* The agent's policy is denoted by $p_\\theta$ (prior to Eq (1)) and then with $\\pi$ later. Comment below Eq (8) makes a reference to $\\theta$ as prior parameters. This presentation is unclear and needs additional information to disambiguate. \n\n* What do the dependencies in Figure 1 right denote? For instance, $p_\\theta$ has a dependence on $p_{env}$ in this figure, but it's not clear what that means.\n\n* How is $\\lambda$ initialized in the inner update loop? The capability to find multiple modes seems dependent upon providing diverse initializations, which seems unspecified. \n\n* The value overestimation section discussion seems to suggest that the iterative scheme does. worse than the direct even with changing $\\beta$. Not clear what the take away from this section is in the context of the rest of the paper.\n\n\n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "R7Ed6Xcgc4P", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\nThe paper studies the popular soft actor-critic (SAC) algorithm, and points out that it does not find the most optimal actor (action distribution) that maximizes the critic (Q value). The reason is that SAC uses “direct amortization” (a feed-forward network from state to Gaussian params for the action distribution), and hopes that simple gradient descent (eqn (13) in the SAC paper) can sufficiently minimize loss (eqn (10) in SAC) over time, which is not true as Fig 6 in this paper suggests. \n\nTo solve this issue, the paper suggests an “iterative amortization” method that optimizes the Gaussian params through a few iterations of 1st order optimization (Alg 2). The optimization is done implicitly with a neuralnet that takes the policy gradient as input (eqn (10)). The resulting algorithm reduces the amortization gap (Fig 6) and improves policy performance over time (Fig 5).\n\nRecommendation\nOverall, I’m slightly inclined to accept this paper. The amortization gap is an interesting and often ignored phenomenon. The paper offers a simple solution and performs convincing experimental study. However, the presentation is very confusing, as the focus is selling this A + B solution (A = iterative amortization, B = SAC), rather than analyzing the “amortization gap” issue in full detail. Consequently, the experiments are not comprehensive enough. \n\nStrengths\n1. The paper reveals that the SAC policy isn’t always fully optimized. It’s common to assume that gradient descent + large enough network + enough data will eventually optimize the policy. But as Fig 6 suggests, it might not be the case.\n2. The paper suggests a simple solution through iterative amortization. This plug-and-play trick outperforms vanilla SAC on two Mujoco tasks (Fig 5).\n\n\nWeaknesses\n1. The paper stops at selling a solution instead of going beyond to analyze the amortization gap. Some natural questions:\n(a) There are multiple ways to fully optimize the policy, such as using a larger SGD coeff and running more SGD steps. Do they work?\n(b) If the amortization gap is an issue, do you see a problem in the expressivity of Gaussian? Does a more general form of policy (e.g. in the original SQL paper) further reduce the gap and improve performance? \n(c) In the case of discrete actions, where the optimal policy can be found directly (simple softmax), does the fully optimized policy outperform direct optimization? Any relevant papers?\n\n2. The presentation is very confusing.\n(a) The motivation reads like “Here’s A and here’s B. Let’s try A + B”. If the paper pointed out SAC’s amortization gap earlier and how it seriously impacts performance, the paper would be easier to understand.\n\n(b) Confusing terminology.\n- p_{\\theta}: Initially described as the agent’s “parametric distribution”, later called “policy prior parameters”, then finally declared a uniform distribution.\n- f_{\\phi}: Should not use the same symbol for both direct and iterative amortization.\n- “policy distribution parameters”: Easily confused with “policy network parameters”. Better called “action distribution parameters”.\n- “are policy networks providing fully-optimized policy estimates”: Estimates could mean the Q value. Better called “policy objectives”.\n\n(c) Lack of experimental rigor.\n- In Fig 1 and 2, which state and which training iteration is being studied? Why should the reader trust the plots?\n- Fig 6 shows that the amortization gap in Ant-v2 does not hurt SAC’s performance. Could you explain why?\n\nOther feedbacks\nI put in a lot of effort in understanding this paper and re-summarizing it in this review. I’d hope the authors understand my confusion and restructure the paper correspondingly. \n", "extra_scores": {}, "extra_text": {}}, {"review_id": "7kdClKSvHaF", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This paper draws an interesting connection to variational inference, categorizing current policy optimization methods with KL regularization as direct amortized optimizers. The paper shows how direct amortized policy optimization can be suboptimal and proposes a new class of method called iterative amortized policy optimization. This ‘iterative’ amortized policy optimization performs iterative optimization before the regular optimization step, offering several advantages: better at reaching the optima, better at sampling multiple modes, and support more rapidly changing policy. The paper shows nice visualization to support these claims and also run benchmark experiments to show its improvement.\n\nThis paper has done a good job introducing ‘iterative amortized policy optimization’, providing support with both good visualizations at a simpler level (to showcase what it is actually doing) and empirical evaluation on complex benchmarks. I think the contribution is novel and introduces a different type of optimization that is shown to improve current policy optimization methods. \nDespite the paper being an overall good paper, I think the experiments should be much better supported and recommend a borderline reject unless the issues below can be addressed.\n\nThe results look convincing overall but I think more runs should be conducted. Experiments on benchmarks only conduct 4 runs. There is high variance in results due to random seeds especially in Mujoco domains. Henderson et al. 2018 (https://arxiv.org/pdf/1709.06560.pdf) show that two different groups of runs with 5 random seeds each can lead to significant difference in performance on the same algorithm. It is hard to say that iterative amortized optimization helped improve performance with only 4 runs. \n\nFurthermore, it is not mentioned at all in the main paper that for HalfCheetah and Walker2D, iterative vs direct use different model architectures (it is mentioned in Appendix A.4 Q-value section). \n\nThese two environments in particular were environments that showed significant difference between the two methods, and makes the reader wonder whether this improvement was due to optimization method difference or architectural change.\nI think at least a justification in this particular design choice is needed or perhaps all results for both models can be shown to help readers understand the general performance.\n\nLastly, it would also be helpful to mention how many runs were conducted for model-based value estimate experiment in Figure 8a. \n\nOther minor details:\nI thought it would be helpful to mention where experiment details can be found in the appendix for each figures or sections in the main paper. \n\nAt the bottom of p.4 last paragraph of section 3.1, should \\nabla_\\phi J be \\nabla_\\lambda J?", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: I think there is a lot to commend in this paper: the general approach for training f_phi in this way is creative and interesting, the discussion of the amortization gap is thought-provoking, and the general idea is not something that I have seen in the literature before. That said, the reviewers raise a number of important concerns about the approach, chief of which is that the paper's explanation for why the method works is questionable. The proposed method can be viewed as simply a more expressive policy architecture. In fact, I suspect strongly that the modest increase in performance is likely explained by this alone. The discussion with Reviewer 2 in particular makes this issue very clear. I don't think the authors offered a very compelling response to this. Therefore, I think there are just too many question marks about this approach to accept the paper for publication at this time.

I do however think that this line of research is very promising, and I would encourage the authors to continue this work and flesh out the evaluation to be more rigorous and complete, understand whether the increase in performance comes down simply to increased expressivity or if the discussion of amortization is closer to reality, and also address other concerns raised by R2 and the other reviewers.
Decision: Reject