Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Policy Gradient with Second Order Momentum

Tianyu Sun

TL;DR

PG‑SOM introduces a lightweight second‑order momentum method for policy gradients by maintaining a diagonal Hessian estimate to precondition gradients. The diagonal curvature estimator is shown to be an unbiased, positive‑definite surrogate that ensures a descent direction in expectation, and the method adds only O$(D)$ memory. Empirical results on standard RL benchmarks demonstrate up to about 2.1x improvements in sample efficiency and reduced variance compared with first‑order and Fisher‑based baselines, with Runge–Kutta style two‑stage updates providing the best early learning and final performance when combined with gradient clipping. The work suggests that coarse curvature information, when paired with simple stabilizers, yields meaningful practical gains without the overhead of full second‑order methods or learned critics.

Abstract

We develop Policy Gradient with Second-Order Momentum (PG-SOM), a lightweight second-order optimisation scheme for reinforcement-learning policies. PG-SOM augments the classical REINFORCE update with two exponentially weighted statistics: a first-order gradient average and a diagonal approximation of the Hessian. By preconditioning the gradient with this curvature estimate, the method adaptively rescales each parameter, yielding faster and more stable ascent of the expected return. We provide a concise derivation, establish that the diagonal Hessian estimator is unbiased and positive-definite under mild regularity assumptions, and prove that the resulting update is a descent direction in expectation. Numerical experiments on standard control benchmarks show up to a 2.1x increase in sample efficiency and a substantial reduction in variance compared to first-order and Fisher-matrix baselines. These results indicate that even coarse second-order information can deliver significant practical gains while incurring only D memory overhead for a D-parameter policy. All code and reproducibility scripts will be made publicly available.

Policy Gradient with Second Order Momentum

TL;DR

PG‑SOM introduces a lightweight second‑order momentum method for policy gradients by maintaining a diagonal Hessian estimate to precondition gradients. The diagonal curvature estimator is shown to be an unbiased, positive‑definite surrogate that ensures a descent direction in expectation, and the method adds only O memory. Empirical results on standard RL benchmarks demonstrate up to about 2.1x improvements in sample efficiency and reduced variance compared with first‑order and Fisher‑based baselines, with Runge–Kutta style two‑stage updates providing the best early learning and final performance when combined with gradient clipping. The work suggests that coarse curvature information, when paired with simple stabilizers, yields meaningful practical gains without the overhead of full second‑order methods or learned critics.

Abstract

We develop Policy Gradient with Second-Order Momentum (PG-SOM), a lightweight second-order optimisation scheme for reinforcement-learning policies. PG-SOM augments the classical REINFORCE update with two exponentially weighted statistics: a first-order gradient average and a diagonal approximation of the Hessian. By preconditioning the gradient with this curvature estimate, the method adaptively rescales each parameter, yielding faster and more stable ascent of the expected return. We provide a concise derivation, establish that the diagonal Hessian estimator is unbiased and positive-definite under mild regularity assumptions, and prove that the resulting update is a descent direction in expectation. Numerical experiments on standard control benchmarks show up to a 2.1x increase in sample efficiency and a substantial reduction in variance compared to first-order and Fisher-matrix baselines. These results indicate that even coarse second-order information can deliver significant practical gains while incurring only D memory overhead for a D-parameter policy. All code and reproducibility scripts will be made publicly available.
Paper Structure (29 sections, 2 theorems, 16 equations, 3 figures, 1 table, 1 algorithm)

This paper contains 29 sections, 2 theorems, 16 equations, 3 figures, 1 table, 1 algorithm.

Table of Contents

  1. Introduction
  2. Motivation.
  3. Contributions.
  4. Paper outline.
  5. Background and related work
  6. Policy Gradient with Second Order Momentum
  7. Second Order Policy Gradient Theorem
  8. Runge Kutta method
  9. Advantage function
  10. Clipping
  11. Entropy Regularization
  12. Algorithmic Overview
  13. Computational Complexity
  14. Experiments
  15. Dataset and implementation details
  16. ...and 14 more sections

Key Result

Lemma 3.1

Let the reward function be defined as The Hessian of the expected return admits the decomposition where $\Psi(\theta;\tau)=\sum_{h=0}^{H} \ln\pi_\theta(a_h\!\mid\!s_h)\,\sum_{t=h}^{H}\gamma^t\,r(s_t,a_t).$

Figures (3)

  • Figure 1: Per‐method learning curves (mean $\pm$ 1 std over 5 seeds). (a) vanilla PG shows a steady rise but retains wide variance throughout training. (b) Hessian PG leverages curvature estimates to boost early returns, yet its confidence band remains large until late episodes. (c) Runge–Kutta combines two‐stage updates to achieve the fastest initial climb, the narrowest bands, and the strongest mid‐training performance.
  • Figure 2: Learning curves (mean $\pm$ 1 std, 5 seeds). Runge--Kutta attains a return of 200 roughly 40 % faster than vanilla Policy Gradient and 55 % faster than Hessian PG, and maintains the steepest slope throughout training.
  • Figure 3: Final-episode return (mean $\pm$ 1 std). For every base method, adding gradient clipping (“Clipping”) yields the best asymptotic performance; RK+clip achieves the overall maximum ( 390).

Theorems & Definitions (4)

  • Lemma 3.1: Diagonal Curvature Estimate
  • Lemma 3.2: Score‐Function Independence
  • proof
  • proof