Randomized Advantage Transformation (RAT): Computing Natural Policy Gradients via Direct Backpropagation
Mingfei Sun 1
arXiv:2605.18591v1 [cs.LG] 18 May 2026
Abstract
Despite their advantages, natural policy gradients are rarely used directly in large-scale deep reinforcement learning due to computational constraints. The Fisher matrix scales with the number of policy parameters, making explicit construction or inversion infeasible. To address this, prior work has largely followed two directions. Hessian-free approaches compute Fisher-vector products and solve the resulting linear systems using conjugate gradient methods, as in TRPO (Schulman et al., 2015). While effective, these methods introduce substantial computational overhead, require careful tuning of inner-loop solvers, and are difficult to apply in shared actor-critic architectures (Schulman et al., 2017b; Wu et al., 2017). Alternatively, structured approximations such as KFAC (Martens & Grosse, 2015) exploit layer-wise factorizations of the Fisher matrix, trading accuracy for efficiency but relying on architecture-dependent assumptions (Benzing, 2022) and nontrivial implementation optimizations (George et al., 2018).
Natural policy gradients improve optimization by accounting for the geometry of distribution space, but their practical use is limited by the cost of estimating and inverting the Fisher matrix. We present Randomized Advantage Transformation (RAT), a method for estimating Tikhonovregularized natural policy gradients via direct backpropagation. By applying the Woodbury formula, we reformulate the regularized natural policy gradients as vanilla policy gradients with a transformed advantage. RAT computes this transformation efficiently via randomized block Kaczmarz iterations on on-policy mini-batches, avoiding explicit Fisher construction, conjugategradient solvers, and architecture-specific approximations. We provide convergence guarantees for RAT and demonstrate empirically that it matches or exceeds established natural-gradient methods across continuous and visual control benchmarks, while remaining simple to implement and compatible with various architectures.
In this work, we show that estimating natural policy gradients can be reduced to a simpler and more general procedure. Our starting point is the observation that Tikhonovregularized natural policy gradients admit an equivalent least-squares formulation. By applying the Woodbury formula, we derive a representation in which the inverse Fisher matrix is absorbed into a transformation of the advantage function. Under this reformulation, the natural policy gradient takes the same form as a vanilla policy gradient, but with a modified advantage, see Figure 1 for an overview.
1. Introduction Natural policy gradients are a foundational tool in deep Reinforcement Learning (RL), offering parameterizationinvariant update directions (Bagnell & Schneider, 2003) by pre-conditioning policy gradients with the inverse Fisher matrix (Amari, 1998; Kakade, 2001). This geometric correction has been shown to significantly improve convergence properties (Agarwal et al., 2021), and underlies several influential RL algorithms, including Natural Actor-Critic (NAC) (Peters & Schaal, 2008), Trust Region Policy Optimization (TRPO) (Schulman et al., 2015), ACKTR (Wu et al., 2017), and connections to PPO-style updates (Schulman et al., 2017b; Hilton et al., 2022). 1
Building on this insight, we introduce Randomized Advantage Transformation (RAT), an algorithm that approximates the transformed advantage using randomized block Kaczmarz iterations (Needell & Tropp, 2014). RAT operates on small subsets of on-policy samples and iteratively refines an estimate of the regularized natural policy gradient. Each iteration requires only standard backpropagation through a surrogate loss, eliminating the need for explicit Fisher construction, Fisher-vector products, or architecture-specific curvature approximations.
The University of Manchester, United Kingdom. Code URL. Correspondence to: <[email protected]>.
We provide a convergence analysis showing that RAT converges linearly to the regularized natural policy gradient under standard assumptions on the function approximation and state-action coverage. We evaluate RAT on a range
Proceedings of the 43 rd International Conference on Machine Learning, Seoul, South Korea. PMLR 306, 2026. Copyright 2026 by the author(s).
1
Randomized Advantage Transformation (RAT) Woodbury Formula Matrix Inverse
Matrix Inverse
Natural Policy Gradients
Vanilla Policy Gradients with Transformed Advantage
Figure 1. Fisher matrix F ∈ R|θ|×|θ| estimated from samples τ (shown in green) is often ill-conditioned and hard to invert directly. Randomized Advantage Transformation (RAT) leverages Woodbury formula to replace the inversion of F to that of a sampled preconditioner of size n × n (shown in grey). The resulting inverse is absorbed into a randomized, block-wise transformation of the advantage function, yielding a surrogate objective whose first-order gradient via direct backpropagation approximates natural policy gradient.
of benchmark reinforcement learning tasks, including continuous control in MuJoCo (Brockman et al., 2016) and high-dimensional visual control in Procgen (Cobbe et al., 2020). Across these settings, RAT matches or outperforms established natural policy gradient methods while offering a simpler implementation and broader applicability, including support for shared actor-critic architectures. Our contributions are as follows:
practice, however, achieving sufficient accuracy often requires many CG iterations at each training iteration, leading to substantial computational overhead. To address this issue, alternative approaches rely on structured approximations of the Fisher matrix, including diagonal (Liu et al., 2024), layer-wise block-diagonal (Martens & Grosse, 2015; George et al., 2018), or low-rank forms (Dangel et al., 2023; Yang et al., 2022). Among these, Kronecker-Factored Approximate Curvature (KFAC) has been widely used in reinforcement learning (Wu et al., 2017; Bae et al., 2022). KFAC approximates each layer’s Fisher matrix as a Kronecker product of two smaller matrices, relying on independence assumptions about the statistics of the gradients (Ba et al., 2017) or their eigen-structure (George et al., 2018). While effective, KFAC depends strongly on the gradient structure of the neural networks. In contrast, our method RAT is oblivious to the particular model architecture and estimates natural policy gradients through standard backpropagation. Another influential line of work estimates natural policy gradients via natural actor-critic (NAC) formulations (Peters & Schaal, 2008; Cayci & Eryilmaz, 2025). In these methods, the advantage function is represented in a compatible form (Sutton et al., 1999), ensuring that the natural policy gradient emerges exactly as the solution to a leastsquares regression problem (Kakade, 2001; Schulman et al., 2017a). Our work is related in spirit, but differs in how the least-squares problem is constructed and solved.
• We derive a Woodbury-based reformulation of Tikhonov-regularized natural policy gradients as vanilla policy gradients with transformed advantages. • We propose Randomized Advantage Transformation (RAT), an efficient algorithm based on randomized block Kaczmarz iterations that estimates NPG using standard backpropagation. • We provide theoretical convergence guarantees and empirical evidence demonstrating the effectiveness of RAT on diverse benchmarks.
2. Related work We review prior work on estimating natural policy gradients and on using the Woodbury formula to reduce the cost of Fisher inversion. Estimating natural policy gradients. The central computational challenge in natural policy gradient methods is to estimate or apply the inverse Fisher matrix efficiently and robustly in the presence of sampling noise and function approximation. A widely adopted strategy avoids forming the Fisher explicitly and instead computes Fisher Vector Products (FVP), also known as Hessian-free methods (Martens, 2010; Pascanu & Bengio, 2014). The resulting linear systems are typically solved using iterative methods such as conjugate gradient (CG) algorithm. This approach underlies trust-region style reinforcement learning algorithms (Schulman et al., 2015) and many modern natural policy gradient implementations (Kuba et al., 2022; Sun et al., 2023). In
Woodbury for inverting Fisher. The Woodbury formula provides an efficient way to compute matrix inverses and solve linear systems involving low-rank updates. Its use in natural gradients dates back to Amari et al. (2000), where the Fisher inverse can be stored explicitly and is estimated directly using the Sherman-Morrison lemma, a special case of the Woodbury formula. More recently, this idea has been extended to rank-1 approximations of natural policy gradients with neural parameterization (Huo et al., 2026). Woodbury formula has also been used to reduce the computational cost 2
Randomized Advantage Transformation (RAT)
Vanilla and Natural Policy Gradients. We consider a parameterized policy π(a|s; θ). Let πk denote the policy induced by parameters θk , and dk denote the corresponding discounted state-action distribution. The vanilla policy gradient is given by (Sutton et al., 1999): ∂ := ∇PG J(θ) E log π(a|s; θ)A (s, a) . π (s,a)∼dπ θ ∂θ
of natural gradient optimization steps. For example, Chen & Heyl (2024) apply Woodbury-based updates to accelerate natural gradient methods, and related work develops a momentum scheme to further improve convergence, e.g., SRPING (Goldshlager et al., 2024). Only very recently has the full Woodbury formula been used to reformulate the Tikhonov-regularized natural gradients. For instance, Wu et al. (2024) employs the push-through identity to analyze per-sample loss reduction under natural gradient updates, while Guzmán-Cordero et al. (2025) uses Woodbury-based transformations to reduce the cost of Fisher inverse, for which convergence guarantees have been provided in Goldshlager et al. (2026). Our approach builds on these insights but differs in how the Woodbury reformulation is exploited. Instead of directly approximating the inverse Fisher, we use the Woodbury identity to transform the advantage function and then estimate the resulting natural policy gradient iteratively through standard backpropagation.
When the parameter space has a non-Euclidean geometry, the vanilla gradient does not correspond to the direction of steepest ascent. To address this, Amari (1998) proposed the natural gradient, which accounts for the information geometry of the parameter manifold. In reinforcement learning, Kakade (2001) introduced the natural policy gradient based on the Fisher matrix ∂ ∂ F (θ) := Edπ (s,a) log π(a|s; θ) ⊤ log π(a|s; θ) . ∂θ ∂θ The resulting natural policy gradient is defined as −1 PG ∇NPG ∇θ J(θ). θ J(θ) := F (θ)
Position of RAT. RAT differs from prior natural gradient methods in that it neither relies on conjugate-gradient solvers nor on structured Fisher approximations. Instead, it leverages a Woodbury-based reformulation to shift curvature information into an advantage transformation, which is approximated via randomized linear solvers. This perspective allows RAT to remain architecture-agnostic, computationally efficient, and compatible with various architectures.
We focus on Kakade’s formulation of the natural policy gradient, which has been shown to be invariant to reparameterization (Bagnell & Schneider, 2003). Following Kunstner et al. (2019), we refer to the Fisher matrix estimated from finite samples as the empirical Fisher, and call the resulting gradients the Empirical Natural Policy Gradients. Randomized block Kaczmarz method. The randomized block Kaczmarz method is an iterative algorithm for solving overdetermined least-squares problems of the form:
3. Preliminaries
2
min ∥Ax − b∥2 ,
Reinforcement learning formulation. We consider a standard reinforcement learning setup in which an agent interacts with an environment over discrete timesteps. At each timestep t, the agent observes a state st ∈ S, selects an action at ∈ A and receives a scalar reward rt ∈ R. The agent’s behavior is defined by a stochastic policy π, which maps states to actionPdistributions π : S 7→ ∆|A| (where n ∆d := {x ∈ Rn+ : i xi = 1} denotes the probability simplex). The environment is modelled as a Markov Decision Process (MDP) {S, A, P, r, p0 }, where P (·|s, a) is the transition kernel, r(s, a) is the reward function, and p0 is the initial state distribution. The (discounted) return PT from time step t is defined as Rt := i=t γ i−t r(si , ai ), where γ ∈ [0, 1) is the discount factor. The objective is to learn a policy that maximizes the expected return from the initial state distribution. The action-value function of a policy π is defined as Qπ (st , at ) := E [Rt |st , at ]. The value function Vπ and advantage h functioni Aπ (s, a) are given := by: Vπ (st ) Eat ∼π(·|st ) Qπ (st , at ) , and Aπ (s, a) :=
x
(1)
where A ∈ Rn×d and b ∈ Rn . Starting from an initial estimate x0 , the method iteratively refines the solution by projecting onto the solution space of randomly selected row blocks of A. Specifically, let T = {τ1 , τ2 , . . . , τm } be a partition of the rows of A. At iteration j, a block τj ∈ T is sampled (typically uniformly at random), and the update is: ⊤ −1 xj = xj−1 + A⊤ bτ − Aτj xj−1 . τj (Aτj Aτj ) Under standard assumptions, this randomized block scheme converges to the least-squares solution, often with favorable convergence properties compared to deterministic variants (Needell & Tropp, 2014).
4. Randomized Advantage Transformation In this section, we first show that Tikhonov-regularized natural policy gradients (NPG) can be formulated as a regularized least-squares problem. We then apply the Woodbury formula to express the resulting NPG update in terms of a transformed advantage, reducing it to a vanilla policy gradient form. Finally, we introduce a randomized Kaczmarz iteration to efficiently solve the corresponding least squares.
Qπ (s, a) − Vπ (s). We define P∞the discounted state distribution as dπ (s) := (1 − γ) t=0 γ t P (st = s|π, p0 ), where the factor (1 − γ) ensures normalization. With slight abuse of notation, we use dπ (s, a) to mean s ∼ dπ (s), a ∼ π(·|s). 3
Randomized Advantage Transformation (RAT)
4.1. Tikhonov Regularized NPG
Thus, the T-NPG corresponds to a vanilla policy gradient with a transformed advantage function,
As shown by Kakade (2001), the natural policy gradient can be obtained as the solution to a least squares problem. We follow the notation of Schulman et al. (2017a). Use n to denote the cardinality of state-action space, i.e., n = |S||A|, and p to denote the number of policy parameters, i.e., p = |θ|. Let Σ ∈ Rn×n be a diagonal matrix with diagonal entries dπ (s)π(a|s), H ∈ Rn×p be the matrix whose rows are ∂θ∂⊤ log π(a|s; θ), y ∈ Rn be the vector with entries Aπ (s, a). Under this notation, the vanilla PG and NPG can be written as: ⊤ PG: ∇PG θ J(θ) := H Σy,
Āπ (s, a) := (λIn + HH ⊤ Σ)−1 y (s,a) .
Importantly, the matrix inversion in the above transformation does not depend on the number of parameters p since λIn + HH ⊤ Σ is of size n × n. This important property is the key to develop our method. 4.3. Randomized Kaczmarz Iteration
(2)
⊤ −1 NPG: ∇NPG H ⊤ Σy. θ J(θ) := (H ΣH)
A major limitation of the above advantage transformation is that n is typically much larger than p, essentially in continuous state-action spaces, making the exact transformation infeasible. Since T-NPG naturally arises from a least-squares problem, we turn to randomized iterative solvers for linear systems (Gower & Richtárik, 2015) and propose to approximate the advantage transformation via randomization.
(3)
The NPG update coincides with the solution to the following least squares. Proposition 1 (Kakade (2001)). The minimizer of 2 minx ∥y − Hx∥Σ is given by x∗ = (H ⊤ ΣH)−1 H ⊤ Σy, provided the inverse exists.
Equation (5) involves weighting by Σ, which is generally unknown. To bypass this, we replace this objective with an unweighted least squares problem constructed from on-policy samples. This corresponds to a Monte Carlo approximation in which state-action pairs are drawn i.i.d from dπ (s)π(a|s), so that expectations over Σ are replaced by empirical averages. Specifically, sampling (s, a) ∼ dπ (s)π(a|s) yields
In practice, the Fisher matrix is estimated from a finite number of samples, (i.e., empirical Fisher), which can render it ill-conditioned or singular. To stabilize inversion, it is common to apply Tikhonov regularization by adding a damping term (Schulman et al., 2015; Martens & Grosse, 2015). T-NPG: ∇T-NPG J(θ) := (λI + H ⊤ ΣH)−1 H ⊤ Σy, (4) θ
2
2
2
(5)
To solve Equation (5), we adopt randomized block Kaczmarz method (Needell & Tropp, 2014). Let Dk denote onpolicy samples at iteration k, partitioned into mini-batches {τ1 , τ2 , . . . , τm }. Starting from an initial estimate g0 , at iteration j, we select a batch τj , and perform the following:
This formulation motivates our use of iterative least-squares solvers in the following section. 4.2. Woodbury Reformulation
2
Interestingly, the introduction of Tikhonov regularization enables the use of the Woodbury formula to transform the inverse (Wu et al., 2024; Guzmán-Cordero et al., 2025). Applying the formula, (I + U V )−1 U = U (I + V U )−1 for any conformable matrices U , V , twice yields ∇T-NPG J(θ) = (λIp + H ⊤ ΣH)−1 H ⊤ Σy θ
(6)
Σy
(7)
⊤
⊤ −1
= H (λIn + ΣHH )
= H ⊤ Σ (λIn + HH ⊤ Σ)−1 y .
2
gj ← arg min yτj − Hτj g 2 + λ ∥g − gj−1 ∥2 . (12) g
Here, Hτj denotes the rows of H indexed by τj . Note that this corresponds to a regularized block update in the randomized block Kaczmarz framework. In the classical (unregularized) formulation, each step projects the current iterate onto the solution space of the sampled linear system Hτj g = yτj . However, enforcing this hard constraint directly can be unstable in our setting due to noise and potential rank deficiency of minibatches. Thus, we consider the regularized version, which can be viewed as a proximal update that balances fitting the current batch with staying close to the previous estimate. Importantly, as shown in Goldshlager et al. (2024), Equation (12) admits a closed-form
(8)
⊤ Compared to vanilla policy gradients ∇PG θ J(θ) = H Σy, the only difference is the advantage term. Specifically, TNPG can be written as ∇T-NPG J(θ) = H ⊤ Σỹ where θ
ỹ = (λIn + HH ⊤ Σ)−1 y.
(11)
Under standard on-policy assumptions, the resulting estimator is unbiased in expectation, and the discrepancy introduced by ignoring Σ vanishes as the batch size increases.
(λI + H ⊤ ΣH)−1 H ⊤ Σy g
2
Eτ [∥yτ − Hτ g∥2 ] ∝ ∥y − Hg∥Σ .
where λ > 0 is a damping coefficient. This Tikhonovregularized NPG (T-NPG) is equivalently the solution to the regularized least-squares:
= arg min ∥y − Hg∥Σ + λ ∥g∥2
(10)
(9) 4
Randomized Advantage Transformation (RAT)
update
terpretation, whereas SPRING can be viewed as a Kaczmarzinspired momentum-based method. Second, RAT performs inner iterations over mini-batches within a single on-policy rollout, whereas SPRING applies a single update per batch. This distinction is crucial as it enables iterative refinement of the NPG estimate without cumulating gj across rollouts.
gj ← gj−1 +Hτ⊤j (λI + Hτj Hτ⊤j )−1 yτj − Hτj gj−1 . | {z } Randomized Advantage Transformation
(13) The bracketed term performs an advantage transformation on the sampled data. We therefore refer to this method as Randomized Advantage Transformation (RAT), i.e., h i Ãj (s, a) := (λI + Hτj Hτ⊤j )−1 yτj − Hτj gj−1 .
As in KFAC (Martens & Grosse, 2015), gradient norm clipping is essential for stability. Instead of clipping in the Fisher norm, we clip ℓ2 -norm of the estimated gradiν ents gj (Zhang et al., 2020): αj := min η, ∥gj ∥ , where 2 ν > 0 is a threshold and η > 0 is the learning rate.
(s,a)
(14) With minibatch size B we have Hτ ∈ RB×p and yτ ∈ RB . When B ≪ p, the matrix inversion is only B × B, in contrast to the original n × n matrix. The main computational cost arises from forming Hτ Hτ⊤ , which scales as O(pB 2 ); this cost can be further reduced using Nyström Approximation (Gittens & Mahoney, 2013). Further, as pointed by Guzmán-Cordero et al. (2025), HH ⊤ is the neural tangent kernel (Jacot et al., 2018), and thus can be approximated efficiently in various ways (Novak et al., 2022). H can be estimated efficiently using the per-sample gradients in PyTorch1 . The advantage can be computed via torch.linalg.solve: which directly solve linear system with matrix (λI + HH ⊤ ) and vector y. It is faster and more numerically stable than explicitly computing the inverse.
Shared actor-critic architectures. We follow Wu et al. (2017) and estimate joint natural policy gradients by modeling the value function as a Gaussian. To apply RAT to the critic, we explicitly introduce a pseudo advantage (e.g., an all-ones vector). RAT is then applied jointly to the policy advantage and the critic pseudo advantage, yielding a unified and stable optimization objective for shared actor-critic networks. We refer readers to Section C.1 for details. This joint application of RAT to both policy and critic updates clearly shows the flexibility of the advantage transformation, and distinguishes our method from prior work on Woodburybased approaches (Guzmán-Cordero et al., 2025). Specifically, while Guzmán-Cordero et al. (2025) can in principle be applied in this setting, they typically require maintaining separate curvature-adjusted gradients for actor and critic and carefully merging them during parameter updates. This merging is inherently architecture-dependent, as it requires explicit knowledge of which parameters are shared and how gradients from different heads should be combined. In contrast, RAT introduces a pseudo-advantage formulation that unifies the actor and critic objectives into a single surrogate loss. The resulting gradient is computed via standard backpropagation. As a result, curvature-adjusted updates are handled implicitly by autograd, without requiring manual gradient partitioning or architecture-specific merging logic. This allows RAT to remain architecture-agnostic in practice, even in shared-network settings. The full RAT is summarized in Algorithm 1 in Appendix.
The natural policy gradients can then be computed via direct backpropagation using the following PPO-like objective: π(a|s; θ) JRAT (θ) := E(s,a)∼Dk Ãj (s, a) , (15) πold (a|s) where πold is the behavior policy used to collect Dk , and remains the same during the inner iterations. Intuition Underlying RAT. RAT can be viewed as an efficient method for constructing a compatible approximation of advantage function by solving linear system Hx = y with Tikhonov regularization. At each iteration, RAT updates the current estimate gj−1 by projecting the residual yτ − Hτ gj−1 onto the solution space of Hτ x = yτ , with Tikhonov regularization. This projection implicitly injects curvature information into the advantage estimates. By iterating over mini-batches, RAT progressively aggregates local curvature information, yielding an accurate approximation of the full natural policy gradient without explicitly forming or inverting the Fisher matrix.
4.4. Convergence Results of RAT RAT is closely related to randomized block Kaczmarz method (Needell & Tropp, 2014) and, more generally, to the class of randomized iterative methods (Gower & Richtárik, 2015). Its convergence analysis can be viewed as an extension of these methods to the setting of Tikhonov-regularized natural policy gradients.
Importantly, RAT differs from SPRING (Goldshlager et al., 2024) in two key aspects. First, RAT has no momentum in1
We begin with a standard assumption on the matrix H. Assumption 1 (Full column rank). The full data matrix H ∈ Rn×p satisfies rank(H) = p and p ≪ n.
https://docs.pytorch.org/tutorials/interme diate/per_sample_grads.html
5
Randomized Advantage Transformation (RAT)
1.0
Assumption 2 (State-action coverage). Let h⊤ i denote i-th row of H, and τ denote a random minibatch sampled from dπ (s, a). For each index i ∈ {1, . . . , n}, P(i ∈ τ ) > 0.
0.8
θ2 = log σ
This assumption is consistent with the common setting in reinforcement learning in which the state-action space is large or continuous, and function approximation is needed.
This assumption ensures that every state-action pair has a non-zero probability of being sampled, which is a standard requirement for defining the natural policy gradient (Bagnell & Schneider, 2003; Kakade, 2001). Under the above assumptions, we obtain the following:
0.4 0.2
Lemma 1. Define Pτ := Hτ⊤ (λI + Hτ Hτ⊤ )−1 Hτ , then µ := λmin (E[Pτ ]) > 0.
0.0
We now present two theorems characterizing the convergence behavior of RAT. We first analyze an idealized case in which the advantage is exactly compatible (Peters & Schaal, 2008). Specifically, for all minibatches τ , yτ = Hτ g ∗ , where g ∗ is the solution to Equation (5).
−1.0
E∥gj − g ∗ ∥22 ≤ (1 − µ)j ∥g0 − g ∗ ∥22 .
0.0 θ1 = µ
0.5
1.0
(16) It is worth noting that Algorithm 1 in Appendix implements multiple inner iterations per batch, resulting in a time-varying sequence of linear systems. We show in Section C.2 that the practical implementation can be interpreted as a contractive solver tracking a slowly varying sequence of systems. Let gt∗ denote the solution of the regularized leastsquares problem defined by the current policy θt , and define the tracking error et := ∥gt − gt∗ ∥. Under our settings (small learning rates and gradient clipping), the tracking error remains small, yielding a bounded steady-state error of order O(maxt ||θt+1 −θt ||). This aligns with standard analyses of stochastic approximation in RL, where updates track a moving target induced by policy changes, and provides a heuristic justification for the implementation.
This theorem, together with Lemma 1, guarantees that the RAT update is contractive, which is essential for establishing linear convergence. The convergence rate of RAT is entirely characterized by the spectrum of E[Pτ ]. In particular, as Pτ = Hτ⊤ Hτ (λI + Hτ⊤ Hτ )−1 , when Hτ⊤ Hτ (empirical Fisher) has a low rank, a smaller λ generally leads to a larger µ, and thus faster convergence. This behavior is also observed in our sensitivity analysis. In addition, Assumption 1 applies to the full matrix H, not to individual minibatches. In practice, minibatch matrices Hτ can be low-rank. The Tikhonov damping term λ > 0 ensures that (λI + Hτ Hτ⊤ ) is always invertible, even when Hτ is rank-deficient. Poor state-action coverage affects the convergence rate through µ = λmin (E[Pτ ]), but does not invalidate the analysis.
5. Experiments
We now consider the more realistic case in which the advantage estimates are noisy: yτ = Hτ g ∗ + ξτ , where ξτ is a zero-mean random variable satisfying E[ξτ | τ ] = 0.
We evaluate Randomized Advantage Transformation (RAT) through a combination of controlled illustrations, continuous control benchmarks, and high-dimensional visual domains. Our goals are to: (1) verify that RAT accurately approximates empirical natural gradients, (2) assess its empirical performance and efficiency relative to established natural policy gradients methods, and (3) analyze the sensitivity to its key design choices.
Theorem 2 (Convergence with error floor). Define η 2 := ⊤ E ∥Hτ (λI + Hτ Hτ⊤ )−1 ξτ ∥22 . Then η2 . µ
−0.5
Figure 2. Univariate Gaussian with θ1 = µ and θ2 = log σ (closed-formed natural gradients, empirical natural gradients, RAT gradients and vanilla gradients; ⋆ for the optimum). RAT closely approximates empirical natural gradients.
Theorem 1 (Linear convergence of RAT). Assume minibatches τj are sampled i.i.d. from dπ (s, a). Then
E∥gj − g ∗ ∥22 ≤ (1 − µ)j ∥g0 − g ∗ ∥22 +
0.6
(17)
η 2 effectively quantifies the norm discrepancy between the true gradient and its stochastic estimate. This motivates the use of gradient norm clipping in practice.
Across all experiments, we apply standard stabilization techniques, including observation normalization (Mnih et al., 2016), advantage normalization (Schulman et al., 2017b), 6
Randomized Advantage Transformation (RAT) halfcheetah
ant
6000
5000
Episodic return
5000
humanoid
4000
4000
160000
6000
140000
5000
3000
1000 0 0
250
500 750 Epoch
1000
3000
2000
FVP+CG PPO KFAC Sophia RAT
120000
FVP+CG PPO KFAC Sophia RAT
4000
3000
2000
humanoidstandup
7000
FVP+CG PPO KFAC Sophia RAT
100000 FVP+CG PPO KFAC Sophia RAT
80000
2000 1000
60000 1000
0
1250
40000
0 0
250
500 750 Epoch
1000
1250
0
250
500 750 Epoch
1000
1250
0
250
500 750 Epoch
1000
1250
Figure 3. Optimizing MLP policies on continuous control tasks with separate actor-critic networks. RAT outperforms KFAC, FVP+CG and Sophia in most tasks. The shaded region denotes the standard error over 5 random seeds. Table 1. Final performance (mean ± stderr over 5 seeds) of different methods on continuous controls with shared actor-critic networks. Ep. Returns S ×A RAT (Ours) ACKTR PPO Sophia
Swimmer ↑ 8×2
±36.3
271.6 59.1±13.0 191.3±32.7 57.9±5.9
Hopper ↑ 11 × 3
±524.9
2334.6 2138.9±171.6 2346.8±202.7 1104.0±90.6
HalfCheetah ↑
Walker2d ↑
±287.4
±293.6
17 × 6
4629.2 3630.9±282.6 4146.0±107.5 899.5±113.2
17 × 6
3156.0 2576.6±154.6 2225.3±303.4 1256.0±129.7
and PopArt for value normalization (Hessel et al., 2019). Unless stated otherwise, all methods use the same network architectures and training pipelines. We report results averaged over five random seeds for a fixed training budget of 1250 epochs, which corresponds to approximately 10 million environment steps (a standard budget in continuous control benchmarks). Additional implementation details are provided in Section C. The code for reproducing our results is available at Code URL2 .
105 × 8 ±353.1
2926.6 23.4±3.2 1373.9±26.0 −7.0±1.4
Humanoid ↑ 376 × 17
±117.3
5382.7 2571.7±838.7 5357.9±150.9 669.4±56.2
HumanoidStandup ↑ 376 × 17
146529.7±2317.6 127928.5±5433.7 130014.2±6463.7 111212.6±13449.9
optimum. The updates produced by RAT closely match the empirical natural gradients, demonstrating the effectiveness and accuracy of RAT with finite samples. 5.2. Continuous Control with MLP Policies We next evaluate RAT on standard continuous control benchmarks from OpenAI Gym (Brockman et al., 2016) implemented in MuJoCo (Todorov et al., 2012). We consider Walker2d-v4 (A ∈ R6 ), HalfCheetah-v4 (A ∈ R6 ), Antv4 (A ∈ R8 ), and Humanoid-v4 (A ∈ R17 ), which span action dimensions from 6 to 17. Policies and value functions are parameterized by two-layer MLPs with 256 hidden units and Tanh activations; the policy outputs the mean of a Gaussian distribution, with a state-independent log standard deviation. We compare RAT against several strong baselines for estimating natural policy gradients, including Fisher-vector products with conjugate gradient (FVP+CG) from (Schulman et al., 2015), Kronecker-Factored Approximate Curvature (KFAC) (Martens & Grosse, 2015), and a diagonal Fisher approximation (i.e., diag(λI + H ⊤ H)) according to Sophia (Liu et al., 2024). All methods are implemented within the same codebase, and baseline hyperparameters are tuned for best performance.
5.1. Illustration of Natural Gradients Estimation We begin with a low-dimensional example that admits an analytic form of the natural gradient, enabling direct visualization of the update directions. Specifically, we consider maximum-likelihood estimation for a univariate Gaussian parameterized by its mean and log-standard deviation: N (x|µ, σ; θ1 , θ2 ), where µ = θ1 and log standard deviation log σ = θ2 . In this setting, the inverse Fisher matrix is available in closed form, allowing us to compute exact natural gradients (see Section A for details). Figure 2 compares the vanilla gradients, the natural gradient, the empirical natural gradients and the gradient estimated by RAT. We also plot the contour lines for log N (x|θ) (loss landscape) to better illustrate how the natural gradient differs from the vanilla gradient. While vanilla gradients follow the steepest ascent direction of log N (x|θ), perpendicular to the contour lines, the natural gradient accounts for the geometry induced by the parameterization and points more directly towards the 2
Ant ↑
Separate actor-critic networks. Figure 3 reports learning curves when actor and critic are optimized separately. RAT consistently matches or outperforms all baselines across all tasks, exhibiting both faster learning and higher final returns. In particular, RAT remains stable on challenging tasks such as Ant-v4 and Humanoid-v4, whereas KFAC frequently
https://github.com/agent-lab/ICML2026-RAT
7
Randomized Advantage Transformation (RAT)
Episodic return
bigfish 20.0 17.5 15.0 12.5 10.0 7.5 5.0 2.5 0.0
bossfight
PPO ACKTR RAT
4 3 2 1 0
0
500
Epoch
1000
1500
jumper
10
0
500
Epoch
1000
1500
PPO ACKTR RAT
4 2 0
500
Epoch
1000
1500
6 4
0
500
PPO ACKTR RAT
4
1000
1500
0
500
0
0 1000
1500
1000
1500
PPO ACKTR RAT
30 25 20 15
4 2
Epoch
starpilot
6
2
Epoch
Epoch
PPO ACKTR RAT
8
6
500
PPO ACKTR RAT
2
10
0
8
miner
8
6
coinrun
10
PPO ACKTR RAT
maze
10
8 Episodic return
caveflyer 7 6 5 4 3 2 1 0
PPO ACKTR RAT
5
10 5 0
500
Epoch
1000
1500
0
0
500
Epoch
1000
1500
Figure 4. Optimizing ResNet policies for discrete controls in ProcGen environments: RAT performs consistently well across all 8 tasks, delivering comparable or higher episodic returns than all baselines. The shaded region denotes the standard error over 5 random seeds. Table 2. Wallclock time per update (in ms) on continuous control tasks with separate actor-critic networks.
Shared
Separate
Time (ms)
HalfCheetah ↓ ±1.49
Ant ↓
±1.35
RAT (Ours) FVP+CG KFAC Sophia PPO
9.83 19.86±1.15 5.60±1.28 3.92±0.71 3.12±1.38
10.04 19.95±1.20 5.61±1.23 3.98±0.73 3.18±1.36
RAT (Ours) ACKTR Sophia PPO
11.53±1.69 6.92±1.63 5.97±0.94 3.70±1.56
11.66±1.55 6.85±1.55 6.03±1.02 3.70±1.46
et al., 2017), an extended KFAC method for shared networks, and Sophia (Liu et al., 2024). Table 1 summarizes final performance. RAT achieves the best overall returns on most tasks, with substantial gains on Ant and Humanoid, highlighting its robustness in challenging shared architectures.
Humanoid ↓ 18.17±3.04 19.81±1.18 6.57±1.47 5.71±0.68 3.22±1.40
Table 2 reports wallclock time per update (in ms; averaged over 124 updates on Xeon(R) w5-2445 GeForce RTX 4090). While the PPO is the fastest, RAT is significantly more efficient than FVP+CG and offers a favorable trade-off between computational cost and performance. Despite higher per-update cost, RAT is most beneficial in regimes where curvature matters (e.g., high-dimensional settings), where PPO often plateaus or requires careful tuning. RAT is not designed to match PPO’s per-step efficiency, but to provide a simple, architecture-agnostic, and principled approximation to natural policy gradients. Compared to existing naturalgradient methods, it offers a stronger performance–compute trade-off while avoiding architecture-specific approximations and complex inner solvers.
19.85±3.11 7.87±1.70 7.58±0.97 3.72±1.49
learns slowly. We also evaluated an enhanced variant of KFAC, i.e., eKFAC (George et al., 2018), and found that the eKFAC method did not yield noticeable improvements over KFAC in our experiments, see Figure 6 in Section E. The Sophia method performs poorly on all tasks, highlighting the importance of capturing parameter correlations. In addition, RAT performs significantly better than PPO on challenging Ant and Humanoid tasks.
5.3. Visual Control with ResNet Policies To assess scalability to high-dimensional visual inputs, we evaluate RAT on the challenging Procgen Benchmark (Cobbe et al., 2020), which features procedurally generated environments with 64x64 RGB observations. We consider 8 representative environments, including BigFish, BossFight, CaveFlyer, Climber, Dodgeball, FruitBot, Heist, and StarPilot. Policies are parameterized using a ResNet-
Shared actor-critic networks. We further evaluate RAT in the shared-network setting, where curvature estimation is more challenging. Since FVP+CG is not directly applicable, we compare against Proximal Policy Optimization (PPO) (Schulman et al., 2017b), a simplified approximation to natural policy gradients (Hilton et al., 2022), ACKTR (Wu 8
Randomized Advantage Transformation (RAT)
based architecture, adapted from Espeholt et al. (2018), and training follows the standard Procgen protocol for evaluating sample efficiency, i.e., training and testing on the same distribution of levels in each environment.
mations may improve scalability. Second, the convergence rate of RAT depends on the minimum singular value of Pτ . Poorly conditioned minibatches may slow convergence, although gradient norm clipping alleviates this issue in practice. Third, our analysis focus on the on-policy setting; extending RAT to full off-policy settings where the advantage function is estimated from replay buffer samples remains an open direction.
Figure 4 presents learning curves comparing RAT with PPO and ACKTR. RAT performs consistently well across all tasks, delivering comparable or higher returns than all baselines while avoiding training instabilities frequently observed in KFAC-based methods. These results demonstrate that RAT scales effectively to visual and complex dynamics. (a) RAT and Grad Clip
(b) Batch sizes
(c) Kaczmarz Iterations
(d) Damping λ
28 29 Batch size
1×8 2×8 3×8 4×8 Iteration
.01 .05 .1 .2 .4 .8 Coefficient λ
Conclusion. We proposed Randomized Advantage Transformation (RAT), an efficient and architecture-agnostic method for estimating natural policy gradients. By using a Woodbury-based reformulation, RAT transforms curvature information into the advantage function and estimates regularized natural policy gradients using standard backpropagation, without explicit Fisher construction or conjugate-gradient solvers. We provided convergence guarantees and demonstrated strong empirical performance on continuous control and high-dimensional visual benchmarks. RAT bridges the gap between principled natural gradient methods and practical deep reinforcement learning, offering a simple and scalable alternative for second-order policy optimization.
6
Episodic return (K)
5 4 Full W/o grad clip
3
W/o RAT 2 1 0 0
200 400 Epoch
27
210
Figure 5. Ablation and sensitivity analysis of RAT on Humanoid.
5.4. Ablations and Sensitivity Analysis
Impact Statement
Finally, we conduct an ablation study and sensitivity analysis to pinpoint the influence of key components and hyperparameters of RAT on performance, focusing on the challenging Humanoid task. Figure 5(a) presents ablations that remove either the advantage transformation or gradient norm clipping. Both components are essential: removing either leads to substantial performance degradation. We further study sensitivity of RAT to batch size, number of Kaczmarz iterations and damping coefficient λ. The results are presented in Figure 5(b), (c) and (d). RAT benefits from sufficiently large batch sizes (210 ) (as the batch size effectively determines the rank of empirical Fisher matrix), while remaining relatively robust to the number of Kaczmarz iterations within a reasonable range (too few iterations result in suboptimal performance). The performance of RAT is also robust across a broad range of damping values (from 0.01 to 0.4), with degradation only occurring at very large values (e.g., λ = 0.8). More analysis on Ant can be found in Figure 8 in Appendix. Overall, these results indicate that RAT is robust and does not require fine-grained tuning.
This work advances scalable optimization methods for reinforcement learning by enabling efficient computation of natural policy gradients without explicit curvature estimation. By simplifying implementation and reducing computational overhead, the proposed method may facilitate broader adoption of principled second-order optimization techniques in practice. As with reinforcement learning methods in general, potential downstream applications span a wide range of domains and should be deployed responsibly, particularly in efficiency-critical settings. This paper focuses on algorithmic contributions and does not involve human subjects or sensitive data.
References Agarwal, A., Kakade, S. M., Lee, J. D., and Mahajan, G. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. J. Mach. Learn. Res., 22:98:1–98:76, 2021. URL https://jmlr.o rg/papers/v22/19-736.html.
6. Limitations and Conclusion
Amari, S. Natural gradient works efficiently in learning. Neural Comput., 10(2):251–276, 1998. doi: 10.1162/08 9976698300017746. URL https://doi.org/10 .1162/089976698300017746.
Limitations. Despite strong theoretical and empirical results, RAT has several limitations. First, RAT requires forming minibatch-level matrices Hτ Hτ⊤ , with cost O(pB 2 ). While substantially cheaper than full Fisher inversion and architecture-agnostic, this can become a bottleneck for large policies or batch sizes; low-rank or sketch-based approxi-
Amari, S., Park, H., and Fukumizu, K. Adaptive method of realizing natural gradient learning for multilayer perceptrons. Neural Comput., 12(6):1399–1409, 2000. doi: 9
Randomized Advantage Transformation (RAT)
10.1162/089976600300015420. URL https: //doi.org/10.1162/089976600300015420.
Learning Research, pp. 2048–2056. PMLR, 2020. URL http://proceedings.mlr.press/v119/cob be20a.html.
Ba, J., Grosse, R. B., and Martens, J. Distributed secondorder optimization using kronecker-factored approximations. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net, 2017. URL https://openreview.net/forum ?id=SkkTMpjex.
Dangel, F., Tatzel, L., and Hennig, P. Vivit: Curvature access through the generalized gauss-newton’s low-rank structure. Trans. Mach. Learn. Res., 2023, 2023. URL https://openreview.net/forum?id=DzJ7 JfPXkE.
Bae, J., Vicol, P., HaoChen, J. Z., and Grosse, R. B. Amortized proximal optimization. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, 2022. URL http://papers.nips.cc/paper_files/pap er/2022/hash/3af25aa3de8b7b02ddbd1b6 be5031be8-Abstract-Conference.html.
Espeholt, L., Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning, I., Legg, S., and Kavukcuoglu, K. IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1406–1415. PMLR, 2018. URL http://proceedings.mlr.press/ v80/espeholt18a.html.
Bagnell, J. A. and Schneider, J. Covariant policy search. In Proceedings of the 18th International Joint Conference on Artificial Intelligence, IJCAI’03, pp. 1019–1024, San Francisco, CA, USA, 2003. Morgan Kaufmann Publishers Inc.
George, T., Laurent, C., Bouthillier, X., Ballas, N., and Vincent, P. Fast approximate natural gradient descent in a kronecker factored eigenbasis. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 9573–9583, 2018. URL https://proceedings.neurips. cc/paper/2018/hash/48000647b315f6f00 f913caa757a70b3-Abstract.html.
Benzing, F. Gradient descent on neurons and its link to approximate second-order optimization. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 1817–1853. PMLR, 2022. URL https://proceedings.mlr.press/v162/b enzing22a.html. Brockman, G., Cheung, V., Pettersson, L., Schneider, J., Schulman, J., Tang, J., and Zaremba, W. Openai gym. CoRR, abs/1606.01540, 2016. URL http://arxiv. org/abs/1606.01540. Cayci, S. and Eryilmaz, A. Recurrent natural policy gradient for pomdps. Trans. Mach. Learn. Res., 2025, 2025. URL https://openreview.net/forum?id=6G01 e0vgIf.
Gittens, A. and Mahoney, M. Revisiting the nystrom method for improved large-scale machine learning. In Dasgupta, S. and McAllester, D. (eds.), Proceedings of the 30th International Conference on Machine Learning, volume 28 of Proceedings of Machine Learning Research, pp. 567– 575, Atlanta, Georgia, USA, 17–19 Jun 2013. PMLR. URL https://proceedings.mlr.press/v28/ gittens13.html. Goldshlager, G., Abrahamsen, N., and Lin, L. A kaczmarzinspired approach to accelerate the optimization of neural network wavefunctions. Journal of Computational Physics, 516:113351, 2024.
Chen, A. and Heyl, M. Empowering deep neural quantum states through efficient optimization. Nature Physics, 20 (9):1476–1481, 2024.
Goldshlager, G., Hu, J., and Lin, L. A sketch-and-project analysis of subsampled natural gradient algorithms, 2026. URL https://arxiv.org/abs/2508.21022.
Cobbe, K., Hesse, C., Hilton, J., and Schulman, J. Leveraging procedural generation to benchmark reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine
Gower, R. M. and Richtárik, P. Randomized iterative methods for linear systems. SIAM J. Matrix Anal. Appl., 36 (4):1660–1690, 2015. doi: 10.1137/15M1025487. URL https://doi.org/10.1137/15M1025487. 10
Randomized Advantage Transformation (RAT)
Guzmán-Cordero, A., Dangel, F., Goldshlager, G., and Zeinhofer, M. Improving energy natural gradient descent through woodbury, momentum, and randomization. CoRR, abs/2505.12149, 2025. doi: 10.48550/ARXIV.2 505.12149. URL https://doi.org/10.48550 /arXiv.2505.12149.
in Neural Information Processing Systems 14 [Neural Information Processing Systems: Natural and Synthetic, NIPS 2001, December 3-8, 2001, Vancouver, British Columbia, Canada], pp. 1531–1538. MIT Press, 2001. URL https://proceedings.neurips.cc/p aper/2001/hash/4b86abe48d358ecf194c5 6c69108433e-Abstract.html.
Haarnoja, T., Zhou, A., Abbeel, P., and Levine, S. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 1856–1865. PMLR, 2018. URL http://procee dings.mlr.press/v80/haarnoja18b.html.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In Bengio, Y. and LeCun, Y. (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: //arxiv.org/abs/1412.6980. Kuba, J. G., Chen, R., Wen, M., Wen, Y., Sun, F., Wang, J., and Yang, Y. Trust region policy optimisation in multiagent reinforcement learning. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022. OpenReview.net, 2022. URL https://openreview.net/forum?id= EcGGFkNTxdJ.
Hessel, M., Soyer, H., Espeholt, L., Czarnecki, W., Schmitt, S., and van Hasselt, H. Multi-task deep reinforcement learning with popart. In Proceedings of the Thirty-Third AAAI Conference on Artificial Intelligence and Thirty-First Innovative Applications of Artificial Intelligence Conference and Ninth AAAI Symposium on Educational Advances in Artificial Intelligence, AAAI’19/IAAI’19/EAAI’19. AAAI Press, 2019. ISBN 978-1-57735-809-1. doi: 10.1609/aaai.v33i01.33013796. URL https://doi.org/10.1609/aaai.v33 i01.33013796.
Kunstner, F., Hennig, P., and Balles, L. Limitations of the empirical fisher approximation for natural gradient descent. In Wallach, H. M., Larochelle, H., Beygelzimer, A., d’Alché-Buc, F., Fox, E. B., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 4158–4169, 2019. URL https://proceedings.neurips.cc/p aper/2019/hash/46a558d97954d0692411c 861cf78ef79-Abstract.html.
Hilton, J., Cobbe, K., and Schulman, J. Batch sizeinvariance for policy optimization. In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K. (eds.), Advances in Neural Information Processing Systems, 2022. URL https: //openreview.net/forum?id=lXuZaxEaI7. Huo, Y., Dash, S. P., Stoican, R., Kaski, S., and Sun, M. Rank-1 approximation of inverse fisher for natural policy gradients in deep reinforcement learning. Transactions on Machine Learning Research, 2026. ISSN 2835-8856. URL https://openreview.net/forum?id= ko8Kn7TS6m.
Liu, H., Li, Z., Hall, D. L. W., Liang, P., and Ma, T. Sophia: A scalable stochastic second-order optimizer for language model pre-training. In The Twelfth International Conference on Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL https://openreview.net/forum?id=3xHD eA8Noi.
Jacot, A., Hongler, C., and Gabriel, F. Neural tangent kernel: Convergence and generalization in neural networks. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pp. 8580–8589, 2018. URL https://proc eedings.neurips.cc/paper/2018/hash/5 a4be1fa34e62bb8a6ec6b91d2462f5a-Abstr act.html.
Martens, J. Deep learning via hessian-free optimization. In Fürnkranz, J. and Joachims, T. (eds.), Proceedings of the 27th International Conference on Machine Learning (ICML-10), June 21-24, 2010, Haifa, Israel, pp. 735–742. Omnipress, 2010. URL https://icml.cc/Conf erences/2010/papers/458.pdf. Martens, J. and Grosse, R. B. Optimizing neural networks with kronecker-factored approximate curvature. In Bach, F. R. and Blei, D. M. (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 2408–2417.
Kakade, S. M. A natural policy gradient. In Dietterich, T. G., Becker, S., and Ghahramani, Z. (eds.), Advances 11
Randomized Advantage Transformation (RAT)
JMLR.org, 2015. URL http://proceedings.ml r.press/v37/martens15.html.
Schulman, J., Abbeel, P., and Chen, X. Equivalence between policy gradients and soft q-learning. CoRR, abs/1704.06440, 2017a. URL http://arxiv.org/ abs/1704.06440.
Mnih, V., Badia, A. P., Mirza, M., Graves, A., Lillicrap, T. P., Harley, T., Silver, D., and Kavukcuoglu, K. Asynchronous methods for deep reinforcement learning. In Balcan, M. and Weinberger, K. Q. (eds.), Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1928–1937. JMLR.org, 2016. URL http://proceedings.mlr.press/v48/mnih a16.html.
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O. Proximal policy optimization algorithms. CoRR, abs/1707.06347, 2017b. URL http://arxiv. org/abs/1707.06347. Sun, M., Devlin, S., Beck, J., Hofmann, K., and Whiteson, S. Trust region bounds for decentralized PPO under non-stationarity. In Agmon, N., An, B., Ricci, A., and Yeoh, W. (eds.), Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS 2023, London, United Kingdom, 29 May 2023 - 2 June 2023, pp. 5–13. ACM, 2023. doi: 10.5555/3545946.3598613. URL https://dl.acm .org/doi/10.5555/3545946.3598613.
Needell, D. and Tropp, J. A. Paved with good intentions: analysis of a randomized block kaczmarz method. Linear Algebra and its Applications, 441:199–221, 2014. Novak, R., Sohl-Dickstein, J., and Schoenholz, S. S. Fast finite width neural tangent kernel. In Chaudhuri, K., Jegelka, S., Song, L., Szepesvári, C., Niu, G., and Sabato, S. (eds.), International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162 of Proceedings of Machine Learning Research, pp. 17018–17044. PMLR, 2022. URL https://proceedings.mlr.press/v162/n ovak22a.html.
Sutton, R. S. and Barto, A. G. Reinforcement learning an introduction, 2nd Edition. MIT Press, 2018. URL http://www.incompleteideas.net/book/t he-book-2nd.html. Sutton, R. S., McAllester, D. A., Singh, S., and Mansour, Y. Policy gradient methods for reinforcement learning with function approximation. In Solla, S. A., Leen, T. K., and Müller, K. (eds.), Advances in Neural Information Processing Systems 12, [NIPS Conference, Denver, Colorado, USA, November 29 - December 4, 1999], pp. 1057–1063. The MIT Press, 1999. URL http: //papers.nips.cc/paper/1713-policy-g radient-methods-for-reinforcement-lea rning-with-function-approximation.
Pascanu, R. and Bengio, Y. Revisiting natural gradient for deep networks. In Bengio, Y. and LeCun, Y. (eds.), 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 1416, 2014, Conference Track Proceedings, 2014. URL http://arxiv.org/abs/1301.3584. Peters, J. and Schaal, S. Natural actor-critic. Neurocomputing, 71(7-9):1180–1190, 2008. doi: 10.1016/J.NEUC OM.2007.11.026. URL https://doi.org/10.1 016/j.neucom.2007.11.026.
Todorov, E., Erez, T., and Tassa, Y. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, IROS 2012, Vilamoura, Algarve, Portugal, October 7-12, 2012, pp. 5026–5033. IEEE, 2012. doi: 10.1109/IROS.2012.6386109. URL https://do i.org/10.1109/IROS.2012.6386109.
Schulman, J., Levine, S., Abbeel, P., Jordan, M. I., and Moritz, P. Trust region policy optimization. In Bach, F. R. and Blei, D. M. (eds.), Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, volume 37 of JMLR Workshop and Conference Proceedings, pp. 1889–1897. JMLR.org, 2015. URL http://proceedings.ml r.press/v37/schulman15.html.
Wu, X., Yu, W., Zhang, C., and Woodland, P. C. An improved empirical fisher approximation for natural gradient descent. In Globersons, A., Mackey, L., Belgrave, D., Fan, A., Paquet, U., Tomczak, J. M., and Zhang, C. (eds.), Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurIPS 2024, Vancouver, BC, Canada, December 10 - 15, 2024, 2024. URL http://papers.nips.cc/paper_files/pap er/2024/hash/f23098fa0cfcdef0e743b13 4d380eeb9-Abstract-Conference.html.
Schulman, J., Moritz, P., Levine, S., Jordan, M. I., and Abbeel, P. High-dimensional continuous control using generalized advantage estimation. In Bengio, Y. and LeCun, Y. (eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings, 2016. URL http://arxiv.org/abs/1506.02438. 12
Randomized Advantage Transformation (RAT)
Wu, Y., Mansimov, E., Grosse, R. B., Liao, S., and Ba, J. Scalable trust-region method for deep reinforcement learning using kronecker-factored approximation. In Guyon, I., von Luxburg, U., Bengio, S., Wallach, H. M., Fergus, R., Vishwanathan, S. V. N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pp. 5279–5288, 2017. URL https://proceedings.neurips.cc/paper /2017/hash/361440528766bbaaaa1901845 cf4152b-Abstract.html. Yang, M., Xu, D., Wen, Z., Chen, M., and Xu, P. Sketchbased empirical natural gradient methods for deep learning. J. Sci. Comput., 92(3):94, 2022. doi: 10.1007/S109 15-022-01911-X. URL https://doi.org/10.1 007/s10915-022-01911-x. Zhang, J., He, T., Sra, S., and Jadbabaie, A. Why gradient clipping accelerates training: A theoretical justification for adaptivity. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https: //openreview.net/forum?id=BJgnXpVYwS.
13
Randomized Advantage Transformation (RAT)
A. Gradient Calculation in the Likelihood Example For the likelihood example, the PG and NPG are given in the closed forms as follows: # " (x−θ )
"
1
∇θ Ex [log N (x|θ)] = Ex |
{z
exp(2θ2 ) (x−θ1 )2 −1 + exp(2θ 2)
Vanilla gradient
,
F
}
|
−1
(θ)∇θ Ex [log N (x|θ)] = Ex
(x − θ1 ) 2 1 1) − 2 + 2(x−θ exp(2θ2 )
{z
#
}
Natural gradient
Specifically, with parameterization θ1 = µ and θ2 = log σ, we have 1 1 (x − µ)2 (x − θ1 )2 √ N (x|µ, σ; θ) = √ = exp − exp − , 2σ 2 2 exp(2θ2 ) 2πσ 2π exp(θ2 ) √ (x − θ1 )2 log N (x|µ, σ; θ) = − log 2π − θ2 − . 2 exp(2θ2 )
(18) (19)
The gradient of log probability is given as # " (x−θ1 ) exp(−2θ2 )(x − θ1 ) 2 σ ∇θ log N (x|θ) = = 2 1) −1 + exp(−2θ2 )(x − θ1 )2 . −1 + (x−θ σ2
Accordingly, the Fisher matrix at point x is given " F (x, θ) =
(20)
#
(x−θ1 )2 σ4 (x−θ1 )3 1) − (x−θ σ4 σ2
(x−θ1 )3 1) − (x−θ σ4 σ2 2(x−θ1 )2 (x−θ1 )4 1− + σ4 . σ2
(21)
Thus, the full Fisher matrix is " F (θ) = Ex∼N (x|µ,σ,θ)
(x−θ1 )2 σ4 (x−θ1 )3 1) − (x−θ 4 σ σ2
h
i
(x−θ1 )3 1) − (x−θ σ4 σ2 2 4 1) 1) 1 − 2(x−θ + (x−θ σ2 σ4
(x−θ1 )2 σ4 i (x−θ1 )3 (x−θ1 ) − σ4 σ2
Ex∼N (x|µ,σ;θ) h
= Ex∼N (x|µ,σ;θ) 1 exp(−2θ2 ) 0 = σ2 = 0 0 2
# (22)
h i 3 (x−θ1 ) 1) Ex∼N (x|µ,σ;θ) (x−θ − 4 2 σ σ h i 2 (x−θ1 )4 1) Ex∼N (x|µ,σ;θ) 1 − 2(x−θ + σ2 σ4
0 . 2
(23)
(24)
The transition from Equation (23) to Equation (24) is based on the fact that if x has a normal distribution N (x|µ, σ), the non-central moments exist for any non-negative integer p and are given as follows: ( 0 if p is odd, p Ex [(x − µ) ] = (25) p σ (p − 1)!! if p is even. For a diagonal Fisher matrix, its inverse exists since exp(2θ2 ) > 0 and is given below 2 σ 0 exp(2θ2 ) 0 F −1 (θ) = = 1 . 0 0 12 2
(26)
Thus, the vanilla gradient direction for the parameter update is " ∇θ Ex [log p(x|θ)] = Ex
(x−θ1 ) exp(2θ2 ) (x−θ1 )2 −1 + exp(2θ 2)
# ,
(27)
and the natural gradient direction is " F
−1
(θ)Ex [∇θ p(x|θ)] = Ex F −1 (θ)∇θ p(x|θ) = Ex 14
(x − θ1 ) 2 1 1) − 2 + 2(x−θ exp(2θ2 )
# (28)
Randomized Advantage Transformation (RAT)
For damped Fisher λI + F (θ), its inverse is given by " 2 −1
(λI + F (θ))
σ 1+λσ 2
=
#
0
" =
1 2+λ
0
exp(2θ2 ) 1+λ exp(2θ2 )
0
0
#
1 2+λ
.
(29)
The regularized natural gradient direction is " −1
(λI + F (θ))
Ex [∇θ p(x|θ)] = Ex
x−θ1 1+λ exp(2θ2 ) (x−θ1 )2 1 + (2+λ) − 2+λ exp(2θ2 )
# (30)
To compute the gradients, we randomly sample 2,000 data points from N (0, 1).
B. Convergence Results of Randomized Advantage Transformation Consider the regularized least-squares objective 1 λ ∥y − Hg∥22 + ∥g∥22 , 2 2
(31)
g ∗ := (H ⊤ H + λI)−1 H ⊤ y.
(32)
minp
g∈R
with unique solution At iteration j, RAT samples a minibatch τj and performs the update gj+1 = gj + Hτ⊤j (λI + Hτj Hτ⊤j )−1 yτj − Hτj gj .
(33)
Define Pτ := Hτ⊤ (λI + Hτ Hτ⊤ )−1 Hτ ,
Sτ := Hτ⊤ (λI + Hτ Hτ⊤ )−1 .
(34)
Let the estimation error be ej := gj − g ∗ .
(35)
Lemma 2. For any minibatch τ , the matrix Pτ satisfies
0 ⪯ Pτ ⪯ I,
Pτ2 ⪯ Pτ .
Proof. Let Hτ = U ΣV ⊤ be the singular value decomposition. Then Pτ = V Σ⊤ (λI + ΣΣ⊤ )−1 ΣV ⊤ = V diag
σi2 λ + σi2
V ⊤.
Each eigenvalue lies in [0, 1), implying 0 ⪯ Pτ ⪯ I. Since x2 ≤ x for x ∈ [0, 1], we also have Pτ2 ⪯ Pτ . Lemma 3. Under the assumptions of full column rank of H and full data coverage, µ := λmin (E[Pτ ]) > 0. Proof. Recall Pτ = Hτ⊤ (λI + Hτ Hτ⊤ )−1 Hτ ⪰ 0,
and
v ⊤ Pτ v = 0 ⇐⇒ Hτ v = 0.
We prove E[Pτ ] ≻ 0 by showing that no nonzero vector v can satisfy Hτ v = 0 almost surely. Fix any v ̸= 0. Since H has full column rank, we have Hv ̸= 0, hence there exists at least one index i such that h⊤ i v ̸= 0. By full data coverage, P(i ∈ τ ) > 0. On the event {i ∈ τ }, the submatrix Hτ contains row h⊤ , and thus H v = ̸ 0 (because its i-th component τ i equals h⊤ v = ̸ 0). Therefore, i P(Hτ v ̸= 0) ≥ P(i ∈ τ ) > 0. Consequently, Hτ v = 0 cannot hold almost surely for any nonzero v.
Since (λI + Hτ Hτ⊤ )−1 ≻ 0, we have v ⊤ Pτ v > 0 whenever Hτ v ̸= 0, and hence v ⊤ E[Pτ ]v = E[v ⊤ Pτ v] > 0
Thus E[Pτ ] ≻ 0, which implies µ = λmin (E[Pτ ]) > 0. 15
for all v ̸= 0.
Randomized Advantage Transformation (RAT)
We first analyze the idealized case where minibatch targets are exact. For all minibatches τ , yτ = Hτ g ∗ . Theorem 3 (Linear convergence of RAT). Assume minibatches τj are sampled i.i.d. from an arbitrary distribution. Define µ := λmin (E[Pτ ]), then E∥gj − g ∗ ∥22 ≤ (1 − µ)j ∥g0 − g ∗ ∥22 . (36) Proof. Using the noise-free assumption, the update Equation (33) becomes gj+1 = gj − Pτj (gj − g ∗ ). Subtracting g ∗ yields ej+1 = (I − Pτj )ej . Conditioned on ej and τj , Since Pτ2j ⪯ Pτj ,
2 ∥ej+1 ∥22 = e⊤ j (I − 2Pτj + Pτj )ej .
I − 2Pτj + Pτ2j ⪯ I − Pτj ,
and therefore ∥ej+1 ∥22 ≤ ∥ej ∥22 − e⊤ j Pτj ej . Taking conditional expectation, E[∥ej+1 ∥22 | ej ] ≤ ∥ej ∥22 − e⊤ j E[Pτ ]ej . 2 Since e⊤ j E[Pτ ]ej ≥ µ∥ej ∥2 ,
E[∥ej+1 ∥22 | ej ] ≤ (1 − µ)∥ej ∥22 .
Iterating proves the claim. We now consider stochastic targets, as in reinforcement learning. For each minibatch τ , yτ = Hτ g ∗ + ξτ , where the noise ξτ satisfies E[ξτ | τ ] = 0. Theorem 4 (Convergence with error floor). Define η 2 := E ∥Sτ ξτ ∥22 , then E∥gj − g ∗ ∥22 ≤ (1 − µ)j ∥g0 − g ∗ ∥22 +
η2 . µ
(37)
Proof. Substituting the stochastic model into Equation (33) yields ej+1 = (I − Pτj )ej + Sτj ξτj . Squaring and taking conditional expectation, the cross term vanishes since E[ξτj | τj ] = 0, giving E[∥ej+1 ∥22 | ej ] ≤ (1 − µ)∥ej ∥22 + η 2 . Unrolling the resulting recursion completes the proof.
C. Implementation Details We implemented RAT using the per-sample gradients feature in PyTorch3 . Specifically, we first compute the per-sample gradients of the policy network’s outputs with respect to its parameters using the built-in function torch.func.grad, torch.func.vmap and torch.func.functional call, and then flat these per-sample gradients to compute H. When computing HH ⊤ , we average over the samples in the mini-batch. We use torch.linalg.solve to solve the linear system involving the damped Fisher matrix (λI + HH ⊤ ), and thus apply the advantage transformation. This is a faster and more numerically stable way than performing the computations separately. Besides, we also incorporated the following training techniques: 3
https://docs.pytorch.org/tutorials/intermediate/per_sample_grads.html
16
Randomized Advantage Transformation (RAT)
Observation normalization. We normalize the observations with running mean and standard deviation as in (Schulman et al., 2017b) for all the MuJoCo tasks and set the clip range to [−5, 5]. For tasks with image observations, we normalize the pixel values to [0, 1] by dividing them by 255 and then normalize each pixel value with 0.5 mean and 0.5 standard deviation (to ensure the pixel values are in the range of [−1, 1]). We also stack the last three frames as the input to the policy network. We found that observation normalization is crucial for stabilizing training, especially for tasks in the Mujoco suite. Note that after observation normalization, we re-evaluate the action distribution’s mean and variance, i.e., µ(s) and σ(s), to ensure that the action distribution is consistent with the normalized observations. Advantage normalization. We use the Generalized Advantage Estimation (GAE) (Schulman et al., 2016) to compute the advantage estimates. We then normalize the advantage estimates to have zero mean and unit standard deviation within each batch as in (Schulman et al., 2017b). We found that advantage normalization is important for stabilizing training, especially when using high learning rates. PopArt value normalization. We use the PopArt normalization technique (Hessel et al., 2019) to normalize the value function targets. Specifically, we maintain running estimates of the mean µ and standard deviation σ of the value function targets and normalize the targets as (Gt − µ)/σ. We also adjust the parameters of the value network to account for the change in normalization following the procedure described in (Hessel et al., 2019). We set the decay rate for the running estimates to 0.99999. One slight improvement we made in our implementation is that we correct the bias in the running estimates of the mean and standard deviation by dividing the estimates by (1 − decayt ) at time step t, similar to Adam (Kingma & Ba, 2015). Gradient clipping. We clip the gradient norm to be at most 0.5 when updating the shared policy and value networks. If the gradient norm exceeds this threshold, we scale down the gradient to have a norm of 0.5. We also tried the Fisher norm clipping technique proposed in (Ba et al., 2017), but found that l2 norm clipping works equally well in our experiments. When the actor and critic networks are separate, we apply gradient clipping to the policy network with a threshold of 0.5, and to the value network with a threshold of 5.0. Action squashing. For environments with bounded action spaces, we apply a squashing function (tanh) to the actions sampled from the Gaussian policy to ensure that the actions lie within the valid range. Different from the procedure described in (Haarnoja et al., 2018), we do not adjust the log-probability of the actions to account for the squashing transformation. This is because the policy ratios, KL divergences, and Fisher matrix are all invariant to such transformations, as long as the transformation is differentiable and invertible. Ratio clamping.
When computing the policy ratios, we clamp ratios to be within [10−1 , 101 ] to avoid numerical instability.
C.1. RAT in Shared Actor-Critic In Actor-Critic methods, the actor and critic often share a common neural architecture (Mnih et al., 2016; Wu et al., 2017). When parameters are shared, we follow Wu et al. (2017) and estimate the joint natural policy gradients for the actor and critic. Specifically, we model the value output as a Gaussian distribution with fixed variance σ 2 , i.e., p(v|s) = N (v; V (s), σ 2 ), where v denotes a target value obtained from Monte-Carlo rollouts or Temporal Difference (TD) methods (Sutton & Barto, 2018). The critic is trained by maximizing the log-likelihood of this distribution, and the Fisher matrix for the critic is defined with respect to the corresponding log-likelihood. In practice we set σ to 1 without loss of generality, yielding 2 log p(v|s) ∝ − ∥v − V (s)∥ . h i max Ev∼q [log p(v|s)] = Ev∼q − ∥v − V (s)∥
2
Under this formulation, the score function for the joint distribution p(a, v|s) factorizes as ∇θ log p(a, v|s, θ) = ∇θ log π(a|s; θ) + ∇θ log p(v|s; θ), which is used to construct the matrix H in RAT. Following Wu et al. (2017), we sample the network outputs independently for the actor and critic, and inject unit-variance Gaussian noise to the value outputs. To apply RAT to the critic, we introduce a pseudo advantage for the value loss, such as an all-ones vector 1 with the same size as the mini-batch. Applying RAT to this pseudo advantage yields w̃(s) for each state, which can be interpreted as the 17
Randomized Advantage Transformation (RAT)
natural gradient update direction for the critic. The resulting joint loss for optimizing shared actor-critic networks is h i π(a|s; θ) 2 L(θ) = −E Ã(s, a) + E w̃(s) ∥v − V (s; θ)∥ πold (a|s) where Ã(s, a) is the transformed advantage for the actor, and w̃(s) is the transformed pseudo advantage for the critic. At each iteration, RAT is applied jointly to transform both the actor advantage for the critic pseudo advantage, enabling a unified and stable natural-gradient update for shared acrtor-critic networks. Algorithm 1 Randomized Advantage Transformation (RAT) 1: Input: Initial policy parameters θ0 , batch size B, total iterations K 2: for k = 0 to K − 1 do 3: Collect on-policy samples: πold (a|s) ← π(a|s; θk ) 4: Randomly partition Dk into batches of size B 5: for each batch τj in Dk do 6: Form Hτj using samples in τj 7: Apply RAT: h
Ãj (s, a) = (λI + Hτj Hτ⊤j )−1 yτj − Hτj gj−1
i (s,a)
Estimate gradient via a single backpropagation of the following loss: π(a|s; θ) ∂ Ãj (s, a) gj = − E ∂θ πold (a|s) 9: Apply gradient clipping αk := min η, ∥gνj ∥ . 2 10: Update policy parameters: θ ← θ + αk gj 11: end for 12: end for 13: Output: Policy parameters θ 8:
C.2. Gap between Theoretical Analysis and Practical Implementation The theorems analyze RAT as a fixed-policy linear system solver, while Algorithm 1 interleaves these updates with policy optimization, resulting in a time-varying sequence of systems. This makes the linear system non-stationary across inner iterations. The resulting algorithm is closer to “multiple PPO-like updates per rollout with curvature-corrected advantages” than to an iterative linear solver converging to a fixed target. To clarify the gap formally, let gt∗ denote the solution of the regularized least-squares problem defined by the current policy θt , and define the tracking error et := ∥gt − gt∗ ∥ .
(38)
At iteration t, RAT performs an update yielding gt+1 . The fixed-system analysis ( Theorem 1) implies a contraction: ∥gt+1 − gt∗ ∥ ≤ ρ ∥gt − gt∗ ∥ = ρet ,
(39)
where ρ = 1 − µ < 1. After the policy update θt 7→ θt+1 , the target solution shifts. Under standard smoothness assumptions on H and y, the solution map θ 7→ g ∗ (θ) is Lipschitz: ∥g ∗ (θt+1 ) − g ∗ (θt )∥ ≤ L ∥θt+1 − θt ∥ .
(40)
et+1 ≤ ρet + L ∥θt+1 − θt ∥
(41)
Combining these yields: Unrolling: et ≤ ρt e0 + L
t−1 X s=0
ρt−1−s ∥θs+1 − θs ∥
18
(42)
Randomized Advantage Transformation (RAT)
This shows that RAT can be interpreted as a contractive solver tracking a slowly varying sequence of systems. Under our settings (small learning rates and gradient clipping), the drift term remains small, yielding a bounded steady-state error of order
O(max ∥θt+1 − θt ∥)
(43)
t
This is a standard tracking bound for contractive iterative methods applied to slowly varying systems and provides a principled justification for the interleaved algorithm. Specifically, while the Lipschitz constant L is not directly measurable, the step size ∥θt+1 − θt ∥ is explicitly controlled by the learning rate and gradient clipping. In particular, we use ℓ2 gradient norm clipping at 0.5 together tiwth a learning rate 0.05 for the policy network, which ensures that each parameter update is bounded by at most 0.5 × 0.05 = 0.025 in ℓ2 norm. This keeps the drift term small throughout training up to the unknown constant L. More broadly, this is precisely why gradient clipping and small learning rates are important in our implementation: they ensure that the target system evolves slowly enough for the tracking interpretation to be meaningful.
D. Hyperparameters We summarize the hyperparameters used in our experiments in Table 3.
Table 3. Common hyperparameters used in our experiments.
Hyperparameter
Value
Discount factor γ GAE parameter λ Damping factor λ Mini-batch size PPO clipping parameter ϵ Number of epochs per update Number of steps per update Gradient clipping threshold (policy) Gradient clipping threshold (value) PopArt decay rate Observation normalization clip range Entropy coefficient
0.99 0.95 10−1 1024 0.2 8 256 * 32 0.5 5.0 0.99999 [−5, 5] 0
Table 4. Hyperparameters for RAT.
Hyperparameter
Value
pi lr vf lr lr for shared network damping for MLP damping for CNN & ResNet
0.05 0.001 0.1 0.1 0.5
19
Randomized Advantage Transformation (RAT) Table 5. Hyperparameters for KFAC (adopted from (George et al., 2018)).
Hyperparameter
Value
lr momentum stat decay damping kl clip weight decay TCov TInv batch averaged
0.001 0.9 0.95 0.001 0.001 0 1 10 True
Table 6. Hyperparameters for PPO.
Hyperparameter
Value
pi lr vf lr lr for shared networks momentum stat decay damping kl clip weight decay TCov TInv batch averaged
0.01 0.001 0.001 0.9 0.95 0.001 0.001 0 1 10 True
Architecture details: • For MuJoCo tasks, we use a two-layer MLP with 256 hidden units per layer and tanh activations for both the policy and value networks. • For Procgen tasks, we use the same ResNet architecture as in (Cobbe et al., 2020): four residual blocks with 16, 32, and 32 filters respectively, followed by a fully connected layer with 256 units. We use ReLU activations after each layer.
E. Additional Experimental Results ant 600 Episodic return
walker2d 4000 3500 3000 2500 2000 1500 1000 500 0
ekfac kfac
400 200 0 0
100
Epoch
200
300
halfcheetah
ekfac kfac
4000
3000 2500
3000
2000
2000
1500 1000
1000
500
0 0
100
Epoch
200
300
hopper
3500
ekfac kfac
5000
0
100
Epoch
200
300
0
ekfac kfac 0
100
Epoch
200
Figure 6. Comparing KFAC and EKFAC on Continuous Control Tasks. EKFAC performs similar to KFAC in most tasks.
20
300
Randomized Advantage Transformation (RAT)
swimmer
250 Episodic return
halfcheetah
PPO ACKTR Sophia RAT
300
4000
200
ant
PPO ACKTR Sophia RAT
5000
3000 2500
2500 2000
2000
3000
150
hopper 3000
PPO ACKTR Sophia RAT
1500
1500
2000
100
1000
1000
500
500
PPO ACKTR Sophia RAT
1000 50 0
0
0 0
200
400
600
0
200
400
Epoch
walker2d 3500
PPO ACKTR Sophia RAT
3000
Episodic return
2500
4000
0 0
200
400
Epoch
Epoch
humanoid
humanoidstandup
PPO ACKTR Sophia RAT
5000
2000
600
0.6
80000
0.4 PPO ACKTR Sophia RAT
1000 40000 0 600
1.0
100000
60000
400
600
0.8
2000
200
400 Epoch
1000
0
200
120000
1500
0
0
140000
3000
500
600
0.2
0.0 0
200
Epoch
400
600
0
Epoch
200
400
600
0.0
0.2
0.4
0.6
0.8
1.0
Epoch
Figure 7. Optimizing MLP Policies on Continuous Control Tasks with Shared Actor-Critic Networks. RAT outperforms KFAC and FVP+CG in most tasks. The shaded region denotes the standard deviation over 5 random seeds.
(a) RAT and Grad Clip
(b) Batch sizes
(c) Kaczmarz Iterations
(d) Damping λ
28 29 Batch size
1×8 2×8 3×8 4×8 Iteration
.01 .05 .1 .2 .4 .8 Coefficient λ
Full W/o grad clip
4 Episodic return (K)
W/o RAT 3
2
1
0 0
200 400 Epoch
27
210
Figure 8. Ablation study and sensitivity analysis of RAT on Ant.
21