Policy Gradient with Second Order Momentum

TL;DR

This paper introduces PG-SOM, a second-order optimization method for reinforcement learning that improves sample efficiency and stability by using a diagonal Hessian approximation to precondition policy gradients.

Contribution

The paper presents PG-SOM, a novel lightweight second-order optimizer for policy gradients that leverages a diagonal Hessian estimate to enhance learning performance.

Findings

Up to 2.1x increase in sample efficiency.

Significant reduction in variance compared to baselines.

Method incurs only D memory overhead for D-parameter policies.

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.