Bi-Level Policy Optimization with Nyström Hypergradients

TL;DR

This paper reframes actor-critic reinforcement learning as a bilevel optimization problem, with the actor as the leader and the critic as the follower. It introduces BLPO, a policy-gradient algorithm that nests the critic update and uses Nyström-based hypergradients to approximate the inverse Hessian–vector product, improving stability. Under a linear critic parametrization, BLPO is shown to converge in polynomial time to a local strong Stackelberg equilibrium with high probability, and empirically BLPO matches or surpasses PPO on discrete and continuous control tasks. The work demonstrates that low-rank IHVP approximations can enable scalable second-order optimization in deep RL and points to broader BLO applications such as RLHF and meta-learning.

Abstract

The dependency of the actor on the critic in actor-critic (AC) reinforcement learning means that AC can be characterized as a bilevel optimization (BLO) problem, also called a Stackelberg game. This characterization motivates two modifications to vanilla AC algorithms. First, the critic's update should be nested to learn a best response to the actor's policy. Second, the actor should update according to a hypergradient that takes changes in the critic's behavior into account. Computing this hypergradient involves finding an inverse Hessian vector product, a process that can be numerically unstable. We thus propose a new algorithm, Bilevel Policy Optimization with Nyström Hypergradients (BLPO), which uses nesting to account for the nested structure of BLO, and leverages the Nyström method to compute the hypergradient. Theoretically, we prove BLPO converges to (a point that satisfies the necessary conditions for) a local strong Stackelberg equilibrium in polynomial time with high probability, assuming a linear parametrization of the critic's objective. Empirically, we demonstrate that BLPO performs on par with or better than PPO on a variety of discrete and continuous control tasks.

Figures (6)

• Figure 1: In continuous control tasks, BLPO either outperforms PPO or performs comparably.
• Figure 2: Performance plots for (a) CartPole and (b) Acrobot.
• Figure 3: We conduct and ablations to show that the hypergradient with Nyström method is key. Here we show examples where BLPO using conjugate gradient is more unstable and under performs the Nyström method. Importantly, the Nyström method never underperforms conjugate gradient in any of our other experiments.
• Figure 4: Here we show examples of where nesting the critic without the hypergradient is not enough. The Nyström method never under performs nesting in any of our other experiments.
• Figure 5: While BLPO-Nyström is slower than PPO, in most cases the simulator is still the bottleneck, so BLPO's clock time is not significantly behind PPO's. In all cases apart from Humanoid Standup, BLPO-Nyström is faster than BLPO-CG.

Theorems & Definitions (31)

• Definition 2.1
• Definition 2.2
• Theorem 3.1
• Lemma 4.4
• Theorem 5.1
• Theorem 5.2
• Corollary 5.3
• Corollary 5.4
• Theorem B.1
• proof