Sample Complexity Analysis for Constrained Bilevel Reinforcement Learning

Several important problem settings within the literature of reinforcement learning (RL), such as meta-learning, hierarchical learning, and RL from human feedback (RL-HF), can be modelled as bilevel RL problems. A lot has been achieved in these domains empirically; however, the theoretical analysis of bilevel RL algorithms hasn’t received a lot of attention. In this work, we analyse the sample complexity of a constrained bilevel RL algorithm, building on the progress in the unconstrained setting. We obtain an iteration complexity of $O(ε^{-2})$ and sample complexity of $\tilde{O}(ε^{-4})$ for our proposed algorithm, Constrained Bilevel Subgradient Optimization (CBSO). We use a penalty-based objective function to avoid the issue of primal-dual gap and hyper-gradient in the context of a constrained bilevel problem setting. The penalty-based formulation to handle constraints requires analysis of non-smooth optimization. We are the first ones to analyse the generally parameterized policy gradient-based RL algorithm with a non-smooth objective function using the Moreau envelope.

💡 Research Summary

The paper tackles a gap in the theoretical understanding of constrained bilevel reinforcement learning (RL), a setting that naturally arises in meta‑learning, hierarchical RL, and especially RL from human feedback (RL‑HF) where safety or ethical constraints must be respected. While many empirical works have demonstrated the practical success of bilevel formulations, prior theoretical analyses have been limited to unconstrained or strongly convex inner problems. This work introduces the Constrained Bilevel Subgradient Optimization (CBSO) algorithm, which extends recent progress on unconstrained bilevel RL to the constrained, non‑convex inner‑level case.

The authors first formalize the constrained bilevel RL problem: an outer objective (f(x,y)) (non‑convex) depends on a reward‑parameter vector (x) and a policy parameter vector (y); the inner problem minimizes a non‑convex loss (g(x,y)) subject to an inequality constraint (h(y)\le c_0) that captures cost or safety limits. Direct primal‑dual methods fail because the inner problem is non‑convex, leading to a non‑zero primal‑dual gap. To circumvent this, the paper adopts a penalty‑based reformulation that aggregates the inner objective and the constraint violation into a single scalar penalty term. The resulting objective is non‑smooth due to the max‑type penalty.

To analyze this non‑smooth landscape, the authors employ the Moreau envelope, which smooths the objective while preserving its global minima. They define Clarke subdifferentials for the envelope and develop a subgradient descent scheme. Crucially, they assume a Kurdyka‑Łojasiewicz (KL) condition for the inner loss; this condition, when applied to the Moreau envelope, yields a Polyak‑Łojasiewicz (PL) inequality, which in turn guarantees a Quadratic Growth (QG) condition. The QG condition provides a lower bound on the distance to the global optimum, enabling convergence guarantees even though the problem is non‑convex.

The algorithm proceeds in two nested loops. In the inner loop, two policy parameters are updated independently: (y) is optimized with respect to a composite objective that includes the outer loss, the inner loss, and the penalty; (z) is optimized only with respect to the inner loss and penalty, serving as a proxy for the exact inner optimum. Both updates use stochastic subgradients estimated from mini‑batches. After (K) inner updates, the outer parameter (x) is updated using a gradient of a constructed function (\phi(x,y_K,z_K)=h_1(x,y_K)-\frac{1}{\sigma_1}h_2(x,z_K)), where (h_1) and (h_2) are the penalized objectives for (y) and (z) respectively. The analysis shows that, under the KL‑condition, ρ‑hypomonotonicity of the subgradients, and appropriate choices of penalty coefficients ((\sigma_1,\sigma_2,\sigma_3)), the algorithm converges to an (\epsilon)-optimal solution with iteration complexity (O(\epsilon^{-2})). Because each iteration requires a stochastic gradient estimate based on a batch of size (B), the total sample complexity becomes (\tilde O(\epsilon^{-4})).

These results are summarized in Table 1, where CBSO matches or improves upon the iteration and sample complexities of prior unconstrained bilevel RL works while handling inequality constraints and non‑convex inner objectives. The paper also notes that the same theoretical guarantees extend to general constrained bilevel optimization problems, making CBSO the first algorithm to provide non‑asymptotic sample complexity for such settings without assuming convexity of the inner problem.

The contribution is threefold: (1) a novel algorithmic framework for constrained bilevel RL that avoids primal‑dual gaps by using a penalty formulation; (2) a rigorous non‑smooth analysis leveraging Moreau envelopes, Clarke subdifferentials, and KL‑to‑PL‑to‑QG condition chaining; (3) concrete iteration and sample complexity bounds that fill a notable gap in the literature.

While the paper does not present empirical experiments, it outlines the practical implementation details (batch size, learning rates, inner‑loop steps) and highlights future directions, such as validating the KL and QG assumptions in realistic RL environments and exploring adaptive penalty schemes. Overall, the work provides a solid theoretical foundation for safe and constrained RL systems, especially those involving large language models trained with human feedback, where respecting cost or safety constraints is paramount.