🤖 machine learning

Hypergradient-based Bilevel Reinforcement Learning with Improved Sample Complexity

This paper proposes a Hessian-free, hypergradient-based bilevel reinforcement learning algorithm that leverages the optimality of Boltzmann policies to achieve state-of-the-art sample complexity of $\tilde{O}(\epsilon^{-2})$ and iteration complexity of $O(\epsilon^{-1})$ without requiring the Polyak-Lojasiewicz condition on the outer-level objective.

Original authors: Naman Saxena, Mudit Gaur, Vaneet Aggarwal

Published 2026-08-03

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Naman Saxena, Mudit Gaur, Vaneet Aggarwal

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to walk, but you don't know exactly what "walking well" looks like. You have a coach (the robot's brain) who decides how to move its legs, and a judge (the reward system) who decides if those moves are good. The tricky part is that the judge's opinion changes based on what the robot does, and the robot's moves change based on what the judge says. This is a bit like a dance where both partners are trying to guess what the other wants to do next. In the world of artificial intelligence, this is called Reinforcement Learning. Usually, we just try to make the robot happy by giving it points for good moves. But sometimes, we want the robot to learn from human feedback, like when a human says, "I liked that path better than this one." This turns the problem into a Bilevel challenge: an "inner" loop where the robot learns to move, and an "outer" loop where we tweak the judge's rules to match human preferences.

The big problem with solving this dance is that it's incredibly expensive to practice. Every time the robot tries a new step, it needs to see thousands of examples to figure out if it's getting better. Previous methods were like trying to solve a massive jigsaw puzzle while wearing thick gloves; they either needed to calculate the shape of every single piece (which is slow and computationally heavy) or they used a rough guess that required way too many practice attempts to get right. Scientists have been looking for a way to teach these robots efficiently without needing supercomputers or millions of tries. This is where the new research comes in, offering a smarter, lighter way to navigate this complex dance.

The Paper: A New Way to Teach Robots Without the Heavy Lifting

This paper introduces a new algorithm called Approximate Hypergradient Optimization (AHO). Think of it as a clever shortcut for teaching a robot to learn from human preferences. The authors, Naman Saxena, Mudit Gaur, and Vaneet Aggarwal from Purdue University, propose a method that is both faster and requires far fewer practice attempts than the current best methods.

To understand their trick, imagine the robot's learning process as a chef trying to perfect a recipe.

The Inner Level: The chef (the robot's policy) is tasting the dish and adjusting the spices to make it delicious.
The Outer Level: The food critic (the reward parameter) is deciding what "delicious" means. If the critic changes their mind, the chef has to start over.

In the past, to figure out how to change the critic's mind to get a better dish, previous methods tried to calculate the "curvature" of the entire kitchen—the exact shape of every possible mistake the chef could make. This is like trying to measure the exact curve of every spice jar on the shelf. It's accurate, but it's so heavy and slow that it breaks the computer (a problem known as needing the Hessian). Other methods tried to guess the answer by penalizing bad guesses, but this was like guessing the recipe by trial and error, requiring the chef to cook the dish thousands of times just to get it right.

The authors' new approach, AHO, uses a different secret ingredient: the Boltzmann policy. Imagine that instead of the chef guessing randomly, they follow a very specific, mathematically perfect "ideal" recipe that naturally balances trying new things (exploration) with sticking to what works (exploitation). The paper shows that even if the robot's brain (the policy class) isn't perfect enough to hold every possible ideal recipe, it can still use the idea of this perfect recipe to skip the heavy lifting.

Here is what they found:

No More Heavy Lifting: By using the properties of this "ideal" recipe, they managed to remove the need to calculate the heavy curvature (the Hessian). This makes the algorithm scalable, meaning it can run on standard computers even when the robot has a huge brain with millions of parameters.
Fewer Tries Needed: The most exciting result is about efficiency. Previous methods needed a massive number of practice attempts (sample complexity) to learn, roughly proportional to $1/\epsilon^3$ (where $\epsilon$ is how close you want to get to the perfect solution). The new AHO algorithm reduces this to roughly $1/\epsilon^2$ . In plain English, if you want to be twice as accurate, the old way might need eight times more practice, while the new way only needs four times more. This is a significant improvement in how quickly the robot can learn.
Dropping Old Assumptions: The paper also proves that you don't need to assume the "judge" (the outer objective) has a very specific, rigid shape (called the Polyak-Łojasiewicz or PL condition) for the math to work. This makes the method more flexible and applicable to real-world problems where things aren't always perfectly shaped.

How sure are they?
The authors provide a rigorous mathematical proof showing that their algorithm converges to a good solution under certain standard conditions. They didn't just guess; they derived the math to show that the error decreases at a predictable rate. They also tested their idea on two specific robot tasks: making a bipedal robot walk and a cheetah-like robot run. In these simulations, their method (AHO) learned to walk and run faster and with higher rewards than the previous best method (Gaur et al., 2025).

What they ruled out:
The paper explicitly argues against the idea that you must use the heavy, slow Hessian calculations to get good results in this setting. They also show that you don't need the strict "unique minimizer" assumption (where there is only one single best answer) or the rigid PL condition on the outer level that other top-tier methods required.

The Bottom Line:
This paper suggests that by using a clever mathematical shortcut based on the "ideal" Boltzmann policy, we can teach robots to learn from human feedback much faster and with less computing power. It's not a magic wand that solves everything instantly, but it removes the heavy weights that were slowing down the dance, allowing the robot to learn its steps with far fewer tries. The authors demonstrate this through both solid math and computer simulations, showing a clear path toward more efficient and scalable AI learning.

Technical Summary: Hypergradient-based Bilevel Reinforcement Learning with Improved Sample Complexity

Problem Formulation

The paper addresses Bilevel Reinforcement Learning (RL), a framework used to formalize problems such as meta-learning, hierarchical task decomposition, and Reinforcement Learning from Human Feedback (RL-HF). In this setting, an outer-level objective $\phi(x, \theta^*(x))$ is optimized with respect to a parameter $x$ (e.g., reward function parameters), while the inner-level objective $J(\theta, x)$ is optimized with respect to policy parameters $\theta$ . The inner-level problem is typically non-convex, and the outer-level objective depends on the optimal solution $\theta^*(x)$ of the inner problem.

Existing approaches face significant limitations:

Penalty-based methods (e.g., Gaur et al., 2025) often rely on approximate gradients and require the Polyak-Łojasiewicz (PL) condition on the outer-level objective, resulting in sub-optimal sample complexity of $\tilde{O}(\epsilon^{-3})$ .
Hessian-based hypergradient methods (e.g., Chakraborty et al., 2024) avoid the PL assumption but suffer from scalability issues due to the need to compute and invert the Hessian matrix, which is computationally prohibitive for large parameter spaces.
Finite-space methods (e.g., Yang et al., 2025) utilize Boltzmann policy optimality to derive Hessian-free gradients but do not directly extend to parameterized, non-realizable policy classes in continuous spaces.

The core challenge is to develop a scalable, Hessian-free bilevel RL algorithm that achieves state-of-the-art sample complexity without assuming the PL condition on the outer-level objective or requiring a realizable policy class (where the policy class can represent all Boltzmann policies).

Methodology

The authors propose the Approximate Hypergradient Optimization (AHO) algorithm. The methodology relies on three key theoretical and algorithmic components:

1. Exploiting Boltzmann Policy Optimality

The inner-level objective is defined as an entropy-regularized discounted reward function. The paper leverages the known optimality of the Boltzmann policy ( $\pi_B$ ) for this specific objective. While the policy class $\Theta$ is generally non-realizable (i.e., it cannot perfectly represent the optimal Boltzmann policy), the authors establish conditions under which the optimality of $\pi_B$ can still be utilized to derive gradients.

2. Gradient-Shifted Value Functions

To avoid the Hessian term typically found in hypergradient expressions (derived via the Implicit Function Theorem), the authors introduce gradient-shifted value functions ( $U^\pi$ and $W^\pi$ ).

$U^\pi(s, \theta, x)$ and $W^\pi(s, a, \theta, x)$ are defined as the long-term discounted summations of the reward gradient $\nabla r(s, a, x)$ under the Boltzmann policy.
These functions act as surrogates for the gradients of the Boltzmann value function ( $\nabla V^{\pi,B}$ ) and Q-value function ( $\nabla Q^\pi$ ).
Approximation Error: Due to the non-realizability of the policy class, there is an approximation error between these surrogate functions and the true gradients. This error is bounded by terms $\epsilon_{kl}$ (related to KL divergence) and $\epsilon_{fd}$ (related to gradient differences), which vanish if the policy class is realizable.

3. Approximate Hypergradient Theorem

The authors derive an approximate hypergradient for the outer-level objective. By substituting the true gradients with the gradient-shifted functions, they obtain a closed-form expression that eliminates the Hessian inverse term. The resulting hypergradient consists of:

The direct gradient of the preference loss with respect to $x$ .
A term involving the difference between the gradient-shifted Q-value and value functions, scaled by the temperature coefficient $\tau$ .

4. Algorithm (AHO)

The AHO algorithm (Algorithm 1) operates in a nested loop:

Inner Loop: Updates the policy parameters $\theta$ using a standard policy gradient method with an estimated Q-value function.
Outer Loop: Updates the reward parameter $x$ using the approximate hypergradient. This involves estimating the gradient-shifted functions ( $U^\pi_{off}$ and $W^\pi_{off}$ ) using off-policy Q-learning type updates (Algorithms 3 and 4) based on collected trajectories.

Key Contributions

Boltzmann Optimality with Non-Realizable Classes: The paper provides conditions (Assumption 3.1) under which the optimality of the Boltzmann policy can be applied to parameterized, non-realizable policy classes, discussing the impact of the resulting approximation errors on convergence.
Hessian-Free Scalable Algorithm: The proposed AHO algorithm is fully Hessian-free, making it scalable to high-dimensional parameter spaces where Hessian-based methods fail.
Improved Sample Complexity: The algorithm achieves an iteration complexity of $O(\epsilon^{-1})$ and a sample complexity of $\tilde{O}(\epsilon^{-2})$ . This improves upon the previous state-of-the-art sample complexity of $\tilde{O}(\epsilon^{-3})$ found in penalty-based methods (Gaur et al., 2025).
Relaxed Assumptions: Unlike prior state-of-the-art work (Gaur et al., 2025), AHO removes the assumption of the PL condition on the outer-level objective function. It also relaxes the requirement for a unique minimizer in the inner-level problem, requiring only a locally isolated and non-degenerate minimizer.

Results and Analysis

Theoretical Guarantees

Under mild regularity conditions (Lipschitz continuity, bounded rewards/log-probabilities, and PL condition on the inner level only), the authors prove that the average squared norm of the outer-level gradient converges to a bound dependent on the approximation errors and sampling parameters:
$\frac{1}{T} \sum_{t=0}^{T-1} \|\nabla \Phi(x_t)\|^2 \leq O(\epsilon) + O(\epsilon_{approx}) + O(\epsilon_{fd}) + O(\epsilon_{kl})$
By setting parameters $T = \Theta(\epsilon^{-1})$ , $B = \Theta(\epsilon^{-1})$ , $K = \Theta(\log(\epsilon^{-1}))$ , and $H = \Theta(\log(\epsilon^{-1}))$ , the sample complexity is established as $\tilde{O}(\epsilon^{-2})$ .

Empirical Evaluation

The authors evaluate AHO on continuous control tasks from the DM Control suite (walker-walk and cheetah-run) formulated as preference-based bilevel RL problems.

Baseline: They compare AHO against the penalty-based bilevel RL algorithm of Gaur et al. (2025).
Comparison Constraint: They explicitly note that they do not compare against PARL (Chakraborty et al., 2024) because PARL's Hessian-based approach is computationally infeasible for the network sizes used (requiring inversion of a Hessian with $\sim 10^{12}$ entries).
Outcome: AHO achieves higher total evaluation reward than the Gaur et al. (2025) baseline on both tasks, empirically supporting the theoretical improvement in sample complexity.

Significance

The paper claims that AHO represents a significant step forward in bilevel RL by resolving the trade-off between scalability and sample efficiency.

Scalability: By removing the Hessian, the method becomes applicable to modern, large-scale neural network policies used in preference-based RL.
Efficiency: By utilizing the structure of the entropy-regularized objective and Boltzmann optimality, the method avoids the gradient approximation vs. sampling error trade-off inherent in penalty methods, leading to a better sample complexity bound ( $\tilde{O}(\epsilon^{-2})$ vs. $\tilde{O}(\epsilon^{-3})$ ).
Theoretical Robustness: The removal of the PL condition on the outer level broadens the applicability of the algorithm to a wider range of non-convex outer objectives common in RL-HF and meta-learning.