KF-LAX: Kronecker-factored curvature estimation for control variate optimization in reinforcement learning

TL;DR

This paper introduces KF-LAX, a method that applies Kronecker-factored curvature estimation to improve sample efficiency and reduce variance in gradient-based reinforcement learning optimization.

Contribution

It combines KFAC with RELAX for control variate optimization, enhancing sample efficiency in reinforcement learning tasks.

Findings

Improved sample efficiency in synthetic and Atari tasks

Reduced variance in gradient estimates

Enhanced performance over baseline methods

Abstract

A key challenge for gradient based optimization methods in model-free reinforcement learning is to develop an approach that is sample efficient and has low variance. In this work, we apply Kronecker-factored curvature estimation technique (KFAC) to a recently proposed gradient estimator for control variate optimization, RELAX, to increase the sample efficiency of using this gradient estimation method in reinforcement learning. The performance of the proposed method is demonstrated on a synthetic problem and a set of three discrete control task Atari games.

Equations

$$f(x){\nabla }_{\theta }\log p(x|\theta )$$

$$\frac{\partial f}{\partial T}\frac{\partial T}{\partial \theta },x=T(ϵ,\theta ),ϵ\sim p(ϵ)$$

$${\hat{g}}_{new}(x)=\hat{g}(x)-c(x)+{\mathbb{E}}_{p(x|\theta )}[c(x)]$$

$$F_l=\mathbb{E}[vec\{{\nabla }_{W_l}L\}vec\{{\nabla }_{W_l}L{\}}^T]=\mathbb{E}[a{a}^T\otimes {\nabla }_{s}L({\nabla }_{s}L{)}^T]$$

$$\approx \mathbb{E}[a{a}^T]\otimes \mathbb{E}[{\nabla }_{s}L({\nabla }_{s}L{)}^T]:=A\otimes S:={\hat{F}}_l$$

$$(X\otimes Y{)}^{-1}=X^{-1}\otimes Y^{-1}$$

$$\Delta W_l={\widehat{F_l}}^{-1}{\nabla }_{W_l}J=A^{-1}{\nabla }_{W_l}J{S}^{-1}$$

$${\hat{g}}_{LAX}={\hat{g}}_{REINFORCE}[f]-{\hat{g}}_{REINFORCE}[c_{\varphi }]+{\hat{g}}_{REPARAM}[c_{\varphi }]$$

$$=[f(x)-c_{\varphi }(x)]\frac{\partial }{\partial \theta }\log p(x|\theta )+\frac{\partial }{\partial \theta }{c}_{\varphi }(x),x=T(\theta ,ϵ),ϵ\sim p(ϵ)$$

$${\hat{g}}_{RELAX}=[f(x)-c_{\varphi }(\tilde{z})]\frac{\partial }{\partial \theta }\log p(x|\theta )+\frac{\partial }{\partial \theta }{c}_{\varphi }(z)-\frac{\partial }{\partial \theta }{c}_{\varphi }(\tilde{z}),$$

$$x=H(z),z\sim p(z|\theta ),\tilde{z}\sim p(z|x,\theta )$$

$$\nabla J(\theta )={\mathbb{E}}_{s\sim {\rho }^{\pi }(s),a,\tau }[Q^{\pi }(s,a){\nabla }_{\theta }\log {\pi }_{\theta }(a|s)]={\mathbb{E}}_{s\sim {\rho }^{\pi }(s),a,\tau }[A^{\pi }(s,a){\nabla }_{\theta }\log {\pi }_{\theta }(a|s)]$$

$${\hat{g}}_{RELAX}^{RL}=\sum _{t=1}^T\frac{\partial }{\partial \theta }\log \pi (a_t|s_t,\theta )[{\hat{Q}}^{\pi }(a_t,s_t)-c_{\varphi }({\tilde{z}}_t,s_t)]-\frac{\partial }{\partial \theta }{c}_{\varphi }({\tilde{z}}_t,s_t)+\frac{\partial }{\partial \theta }{c}_{\varphi }(z_t,s_t)$$

$$\tilde{w}=\tilde{\nabla }J(\theta )$$

$$\sum _{s,a}{\rho }^{\pi }(s)\pi (a|s,\theta )\psi (s,a)(\psi (s,a{)}^T\tilde{w}-A^{\pi }(s,a))=0$$

$$\sum _{s,a}{\rho }^{\pi }(s)\pi (a|s,\theta )\psi (s,a)\psi (s,a{)}^T\tilde{w}=\sum _{s,a}{\rho }^{\pi }(s)\pi (a|s,\theta )\psi (s,a)A^{\pi }(s,a)$$

$$\nabla J(\theta )={\mathbb{E}}_{s\sim {\rho }^{\pi }(s),a,\tau }[Q^{\pi }(s,a){\nabla }_{\theta }\log {\pi }_{\theta }(a|s)]={\mathbb{E}}_{s\sim {\rho }^{\pi }(s),a,\tau }[A(s,a)\nabla \log \pi (a|s)]$$

$$F(\theta )\tilde{w}=\nabla J(\theta )$$

$$\pi (a|s,{\theta }^{\prime })=\pi (a|s,\theta )(1+f^{\pi }(s,a;w))+O({\alpha }^2)$$

$$\pi (a|s,{\theta }^{\prime })=\pi (a|s,\theta )+\frac{\partial \pi (a|s,\theta )}{\partial \theta }\Delta \theta +O(\Delta {\theta }^2)$$

$$=\pi (a|s,\theta )(1+\alpha \psi (s,a{)}^T\tilde{w})+O(\Delta {\theta }^2)$$

$$=\pi (a|s,\theta )(1+\alpha f^{\pi }(s,a;w))+O({\alpha }^2)$$

Full text

1 Introduction

Latent variable models are widely used in machine learning. Specifically, discrete latent variable models have been studied in many areas such as reinforcement learning (RL) and natural language processing. In the RL applications, for the optimization of both continuous and discrete variable models, policy gradient methods (Williams, 1992, Glynn, 1990) have gained high attention and have been applied in a wide variety of tasks. However, they mostly suffer from having a high variance for the gradient estimation.

Various works have focused on reducing the variance of the estimated gradient in the policy gradient methods. A subset of these approaches try to reduce the variance through use of a well designed control variate or a baseline. For example, Q-prop (Gu et al., 2016) and MuProp (Gu et al., 2015) use control variate based on the first order Taylor expansion of the target function around a fixed point, and (Wu et al., 2018) uses a factorized action-dependent baseline. In many applications, despite the estimated gradient being unbiased, it still suffers from high variance even when it is aligned with a control variate. In cases that the target function is continuous, reparameterization trick (Kingma and Welling, 2013, Rezende et al., 2014) could be used as a way of obtaining a low variance gradient estimator.

Particularly, in discrete domains, concrete relaxation (Maddison et al., 2016, Jang et al., 2016) uses this technique to obtain a low variance but biased gradient estimator. More recent gradient estimators such as REBAR (Tucker et al., 2017) and RELAX (Grathwohl et al., 2017) leverage both the control variate and the reparameterization trick to obtain a low variance and unbiased estimator. To estimate the gradient and apply the reparameterization trick, REBAR uses two different uniform random variables. In general, REBAR is an unbiased gradient estimator. RELAX gradient estimation approach is further discussed in the next sections.

Natural gradient descent is an approach that leverages the local curvature structures of a target function and can provide gradient updates invariant to the parameterization. In problems that parameter space has an underlying structures, the ordinary gradient does not represent the steepest descent direction, but the natural gradient does (Amari, 1998). Therefore, an efficient estimation of the natural gradient can provide a practical way to build a faster training procedure. In particular, in reinforcement learning, natural policy gradient method (Kakade, 2002) was introduced over a decade ago, and works such as (Wu et al., 2017) and (Schulman et al., 2015) successfully applied it in practice for discrete and continuous control tasks. As an another practical usage of the natural gradient, authors in (Hoffman et al., 2013) used a noisy computation of natural gradient to update variational parameters in mean field variational inference, and applied it for topic modeling of large data sets.

KFAC (Martens and Grosse, 2015) is a technique to approximate the local curvature structure of a given function with an efficient way of using mini-batch samples of data. Sample efficiency and fast optimization procedure are critical challenges in deep RL. As a second order optimization technique, leveraging KFAC to estimate the natural gradient can be helpful to train an agent faster. In this paper, inspired by KFAC and a recently introduced gradient estimator, RELAX (Grathwohl et al., 2017), we develop a gradient optimization procedure for black box functions which approximates the natural gradient for a differentiable surrogate function as a control variate. Black box functions are particularly important in RL where the reward function is not necessarily known. Section 2 describes the preliminaries. Section 3 is devoted for the proposed method. Section 4 provides a summary of related works. The proposed method is applied to a synthetic problem and a few discrete control task Atari games in section 5, and the conclusion and future work are provided in the last sections.

2 Preliminaries

Many loss functions in machine learning such as variational objective in variational inference, expected reward in reinforcement learning, and empirical risk in supervised learning can be written as an expectation of a function over the data distribution. Consider the optimization problem of a function in the form of $\mathbb{E}_{p(x|\theta)}[f(x)]$ with respect to $\theta$. Some of the popular techniques to obtain a gradient estimator for the parameters of this loss function are described in sections 2.1-2.3. In section 2.4 the natural gradient and KFAC are explained in more details.

2.1 Score function estimator

A general approach to estimate the gradient of $\mathbb{E}_{p(x|\theta)}[f(x)]$ with respect to $\theta$ is known as the score function estimator (also known as REINFORCE, and likelihood ratio estimator) (Williams, 1992, Glynn, 1990) which uses the derivatives of $\log p(x|\theta)$ to estimate the gradient as,

[TABLE]

Where $x$ is sampled from $p(x|\theta)$. A straight forward calculation shows that the score function is an unbiased gradient estimator (Williams, 1992). Moreover, this estimator assumes no special restriction for the function $f$ and is generally applicable for black box functions. In situations that $f$ directly depends on $\theta$ and the dependency is known, the derivatives of $f$ can be added to the above estimator to come up with an unbiased estimator.

2.2 Reparameterization trick

Reparameterization trick (Kingma and Welling, 2013, Rezende et al., 2014) is a technique for cases that $f$ is a continuous and differentiable function, and $x$ can be written as a continuous function of $\theta$ and a random noise $\epsilon$ with a known distribution. The reparameterized gradient estimator has the form,

[TABLE]

Reparameterization gradient estimate is generally unbiased and has low variance. In practice, the standard Gaussian and Gumbel distributions are popular distributions for reparameterizing the continuous and discrete $x$ values, respectively (Maddison et al., 2016, Rezende et al., 2014, Kingma and Welling, 2013, Jang et al., 2016).

2.3 Control variates

For a given gradient estimator $\hat{g}$, using a well designed function $c(x)$ which $\mathbb{E}_{p(x|\theta)}[c(x)]$ can be calculated analytically can result a lower variance gradient estimator without changing the bias of the initial estimator as follow,

[TABLE]

Stronger positive correlation between $c(x)$ and $\hat{g}(x)$ results more reduction in the variance of the new estimator.

2.4 Natural gradient and Kronecker-factored curvature approximation

Consider the problem of minimizing $J(\theta)$ with respect to $\theta\in\Theta$. A Conventional gradient descent approach uses Euclidean distance over the parameter space $\Theta$. However, Euclidean distance does not benefit from the curvature information of the function and may result a slower convergence. Natural gradient (Amari, 1998) uses particular Mahalanobis distance as a local metric and gives an update rule of the form $\Delta\theta\propto-F^{-1}\nabla_{\theta}J$, where $F$ is the Fisher information matrix which is a function of $\theta$.

Modern neural networks could have millions of parameters. In cases that the optimization is with respect to a neural network parameters, computing the Fisher information matrix and its inverse can be computationally impractical and inefficient. For example, for a neural network with only one fully-connected layer with 1000 input and 1000 outputs, direct estimation of the entries of the Fisher information matrix which is a square matrix with a million rows might be inefficient. Kronecker-factored approximate curvature (KFAC) (Martens and Grosse, 2015) leverages the properties of Kronecker product to approximate the Fisher information matrix and its inverse efficiently. Consider a multi-layer neural network. Let the output of the neural network be $p(y|x)$ where $x$ denotes the input, and let $L=\log p(y|x)$ indicates the log likelihood. For the $l$-th layer of the neural network, let $d_{in}$, $d_{out}$ denote the input and output size of the layer (the index $l$ is dropped for simplicity), $W_{l}\in\mathbb{R}^{d_{in}\times d_{out}}$ the weights, $a\in\mathbb{R}^{d_{i}n}$ input vector to the layer, and $s=W_{l}a$ as the pre-activation vector for the next layer. We note that $\nabla_{W_{l}}L=(\nabla_{s}L)a^{T}$. KFAC uses this relationship to approximate the Fisher information matrix in a block diagonal manner, where each block corresponds to a layer. For the layer $l$, denote the corresponding block as $F_{l}$. The approximate Fisher information for this layer can be written as,

[TABLE]

[TABLE]

Where $A:=\mathbb{E}[aa^{T}]$, and $S:=\mathbb{E}[\nabla_{s}L(\nabla_{s}L)^{T}]$. $A$ and $S$ can be approximated during training using mini-batches of the data, or through the states and actions resulting from the agent experience of different trajectories. Having these approximations, the inverse of the block could be obtained using the following property of Kronecker product,

[TABLE]

Therefore, the natural gradient update for the layer $l$ can be approximated as,

[TABLE]

For more details about the derivation refer to the (Martens and Grosse, 2015).

3 Proposed Method

This section uses KFAC for the LAX/RELAX (Grathwohl et al., 2017) gradient estimators which can be applied in both continuous and discrete domains. LAX (continuous cases) and RELAX (discrete cases) leverage the score function and the reparameterization trick to obtain a low variance estimator without adding bias to the initial gradient estimator.

Let $x$ be a continuous random variable and consider the optimization of $\mathbb{E}_{p(x|\theta)}[f(x)]$ with respect to $\theta$. The LAX estimator (Grathwohl et al., 2017) is given by,

[TABLE]

[TABLE]

where $c_{\phi}$ is the surrogate function. Parameters of the surrogate function are trained in order to minimize the variance of the gradient estimator. In case that that $c_{\phi}$ is equal to $f$, LAX and the reparameterization gradient estimator are the same. By applying KFAC to estimate the natural gradient for the surrogate function parameters, KF-LAX procedure is obtained. The procedure is described in the algorithm 1.

Similarly, in case that $x$ is a discrete random variable, RELAX can be written as,

[TABLE]

[TABLE]

where $H$ represents the heaviside function222$H(z)=1$ if $z\geq 0$ and $H(z)=0$ if $z<0$. Following the same approach for the LAX, we get another optimization algorithm for RELAX which we call KF-RELAX. This algorithm is provided in Appendix A.

4 Experiments

4.1 Synthetic problem

KF-RELAX is applied to a synthetic problem. Consider the minimization of $\mathbb{E}_{p(b|\sigma(\theta))}[(b-t)^{2}]$ where $t\in(0,1)$ is a constant and $b$ is a binary random variable sampled from a Bernoulli distribution with parameter $\sigma(\theta)$, where $\sigma(.)$ indicates the sigmoid function. According to the selection of $t$, the optimal distribution is a deterministic Bernoulli distribution with parameter [math] or $1$. Figure 1 compares the convergence of REINFORCE, KF-RELAX and RELAX for $t=0.499$, using only a single sample of $x$ per iteration. Also, in appendix B, it is empirically shown that all gradient estimators are unbiased.

4.2 Reinforcement learning

Consider a finite discrete-time Markov Decision Process (MDP) which is a tuple $(\chi,\rho_{0},\gamma,A,r,P)$. At time $t=0$, an agent is in the state $s_{0}$ which is sampled from the initial distribution $\rho_{0}:\chi\rightarrow\mathbb{R}$. At each time $t$, the agent chooses an action $a_{t}\in A$ from its policy distribution $\pi_{\theta}(a_{t}|s_{t})$, gets a reward according to the reward function $r(s_{t},a_{t})$ and transitions to the next state $s_{s+1}$ according to the Markov transition probabilities $P(s_{t+1}|s_{t},a_{t})$. Final goal of the agent is to maximize its expected reward $J(\theta)=\mathbb{E}_{s_{0},a_{0},\tau}[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})]=\mathbb{E}_{s\sim\rho^{\pi}(s),a,\tau}[\sum_{t=0}^{\infty}\gamma^{t}r(s_{t},a_{t})]$, where $\tau$ indicate the trajectories, $\gamma$ is the discount factor and $\rho^{\pi}=\sum_{t=0}^{\infty}\gamma^{t}P(s_{t}=s)$ is the unnormalaized state visitation frequency. Using the policy gradient theorem (Sutton et al., 1998), the gradient of $J(\theta)$ can be written in the following expectation form,

[TABLE]

Where, $Q^{\pi}(s,a)$ is the state-action value function, $V^{\pi}(s)$ is the value function, and $A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s)$ is the advantage function. Due to simplicity of the selected tasks, same as (Grathwohl et al., 2017), we use the $Q$ function approximation, instead of directly estimating the advantage function.

4.2.1 Discrete Control Tasks

KF-RELAX and RELAX performance are compared in three discrete control RL tasks. According to appendix C of the (Grathwohl et al., 2017), using a sampled trajectory, the RELAX gradient estimator can be written as,

[TABLE]

KF-RELAX gradient estimator has a similar form with a difference that at each step, the natural gradient of the surrogate function is estimated.

The learning rate for both the agent and the surrogate function were chosen from the set $\{0.03,0.01,10^{-3},10^{-4}\}$. $c_{\phi}$ is chosen as a three fully-connected layer neural network. For the RELAX, all models were trained using Adam (Kingma and Ba, 2014). Entropy regularization with a weight of $0.01$ was used to increase the exploration, and the discount factor was set to $0.99$. Tikhonov damping technique described by (Martens and Grosse, 2015) with the value in $\{0.1,0.01,10^{-3},5\times 10^{-4}\}$ was used to further control the training stability. The same value of mini-batch size was used for both gradient estimators. In addition, as in (Wu et al., 2017), a trust region approach was used for KF-RELAX, whereby the parameter updates capped at a specific upper bound to prevent the agent from converging to a near-deterministic policy during the early training stages. The value of the upper bound was chosen from $\{10^{-3},10^{-4},10^{-5},10^{-6}\}$. The experiments are done on Cart-pole, Lunar-Lander, and Acrobot which were selected from the OpenAI gym environments (Brockman et al., 2016). The comparison of KF-RELAX and RELAX in these environments are shown in the figure 2.

5 Related Works

Natural policy gradient method was first introduced in (Kakade, 2002). In particular case where the state-value function is approximated with a compatible function approximator, the natural gradient tends toward the optimal greedy policy improvement (Appendix C shows the same fact holds for the advantage function which is frequently used for policy optimization). Natural gradient was used for actor-critic methods by (Peters and Schaal, 2008). More recently, Trust Region Policy Optimization (TRPO) (Schulman et al., 2015) performs an optimization for each step of the policy update and the natural gradient policy update is a particular case of TRPO. However, TRPO does not scale to neural networks with larger architectures.

Kronecker-factored curvature approximation (KFAC) technique is also applied in practice to improve the training convergence speed. For example, (Grosse and Martens, 2016) used KFAC to estimate the Fisher information of the convolutional layers in neural networks. Recently, (Martens et al., 2018) extended KFAC to recurrent neural networks. (Wu et al., 2017) applies KFAC to optimize both the actor and the critic of an actor-critic model for deep RL and proposed the ACKTR method. Inspired by ACKTR, this work uses the KFAC technique for the surrogate function in RELAX. In the continuous deep RL applications, a vanilla Gauss-Newton method can be used to leverage KFAC for the agent as well (Wu et al., 2017, Nocedal and Wright, 2006).

6 Conclusion and Future Work

Even though the KF-LAX (KF-RELAX) could provide stronger update rules for optimization and reduce the number of samples for convergence, tuning the hyperparameters for KFAC can be challenging in the deep RL framework as the estimation of the gradient may exhibit a high variance. Ill specification of parameters like the damping factor, and the scaling factor may cause a training collapse or no improvement in sample size needed for training. Fortunately, these hyperparameters are studied for practical usage and these studied are beneficial for the tuning process (Martens and Grosse, 2015, Wu et al., 2017).

Choosing a complicated architecture for the surrogate function can negatively influence the training progress or even add to the variance of the gradient estimation, hence degrading the convergence. In our experiments, a surrogate function with more than three layers increases the gradient estimation variance. This is considered a limitation of using this approach. Stabilizing the training process for more complicated surrogate functions is considered as a future work.

In this work we leveraged two recently introduced optimization techniques, KFAC and RELAX, to get more powerful gradient steps. We experimented the proposed approach on a few discrete control tasks in reinforcement learning. We believe this approach can be promising in improving the sample efficiency of RELAX in reinforcement learning. Further experiments for discrete and continuous tasks are considered as the future works.

Appendix A KF-RELAX

Using the KFAC technique for the RELAX gradient estimator results the KF-RELAX procedure for discrete cases which is explained in the Algorithm 2.

[TABLE]

Appendix B RELAX and KF-RELAX are unbiased

Figure 3 shows that same as the REIFNORCE, RELAX and KF-RELAX are unbiased and converge to the true optimal solution in the synthetic problem which is described in section 4.1. Since for $t=0.499$ the REINFORCE gradient estimator does not necessarily converge, a simpler case $t=0.49$ is used in figure 3.

Appendix C Use of the advantage function in the natural policy gradient

In this section we rewrite a similar results to (Kakade, 2002) that with an approximation of the advantage function by some compatible function approximator, the natural gradient update rule tends to move toward the best action.

Lemma: Suppose an approximation of the advantage function with a function $f^{\pi}(s,a;w)$ is desired. Let $f^{\pi}$ be in the form of $w^{T}\nabla\log\pi(a|s,\theta)$ and let $\tilde{w}$ be the minimizer of the square error $\epsilon(w,\pi)=\mathbb{E}_{s\sim\rho^{\pi},a}[(A^{\pi}(s,a)-f^{\pi}(s,a;w))^{2}]$. Then

[TABLE]

where $J(\theta)=\mathbb{E}_{s,a,\tau}[\sum_{t=0}^{\infty}\gamma^{t}r_{t}]$, and $\tilde{\nabla}$ indicates the natural gradient.

Proof: Since $\tilde{w}$ minimizes the $\epsilon(w,\pi)$, we have $\frac{\partial\epsilon}{\partial w_{i}}=0$. For simplicity in the notation we set $\psi(s,a)=\nabla\log\pi(a|s,\theta)$. Therefore,

[TABLE]

Rearranging this equation,

[TABLE]

The policy gradient theorem (Sutton et al., 1998) states that

[TABLE]

Also, by using the fact that $\nabla\pi(a|s,\theta)=\pi(a|s,\theta)\psi^{\pi}(s,a)$ and the definition of the Fisher information matrix $F(\theta)=\sum_{s,a}\rho^{\pi}(s)\pi(a|s,\theta)\psi(s,a)\psi(s,a)^{T}$ both sides can be rewritten as,

[TABLE]

Solving for $\tilde{w}$ gives $\tilde{w}=F^{-1}(\theta)\nabla J(\theta)$ which is the definition of the natural gradient. $\square$

The following lemma states that based on our approximation of the $A^{\pi}(a|s,\theta)$, the natural gradient locally tends toward the best action. An action $a$ is best if $a\in argmax_{a^{\prime}}f^{\pi}(s,a^{\prime};\tilde{w})$. Intuitively, the best action is defined by performing greedy policy with regard to our approximation of the advantage function.

Lemma: Let $\theta^{\prime}=\theta+\Delta\theta$ be the natural gradient update of the policy parameters where $\Delta\theta=\alpha\tilde{\nabla}J(\theta)=\alpha\tilde{w}$, then

[TABLE]

Proof: To the first order Taylor approximation,

[TABLE]

[TABLE]

[TABLE]

$\square$