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

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 used with entropy or KL regularization, are a form of amortized optimization, optimizing network parameters rather than the policy distributions directly. However, direct amortized mappings can yield suboptimal policy estimates and restricted distributions, limiting performance and 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 on benchmark continuous control tasks.

Iterative Amortized Policy Optimization
Joseph Marino" California Institute of Technology
Alexandre Piché Mila, Université de Montréal
Alessandro Davide lalongo University of Cambridge
Yisong Yue California Institute of Technology
Abstract
Policy networks 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 used with entropy or KI. regularization, are form of amortized optimization, opti- mizing network parameters rather than the policy distributions directly. However, direct amortized mappings can yield suboptimal policy estimates and restricted distributions, limiting performance and exploration. Given this perspective, we con- sider the more flexible class of iterative amortized optimizers. We demonstrate that the resulting technique, iterative amortized policy optimization, yields performance improvements over direct amortization on benchmark continuous control tasks. Accompanying code: github. com/ hommmjoeouuieee ina.noooaarraaiinnn
RI., policy networks, when
Introduction
Reinforcement learning (RL) algorithms involve policy evaluation and policy optimization Given a policy, one can estimate the value for each state or state-action pair following that policy, and given value estimate, one can improve the policy 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 often parameterized by deep networks, such issues are typically tackled using regularization from previous policies [67] 668 or by maximizing policy entropy (57 23). These techniques can be interpreted as variational inference [51], using optimization to infer policy that yields high expected return while satisfying prior policy constraints. This smooths the optimization landscape, improving stability and performance [3].
73
However, one subtlety arises: when used with entropy or KL regularization, policy networks perform amortized optimization (26). That is, rather than optimizing the action distribution, e.g., mean and variance, many deep RL algorithms, such as soft actor-critic (SAC) [1 32), instead optimize a network to output these parameters, learning optimize the policy. Typically, this implemented as a direct mapping from states to action distribution parameters. While such direct amortization schemes have improved the efficiency of variational inference as "encoder" networks they also suffer from several drawbacks: 1) they tend to provide suboptimal estimates 20 43 55
44 56].
yielding a so-called "amortization gap" in performance [20). 2) they are restricted a single estimate [27]. thereby limiting exploration, and 3) they cannot generalize to new objectives, unlike, e.g.. gradient-based (36) or gradient- free optimizers (66).
Inspired by techniques and improvements from variational inference, we investigate iterative amor- tized policy optimization. Iterative amortization [55] uses gradients errors iteratively update the parameters of distribution. Unlike direct amortization, which receives gradients only after
Now at DeepMind, London, UK. Correspondence to jeeeppharrinoot com.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

outputting the distribution, iterative amortization uses these gradients online, thereby learning to iteratively optimize. generative modeling settings, iterative amortization empirically outperforms direct amortization 55 54
and can find multiple modes of the optimization landscape [27
The contributions of this paper
follows:
We propose iterative amortized policy optimization. exploiting a new, fruitful connection between amortized variational inference and policy optimization.
Using the suite of MuJoCo environments [78 [2]
we demonstrate performance improve-
ments over direct amortized policies, as well more complex flow-based policies.
We demonstrate novel benefits of this amortization technique: improved accuracy, providing multiple policy estimates, and generalizing to new objectives.
Background
2.1 Preliminaries
consider Markov decision processes (MDPs), where S; S and a, E A are the state and action at time t, resulting in reward тı r(s;,, Environment state transitions given by S:+1 ~ Peny (S++1 Ss, ay), and the agent is defined by a parametric distribution, Po(ar(s)) with parameters 0 The discounted sum of rewards is denoted as R(T) E. 2're. where (0, 1] is the discount factor, and r = ($1, al,...) is trajectory. The distribution over trajectories is:
T-1
p(r) = p(s)) II Pew (s (sc, a (a ss:),
where the initial state is 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 the undiscounted setting (y 1), though the formulation can be applied with any valid y.
2.2
KL-Regularized Reinforcement Learning
Various works have formulated RL, planning, and control problems in terms of probabilistic inference 21 8 79 77 11 51. These approaches consider the encer-enviooment interaction as graphical model, then convert reward maximization into maximum marginal likelihood estimation, learning and inferring policy that results in maximal reward. This conversion is accomplished by introducing one or more binary observed variables
19] denoted as for "optimality" [51] with
(0 = 1/T) x exp (R(r)/a),
where is temperature hyper- parameter. We would like to infer latent variables, and learn parameters, 8, that yield the maximum log-likelihood of optimality, i.., log = 1). Evaluating 1)dr. This
this likelihood requires marginalizing the joint distribution, p(O 1) p p(r, involves averaging over all trajectories, which is intractable high-dimensional spaces. Instead, we use variational inference lower bound this objective, introducing structured approximate posterior distribution:
T-1
क (r|()) = p(s1) II
Peav (St+ 11,, ac) O).
This provides the following lower bound on the objective:
log p(O = 1) log
plO = -1/))((T)d
(3)
2 f n(1|0) log p(0 = 1|7)p(T)
dr
(4)
= R(T/aa) (tt1:///)lDD
(5)
"In this paper, we consider the entropy-regularized case. where ps (a.ls)) = u(-1,1), iee. uniform. However, we present the derivation for the KL-regularized case for full generality.

---

204
E. [Q)(
(s.. ac)
teabl
505
-aiSe log Zolax 1r.
Te (ar//e8 O)
po(au (s.)
Polae
-05
Polky Natwoak
Tueettve Pairry Neework
Rea 1:a x]
Opeinall Exiiadi
Direct
Iterative
Lanlig))
Figure Amortization. Left: Optimization over two dimension of the policy mean, fi and μa, for particular state. A direct amortized policy network outputs suboptimal estimate, yielding 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, J (Eq. (). Dashed arrows denote amortization. Iterative amortization uses gradient feedback during optimization, while direct amortization does not.
amortization gap performance.
Equivalently, we can multiply by a, defining the variational RL objective as:
J 0) = E, [R((T] (
(6)
This objective consists of the expected return (i.e., the standard RL objective) and KL divergence between क (T|O) and p(r). In terms states and actions, this objective is:
T Σ r-z-a a log
T ( S.,, O) po (a, (Se)
I [ 0) = Exs a peppe
At a given timestep, L, one optimize this objective by estimating the future terms in the sum using a "soft" action-value (Q+) network (30) or model [62]. For instance, sampling St ~ Penv, slightly abusing notation, we can write the objective at time / as:
[ E, [Qx at) aDKu (=(al( O)llpa (a. s:)).
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, e.g., Gaussian mean, μu, and variance, o². In this case, policy optimization corresponds to maximizing:
X
arg max J (T,
(9)
Optionally, can then also learn the policy prior parameters, 0 1
2.3 Entropy & KL Regularized Policy Networks Perform Direct Amortization
Policy-based approaches RL typically do directly optimize the action distribution parameters, e.g., through gradient- optimization. Instead, the action distribution parameters are output by function approximator (deep network), fa. which is trained using deterministic 770 52 stochastic gradients [83] When combined with entropy or KL regularization, this policy network is form of amortized optimization [26), learning to estimate policies. Again, denoting the action distribution parameters, e.g., mean and variance, as X, for given state, s, we can express this direct mapping as
X < fo(s),
(direct amortization)
(10)
denoting the corresponding policy as (aja,,, O; A). Thus, fo attempts to learn to optimize Eq. 9 This setup is shown in Figure I (Right). Without entropy KL regularization, i.e. (als)) po (als). we can instead interpret the network as directly integrating the LHS of Eq. which is less efficient and more challenging. Regularization smooths the optimization landscape, yielding more stable improvement and higher asymptotic performance B].
Viewing policy networks as form of direct amortized variational optimizer (Eq. 10 allows us to see that they are similar "encoder" networks in variational autoencoders (VAEs) there several drawbacks to direct amortization.
However,

---

Algorithm 2 Iterative Amortization
Algorithm 1 Direct Amortization
Initialize ø
for each environment step do Initialize X
Initialize ø
for each environment step do
for each policy optimization iteration do
X 4 (st)
À f- (st X, VxJ) end for
a, ~ (a S,, O:
S Penv (St+1 St, ar)
3, ~ (ar 0; A)
end for
St ~ Pem (Sr+1 Sr, as)
for each training step do
end for for each training step do
Φ - 0 NV
end for
Φ + V
end for
Amortization Gap. Direct amortization results in suboptimal approximate posterior estimates, with the resulting gap in the variational bound referred to as the amortization gap (0]. Thus, in the RL setting, an amortized policy, Tg. results in worse performance than the optimal policy within the parametric policy class, denoted as î. The amortization gap is the gap following inequality:
[Fo, 0) S O).
Because is variational bound on the RL objective, i.e., expected return, a looser bound, due to amortization, prevents one from more completely optimizing this objective.
This is shown in Figure I (Left) where I is plotted over two dimensions of the policy mean at particular state in the MuloCo environment Hopper- vv.. The estimate of a direct amortized policy is suboptimal, far from the optimal estimate (*) While the relative difference in the objective is relatively small, suboptimal estimates prevent sampling and exploring high-value regions of the action-space. That is, suboptimal estimates have only minor impact evaluation performance (see Appendix B.6) but hinder effective data collection.
Single Estimate. Direct amortization limited single, static estimate. In other words, if there are multiple high-value regions of the action-space, uni-modal (e.g., Gaussian) direct amortized policy restricted to only one region, thereby limiting exploration. Note that this an additional restriction beyond simply considering uni-modal distributions, generic optimization procedure may arrive at multiple uni-modal estimates depending on initialization and stochastic sampling (see Section 3.2). While multi-modal distributions reduce the severity of this restriction 774 29]] the other limitations of direct amortization still persist.
Inability to Generalize Across Objectives. Direct amortization is feedforward procedure, receiv- ing gradients from the objective only after estimation. This is contrast to other forms of optimization, which receive gradients (feedback) during estimation. Thus, unlike other optimizers, direct amortiza- tion is incapable of generalizing to new objectives, e.g., if Qx (s, a) or (a)s) change, which is a desirable capability for adapting to new tasks or environments.
To improve upon this scheme and overcome these drawbacks, in Section we turn to a technique developed in generative modeling, iterative amortization [55], retaining the efficiency of amortization while employing more flexible iterative estimation procedure.
2.4 Related Work
Previous works have investigated methods for improving policy optimization. QT-Opt 41 uses
the cross-entropy method (CEM) [66], an iterative derivative-free optimizer, to optimize Q-value estimator for robotic grasping. CEM and related methods are also used in model-based RL. for
performing model-ppeeicctvvv control 6 [4 62 33
Gradient-based policy optimization [36] 71 10
contrast, less common, however, gradient- based optimization can also be combined with CEM
Additional 2D plots shown in Figure 19 in the Appendix.

---

0,5
Fulals 07)
0.5
3.5
1.0
63
CO tanilpp)
1L.
10
Dim.
Dim.
Figure 2: Estimating Multiple Policy Modes. Unlike direct amortization, which restricted to a single estimate, iterative amortization effectively sample from multiple high-value action modes. This is shown for a particular state 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.
Most policy-based methods use direct amortization, either using feedforward 31 or recurrent 28 network. Similar approaches have also been applied to model-based value estimates 13
16
4
as well as combining direct amortization with model predictive control [50) and planning
665
separate line of work has explored improving the policy distribution, using normalizing flows
29 74
and latent variables [6].. principle, iterative amortization can perform policy optimization each
of these settings.
Iterative amortized policy optimization is conceptually similar negative feedback control 77 using eπrors to update policy estimates. However, while conventional feedback control methods are often restricted in their applicabilityy e.g., linear systems and quadratic cost, iterative amortization is generally applicable to any differentiable control objective. This is analogous to the generalization of Kalman filtering
42 to amortized filtering [54] for state estimation.
3 Iterative Amortized Policy Optimization
3.1 Formulation
Iterative amortization (55 utilizes errors or gradients update the approximate posterior distribution parameters. While various forms exist, we consider gradient-encoding models [6] due to their generality. Compared with direct amortization (Eq. 10), use iterative amortized optimizers of the general form
A + fo A, VAJ))
(iterative amortization)
(11)
also shown in Figure (Right), where fo is a deep network and À are the action distribution parameters. For example, if = Na μt, diag(o²)). then X x= [J,, Technically, is redundant, as the state dependence is already captured in J. but this empirically improve performance [55). In
practice, the update is carried out using a "highway" gating operation [38] [2]. Denoting wa E [0. 1] as the gate and dg the update, both of which are output fo. the gating operation is expressed
A wp 3 X A (1 wws) @ 85.
(12)
where denotes dlement-wise multiplication. This update is typically run for fixed number of steps, and, as with a direct policy, the iterative optimizer is trained using stochastic gradient estimates of VaJ, obtained through the path-wise derivative estimator 44
64 35 Because the gradients
VAJ must be estimated online, i.e., during policy optimization, this scheme requires some way of
estimating I online through parameterized Q-value network (58 differentiable model 35
3.2 Benefits of Iterative Amortization
Reduced Amortization Gap. Iterative amortized optimizers are more flexible than their direct
counterparts, incorporating feedback from the objective during policy optimization (Algorithm) 2
rather than only after optimization (Algorithm [T]. Increased flexibility improves the accuracy of
optimization, thereby tightening the variational bound 55 54
We see this flexibility in Figure [
(Left), where an iterative amortized policy network iteratively refines the policy estimate (+), quickly arriving near the optimal estimate.
5

---

Multiple Estimates. Iterative amortization, 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, referred to as "multi-stability" in the genera- tive modeling literature (77). We illustrate this capability across two action dimensions in Figure 2 for a state in the Ant-v2 MuloCo 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, despite only using uni-modal policy distribution.
Generalization Across Objectives. Iterative amortization uses the gradients of the objective during optimization. i.e., feedback, allowing it to potentially generalize to new or updated objectives. We see this in Figure ( (Left). where iterative amortization, despite being trained with a different value estimator, is capable of generalizing to this new objective. We demonstrate this capability further in Section 4 This opens the possibility of accurately and efficiently performing policy optimization in new settings, e.g., a rapidly changing model or new tasks.
3.3 Consideration: Mitigating Value Overestimation
Why are more powerful policy optimizers typically not used in practice? Part of the issue stems from value overestimation. Model-free approaches generally estimate Qx using function approximation and temporal difference learning. However, this has the pitfall of value overestimation, i.c., positive bias in the estimate, @ [75]. 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 flexible policy optimizers exacerbate the problem, exploiting this uncertainty to greater degree. Further, a rapidly changing policy increases the difficulty of value estimation 63
Various techniques have been proposed for mitigating value overestimation in deep RL. The most prominent technique, double deep Q-network [81] maintains two Q-value estimates [80), attempting to decouple policy optimization from value estimation. Fujimoto et al. (25) apply and improve upon this technique for actor-critic settings, estimating the target Q-value as the minimum of two Q-networks, Qu and Qu:i
(s. a) = min (s, 11122
where vig denotes the target network parameters. As noted by Fujimoto al. 25 this not only counteracts value overestimation, but also penalizes high-variance value estimates, because the minimum decreases with the variance of the estimate. Ciosek al. [15) noted that, for bootstrapped ensemble of two Q-networks, the minimum operation can be interpreted as estimating
(s. a) (s a) (s,
(13)
with mean a) as) 4 (s a). standard deviation ( a) 1. (s,
Mo(s, a))3)1 and B = 1. Thus, to 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 3 = results in increased value bias and more rapidly changing policy as compared with direct amortization (Figure 3). Further pe- nalizing high-variance target values (3 2.5) reduces value overestimation and improves sta- bility. For details, see Appendix A.2 Recent techniques for mitigating overestimation have been proposed, such as adjusting a 2 In of- fline RL, this issue has been tackled through the action prior [44 48 84) or by altering Q- network training 2 49
Lavm J=1
300
Dua (ooeell
Millon Skope
Millan Stepe
(a)
(b)
Figure 3: Mitigating Value Overestimation.
While such techniques could be used here, increasing provides sim- ple solution with no additional computational overhead. This is a meaningful insight toward applying more powerful policy optimizers.
With 3 1, iterative amortization results in (a)
higher value overestimation and (b) a more rapidly changing policy as compared with direct amortiza- tion. Increasing ,9 helps to mitigate these issues.

---

Ope. Ihr.
Oyt. Ibe.
Dan.
11.5
0.0
fr2
Adiun
-0.5
550
Amaot.
Envioonmen
10
Time Step
21
Environment Time Step
10
15
20
25
Opt.
100
200
(a) Policy
(b) Improvement
(c) Comparison
Figure 4: Policy Optimization. Visualization over time steps of (a) one dimension of the policy distribution and (b) the improvement in the objective, AJ. across policy optimization iterations. (c) Comparison of iterative amortization with Adam [45] grradient-basedd and CEM [66] (gradient-free).
Iterative amortization is an order of magnitude more efficient.
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 332 This uses two Q-networks, uniform action prior, Pa(a)s) = 4(-1, 1), and a tuning scheme for the temperature, In our experiments, "direct" refers 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.1b. Unless otherwise stated, we use 5 iterations per time step for iterative amortization,
following 55 For details, refer to Appendix and Haarnoja et
[31 32)
4.2 Analysis
4.2.1
Visualizing Policy Optimization
We provide 2D visualizations of iterative amortized policy optimization in Figures [ & 2 with additional 2D plots in Appendix [B.5]] In Figure we visualize iterative refinement using a single action dimension from Ant-v2 across time steps. The refinements in Figure 44a give rise to the objective improvements in Figure 4b We compare with Adam [45] (gradient-based) and CEM [66]
(gradient-free) Figure 4c where iterative amortization an order of magnitude more efficient. This trend is consistent across environments, as shown in Appendix B.4
4.2.2 Performance Comparison
We evaluate iterative amortized policy optimization on the suite of MuJoCo (78 continuous control tasks from OpenAl gym [12). In Figure we compare the cumulative reward of direct and iterative amortized policy optimization across environments. Each curve shows the mean and + standard deviation of random seeds. In all cases, iterative amortized policy optimization matches or outperforms the baseline direct amortized method, both in sample efficiency and final performance. Iterative amortization also yields more consistent, lower variance performance.
4.2.3 Improved Exploration: Multiple Policy Modes
As described in Section 3.2 iterative amortization is capable of obtaining multiple estimates, i.c., multiple modes of the optimization objective. To confirm that iterative amortization has captured multiple modes, at the end of training, we take an iterative agent trained on Walker2d-v2 and histogram the distances between policy means across separate runs of policy optimization per state (Fig. 7a). For the state with the largest distance, we plot 2D projections of the optimization objective, J. across action dimensions in Figure 7b well the policy density across 10 optimization (Fig. (T). The modal policy optimization surface shown 7 resuuls in the multi-modal policy Figure 7c Additional results on other environments are presented in Appendix B.7
To better understand whether the performance benefits of iterative amortization are coming purely from improved exploration via multiple modes, we also compare with direct amortization with
7

---

Clanixbee Rovved
6030
Cumllabee Rosunl
230.01
Cummaleee: Desunl
-300
Cumbatae
100.0
b000
-00
63
04
3t1lio
Steps
Reacher-v2
(b) Hopper-v2
HalfCheetah--2
Swinmer-92
Cantattiv Rooael
Oanbaag Roodd
Onnaattee Roward
Canulaaton Rowaed
2000
(v
1000
2000
dest I8AC)
Renrhhe
Milbon
Millon
Millon Strpe
3011oo
(e) Valzer2d-r2
(1) Ant-v2
Hunanoid-v2
(h) HumenoidStandduprr
Figure 5: Performance Comparison. Iterative amortized policy optimization performs comparably
with or better than direct amortization across MuJoCo environments. Curves show the mean +
std. dev. of performance over random seeds.
Auseiaatte Cap
Auseiiattia Cap
Aovettuaion sasstaatio Cap 101
Auoettattie Cap
(SAC)
Benohive libes
Stepe
Stepe
(a) Hopper-v2
(b) BalrCheatab-v2
Valker2d-v2
(d)
Ant-v2
Figure 6: Amortization Gap. Iterative amortization achieves similar or lower amortization gaps
than direct amortization. Gaps estimated using stochastic gradient-bart based optimization over 100 random states. Curves show the mean and + std. dev. over random seeds.
multi-modal policy distribution. This formed using inverse autoregressive flows 46 type of normalizing flow (NF). Results are presented in Appendix B.2] Using multi-modal policy reduces performance deficiencies on Hopper-v2 and Walker2d-v2, indicating that much of the benefit of iterative amortization is due to lifting direct amortization's 's restriction to single, uni-modal policy estimate. Yet, direct + NF still struggles on HalfCheetah-v2 compared with iterative amortization, suggesting that more complex, multi-modal distributions are not the only consideration.
4.2.4 Improved Optimization: Amortization Gap
To evaluate policy optimization accuracy, we estimate per-step amortization gaps, performing ad- ditional iterations of gradient ascent on J w.r.t. the policy parameters, À = [14, or] (see Appendix A.3 To analyze generalization, we also evaluate the iterative agents trained with iterations for an additional 5 amortized iterations. Results are shown in Figure 6 We emphasize that it is challenging 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 -eetworks, which may be biased (Figure Nevertheless, we find that iterative amortized policy optimization achieves, on average, lower amortization gaps than direct amortization across all environments. Additional amortized iterations at evaluation yield further estimated improvement, demonstrating generalization beyond the optimization horizon used during training.
The amortization gaps are small relative the objective, playing a negligible role in evaluation performance (see Appendix B.6). Rather, improved policy optimization is helpful for training. allowing the agent to explore states where value estimates are highest. To probe this further, we train iterative amortized policy optimization while varying the number of iterations per step in {1, 21 5}. yielding optimizers with varying degrees of accuracy. Note that each optimizer is, in theory, capable of finding multiple modes. In Figure 8 we see that training with additional iterations improves performance and optimization accuracy. We stress that the exact form of this relationship depends on the Q- value estimator and other factors. We present additional results in Appendix B.6

---

10
T,AKK Ch
0.0
banhs (1-0-- bath 1e0l
0.7
Das
Dec.
Das.
Dixx.
1im.
Din.
Daa
(a)
(b)
(c)
Figure 7: Multiple Policy Modes. (a) Histogram of distances between policy means (1)) across optimization runs (i and j) over seeds and states on Walker2d-v2 at 3 million environment steps. For the state with the largest distance, (b) shows the projected optimization surface on each pair of action dimensions, and (c) shows the policy density for 10 optimization runs.
Convlettves Reword
L3000
Cap
12000)
10000)
ds
Creallatiee
MP MB
5000
2008
MB
Million Stopo
Million Stope
Million
Figure Varying Iterations During Training. Performance (Left) and estimated amortization gap (Right) for varying num- bers of policy optimization iterations per step during training on SalfCheetah-v2. Increasing the iterations generally improves performance and decreases the estimated amortization gap.
Figure 9: Zero-shot general- ization of iterative amortiza- tion from model-free (MF) to model-based (MB) value esti- mates.
4.2.5 Generalizing to Model-Based Value Estimates
Direct amortization is a purely feedforward process and therefore incapable of generalizing to new objectives. In contrast, because iterative amortization is formulated through gradient-basee feedback, such optimizers may be capable of generalizing to new objective estimators, as shown in Figure To demonstrate this capability further, we apply iterative amortization with model-based value estimators. using a learned deterministic model on SalfCheetahhv2 (see Appendix A.5). We evaluate the generalizing capabilities in Figure 9 by transferring the policy optimizer from model- free agent to model-based agent. Iterative amortization generalizes to these new value estimates, instantly recovering the performance the model- basel agent. This highlights the opportunity for instantly incorporating new tasks, goals, or model estimates into policy optimization.
5 Discussion
We have introduced iterative amortized policy optimization, flexible and powerful policy optimiza- tion technique. In so doing, we have identified KL-regularized policy networks as form of direct amortization, highlighting several limitations: 1) limited accuracy, quantified by the amortization gap. 2) restriction to single estimate, limiting exploration, and 3) inability to generalize to new objectives, limiting the transfer of these policy optimizers. As shown through our empirical analysis, iterative amortization provides step toward improving each of these restrictions, with accompanying improvements in performance over direct amortization. Thus, iterative amortization can serve as drop-in replacement and improvement over direct policy networks in deep RL.
This improvement, however, is accompanied by added challenges. As highlighted in Section 3.3 improved policy optimization can exacerbate issues Q-value estimation stemming from positive bias. Note that this is not unique to iterative amortization, but applies broadly to any improved optimizer. We have provided simple solution that involves adjusting a factor, 3, counteract this bias. Yet, we see this as an area for further investigation, perhaps drawing on insights from the offline RL community 49 In addition to value estimation issues, iterative amortized policy optimization incurs computational costs that scale linearly with the number of iterations. This is in comparison with direct amortization, which has constant computational cost. Fortunately, unlike standard optimizers, iterative amortization adaptively tunes step sizes. Thus, relative improvements

---

rapidly diminish with each additional iteration, enabling accurate optimization with exceedingly few iterations. practice, even a single iteration per time step can work surprisingly well.
Although we have discussed three separate limitations of direct amortization, these factors are highly interconnected. By broadening policy optimization to iterative procedure, automatically obtain a potentially more accurate and general policy optimizer, with the capability of obtaining multiple modes. While our analysis suggests that improved exploration resulting from multiple modes the primary factor affecting performance, future work could tease out these effects further and assess the relative contributions of these improvements in additional environments. We are hopeful that iterative amortized policy optimization, by providing a more powerful, exploratory, and general optimizer, will enable a range of improved RL algorithms.
Acknowledgments and Disclosure of Funding
JM acknowledges Scott Fujimoto for helpful discussions. This work was funded in part by NSF # #1918839 and Beyond Limits. JM is currently employed by Google DeepMind. The authors declare no other competing interests related to this work.
References
[1] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. In International Conference on Learning Representations, 2018
[2] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective on offline reinforcement learning. In International Conference Machine Learning, pages 104-114, 2020.
[3] Zafarali Ahmed, Nicolas Le Roux. Mohammad Norouzi. and Dale Schuurmans. Understanding the impact of entropy policy optimization. In International Conference on Machine Learning. pages 151-160, 2019.
[4] Brandon Amos, Samuel Stanton, Denis Yarats, and Andrew Gordon Wilson. On the model-based stochastic value gradient for continuous reinforcement learning. In Learning for Dynamics and Control, pages 6-20. PMLR, 2021.
[5] Brandon Amos and Denis Yarats. The differentiable cross-entropy method. In International Conference on Machine Learning, 2020.
[6] Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffman, David Pfau, Tom Schaul, and Nando de Freitas. Learning to learn by gradient descent by gradient descent. In Advances in Neural Information Processing Systems, pages 398 3981-3989, 2016.
[7] Karl Johan Astrom and Richard M Murray. Feedback Systems: An Introduction for Scientists and Engineers. Princeton University Press, 2008
[8] Hagai Attias. Planning by probabilistic inference. In AISTATS. Citeseer, 2003.
[9] Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. NewrIPS Deep Learning Symposium, 2016.
[10] Homanga Bharadhwaj,, Kevin Xie, and Florian Shkurti, Model-predictive planning via cross- entropy and gradient-based optimization. In Learning for Dynamics and Control, pages 277-286. 2020.
[11] Matthew Botvinick and Mare Toussaint. Planning as inference. Trends in cognitive sciences, 16(10):485-88, 2012.
[12] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. arXiv preprint arXiv:1606.015440 2016.
10

---

[13] Arunkumar Byravan, Jost Tobias Springenberg. Abbas Abdolmaleki, Roland Hafner, Michael Neunert, Thomas Lampe, Noah Siegel, Nicolas Heess, and Martin Riedmiller. Imagined value gradients: Model-based policy optimization with tranferable latent dynamics models. In Conference on Robot Learning, pages 566-589, 2020.
[14) Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems, pages 4754-4765, 2018.
[15] Kamil Ciosek, Quan Vuong, Robert Loftin, and Katja Hofmann. Better exploration with optimistic actor critic. In Advances in Neural Information Processing Systems, pages 1787- 1798, 2019.
[16] Ignasi Clavera, Yao Fu, and Picter Abbeel. Model-augmented actor-critic: Backpropagating through paths. In International Conference on Learning Representations, 2020.
[17] Djork-Amé Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep net- work learning by exponential linear units (clus). In International Conference on Learning Representations, 2016.
[18] Comet.ML. Co home page, 2021.
[19] Gregory F Cooper. A method for using belicf networks influence diagrams. In Fourth Workshop on Uncertainty in Artificial Intelligence, 1988.
[20] Chris Cremer, Xuechen Li, and David Duvenaud. Inference suboptimality in variational autoencoders. In International Conference on Machine Learning, pages 1078-1086, 2018.
[21] Peter Dayan and Geoffrey E. Hinton. Using expectation-maxiiizatiie for reinforcement learning. Neural Computation, 9())2711-278, 1997.
[22] Roy Fox. Toward provably unbiased temporal-difference value estimation. Optimization Foundations for Reinforcement Learning Workshop at NeurlPS, 2019.
[23] Roy Fox, Ari Pakman, and Naftali Tishby. Taming the noise in reinforcement learning via soft updates. In Proceedings of the Thirty-Second Conference on Uncertainty in Artificial Intelligence, pages 202-211. AUAI Press, 2016.
[24] Scott Fujimoto, David Meger, and Doina Precup. Off-pooicy deep reinforcement learning without exploration. In International Conference on Machine Learning, pages 2052-2062, 2019.
[25] Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In International Conference on Machine Learning, pages 1587-1596, 2018.
[26] Samuel Gershman and Noah Goodman. Amortized inference in probabilistic reasoning. In Proceedings of the Cognitive Science Society, volume 36, 2014.
[27] Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner. Multi-object repre- sentation learning with iterative variational inference. In International Conference on Machine Learning, pages 2424-2433, 2019.
[28] Arthur Guez, Mehdi Mirza, Karol Gregor, Rishabh Kabra, Sebastien Racaniere, Theophane Weber, David Raposo, Adam Santoro, Laurent Orseau, Tom Eccles, al. An investigation of model-free planning. In International Conference on Machine Learning, pages 2464-244 2019.
[29] Tuomas Haarnoja, Kristian Hartikainen, Pieter Abbeel, and Sergey Levine. Latent space policies for hierarchical reinforcement learning. International Conference on Machine Learning, pages 1851-1860, 2018.

---

[30] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep mueey--based policies. In International Conference on Machine Learning, pages 1352-1361. JMLR. org, 2017.
[31] 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 Conference Muchine Leurning, pages 1856-1865, 2018.
[32] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and applications. arXiv preprint 1rXfv:1812.05900. 2018.
[33] 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, pages 2555-2565. PMLR, 2019.
[34] Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor. Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane. Jaime Fer- nández del Río, Mark Wiebe, Pearu Peterson, Pierre Kérard-Marchant Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-362 September 2020.
Kevin Sheppard, Tyler
[35] Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems, pages 2944-2952, 2015.
[36] Mikael Henaff, William F Whitney, and Yann LeCun. Model-based planning with discrete and continuous actions. arXiv preprint 170vv1705..0177, 2017.
[37] Devon Hjelm, Ruslan R Salakhutdinov, Kyunghyun Cho, Nebojsa Jojic, Vince Calhoun, and Junyoung Chung. Iterative refinement of the approximate posterior for directed belief networks. In Advances in Neural Information Processing Systems, pages 4691-4699, 2016.
[38] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 18):1735- 1997.
[39] D. Hunter. Matplotli: A 2d graphics environment. Computing in Science & Engineering, 9(3):90 95, 2007.
[40] 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, pages 12519-12538 12530, 2019.
[41] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, 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. In Conference on Robot Learning, 2018
[42] Rudolph Emil Kalman. A new approach to linear filtering and prediction problems. Journal of Basic Engineering, 82(1):335-55 45, 1960.
[43] Yoon Kim, Sam Wiseman, Andrew C Miller, David Sontag, and Alexander M Rush. Semi- amortized variational autoencoders. In International Conference on Machine Learning, 2018.
[44] Diederik P Kingma and Max Welling, Stochastic gradient vb and the variational auto-encoder. In International Conference on Learning Representations, 2014.
[45] Durk P Kingma and Jimmy Ba. Adam: method for stochastic optimization. In International Conference on Learning Representations, 2015.
[46] Durk P Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling. Improved variational inference with inverse autoregressive flow. In Advances Neural Infor- mation Processing Systems, pages 4743-4751, 2016.
12

---

[47] Rahul G Krishnan, Dawen Liang, and Matthew Hoffman. On the challenges of learning with inference networks on sparse, ligh-dimensional data. In International Conference on Artificial Intelligence and Statistics, pages 143-151, 2018.
[48] 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 Systems, pages 11784-11794, 2019.
[49] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems, pages 1179-1191, 2020.
[50] Keuntaek Lee. Kamil Saigol, and Evangelos A Theodorou. Safe end-to-end imitation learning for model predictive control. In International Conference on Robotics and Automation, 2019.
[51] Sergey Levine. Reinforcement learning and control as probabilistic inference: Tutorial and review. arXiv preprint arXiv: htt.:::800000000 2018.
[52] Timothy P Lillicrap, Jonathan Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning, In International Conference on Learning Representations, 2016.
[53] Long-Ji Lin. Self-improving reactive agents based on reinforcement learning, planning and teaching. Machine learning, 8(3-4):2933321, 1992.
[54] Joseph Marino, Milan Cvitkovic, and Yisong Yue. general method for amortizing variational filtering. In Advances in Neural Information Processing Systems, 2018.
[55] Joseph Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Conference on Machine Learning, pages 3403-3412, 2018.
[56] Andriy Mnih and Karol Gregor. Neural variational inference and learning in belief networks. In International Conference on Machine Learning, pages 1791-1799, 2014.
[57] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep rein- forcement learning. In International Conference Machine Learning, pages 1937, 2016.
[58] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning. In NewrIPS Deep Learning Workshop, 2013.
[59] Rémi Munos, Tom Stepleton, Anna Harutyunyan, and Mare Bellemare. Safe and efficient off-policy reinforcement learning. In Advances in Neural Information Processing Systems, pages 1054-1062, 2016.
[60] 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 Interna- tional Conference on Robotics and Automation, pages 7559-7566. IEEE, 2018.
[61] Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in Neural Information Processing Systems, pages 8026-8037, 2019.
[62] Alexandre Piché, Valentin Thomas, Cyril Ibrahim, Yoshua Bengio, and Chris Pal. Probabilistic planning with sequential monte carlo methods. In International Conference Learning Representations, 2019. 2019.
[63] Aravind Rajeswaran, Igor Mordatch, and Vikash Kumar. A game theoretic framework for model based reinforcement learning. International Conference on Machine Learning, pages 7953-7963, 2020.
13

---

[64] Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic backpropagation and approximate inference in deep generative models. In International Conference on Machine Learning, pages 1278-1286, 1286, 2014.
[65] Benjamin Rivière, Wolfgang Hönig, Yisong Yue, and Soon-Jo Chung. Glas: Global-to-locu safe autonomy synthesis for multi-robot motion planning with end-to-end learning. IEEE Robotics and Automation Letters, (3)):4249-45555 2020.
[66] Reuven Y Rubinstein and Dirk P Kroese. The cross-entropy method: a unified approach to combinatorial optimization, Monte-Carlo simulation and machine learning. Springer Science & Business Media, 2013.
[67] John Schulman, Sergey Levine, Pieter Abbeel, Michael Jordan, and Philipp Moritz. Trust region policy optimization. In International Conference on Machine Learning, pages 1889-1897 2015.
[68] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:/707.006 06647, 2017.
[69] David Silver, Aja Huang, Chris J Maddison, Arthur Guez, Laurent Sifre, George Van Den Driess- che, Julian Schrittwieser, Ioannis Antonoglou, Veda Panneershelvam, Mare Lanctot, et al. Mas- tering the game of go with deep neural networks and tree search. Nature, 529(7587):484 489, 2016.
[70] David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In International Conference Machine Learning, pages 387-395, 2014.
[71] 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, pages 4732-4741, 2018.
[72] Rupesh K Srivastava, Klaus Greff, and Jürgen Schmidhuber. Training very deep networks. In Advances in Neural Information Processing Systems, pages 2377-2385, 2015.
[73] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
[74] Yunhao Tang and Shipra Agrawal. Boosting trust region policy optimization by normalizing flows policy. In NeurlPS Deep Reinforcement Learning Workshop, 2018.
[75] 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.
[76] 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.07438, 2019.
[77] Emanuel Todorov. General duality between optimal control and estimation. In Decision and Control, 2008 CDC 2008. 47th IEEE Conference on, pages 4286-4292. IEEE, 2008
[78] Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In International Conference on Intelligent Robots and Systems, pages 5026-5033. IEEE, 2012.
[79] Mare Toussaint and Amos Storkey. Probabilistic inference for solving discrete and continuous state markov decision processes. In International Conference on Machine Learning, pages 945- 2006.
[80] Hado Van Hasselt. Double q-learning. In Advances in Neural Information Processing Systems, pages 2613-2621, 2010.
[81] Hado Van Hasselt, Arthur Guez, and David Silver. Deep reinforcement learning with double g-learning In AAAI Conference on Artificial Intelligence, 2016.
14

---

[82] Michael L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software, 6(60):3021, 2021.
[83] Ronald J Williams. Simple statistical gradient-following algorithms for connectionist reinforce- ment learning. Reinforcement Learning. pages 5-32. Springer, 1992.
[84] Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11611, 2019.
Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope? [Yes] We demonstrate performance improvements and novel benefits of iterative amortization in Section Each claim is supported by empirical evidence.
(b) Did you describe the limitations of your work? [Yes] We discuss the additional computation requirements and the necessity of mitigating value overestimation.
(c) Did you discuss any potential negative societal impacts of your work? [No] We do not see any immediate societal impacts of this work beyond the general impacts of improvements in machine learning.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them? [Yes]
2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [N/A]
(b) Did you include complete proofs of all theoretical results? [N/A]
3. If you experiments..
(a) Did you include the code, data, and instructions needed to reproduce the main ex- perimental results (either in the supplemental material as URL)? [Yes] We have included code the supplementary material with an accompanying README file.
(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)? [Yes] All details and hyperparameters
provided in the Appendix. (c) Did you report error bars (e.g., with respect to the random seed after running exper- iments multiple times)? [Yes] We report error bars on each of our main quantitative results comparing performance.
(d) Did you include the total amount of compute and the type of resources used (c.g., type of GPUs, internal cluster, or cloud provider)? [Yes] We provide these details in the supplementary material.
4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [Yes] We cite the creators of the environments and software packages used in this paper.
(b) Did you mention the license of the assets? [Yes] We mention cite the license for
each asset.
(c) Did you include any new assets either in the supplemental material or as a URL? [N/A]
(d) Did you discuss whether and how consent was obtained from people whose data you're using/curating? [N/A]
(c) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if applicable? [N/A]
(b) Did you describe any potential participant risks, with links Institutional Review Board (IRB) approvals, if applicable? [N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amount spent participant compensation? [N/]]
15


Reviews:
[{"review_id": "qX4uPo4ByVV", "reviewer": "Reviewer_fdJ1", "summary": "This paper proposes iterative amortized policy optimization, a class of methods that use the gradients online to optimize the policy iteratively.", "questions": "", "limitations": "", "rating": 6, "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": "This paper proposes iterative amortized policy optimization, a class of methods that use the gradients online to optimize the policy iteratively. To the best of my knowledge, the idea is novel and original. The paper is generally well-written and easy to follow and in particular, I appreciate the authors making a detailed, aligned, and symmetric comparison of direct amortization and iterative amortization (the one-to-one pros and cons, as well as the graphical demonstration, and the algorithm, etc...). The empirical experiments are quite comprehensive, which includes quite some qualitative analysis of the claimed properties. And indeed, the proposed method is validated by the learning curves in MuJoCo continuous control tasks benchmark, as the proposed method outperforms and matches direct amortization in most continuous control tasks. Some questions and/or suggestions: 1. could you talk more about what's the added complexity of the proposed method? 2. In Fig.5 (c), while direct amortization seems to be very stable, iterative amortization seems to have two moments with sudden performance collapse, is there any reason why? 3. the increase on performance gains on MuJoCo tasks (Fig.5) and the corresponding decrease in amortization gap (Fig.6) do not seem to be be consistent, e.g. there is a significant decrease in amortization gap on Ant, while the expected reward is almost the same. So maybe it's good to investigate what brings the performance gains? 4. for exploration, it's also helpful to test the proposed algorithm in MuJoCo tasks with sparse rewards.\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See the main review.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "10"}, "extra_text": {}}, {"review_id": "65v7oyiXRbt", "reviewer": "Reviewer_Cxw1", "summary": "This work analyzes policy optimization from an inference perspective. It first shows that policy optimization with entropy or KL regularization in DRL with a policy network is a form of amortized optimization where instead of directly optimizing the action distribution parameters, we optimize a network which outputs the action distribution parameters. This results in an amortization gap which negatively impacts performance and exploration. Next the authors propose an alternative iterative amortization method for training the policy network. This new approach is shown to have various benefits which the authors demonstrated through a series of experiments on the MuJoCo environments.", "questions": "", "limitations": "", "rating": 8, "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": "Overall, I found this paper to be very insightful. It brings forward many interesting ideas which helps to bridge the gap between inference and RL. The paper was very easy to read. The authors did a good job of explaining the issues of direct amortization and presenting their novel technique. The diagrams were also very helpful! The discussion on amortization gap is particularly interesting and to the best of my knowledge not something that has previously been explored by the RL community. The experiments in this paper are very well designed in explaining the iterative amortization approach. I have two questions for the authors which I hope they can clarify in their response:\n\n- In the section on value overestimation (Section 3.3), you mentioned that “If the policy can exploit regions of high uncertainty, the resulting target values will introduce positive bias into the estimate. More flexible policy optimizers exacerbate the problem, exploiting this uncertainty to a greater degree.” Could you elaborate on this claim a little further? Why are more powerful policies more prone to overestimation of the Q function? This seems to be a fairly central point to Section 3.3.\n- Empirical performance of the iterative amortization approach does not seem to be that strong. While there are some improvements on some of the simpler environments. Performance on the most challenging MuJoCo environments (Ant, Humanoid, HumanoidStandup) appear to be on par with the direct amortization method. Any sense why the iterative approach does not result in any improvement on these environments?\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Yes", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4-5"}, "extra_text": {}}, {"review_id": "5faDgw3kUQ", "reviewer": "Reviewer_CT92", "summary": "This paper introduces a novel approach to policy optimization which leverages the connection between amortized variational inference and policy-based RL with function approximation. Specifically, rather than rely on a \"direct\" parameterization of the state-conditioned distribution over actions, in which an input state is mapped directly to distribution parameters (e.g., mean and variance), the approach adds an additional optimization procedure over the distribution parameters themselves, rather than those of the policy network. Experimental results are then used to demonstrate the benefits of this approach for performance, accuracy, and flexibility in the optimization process. ", "questions": "", "limitations": "", "rating": 7, "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": "Originality: The application of iterative amortization to policy optimization for RL is interesting and novel to the best of my knowledge. The connections between variational inference and RL with function approximation are deep yet not fully explored, and this is an excellent example of the interesting insights that can be gained through the lens of this relationship. The relationship to previous work and associated citations appear thorough. \n\nQuality: The submission appears technically sound, and the claims seem to be supported. However, it would be helpful if the authors discussed the tradeoffs inherent in iterative amortization (e.g., how is performance improvement balanced by increased computational expense?). Additionally, iterative amortization appears to help on some tasks and not at all on others. What properties of an environment make iterative amortization more or less useful? How do these properties vary among the tasks used in the paper? These are important questions from both a practical and theoretical standpoint. The analysis of the contribution of multimodality to performance is likely a good first step in this direction. There are also a number of supporting results in the appendix. Minor: For eq. 2, should there be a $\\rho(s_1)$ prefixed to the righthand side of the equation? And shouldn't the upper limit of the products in both eqs. 1 and 2 be $T-1$? \n\nClarity: The paper is mostly clear and well-organized. The background section and discussion of value overestimation were particularly enjoyable to read. Minor: Figure 7b is difficult to parse--I'm unclear as to what takeaway is intended. It also would be helpful to mark the exact section of the appendix in which supplementary results appear. \n\nSignificance: The magnitude of the performance gain on the majority of Mujoco tasks is not significant. It's difficult to gauge the significance of reducing the amortization gap--why this gap appear to grow over the course of training (Figure 8, right)?  While the raw performance improvement due to the proposed approach is not generally notable, I think the main significance of this work is captured by results like the ability to flexibly adapt to new objectives, such as to model-based value estimation. Namely, its significance lies in its novelty and the potential for follow-up work. Some theoretical analysis would be helpful in developing a deeper understanding. \n\n\n\nGiven the above factors, I recommend acceptance. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "See the main review for a discussion of limitations and suggestions for improvement. I do not believe there are obvious societal impacts for this work. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "7 hours"}, "extra_text": {}}, {"review_id": "4Vfdb53YyZD", "reviewer": "Reviewer_Embw", "summary": "The authors, via a standard reinterpretation of reinforcement learning as an inference problem, discusses the resulting “amortization gap” that results from using policy networks to optimize the variational lower bound. In order to close this gap, the authors propose to train an iterative procedure parameterized by $\\phi$, which uses estimated gradients of the return to iteratively refine $\\pi(\\cdot | s_t)$ for the current state $s_t$. They then provide an extensive experimental evaluation demonstrating their proposed method and several desirable properties, including adaptability to gradients estimated by different methods.", "questions": "", "limitations": "", "rating": 7, "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": "**Originality:**  \nThe authors leverage a known connection between variational inference and reinforcement learning to motivate their problem (which is known in the VI literature). Although the techniques are not novel, they are novel in their application to the RL setting.\n\n**Quality:**  \nThe submission is technically sound, with authors leveraging known techniques from the generative modeling literature. Furthermore, the authors present a very thorough experimental evaluation of their claims, with very interesting results. Although the performance gains in some of the tasks are not substantial, the demonstrated reduction in the amortization gap, inherent multi-modal behavior, and transferability between different value estimation methods are worthy of note.\n\nAs an additional experiment, it would be interesting to see whether the learned policy is robust to slight changes in the environment (e.g. changing weight/friction parameters of the MuJoCo simulation) as an extension of the generalization experiment.\n\n**Clarity:**  \nThe paper is mostly clear, although as someone who is not as familiar with the generative modeling literature, a few parts of the paper were a bit confusing from my perspective.\nFor example, in the context of RL, what is estimation referring to in lines 112-116 (I assume in this case that it is the process of choosing an action distribution)? I think it would be really helpful to have these connections explicitly mentioned.\n\n**Significance:**  \nThe paper introduces really interesting ideas, including having a trained policy which includes a trained optimizer that improves upon an initial action based on predicted returns. The work suggests several venues for future work, including across-task generalization, better policies for existing algorithms, among others.\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The authors have adequately addressed the limitations of their approach. I particularly appreciated that the authors spent time addressing the value overestimation problem that is inevitably exacerbated by having a more expressive policy class.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}]

Author rebuttal:

Meta-review: All the reviewers think introduce the advanced variational inference for making the policy flexible by exploiting the connection between RL and variational inference is interesting. 

Although I recommend the acceptance for this submission based on the reviewers feedback, there are several issues should be extensively discussed:

- EBM closed-form for policy parametrization: in fact, with entropy-regularization, the policy will lie in energy-based model. Therefore, a natural competitor is advanced sampling algorithms, e.g., Langevin, HMC, SVGD [1], Wasserstein flow [2] etc, based on learned EBM. 

- Performance tradeoff in computation, sample complexity, vs flexibility:  as the complicated parametrization introduced, the computation cost and sample complexity will increase. This should be carefully discussed, as almost all reviewers pointed out. From the empirical study, it seems the flexible policy family does not provide significant benefits, which diminishes the significance of this paper. 

In sum, I think this paper is interesting and could be better if the authors carefully address these questions. 
 
[1] Haarnoja, Tuomas, Haoran Tang, Pieter Abbeel, and Sergey Levine. "Reinforcement learning with deep energy-based policies." In International Conference on Machine Learning, pp. 1352-1361. PMLR, 2017.
[2] Zhang, Ruiyi, Changyou Chen, Chunyuan Li, and Lawrence Carin. "Policy optimization as wasserstein gradient flows." In International Conference on Machine Learning, pp. 5737-5746. PMLR, 2018.
Decision: Accept (Poster)