Original: B ACKPROPAGATION THROUGH THE V OID : O PTIMIZING CONTROL VARIATES FOR BLACK - BOX GRADIENT ESTIMATION
Revised: B ACKPROPAGATION THROUGH THE V OID : O PTIMIZING CONTROL VARIATES FOR BLACK - BOX GRADIENT ESTIMATION

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Gradient-based optimization is the foundation of deep learning and reinforcement learning. Even when the mechanism being optimized is unknown or not differentiable, optimization using high-variance or biased gradient estimates is still often the best strategy.
Revised: Gradient-based optimization is the foundation of deep learning and reinforcement learning, but is difﬁcult to apply when the mechanism being optimized is unknown or not differentiable.

Original: We introduce a general framework for learning low-variance, unbiased gradient estimators for black-box functions of  random variables, based on gradients of a learned function. These estimators can be jointlytrained with model parameters or policies, and are applicable in both discrete and continuous settings.
Revised: We introduce a general framework for learning low-variance, unbiased gradient estimators, applicable to black-box functions of discrete or continuous random variables. Our method uses gradients of a surrogate neural network to construct a control variate, which is optimized jointly with the original parameters.

Original: Wegive unbiased, adaptive analogs of state-of-the-art reinforcement learning methods such as advantage actor-critic. We also demonstrate this framework for training discrete latent-variable models. 
Revised: We demonstrate this framework for training discrete latent-variable models. We also give an unbiased, action-conditional extension of the advantage actor-critic reinforcement learning algorithm.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Gradient-based optimization has been key to most recent advances in machine learning and reinforcement learning.
Revised: Gradient-based optimization has been key to most recent advances in machine learning and reinforcement learning.

Original: The back-propagation algorithm (Rumelhart & Hinton, 1986), also known as reverse-mode automatic differentiation (Speelpenning, 1980; Rall, 1981) computes exact gradients of deterministic, differentiable objective functions.
Revised: The back-propagation algorithm (Rumelhart & Hinton, 1986), also known as reverse-mode automatic differentiation (Speelpenning, 1980; Rall, 1981) computes exact gradients of deterministic, differentiable objective functions.

Original: The reparameterization trick (Williams, 1992;
Revised: The reparameterization trick (Williams, 1992;

Original: Kingma & Welling, 2014; Rezende et al., 2014) allows backpropagation to give unbiased, lowvariance estimates of gradients of expectations of continuous random variables.
Revised: Kingma & Welling, 2014; Rezende et al., 2014) allows backpropagation to give unbiased, lowvariance estimates of gradients of expectations of continuous random variables.

Original: This has allowed effective stochastic optimization of large probabilistic latent-variable models.
Revised: This has allowed effective stochastic optimization of large probabilistic latent-variable models.

Original: Unfortunately, there are many objective functions relevant to the machine learning community for which backpropagation cannot be applied.
Revised: Unfortunately, there are many objective functions relevant to the machine learning community for which backpropagation cannot be applied.

Original: In reinforcement learning, for example, the function being optimized is unknown to the agent and is treated as a black box (Schulman et al., 2015).
Revised: In reinforcement learning, for example, the function being optimized is unknown to the agent and is treated as a black box (Schulman et al., 2015a).

Original: Similarly, when ﬁtting probabilistic models with discrete latent variables, discrete sampling operations create discontinuities giving the objective function zero gradient with respect to its parameters.
Revised: Similarly, when ﬁtting probabilistic models with discrete latent variables, discrete sampling operations create discontinuities giving the objective function zero gradient with respect to its parameters.

Original: Much recent work has been devoted to constructing gradient estimators for these situations.
Revised: Much recent work has been devoted to constructing gradient estimators for these situations.

Original: In reinforcement learning, advantage actor-critic methods (Sutton et al., 2000) give unbiased gradient estimates with reduced variance obtained by jointly optimizing the policy parameters with an estimate of the value function.
Revised: In reinforcement learning, advantage actor-critic methods (Sutton et al., 2000) give unbiased gradient estimates with reduced variance obtained by jointly optimizing the policy parameters with an estimate of the value function.

Original: In discrete latent-variable models, low-variance but biased gradient estimates can be given by continuous relaxations of discrete variables (Maddison et al., 2016; Jang et al., 2016).
Revised: In discrete latent-variable models, low-variance but biased gradient estimates can be given by continuous relaxations of discrete variables (Maddison et al., 2016; Jang et al., 2016).

Original: A recent advance by Tucker et al.
Revised: A recent advance by Tucker et al.

Original: (2017) used a continuous relaxationto construct a control variate for functions of discrete random variables. Low-variance estimates of the expectation of the control variate can be computed using the reparameterization trick to produce an unbiased estimator with lower variance than previous methods.
Revised: (2017) used a continuous relaxation of discrete random variables to build an unbiased and lower-variance gradient estimator, and showed how to tune the free parameters of these relaxations to minimize the estimator’s variance during training.

Original: Furthermore, Tucker et al.

Original: (2017) showed how to tune the free parameters of these relaxations to minimize the estimator’s variance during training.

Original: In this work we generalize the method of Tucker et al.
Revised: We generalize the method of Tucker et al.

Original: (2017) to learn a free-form control variate parameterized by a neural network, giving a lower-variance, unbiased gradient estimator which can be applied to a wider variety of problems with greater ﬂexibility.
Revised: (2017) to learn a free-form control variate parameterized by a neural network. This gives a lower-variance, unbiased gradient estimator which can be applied to a wider variety of problems.

Original: Most notably, our method is applicable even when no continuous relaxation is available, as in reinforcement learning or black box function optimization.
Revised: Most notably, our method is applicable even when no continuous relaxation is available, as in reinforcement learning or black-box function optimization.

Original: Furthermore, we derive improved variants of popular reinforcement learning methods with unbiased, action-dependent gradient estimates and lower variance.

Original: [Figure]

Original:  B ACKGROUND : G RADIENT ESTIMATORS
Revised:  B ACKGROUND : G RADIENT ESTIMATORS

Original: How can we choose the parameters of a distribution to maximize an expectation?
Revised: How can we choose the parameters of a distribution to maximize an expectation?

Original: This problem comes up in reinforcement learning, where we must choose the parameters θ of a policy distribution π ( a | s, θ ) to maximize the expected reward E τ ∼ π
Revised: This problem comes up in reinforcement learning, where we must choose the parameters θ of a policy distribution π ( a | s, θ ) to maximize the expected reward E τ ∼ π

Original: [ R ] over state-action trajectories τ .
Revised: [ R ] over state-action trajectories τ .

Original: It also  comes up in ﬁtting latent-variable models, when we wish to maximize the marginal probability p ( x | θ ) =
Revised: It also [Figure] comes up in ﬁtting latent-variable models, when we wish to maximize the marginal probability p ( x | θ ) =

Original: (cid:80) z p ( x | z ) p ( z | θ ) =
Revised: (cid:80) z p ( x | z ) p ( z | θ ) =

Original: E p ( z | θ )
Revised: E p ( z | θ )

Original: [ p ( x | z )] .
Revised: [ p ( x | z )] .

Original: In this paper, we’ll consider the general problem of optimizing
Revised: In this paper, we’ll consider the general problem of optimizing

Original: [Equation]
Revised: [Equation]

Original: When the parameters θ are high-dimensional, gradient-based optimization is appealing because it provides information about how to adjust each parameter individually.
Revised: When the parameters θ are high-dimensional, gradient-based optimization is appealing because it provides information about how to adjust each parameter individually.

Original: Stochastic optimization is essential for scalablility. However, it is only guaranteed to converge to a ﬁxed point of the objective when the stochastic gradients ˆ g are unbiased, i.e. E [ˆ g ] =
Revised: Stochastic optimization is essential for scalablility, but is only guaranteed to converge to a ﬁxed point of the objective when the stochastic gradients ˆ g are unbiased, i.e. E [ˆ g ] =

Original: ∂∂θ E p (b | θ )[ f ( b )] (Robbins & Monro, 1951).
Revised: ∂∂θ L ( θ ) (Robbins & Monro, 1951).

Original: How can we build unbiased, stochastic estimators of ∂∂θ L ( θ ) ?
Revised: How can we build unbiased, stochastic gradient estimators?

Original: There are several standard methods:
Revised: There are several standard methods:

Original: The score-function gradient estimator One of the most generally-applicable gradient estimators is known as the score-function estimator, or REINFORCE (Williams, 1992):
Revised: The score-function gradient estimator One of the most generally-applicable gradient estimators is known as the score-function estimator, or REINFORCE (Williams, 1992):

Original: [Equation]
Revised: [Equation]

Original: This estimator is unbiased, but in general has high variance.
Revised: This estimator is unbiased, but in general has high variance.

Original: Intuitively, this estimator is limited by the fact that it doesn’t use any information about how f depends on b , only on the ﬁnal outcome f ( b ) .
Revised: Intuitively, this estimator is limited by the fact that it doesn’t use any information about how f depends on b , only on the ﬁnal outcome f ( b ) .

Original: The reparameterization trick When f is continuous and differentiable, and the latent variables b can be written as a deterministic, differentiable function of a random draw from a ﬁxed distribution, the reparameterization trick (Williams, 1992; Kingma & Welling, 2014; Rezende et al., 2014) creates a low-variance, unbiased gradient estimator by making the dependence of b on θ explicit through a reparameterization function b = T ( θ, (cid:15) ) :
Revised: The reparameterization trick When f is continuous and differentiable, and the latent variables b can be written as a deterministic, differentiable function of a random draw from a ﬁxed distribution, the reparameterization trick (Williams, 1992; Kingma & Welling, 2014; Rezende et al., 2014) creates a low-variance, unbiased gradient estimator by making the dependence of b on θ explicit through a reparameterization function b = T ( θ, (cid:15) ) :

Original: [Equation]
Revised: [Equation]

Original: This gradient estimator is often used when training high-dimensional, continuous latent-variable models, such as variational autoencoders or GANs (Goodfellow et al., 2014).
Revised: This gradient estimator is often used when training high-dimensional, continuous latent-variable models, such as variational autoencoders.

Original: One intuition for why this gradient estimator is preferable to REINFORCE is that it depends on ∂f∂b , which exposes the dependence of f on b .
Revised: One intuition for why this gradient estimator is preferable to REINFORCE is that it depends on ∂f/∂b , which exposes the dependence of f on b .

Original: Control variates Control variates are a general method for reducing the variance of a Monte Carlo estimator.
Revised: Control variates Control variates are a general method for reducing the variance of a stochastic estimator.

Original: Given an estimator ˆ g ( b ) , a control variate is a function c ( b ) with a known mean E p ( b ) [ c ( b )] .
Revised: A control variate is a function c ( b ) with a known mean E p ( b ) [ c ( b )] .

Original: Subtracting the control variate from our estimator and adding its mean gives us a new estimator:[Equation] [Equation] This new estimator has the same expectation as the old one:
Revised: Given an estimator ˆ g ( b ) , subtracting the control variate from this estimator and adding its mean gives us a new estimator:

Original: [Equation]
Revised: [Equation]

Original: Importantly, the new estimator has  lower variancethan ˆ g ( b ) if c ( b ) is positively correlated with f ( b ) .
Revised: This new estimator has the same expectation as the old one, but has lower variance if c ( b ) is positively correlated with ˆ g ( b ) .

Original:  C ONSTRUCTING AND OPTIMIZING A DIFFERENTIABLE SURROGATE
Revised:  C ONSTRUCTING AND OPTIMIZING A DIFFERENTIABLE SURROGATE

Original: In this section, we introduce a gradient estimator for the expectation of a function ∂∂θ E p ( b | θ )
Revised: In this section, we introduce a gradient estimator for the expectation of a function ∂∂θ E p ( b | θ )

Original: [ f ( b )] that can be applied even when f is unknown, or not differentiable, or when b is discrete.
Revised: [ f ( b )] that can be applied even when f is unknown, or not differentiable, or when b is discrete.

Original: Our estimator combines the score function estimator, the reparameterization trick, and control variates.
Revised: Our estimator combines the score function estimator, the reparameterization trick, and control variates.

Original: We obtain an unbiased estimator whose variance can potentially be as low as the reparameterization-trick estimator, even when f is not differentiable or not computable.

Original: First, we consider the case where b is continuous, but that f cannot be differentiated.
Revised: First, we consider the case where b is continuous, but that f cannot be differentiated.

Original: Instead of differentiating through f , we build a surrogate of f using a neural network c φ , and differentiate c φ instead.
Revised: Instead of differentiating through f , we build a surrogate of f using a neural network c φ , and differentiate c φ instead.

Original: Since the score-function estimator and reparameterization estimator have the same expectation, we can simply subtract the score-function estimator for c φ and add the reparameterization estimator for c φ , to produce a gradient estimator which we call LAX:[Equation]
Revised: Since the score-function estimator and reparameterization estimator have the same expectation, we can simply subtract the score-function estimator for c φ and add back its reparameterization estimator. This gives a gradient estimator which we call LAX:

Original: [Equation]
Revised: [Equation]

Original: This estimator is unbiased for any choice of c φ and when c φ = f , our estimator becomes the reparameterization estimator for c φ .
Revised: This estimator is unbiased for any choice of c φ . When c φ = f , then LAX becomes the reparameterization estimator for f .

Original: Thus our estimator can have variance at least as low as the reparameterization estimator.
Revised: Thus LAX can have variance at least as low as the reparameterization estimator.

Original: 3. O PTIMIZING THEGRADIENT CONTROL VARIATEWITH GRADIENTS
Revised: 3. G RADIENT - BASED OPTIMIZATION OF THE CONTROL VARIATE

Original: Since ˆ g LAX is unbiased for any choice of the surrogate c φ , the only remaining problem is to choose a c φ that gives low variance to ˆ g LAX .
Revised: Since ˆ g LAX is unbiased for any choice of the surrogate c φ , the only remaining problem is to choose a c φ that gives low variance to ˆ g LAX .

Original: How can we ﬁnd a φ which gives our estimator low variance?
Revised: How can we ﬁnd a φ which gives our estimator low variance?

Original: We simply optimize c φ using stochastic gradient descent, at the same time as we optimize the parameters  of our model or policy.
Revised: We simply optimize c φ using stochastic gradient descent, at the same time as we optimize the parameters θ of our model or policy.

Original: To optimize c φ , we require the gradient of the variance of ourgradient estimator.
Revised: To optimize c φ , we require the gradient of the variance of our estimator.

Original: To estimate these gradients, we could simply differentiate through the empirical variance over each mini-batch.
Revised: To estimate these gradients, we could simply differentiate through the empirical variance over each mini-batch.

Original: Or, following Ruiz et al.
Revised: Or, following Ruiz et al.

Original: (2016) and Tucker et al.
Revised: (2016) and Tucker et al.

Original: (2017), we can construct an unbiased, single-sample estimator using the fact that our gradient estimator is unbiased.
Revised: (2017), we can construct an unbiased, single-sample estimator using the fact that our gradient estimator is unbiased.

Original: For any unbiased gradient estimator ˆ g with parameters φ :
Revised: For any unbiased gradient estimator ˆ g with parameters φ :

Original: [Equation]
Revised: [Equation]

Original: Thus, an unbiased single-sample estimate of the gradient of the variance of ˆ g is given by2ˆ g ∂ ˆ g∂φ .
Revised: Thus, an unbiased single-sample estimate of the gradient of the variance of ˆ g is given by ∂ ˆ g 2 /∂φ .

Original: This method of directly minimizing the variance of the gradient estimator stands in contrast to other methods such as Q-Prop (Gu et al., 2016) and advantage actor-critic (Sutton et al., 2000), which train the control variate to minimize the squared error ( f ( b ) − c φ ( b )) 2 .
Revised: This method of directly minimizing the variance of the gradient estimator stands in contrast to other methods such as Q-Prop (Gu et al., 2016) and advantage actor-critic (Sutton et al., 2000), which train the control variate to minimize the squared error ( f ( b ) − c φ ( b )) 2 .

Original: Our algorithm, which jointly optimizes the parameters θ
Revised: Our algorithm, which jointly optimizes the parameters θ

Original: and the surrogate c φ is given in Algorithm 1.
Revised: and the surrogate c φ is given in Algorithm 1.

Original: 3.1. O PTIMAL SURROGATE
Revised: 3.1. O PTIMAL SURROGATE

Original: What is the form of the variance-minimizing c φ ?
Revised: What is the form of the variance-minimizing c φ ?

Original: Inspecting the square of (6), we can see that this loss encourages c φ ( b ) to approximate f ( b ) , but with a weighting based on ∂∂θ log p ( b  ) .
Revised: Inspecting the square of (5), we can see that this loss encourages c φ ( b ) to approximate f ( b ) , but with a weighting based on ∂∂θ log p ( b | θ ) .

Original: Moreover, as c φ → f then ˆ g LAX → ∂∂θ c φ .
Revised: Moreover, as c φ → f then ˆ g LAX → ∂∂θ c φ .

Original: Thus, this objective encourages a balance between the variance of the reparameterization estimator and the variance of the REINFORCE estimator.
Revised: Thus, this objective encourages a balance between the variance of the reparameterization estimator and the variance of the REINFORCE estimator.

Original: Figure 2 shows the learned surrogate on a toy problem.
Revised: Figure 2 shows the learned surrogate on a toy problem.

Revised: 3. D ISCRETE RANDOM VARIABLES AND CONDITIONAL REPARAMETERIZATION

Original: Algorithm
Revised: We can adapt the LAX estimator to the case where b is a discrete random variable by introducing a “relaxed” continuous variable z . We require a continuous, reparameterizable distribution p ( z | θ ) and a deterministic mapping H ( z ) such that H ( z ) = b ∼ p ( b | θ ) when z ∼ p ( z | θ ) . In our implementation, we use the Gumbel-softmax trick, the details of which can be found in appendix B.

Revised: The discrete version of the LAX estimator is given by:

Revised: [Equation]

Original: Require: f ( · ) , log p ( b | θ ) , reparameterized sampler
Revised: Require: f ( · ) , log p ( b | θ ) , reparameterized sampler

Original: b
Revised: b

Original: =
Revised: =

Original: T ( θ, (cid:15) ) , neural network c φ ( · ) while not converged do (cid:15)i ∼ p ( (cid:15) )
Revised: T ( θ, (cid:15) ) , neural network c φ ( · ) while not converged do (cid:15) ∼ p ( (cid:15) )

Original: (cid:46)
Revised: (cid:46)

Original: Sample noise bi ← T ( (cid:15) i , θ ) (cid:46)Compute input g θ ← [ f ( b i ) − c φ ( b i )] ∇ θ log p + ∇ θ c φ ( b i ) (cid:46) Estimate gradient g φ ← 2 g θ ∂g θ ∂φ (cid:46) Estimate gradient of variance of gradient θ ← θ + α 1 g θ (cid:46) Update parameters φ ← φ + α 2 g φ (cid:46) Update control variate end while return θ 3. D ISCRETE RANDOM VARIABLES AND CONDITIONAL REPARAMETERIZATION
Revised: Sample noise b ← T ( (cid:15), θ ) (cid:46)

Original: We can adapt the LAX estimator to the case where b is a discrete random variable by introducing a “relaxed” continuous variable z . We require a continuous, reparameterizable distribution p ( z | θ ) and a deterministic mapping H ( z ) such that H (z ) = b ∼ p ( b | θ ) when z ∼ p ( z | θ  ) . In our implementation, we use the Gumbel-softmax trick, the details of which can be found in appendix B.
Revised: Compute input ˆ g θ ← [ f ( b ) − c φ ( b )] ∇ θ log p ( b | θ ) + ∇ θ c φ ( b ) (cid:46) Estimate gradient of objective ˆ g φ ← ∂ ˆ g 2 θ /∂φ

Original: The discrete version of the LAX estimator is given by: [Equation]
Revised: (cid:46) Estimate gradient of variance of gradient θ ← θ + α 1 g θ (cid:46) Update parameters φ ← φ + α 2 g φ (cid:46) Update control variate end while return θ

Original: This estimator is simple to implement and general.
Revised: This estimator is simple to implement and general.

Original: However, when f = c φ we do not recover the reparameterization estimator as we do with LAX.
Revised: However, when f = c φ we do not recover the reparameterization estimator as we do with LAX.

Original: To achieve this, we must be able to replace the ∂∂θ log p ( z | θ ) in the control variate with ∂∂θ log p ( b | θ ) .
Revised: To achieve this, we must be able to replace the ∂∂θ log p ( z | θ ) in the control variate with ∂∂θ log p ( b | θ ) .

Original: This is the motivation behind our next estimator which we call RELAX.
Revised: This is the motivation behind our next estimator, which we call RELAX.

Original: To construct a more powerful gradient estimator, we incorporate a further reﬁnement due to Tucker et al.
Revised: To construct a more powerful gradient estimator, we incorporate a further reﬁnement due to Tucker et al.

Original: Speciﬁcally, we evaluate our control variate both at a relaxed input z ∼ p ( z | θ ) , and also at a relaxed input conditioned on the discrete variable b , denoted ˜
Revised: Speciﬁcally, we evaluate our control variate both at a relaxed input z ∼ p ( z | θ ) , and also at a relaxed input conditioned on the discrete variable b , denoted ˜

Original: z ∼ p ( z | b, θ ) .
Revised: z ∼ p ( z | b, θ ) .

Original: Thus we deﬁne our estimator as
Revised: Doing so gives us:

Original: [Equation]
Revised: [Equation]

Original: This estimator is unbiased for any c φ .
Revised: This estimator is unbiased for any c φ .

Original: A proof and a detailed algorithm can be found in appendix A.
Revised: A proof and a detailed algorithm can be found in appendix A.

Original: We note that the distribution p ( z | b, θ ) must also be reparameterizable.
Revised: We note that the distribution p ( z | b, θ ) must also be reparameterizable.

Original: We demonstrate how to perform this conditional reparameterization for Bernoulli and categorical random variables in appendix B.
Revised: We demonstrate how to perform this conditional reparameterization for Bernoulli and categorical random variables in appendix B.

Original: 3. C HOOSING THE CONTROL VARIATE ARCHITECTURE
Revised: C HOOSING THE CONTROL VARIATE ARCHITECTURE

Original: The variance-reduction objective introduced above allows us to use any differentiable, parametric function as our control variate c φ .
Revised: The variance-reduction objective introduced above allows us to use any differentiable, parametric function as our control variate c φ .

Original: How should we choose the architecture of c φ ?
Revised: How should we choose the architecture of c φ ?

Original: Ideally, we will take advantage of any known structure in f .
Revised: Ideally, we will take advantage of any known structure in f .

Original: If f is a known, differentiable function of discrete random variables, we can use the concrete relaxation (Jang et al., 2016; Maddison et al., 2016) and let c φ ( z ) =
Revised: In the discrete setting, if f is known and happens to be differentiable, we can use the concrete relaxation (Jang et al., 2016; Maddison et al., 2016) and let c φ ( z ) =

Original: f ( σ λ ( z )) .
Revised: f ( σ λ ( z )) .

Original: In this special case, our estimator is exactly the REBAR estimator.
Revised: In this special case, our estimator is exactly the REBAR estimator.

Original: We are also free to add a learned component to the concrete relaxation and let c φ ( z ) =
Revised: We are also free to add a learned component to the concrete relaxation and let c φ ( z ) =

Original: f ( σ λ ( z ))
Revised: f ( σ λ ( z ))

Original: + r ρ ( z ) where r ρ is a neural network with parameters ρ  .
Revised: + r ρ ( z ) where r ρ is a neural network with parameters ρ making φ = { ρ, λ } .

Original: We took this approach in our experiments training discrete variational auto-encoders.
Revised: We took this approach in our experiments training discrete variational autoencoders.

Original: If f is unknown, we can simply let c φ be a generic function approximator such as a neural network.
Revised: If f is unknown, we can simply let c φ be a generic function approximator such as a neural network.

Original: We took this simpler approach in our reinforcement learning experiments.
Revised: We took this simpler approach in our reinforcement learning experiments.

Original: 3. R EINFORCEMENT LEARNING
Revised: 3. R EINFORCEMENT LEARNING

Original: We now describe how we apply the LAX estimator in the reinforcement learning (RL) setting.
Revised: We now describe how we apply the LAX estimator in the reinforcement learning (RL) setting.

Original: By reinforcement learning, we refer to the problem of optimizing the parameters θ of a policy distribution π ( a | s, θ ) to maximize the sum of rewards.
Revised: By reinforcement learning, we refer to the problem of optimizing the parameters θ of a policy distribution π ( a | s, θ ) to maximize the sum of rewards.

Original: In this setting, the random variable being integrated over is τ , which denotes a series of  actions and states [( s 1 , a 1 ) , ( s 2 , a 2 ) , ..., ( s T , a T )] .
Revised: In this setting, the random variable being integrated over is τ , which denotes a series of T actions and states [( s 1 , a 1 ) , ( s 2 , a 2 ) , ..., ( s T , a T )] .

Original: The function whose expectation is being optimized, R , maps τ to the sum of rewards R ( τ )
Revised: The function whose expectation is being optimized, R , maps τ to the sum of rewards R ( τ )

Original: =
Revised: =

Original: (cid:80) Tt =1 r t ( s t , a t ) .
Revised: (cid:80) Tt =1 r t ( s t , a t ) .

Original: Again, we want to estimate the gradient of an expectation of a black-box function: ∂∂θ E p ( τ | θ )
Revised: Again, we want to estimate the gradient of an expectation of a black-box function: ∂∂θ E p ( τ | θ )

Original: [ R ( τ )] .
Revised: [ R ( τ )] .

Original: The de facto standard approach is the advantage actor-critic estimator (A2C)
Revised: The de facto standard approach is the advantage actor-critic estimator (A2C)

Original: (Sutton et al., 2000):
Revised: (Sutton et al., 2000):

Original: [Equation]
Revised: [Equation]

Original: Where c φ ( s t ) is an estimate of the state-value function, c φ ( s )
Revised: Where c φ ( s t ) is an estimate of the state-value function, c φ ( s )

Original: ≈ V π
Revised: ≈ V π

Original: ( s ) =
Revised: ( s ) =

Original: E τ [ R | s 1 = s ] .
Revised: E τ [ R | s 1 = s ] .

Original: This estimator is unbiased when c  does not depend on a t .
Revised: This estimator is unbiased when c φ does not depend on a t .

Original: The main limitations of A2C are that c  does not depend on a t , and that it’s not obvious how to optimize c  .
Revised: The main limitations of A2C are that c φ does not depend on a t , and that it’s not obvious how to optimize c φ .

Original: Using the LAX estimator addresses both of these problems.
Revised: Using the LAX estimator addresses both of these problems.

Original: First, we assume π ( a t | s t  ) is reparameterizable, meaning that we can write a t = a ( (cid:15) t , s t , θ ) , where (cid:15) t does not depend on θ .
Revised: First, we assume π ( a t | s t , θ ) is reparameterizable, meaning that we can write a t = a ( (cid:15) t , s t , θ ) , where (cid:15) t does not depend on θ .

Original: We again introduce a differentiable surrogate c φ ( a, s ) .
Revised: We again introduce a differentiable surrogate c φ ( a, s ) .

Original: Crucially, this surrogate is a function of the action as well as the state.
Revised: Crucially, this surrogate is a function of the action as well as the state.

Original: Our estimator is deﬁned as: [Equation] t | s t , θ ) ∂θ (cid:34) ∞ (cid:88) t (cid:48) =
Revised: The extension of LAX to Markov decision processes is:

Original: t

Original: [Equation]

Original: t (cid:48)

Original: [Equation]

Original: t ) (cid:35)

Original: +

Original: [Equation]

Original: ∂θc

Original: [Equation]

Original: ( a t

Original: [Equation]

Original: s t )

Original: [Equation]

Original: (11) a t = a ( (cid:15) t

Original: [Equation]

Original: s t

Original: [Equation]
Revised: [Equation]

Original: θ ) (cid:15) t ∼ p ( (cid:15) t ) .

Original: This estimator is unbiased if the true dynamics of the system are Markovian w.r.t.
Revised: This estimator is unbiased if the true dynamics of the system are Markovian w.r.t.

Original: the state s t .
Revised: the state s t .

Original: When T = 1 , we recover the special case ˆ g RLLAX =
Revised: When T = 1 , we recover the special case ˆ g RLLAX =

Original: ˆ g LAX .
Revised: ˆ g LAX .

Original: Comparing ˆ g RLLAX to the standard advantage actor-critic estimator in (10), the main difference is that our baseline c φ ( a t , s t ) is action-dependent while still remaining unbiased.
Revised: Comparing ˆ g RLLAX to the standard advantage actor-critic estimator in (9), the main difference is that our baseline c φ ( a t , s t ) is action-dependent while still remaining unbiased.

Original: To optimize the parameters φ of our control variate c
Revised: To optimize the parameters φ of our control variate c

Original: φ ( a t , s t ) , we can again use the single-sample estimator of the gradient of our estimator’s variance given in (7).
Revised: φ ( a t , s t ) , we can again use the single-sample estimator of the gradient of our estimator’s variance given in (6).

Original: This approach avoids unstable training dynamics, and doesn’t require storage and replay of previous rollouts.
Revised: This approach avoids unstable training dynamics, and doesn’t require storage and replay of previous rollouts.

Original: Details of this derivation, as well as the discrete and conditionally reparameterized version of this estimator can be found in appendix C.
Revised: Details of this derivation, as well as the discrete and conditionally reparameterized version of this estimator can be found in appendix C.

Original:  S COPE AND L IMITATIONS
Revised:  S COPE AND L IMITATIONS

Original: The work most related to ours is the recently-developed REBAR method (Tucker et al., 2017), which inspired our work.
Revised: The work most related to ours is the recently-developed REBAR method (Tucker et al., 2017), which inspired our work.

Original: The REBAR estimator is a special case of the RELAX estimator, when the surrogate is set to c φ ( z ) =
Revised: The REBAR estimator is a special case of the RELAX estimator, when the surrogate is set to c φ ( z ) =

Original: η · f ( softmax λ ( z )) .
Revised: η · f ( softmax λ ( z )) .

Original: The only free parameters of the REBAR estimator are the scaling factor η , and the temperature λ , which gives limited scope to optimize the surrogate.
Revised: The only free parameters of the REBAR estimator are the scaling factor η , and the temperature λ , which gives limited scope to optimize the surrogate.

Original: REBAR can only be applied when f is known and differentiable.
Revised: REBAR can only be applied when f is known and differentiable.

Original: Furthermore, it depends on essentially undeﬁned behavior of the function being optimized, since it evaluates the discrete loss function at continuous inputs.
Revised: Furthermore, it depends on essentially undeﬁned behavior of the function being optimized, since it evaluates the discrete loss function at continuous inputs.

Original: Because LAX and RELAX can construct a surrogate from scratch, they can be used for optimizing black-box functions, as in reinforcement learning settings where the reward is an unknown function of the environment.
Revised: Because LAX and RELAX can construct a surrogate from scratch, they can be used for optimizing black-box functions, as in reinforcement learning settings where the reward is an unknown function of the environment.

Original: LAX and RELAX only require that we can query the function being optimized, and can sample from and differentiate p ( b | θ ) .
Revised: LAX and RELAX only require that we can query the function being optimized, and can sample from and differentiate p ( b | θ ) .

Original: Can RELAX be used to optimize deterministic black-box functions?

Original: The answer is yes, with the caveat that one must introduce stochasticity to the inputs.

Original: Thus, RELAX is most suitable for problems where one is already optimizing a distribution over inputs, such as in inference or reinforcement learning.

Original: Direct dependence on parameters Above
Revised: Direct dependence on parameters Above

Original: , we assumed that the function f being optimized does not depend directly on θ , which is usually the case in black-box optimization settings.
Revised: , we assumed that the function f being optimized does not depend directly on θ , which is usually the case in black-box optimization settings.

Original: However, a dependence on θ can occur when training probabilistic models, or when we add a regularizer to a black-box optimization problem.
Revised: However, a dependence on θ can occur when training probabilistic models, or when we add a regularizer.

Original: In both these settings, if the dependence on θ is known and differentiable, we can use the fact that
Revised: In both these settings, if the dependence on θ is known and differentiable, we can use the fact that

Original: [Equation]
Revised: [Equation]

Original: and simply addthe term ∂∂θ f ( b, θ ) to our gradient estimate.
Revised: and simply add ∂∂θ f ( b, θ ) to any of the gradient estimators above to recover an unbiased estimator.

Original:  R ELATED WORK
Revised:  R ELATED WORK

Original: Miller et al.
Revised: Miller et al.

Original: (2017) reduce the variance of reparameterization gradients in an orthogonal way to ours by approximating the gradient-generating procedure with a simple model and using that model as a control variate.
Revised: (2017) reduce the variance of reparameterization gradients in an orthogonal way to ours by approximating the gradient-generating procedure with a simple model and using that model as a control variate.

Original: NVIL (Mnih & Gregor, 2014) and VIMCO (Mnih & Rezende, 2016) provide reduced variance gradient estimation in the special case of discrete latent variable models and discrete latent variable models with Monte-Carlo objectives.
Revised: NVIL (Mnih & Gregor, 2014) and VIMCO (Mnih & Rezende, 2016) provide reduced variance gradient estimation in the special case of discrete latent variable models and discrete latent variable models with Monte Carlo objectives.

Original: Salimans et al.
Revised: Salimans et al.

Original: (2017) estimate gradients using a form of ﬁnite differences, evaluating hundreds of different parameter values in parallel to construct a gradient estimate.
Revised: (2017) estimate gradients using a form of ﬁnite differences, evaluating hundreds of different parameter values in parallel to construct a gradient estimate.

Original: In contrast, our method is a single-sample estimator.
Revised: In contrast, our method is a single-sample estimator.

Original: Staines & Barber (2012) address the general problem of developing gradient estimators for deterministic black-box functions or discrete optimization.
Revised: Staines & Barber (2012) address the general problem of developing gradient estimators for deterministic black-box functions or discrete optimization.

Original: They introduce a sampling distribution, and optimize an objective similar to ours.
Revised: They introduce a sampling distribution, and optimize an objective similar to ours.

Original: Wierstra et al.
Revised: Wierstra et al.

Original: (2014) also introduce a sampling distribution to build a gradient estimator, and consider optimizing the sampling distribution.
Revised: (2014) also introduce a sampling distribution to build a gradient estimator, and consider optimizing the sampling distribution.

Original: In the reinforcement learning setting, the work most similar to ours is Q -prop (Haarnoja et al., 2017).
Revised: In the reinforcement learning setting, the work most similar to ours is Q -prop (Haarnoja et al., 2017).

Original: Like our method, Q -prop reduces the variance of the policy gradient with an learned, action-dependent control variate whose expectation is approximated via a monte-carlo sample from a taylor series expansion of the control variate.
Revised: Like our method, Q -prop reduces the variance of the policy gradient with a learned, action-dependent control variate whose expectation is approximated via a Monte Carlo sample from a Taylor series expansion of the control variate.

Original: Unlike our method, their control variate is trained off-policy.
Revised: Unlike our method, their control variate is trained off-policy.

Original: While our method is applicable in both the continuous and discrete action domain, Q -prop is only applicable to continuous actions.
Revised: While our method is applicable in both the continuous and discrete action domain, Q -prop is only applicable to continuous actions.

Original:  A PPLICATIONSREINFORCE
Revised:  A PPLICATIONS

Revised: REINFORCE

Original: [Figure]
Revised: [Figure]

Original: ( u )
Revised: function, using different gradient estimators. Because REBAR uses the concrete relaxation of f , which happens to be implemented as a quadratic function, the optimal relaxation is constrained to be a warped quadratic.

Original: [Figure] [Figure] [Figure]
Revised: In contrast, RELAX can choose a free-form relaxation.

Original: We demonstrate the effectiveness of our estimator on a number of challenging optimization problems.
Revised: We demonstrate the effectiveness of our estimator on a number of challenging optimization problems.

Original: Following Tucker et al.
Revised: Following Tucker et al.

Original: (2017) we begin with a simple toy example to illuminate the potential of our method and then continue to the more relevant problems of optimizing
Revised: (2017) we begin with a simple toy example to illuminate the potential of our method and then continue to the more relevant problems of optimizing

Original: binary
Revised: binary

Original: VAE’s and reinforcement learning.
Revised: VAE’s and reinforcement learning.

Original: 6. T OY EXPERIMENT
Revised: 6. T OY EXPERIMENT

Original: As a simple example, we follow Tucker et al.
Revised: As a simple example, we follow Tucker et al.

Original: (2017) in minimizing E p ( b | θ )
Revised: (2017) in minimizing E p ( b | θ )

Original: [( b − t ) 2 ] as a function of the parameter θ where p ( b | θ ) =
Revised: [( b − t ) 2 ] as a function of the parameter θ where p ( b | θ ) =

Original: Bernoulli ( b | θ ) .
Revised: Bernoulli ( b | θ ) .

Original: Tucker et al.
Revised: Tucker et al.

Original: (2017) set the target t = .
Revised: (2017) set the target t = .

Original: We focus on the more challenging case where t = .
Revised: We focus on the more challenging case where t = .

Original: Figures 1a and 1b show the relative performance and gradient log-variance of REINFORCE, REBAR, and RELAX.
Revised: Figures 1a and 1b show the relative performance and gradient log-variance of REINFORCE, REBAR, and RELAX.

Original: Figure 2 plots the learned surrogate c φ for a ﬁxed value of θ .
Revised: Figure 2 plots the learned surrogate c φ for a ﬁxed value of θ .

Original: We can see that c φ is near f for all z , keeping the variance of the REINFORCE part of the estimator small.
Revised: We can see that c φ is near f for all z , keeping the variance of the REINFORCE part of the estimator small.

Original: Moreover the derivative of c φ is positive for all z meaning that the reparameterization part of the estimator will produce gradients pointing in the correct direction to optimize the expectation.
Revised: Moreover the derivative of c φ is positive for all z meaning that the reparameterization part of the estimator will produce gradients pointing in the correct direction to optimize the expectation.

Original: Conversely, the concrete relaxation of REBAR is close to f only near 0 and 1 and its gradient points in the correct direction[Figure] only for values of z > log( 1 − t t ) .
Revised: Conversely, the concrete relaxation of REBAR is close to f only near 0 and 1 and its gradient points in the correct direction only for values of z > log( 1 − t t ) .

Original: These factors together result in the RELAX estimator achieving the best performance.
Revised: These factors together result in the RELAX estimator achieving the best performance.

Original: 6. D ISCRETE VARIATIONAL AUTOENCODER
Revised: 6. D ISCRETE VARIATIONAL AUTOENCODER

Original: Next, we evaluate the RELAX estimator on the task of training a variational autoencoder (Kingma &
Revised: Next, we evaluate the RELAX estimator on the task of training a variational autoencoder (Kingma &

Original: Welling, 2014; Rezende et al., 2014) with Bernoulli latent variables.
Revised: Welling, 2014; Rezende et al., 2014) with Bernoulli latent variables.

Original: We reproduceda subset of the  experiments from Tucker et al.
Revised: We reproduced the variational autoencoder experiments from Tucker et al.

Original: (2017), training models with 1 and 2 layers of200 Bernoulli random variables with linear  mappings between them, on both the MNIST and Omniglot (Lake et al., 2015) datasets.
Revised: (2017), training models with 1 or 2 layers of Bernoulli random variables with linear or nonlinear mappings between them, on both the MNIST and Omniglot (Lake et al., 2015) datasets.

Original: Details of these models and our experimental procedure can be found in appendix E.1.
Revised: Details of these models and our experimental procedure can be found in Appendix E.1.

Original: To take advantage of the available structure in the loss function, we choose the form of our control variate to be c φ ( z ) =
Revised: To take advantage of the available structure in the loss function, we choose the form of our control variate to be c φ ( z ) =

Original: f ( σ λ ( z ))
Revised: f ( σ λ ( z ))

Original: + ˆ r ρ ( z ) where ˆ r ρ is a neural network with parameters ρ
Revised: + ˆ r ρ ( z ) where ˆ r ρ is a neural network with parameters ρ

Original: and f ( σ λ ( z )) is the discrete loss function (the evidence lower-bound) evaluated at continuously relaxed inputs as in REBAR.
Revised: and f ( σ λ ( z )) is the discrete loss function, the evidence lower-bound (ELBO), evaluated at continuously relaxed inputs as in REBAR.

Original: In all experiments, the learned control variate improved the trainingand validation performance, over the state-of-the-art baseline of REBAR. 
Revised: In all experiments, the learned control variate improved the training performance, over the state-of-the-art baseline of REBAR. In both linear models, we achieved improved validation performance as well increased convergence speed. We believe the decrease [Figure]

Revised: in validation performance for the nonlinear models was due to overﬁtting caused by improved optimization of an under-regularized model.

Revised: We leave exploring this phenomenon to further work.

Original: Dataset
Revised: Dataset

Original: Model
Revised: Model

Original: Concrete
Revised: Concrete

Original: NVIL
Revised: NVIL

Original: MuProp
Revised: MuProp

Original: REBAR
Revised: REBAR

Original: RELAX
Revised: RELAX

Original: Nonlinear − 102 .− 101 .-101.
Revised: Nonlinear − 102 .− 101 .-101.

Original: MNIST linear 1 layer -111.
Revised: MNIST linear 1 layer -111.

Original: − 112 .− 111 .-111.
Revised: − 112 .− 111 .-111.

Original: [Table]
Revised: linear 2 layer -99. − 99 .− 99 .-98. Nonlinear − 110 .− 109 .-108. Omniglot linear 1 layer -117.

Original:  To obtain training curves we created our own implementation of REBAR, which gave identical or slightly improved performance compared to the implementation of Tucker et al.
Revised: − 117 .− linear 2 layer − 109 .− 109 .-108. To obtain training curves we created our own implementation of REBAR, which gave identical or slightly improved performance compared to the implementation of Tucker et al.

Original: While we obtained a modest improvement in training and validation scores (tables 1 and 3), the most notable improvement provided by RELAX is in its rate of convergence.
Revised: While we obtained a modest improvement in training and validation scores (tables 1 and 3), the most notable improvement provided by RELAX is in its rate of convergence.

Original: Training curves for the linear models can be seen in ﬁgure 3 and in appendix D. In table 4 we compare the number of training epochs that are required to match the best validation score of REBAR.
Revised: Training curves for all models can be seen in Figure 3 and in Appendix D. In Table 4 we compare the number of training epochs that are required to match the best validation score of REBAR.

Original: In all experiments, RELAX provides an increase in rate of convergence.
Revised: In both linear models, RELAX provides an increase in rate of convergence.

Original: 6. R EINFORCEMENT LEARNING
Revised: 6. R EINFORCEMENT LEARNING

Original: We apply our gradient estimator to a few simple reinforcement learning environments with discrete and continuous actions.
Revised: We apply our gradient estimator to a few simple reinforcement learning environments with discrete and continuous actions.

Original: We use the RELAX and LAX estimators for discrete and continuous actions, 
Revised: We use the RELAX and LAX estimators for discrete and continuous actions, respectively.

Original: Cart-pole
Revised: We compare with the advantage actor-critic algorithm (A2C) (Sutton et al., 2000) as a baseline. As our control variate does not have the same interpretation as the value function of A2C, it was not directly clear how to add reward bootstrapping and other variance reduction techniques common in RL into our model. For instance, to do reward bootstrapping, we would need to use the statevalue function.

Revised: In the discrete experiments, due to the simplicity of the tasks, we chose not to use reward bootstrapping, and therefore omitted the use of state-value function.

Revised: However, with the more complicated continuous tasks, we chose to use the value function to enable bootstrapping.

Revised: In this case, the control variate takes the form: c φ ( a, s )

Revised: =

Revised: V ( s ) + ˆ c ( a, s ) , where V ( s ) is trained as it would be in A2C. Full details of our experiments can be found in Appendix E.

Revised: In the discrete action setting, we test our approach on the Cart Pole and Lunar Lander environments as provided by the OpenAI gym (Brockman et al., 2016).

Revised: In the continuous action setting, we test on the MuJoCo-simulated (Todorov et al., 2012) environments Inverted Pendulum and Inverted

Revised: Cart-pole

Original: Lunar lander
Revised: Lunar lander

Original: Inverted pendulum
Revised: Inverted pendulum

Original: [Figure]

Original:  pendulum
Revised: Inverted double pendulum

Original: [Figure]
Revised: [Figure]

Original: L ogV a r i a n ce
Revised: L ogV a r i a n ce

Original: [Figure]
Revised: [Figure]

Original: respectively.

Original: We compare with the advantage actor-critic algorithm (A2C) (Sutton et al., 2000) as a baseline. Full details of our experiments can be found in Appendix E. 6.3. E XPERIMENTS
Revised: [Table]

Original: In the discrete action setting, we test our approach on the Cart Pole and Lunar Lander environments as provided by the OpenAI gym (Brockman et al., 2016).

Original: In the continuous action setting, we test on the MuJoCo-simulated (Todorov et al., 2012) environments Inverted Pendulum and Inverted

Original: Double Pendulum also found in the OpenAI gym.
Revised: Double Pendulum also found in the OpenAI gym.

Original: In all tested environments we observe improved performance and sample efﬁciency using our method.
Revised: In all tested environments we observe improved performance and sample efﬁciency using our method.

Original: The results of our experiments can be seen in ﬁgure 4, and table 2.
Revised: The results of our experiments can be seen in Figure 4, and Table 2.

Original: We found that our estimator produced policy gradients with drastically reduced variance (see ﬁgure ?? ) allowing for larger learning rates to be used while maintaining stable training.
Revised: We found that our estimator produced policy gradients with drastically reduced variance (see Figure 4) allowing for larger learning rates to be used while maintaining stable training.

Original: In both discrete environments our estimator achieved great than a 2-times speedup in convergence over the [Table]
Revised: In both discrete environments our estimator achieved great than a 2-times speedup in convergence over the baseline.

Original:  C ONCLUSIONS AND FUTURE WORK
Revised:  C ONCLUSIONS AND FUTURE WORK

Original: In this work we synthesized and generalized many of the standard approaches for constructing gradient estimators.
Revised: In this work we synthesized and generalized several standard approaches for constructing gradient estimators.

Original: We proposed asimple and generic gradient estimator that can be applied to expectations of known or black-box functions of discrete or continuous random variables.
Revised: We proposed a generic gradient estimator that can be applied to expectations of known or black-box functions of discrete or continuous random variables, and adds little computational overhead.

Original: We also derive a simple extension toapply our method to reinforcement learning in both discrete- and continuous-action domains.
Revised: We also derived a simple extension to reinforcement learning in both discrete and continuous-action domains.

Original: This approach is relatively simple to implement and adds little computational overhead. The scope and generality of our estimator opens up many new possibilities for models which can now be trained via gradient decent.
Revised: Future applications of this method could include training models with hard attention or memory indexing (Zaremba & Sutskever, 2015).

Original: For example, we could  apply our estimator totrain a VAE with continuous latent variables whose generative model is non-differentiable (a rendering engine perhaps).
Revised: One could also apply our estimators to continuous latentvariable models whose likelihood is non-differentiable, such as a 3D rendering engine.

Original: We also feel that there is much room to explore model design choices for the controlvariate and to better understand the properties of the optimal control variate.
Revised: There is also room to explore other architecture choices for the control variate.

Original: We believe our results in reinforcement learning are promising and should motivate further research into using action-dependent control-variates for policy-gradient methods. We are interested in combining our approach with other popular variance reduction techniques such as generalized advantage estimation (Kimura et al., 2000).
Revised: In the reinforcement learning setting, our method could be combined with other variance-reduction techniques such as generalized advantage estimation (Kimura et al., 2000; Schulman et al., 2015b), or other optimization methods, such as KFAC (Wu et al., 2017).

Original: We are alsointerested in ways to train our control variate off-policy as in Q -prop
Revised: One could also train our control variate off-policy, as in Q -prop

Original: (Gu et al., 2016).
Revised: (Gu et al., 2016).

Original: We also feel that the relationship between our learned control variate and the action-value function (commonly denoted as Q ) is worth exploring and understanding in greater detail.

Original: A PPENDICES A T HE RELAX A LGORITHM
Revised: A PPENDICES A T HE RELAX A LGORITHM

Original: We prove that ˆ g RELAX is unbiased.
Revised: We prove that ˆ g RELAX is unbiased.

Original: Following Tucker et al.
Revised: Following Tucker et al.

Original: [Equation]
Revised: [Equation]

Original: [Equation]

Original: [Equation]
Revised: [Equation]

Original: Algorithm 2 RELAX: Low-variance control variate optimization for black-box gradient estimation.
Revised: Algorithm 2 RELAX: Low-variance control variate optimization for black-box gradient estimation.

Original: Require: f ( · ) , log p ( b | θ ) , reparameterized samplers b
Revised: Require: f ( · ) , log p ( b | θ ) , reparameterized samplers b

Original: = H ( z ) , z = S ( (cid:15), θ ) and ˜ z
Revised: = H ( z ) , z = S ( (cid:15), θ ) and ˜ z

Original: = S ( (cid:15), θ | b ) , neural network c φ ( · ) while not converged do (cid:15) i , (cid:101)
Revised: = S ( (cid:15), θ | b ) , neural network c φ ( · ) while not converged do (cid:15) i , (cid:101)

Original: (cid:15) i ∼ p ( (cid:15) )
Revised: (cid:15) i ∼ p ( (cid:15) )

Original: (cid:46)
Revised: (cid:46)

Original: Sample noise z
Revised: Sample noise z

Original: i ← S ( (cid:15) i , θ ) (cid:46) Compute unconditional relaxed input b
Revised: i ← S ( (cid:15) i , θ ) (cid:46) Compute unconditional relaxed input b

Original: i ← H ( z i ) (cid:46) Compute input (cid:101)
Revised: i ← H ( z i ) (cid:46) Compute input (cid:101)

Original: z
Revised: z

Original: i ← S ( (cid:101) (cid:15)
Revised: i ← S ( (cid:101) (cid:15)

Original: i , θ | b i ) (cid:46)
Revised: i , θ | b i ) (cid:46)

Original: Compute conditional relaxed input g θ ←
Revised: Compute conditional relaxed input g θ ←

Original: [Equation]
Revised: [Equation]

Original: (cid:46) Estimate gradient g φ ← 2 g θ ∂g
Revised: Estimate gradient g [Equation] ← [Equation]

Revised: g

Original: θ ∂φ (cid:46) Estimate gradient of variance of gradient θ ← θ + α 1 g θ (cid:46) Update parameters φ ← φ + α 2 g φ (cid:46) Update control variate end while return θ
Revised: [Equation] (cid:46) Estimate gradient of variance of gradient θ ← θ + α 1 g [Equation] (cid:46) Update parameters φ ← φ + α 2 g φ (cid:46) Update control variate end while return θ

Original: B C ONDITIONAL R E - SAMPLING FOR D ISCRETE R ANDOM V ARIABLES
Revised: B C ONDITIONAL R E - SAMPLING FOR D ISCRETE R ANDOM V ARIABLES

Original: When applying the RELAX estimator to a function of discrete random variables b ∼ p ( b | θ ) , we require that there exists a distribution p ( z | θ ) and a deterministic mapping H ( z ) such that if z ∼ p ( z | θ ) then H ( z ) =
Revised: When applying the RELAX estimator to a function of discrete random variables b ∼ p ( b | θ ) , we require that there exists a distribution p ( z | θ ) and a deterministic mapping H ( z ) such that if z ∼ p ( z | θ ) then H ( z ) =

Original: b ∼ p ( b | θ ) .
Revised: b ∼ p ( b | θ ) .

Original: Treating both b and z as random, this procedure deﬁnes a probabalistic model p ( b, z | θ ) =
Revised: Treating both b and z as random, this procedure deﬁnes a probabilistic model p ( b, z | θ ) =

Original: p ( b | z ) p ( z | θ ) .
Revised: p ( b | z ) p ( z | θ ) .

Original: The RELAX estimator requires reparameterized samples from p ( z | θ ) and p ( z | b, θ ) .
Revised: The RELAX estimator requires reparameterized samples from p ( z | θ ) and p ( z | b, θ ) .

Original: We describe how to sample from these distributions in the common cases of p ( b | θ ) =
Revised: We describe how to sample from these distributions in the common cases of p ( b | θ ) =

Original: Bernoulli ( θ ) and p ( b | θ ) =
Revised: Bernoulli ( θ ) and p ( b | θ ) =

Original: Categorical ( θ ) .
Revised: Categorical ( θ ) .

Original: Bernoulli When p ( b | θ ) is Bernoulli distribution
Revised: Bernoulli When p ( b | θ ) is Bernoulli distribution

Original: we let H ( z ) =
Revised: we let H ( z ) =

Original: I ( z > 0) and we sample from p ( z | θ ) with z =log θ [Equation] − [Equation]
Revised: I ( z > 0) and we sample from p ( z | θ ) with z =

Original: +

Original: [Equation]
Revised: [Equation]

Original: u 1 − u, u ∼ uniform [0 , 1] .

Original: We can sample from p ( z | b, θ ) with˜ z
Revised: We can sample from p ( z | b, θ ) with

Original: = (cid:26) v

Original: [Equation]
Revised: [Equation]

Original: b =v ( [Equation] b =where v ∼ uniform [0 , 1] . Categorical When p ( b | θ ) is a Categorical distribution where θ i =
Revised: Categorical When p ( b | θ ) is a Categorical distribution where θ i =

Original: p ( b = i | θ ) , we let H ( z ) =
Revised: p ( b = i | θ ) , we let H ( z ) =

Original: argmax ( z ) and we sample from p ( z |
Revised: argmax ( z ) and we sample from p ( z |

Original: θ ) with
Revised: θ ) with

Original: [Equation]
Revised: [Equation]

Original: where k is the number of possible outcomes.
Revised: where k is the number of possible outcomes.

Original: Intuitively, to sample from p ( z | b, θ )we should ﬁrst sample v ∼ uniform [0 , 1] k , then compute g b = log θ  b  − log( − log( v b )) .
Revised: To sample from p ( z | b, θ ) , we note that the distribution of the largest ˆ z b is independent of θ , and can be sampled as ˆ z b = − log( − log v b ) where v b ∼ uniform [0 , 1] .

Original: Then we must determine how to scale each v i (cid:54) =
Revised: Then, the remaining v i (cid:54) =

Revised: b can be sampled as before but with their underlying noise truncated so

Original: b such that g i (cid:54) =
Revised: ˆ z i (cid:54) =

Original: b < g b .
Revised: b < ˆ z b .

Original: We can deﬁne v (cid:48) such that
Revised: As shown in the appendix of Tucker et al. (2017), we can then sample from p ( z | b, θ ) with:

Original: [Equation]
Revised: [Equation]

Original: and then ˜ z = log θ − log( − log v (cid:48) ) which is our sample from p ( z | b, θ ) .
Revised: where v i ∼ uniform [0 , 1] .

Original: C D ERIVATIONS OF ESTIMATORS USED IN R EINFORCEMENT LEARNING
Revised: C D ERIVATIONS OF ESTIMATORS USED IN R EINFORCEMENT LEARNING

Original: We give the derivation of the LAX estimator used for continuous RL tasks.
Revised: We give the derivation of the LAX estimator used for continuous RL tasks.

Original: Theorem C.1.
Revised: Theorem C.1.

Original: The LAX estimator,
Revised: The LAX estimator,

Original: [Equation]
Revised: [Equation]

Original: is unbiased.
Revised: is unbiased.

Original: Proof.
Revised: Proof.

Original: Note that by using the score-function estimator, for all t , we have
Revised: Note that by using the score-function estimator, for all t , we have

Original: [Equation]
Revised: [Equation]

Original: Then, by adding and subtracting the same term, we have ∂
Revised: Then [Equation] by adding and subtracting the same term [Equation]

Revised: we have

Original: [Equation]
Revised: [Equation]

Original: In the discrete control setting, our policy parameterizes a soft-max distribution which we use to sample actions.
Revised: In the discrete control setting [Equation] our policy parameterizes a soft-max distribution which we use to sample actions.

Original: We deﬁne z t ∼ p ( z t | s t ) , which is equal to σ (log π − log( − log( u ))) where u ∼ Unif [0 , 1] , a t = argmax ( z t ) , σ is the soft-max function.
Revised: We deﬁne z t ∼ p ( z t | s t ) , which is equal to σ (log π − log( − log( u ))) where u ∼ uniform [0 , 1] , a t = argmax ( z t ) , σ is the soft-max function.

Original: We also deﬁne ˜
Revised: We also deﬁne ˜

Original: z t ∼ p ( z t | a t , s t ) and uses the same reparametrization trick for sampling ˜ z t as explicated in Appendix B.
Revised: z t ∼ p ( z t | a t , s t ) and uses the same reparametrization trick for sampling ˜ z t as explicated in Appendix B.

Original: Theorem C.2.
Revised: Theorem C.2.

Original: The RELAX estimator,
Revised: The RELAX estimator,

Original: [Equation]
Revised: [Equation]

Original: is unbiased.
Revised: is unbiased.

Original: Proof.
Revised: Proof.

Original: Note that by using the score-function estimator, for all t , we have
Revised: Note that by using the score-function estimator, for all t , we have

Original: E p ( a 1: t ,s 1: t )

Original: [Equation]
Revised: [Equation]

Original: [Figure]

Original: Then, by adding and subtracting the same term, we have
Revised: Then, by adding and subtracting the same term [Equation] we have

Original: [Equation]
Revised: [Equation]

Revised: t )]

Revised: (cid:33)(cid:35) + (cid:88) t

Original: [Table]
Revised: E p ( a 1: t − 1 ,s 1: t ) (cid:104) ∂ ∂θ E p ( z t | s t ) [ c φ ( z t , s t )] (cid:105) =

Revised: E p ( τ ) (cid:104) ∞ (cid:88) t =

Revised: ∂ log π ( a t | s t , θ ) ∂θ (cid:32) ∞ (cid:88) t (cid:48)

Revised: =

Revised: t r t (cid:48) − E p ( z t | a t ,s t )

Revised: [ c φ ( z t , s t ) (cid:33) − ∂ ∂θ E p ( z t | a t ,s t )

Revised: [ c φ ( z t , s t )]

Revised: + ∂ ∂θ E p ( z t | s t )

Revised: [ c φ ( z t , s t )] (cid:105)

Original: Since p ( z t | s t ) is reparametrizable
Revised: Since p ( z t | s t ) is reparametrizable

Original: [Table]
Revised: [Table]

Original: we obtain the estimator in Eq.(16)
Revised: we obtain the estimator in Eq.(16)

Original: [Table]
Revised: [Table]

Original: D F URTHER RESULTS ON DISCRETE VARIATIONAL AUTOENCODERS
Revised: D F URTHER RESULTS ON DISCRETE VARIATIONAL AUTOENCODERS

Original: [Table] [Table]
Revised: [Table] [Table]

Original: In the one layer models we optimize the evidence lower bound (ELBO):
Revised: [Figure]

Original: [Equation]
Revised: [Figure]

Original: where q ( b 1 | x ) = σ ( x · W q + β q ) and p ( x | b 1 ) = σ ( b 1 · W p + β p ) with weight matrices W q , W p and bias vectors β q , β p . The parameters of the prior p ( b ) are also learned.
Revised: E E XPERIMENTAL D ETAILS E. D ISCRETE VAE

Original: We run all models for 2 , 000 , 000 iterations with a batch size of 24 .
Revised: We run all models for 2 , 000 , 000 iterations with a batch size of 24 .

Original: For the REBAR models, we tested learning rates in { . 005 , .
Revised: For the REBAR models, we tested learning rates in { . 005 , .

Original: RELAX adds more hyperparameters.
Revised: RELAX adds more hyperparameters.

Original: These are the depth of the neural network component of our control variate r ρ , the weight decay placed on the network, and the scaling on the learning rate for the control variate.
Revised: These are the depth of the neural network component of our control variate r ρ , the weight decay placed on the network, and the scaling on the learning rate for the control variate.

Original: We tested neural network models with l layers of 200 units using the ReLU nonlinearity with l ∈ { 2 , 4 } .
Revised: We tested neural network models with l layers of 200 units using the ReLU nonlinearity with l ∈ { 2 , 4 } .

Original: We trained the control variate with weight decay in { . 001 , . 0001 } .
Revised: We trained the control variate with weight decay in { . 001 , . 0001 } .

Original: We trained the control variate with learning rate scaling in { 1 , 10 } .
Revised: We trained the control variate with learning rate scaling in { 1 , 10 } .

Original: To limit the size of hyperparameter search for the RELAX models, we only test the best performing learning rate for the REBAR baseline and the next largest learning rate in our search set.
Revised: To limit the size of hyperparameter search for the RELAX models, we only test the best performing learning rate for the REBAR baseline and the next largest learning rate in our search set.

Original: In many cases, we found that RELAX allowed our model to converge at learning rates which made the REBAR estimators diverge.
Revised: In many cases, we found that RELAX allowed our model to converge at learning rates which made the REBAR estimators diverge.

Original: We believe further improvement could be achieved by tuning this parameter.
Revised: We believe further improvement could be achieved by tuning this parameter.

Revised: It should be noted that in our experiments, we found the RELAX method to be fairly insensitive to all hyperparameters other than learning rate.

Revised: In general, we found the larger (4 layer) control variate architecture with weight decay of .

Revised: 001 and learning rate scaling of 1 to work best, but only slightly outperformed other conﬁgurations.

Original: All presented results are from the models which achieve the highest ELBO on the validation data.
Revised: All presented results are from the models which achieve the highest ELBO on the validation data.

Original: E.1. T WO LAYER  MODEL
Revised: E.1. O NE LAYER LINEAR MODEL

Original: In the two layer  models we optimize the ELBO
Revised: In the 1 layer linear models we optimize the evidence lower bound (ELBO):

Original: [Equation]
Revised: [Equation]

Original: ) − log q ( b 2 |  b 1 )]
Revised: ] where q ( b 1 | x ) = σ ( x · W q + β q ) and p ( x | b 1 ) = σ ( b 1 · W p + β p ) with weight matrices W q , W p and bias vectors β q , β p . The parameters of the prior p ( b ) are also learned. E.1. T WO LAYER LINEAR MODEL

Revised: In the two layer linear models we optimize the ELBO

Original:  where q ( b 1 | x ) = σ ( x · W q 1 + β q 1 ) , q ( b 2 | b 1 )
Revised: [Equation] ] where q ( b 1 | x ) = σ ( x · W q 1 + β q 1 ) , q ( b 2 | b 1 )

Original: =
Revised: =

Original: σ ( b 1 · W q 2 + β q 2 ) , p ( x | b 1 )
Revised: σ ( b 1 · W q 2 + β q 2 ) , p ( x | b 1 )

Original: =
Revised: =

Original: σ ( b 1 · W p 1 + β p 1 ) , and p ( b 1 | b 2 )
Revised: σ ( b 1 · W p 1 + β p 1 ) , and p ( b 1 | b 2 )

Original: =
Revised: =

Original: σ ( b 2 · W p 2 + β p 2 ) with weight matrices W q 1 , W q 2 , W p 1 , W p 2 and biases β q 1 , β q 2 , β p 1 , β p 2 .
Revised: σ ( b 2 · W p 2 + β p 2 ) with weight matrices W q 1 , W q 2 , W p 1 , W p 2 and biases β q 1 , β q 2 , β p 1 , β p 2 .

Original: As in the one layer model, the prior p ( b 2 ) is also learned.
Revised: As in the one layer model, the prior p ( b 2 ) is also learned.

Revised: E.1. N ONLINEAR MODEL

Revised: In the one layer nonlinear model, the mappings between random variables consist of 2 deterministic layers with 200 units using the hyperbolic-tangent nonlinearity followed by a linear layer withunits.

Original: We run an identical hyperpameter search in the 2 layer model as we do in the 1 layer model.
Revised: We run an identical hyperpameter search in all models.

Original: E. D ISCRETE RL
Revised: E. D ISCRETE RL

Original: In both the baseline A2C and RELAX models, the policy and control variate (value function in the baseline model) were 2 layer neural networks with 10 units per layer.
Revised: In both the baseline A2C and RELAX models, the policy and control variate (value function in the baseline model) were 2 layer neural networks with 10 units per layer.

Original: The ReLU non linearity was used on all layers except for the output layer.
Revised: The ReLU non linearity was used on all layers except for the output layer which was linear.

Original: For these tasks we estimate the policy gradient with a single Monte Carlo sample.
Revised: For these tasks we estimate the policy gradient with a single Monte Carlo sample.

Original: We run one episode of the environment to completion, compute the discounted rewards, and run one iteration of gradient decent.
Revised: We run one episode of the environment to completion, compute the discounted rewards, and run one iteration of gradient descent.

Original: We believe using larger batches will improve performance but would less clearly demonstrate the potential of our method.
Revised: We believe using larger batches will improve performance but would less clearly demonstrate the potential of our method.

Original: As our control variate does not have the same interpretation as the value function of A2C, it was not directly clear how to add reward bootstrapping and other variance reduction techniques common in

Original: RL into our model. We leave the task of incorporating these and other variance reduction techniques to future work. Both models were trained with the RMSProp (Tieleman & Hinton, 2012) optimizer and a reward discount factor of . 99 was used. 
Revised: Both models were trained with the RMSProp (Tieleman & Hinton, 2012) optimizer and a reward discount factor of . 99 was used. Entropy regularization with a weight of . 01 was used to encourage exploration.

Original: Both models have 2 hyperparameters to tune; the global learning rate and the scaling factor on the learning rate for the control variate (or value function).
Revised: Both models have 2 hyperparameters to tune; the global learning rate and the scaling factor on the learning rate for the control variate (or value function).

Original: We complete a grid search for both parameters in { 0 . 01 , 0 . 003 , 0 . 001 } and present the model which “solves” the task in the fewest number of episodes averaged over 5 random seeds.
Revised: We complete a grid search for both parameters in { 0 . 01 , 0 . 003 , 0 . 001 } and present the model which “solves” the task in the fewest number of episodes averaged over 5 random seeds.

Original: “Solving” the tasks was deﬁned by the creators of the OpenAI gym (Brockman et al., 2016).
Revised: “Solving” the tasks was deﬁned by the creators of the OpenAI gym (Brockman et al., 2016).

Original: The Cart Pole task is considered solved if the agent receives an average reward greater than 195 over 100 consecutive episodes.
Revised: The Cart Pole task is considered solved if the agent receives an average reward greater than 195 over 100 consecutive episodes.

Original: The Lunar Lander task is considered solved if the agent receives an average reward greater than 200 over 100 consecutive episodes.
Revised: The Lunar Lander task is considered solved if the agent receives an average reward greater than 200 over 100 consecutive episodes.

Original: The Cart Pole experiments were run for 250,000 frames.
Revised: The Cart Pole experiments were run for 250,000 frames.

Original: The Lunar Lander experiments were run for 5,000,000 frames.
Revised: The Lunar Lander experiments were run for 5,000,000 frames.

Revised: The results presented for the CartPole and LunarLander environments were obtained using a slightly biased sampler for p ( z | b, θ ) .

Original: E. C ONTINUOUS RL
Revised: E. C ONTINUOUS RL

Original: The continuous tasks uses both the value function and the control variate to enable bootstrapping, which is needed due to the increased complexity of the problem.

Original: The three models- policy, value, and control variate, are 2 layer neural networks with 64 hidden units per layer.
Revised: The three models- policy, value, and control variate, are 2 layer neural networks with 64 hidden units per layer.

Original: The value and control variate networks are identical, with the ELU(Djork-Arn´e Clevert & Hochreiter, 2016) nonlinearity in each hidden layer.
Revised: The value and control variate networks are identical, with the ELU (Djork-Arn´e Clevert & Hochreiter, 2016) nonlinearity in each hidden layer.

Original: The policy network has tanh nonlinearity.
Revised: The policy network has tanh nonlinearity.

Original: The policy network, which parameterizes the Gaussian policy comprises of a network (with the architecture mentioned above) that outputs the mean, and a separate, trainable log standard deviation value that is not input dependent.
Revised: The policy network, which parameterizes the Gaussian policy comprises of a network (with the architecture mentioned above) that outputs the mean, and a separate, trainable log standard deviation value that is not input dependent.

Original: All three networks have a linear output layer.
Revised: All three networks have a linear output layer.

Original: We selected the batch size to be 2500, meaning for a ﬁxed timestep (2500) we collect multiple rollouts of a task and update the networks’ parameters with the batch of episodes.
Revised: We selected the batch size to be 2500, meaning for a ﬁxed timestep (2500) we collect multiple rollouts of a task and update the networks’ parameters with the batch of episodes.

Original: Per one policy update, we optimize both the value and control variate network multiple times.
Revised: Per one policy update, we optimize both the value and control variate network multiple times.

Original: The number of times we train the value network is ﬁxed to 25, while for the control variate, it was chosen to be a hyperparameter.
Revised: The number of times we train the value network is ﬁxed to 25, while for the control variate, it was chosen to be a hyperparameter.

Original: All models were trained using ADAM (Kingma & Ba, 2015), with β 1 = 0 .
Revised: All models were trained using ADAM (Kingma & Ba, 2015), with β 1 = 0 .

Original: 9 , β 2 = 0 .
Revised: 9 , β 2 = 0 .

Original: 999 , and (cid:15) =
Revised: 999 , and (cid:15) =

Original: 1 e − 08 .
Revised: 1 e − 08 .

Original: The baseline A2C case has 2 hyperparameters to tune: the learning rate for the optimizer for the policy and value network.
Revised: The baseline A2C case has 2 hyperparameters to tune: the learning rate for the optimizer for the policy and value network.

Original: A grid search was done over the set: { 0 . 03 , 0 . 003 , 0 . 003 } .
Revised: A grid search was done over the set: { 0 . 03 , 0 . 003 , 0 . 0003 } .

Original: RELAX has 4 hyperparameters to tune: 3 learning rates for the optimizer per network, and the number of training iterations of the control variate per policy gradient update.
Revised: RELAX has 4 hyperparameters to tune: 3 learning rates for the optimizer per network, and the number of training iterations of the control variate per policy gradient update.

Original: Due to the large number of hyperparameters, we restricted the size of the grid search set to { 0 . 003 , 0 . 0003 } for the learning rates, and { 10 , 25 , 50 } for the control variate training iteration number.
Revised: Due to the large number of hyperparameters, we restricted the size of the grid search set to { 0 . 003 , 0 . 0003 } for the learning rates, and { 10 , 25 , 50 } for the control variate training iteration number.

Original: We chose the hyperparameter setting that yielded the shortest episode-to-completion time averaged over 5 random seeds.
Revised: We chose the hyperparameter setting that yielded the shortest episode-to-completion time averaged over 5 random seeds.

Original: As with the discrete case, we used the deﬁnition of completion deﬁned by OpenAI gym
Revised: As with the discrete case, we used the deﬁnition of completion deﬁned by OpenAI gym

Original: (Brockman et al., 2016) for each task.
Revised: (Brockman et al., 2016) for each task.

Original: The Inverted Pendulum experiments were run for 1,000,000 frames. The Inverted Double Pendulum experiments were run for 50,000,000 frames.
Revised: The Inverted Pendulum experiments were run for 1,000,000 frames. The Inverted Double Pendulum experiments were run for 5,000,000 frames.