PaperlayerView exact revision on arXiv

Source-linked AI summary

Principled Penalty-based Methods for Bilevel Reinforcement Learning and RLHF

Han Shen, Zhuoran Yang, Tianyi Chen

arXiv:2402.06886v3cs.LGmath.OCstat.ML

TL;DR

Bilevel RL is difficult to address with existing bilevel methods because its discounted-return objective lacks their usual structural assumptions. The paper develops penalty reformulations and a first-order policy-gradient algorithm, with convergence results for standard and zero-sum game constraints.

  • Problem

    Existing bilevel methods rely on strong convexity or uniform error bounds that the nonconvex discounted-return objective in bilevel RL does not satisfy.

    Source: §1 · p. 2, §1.1 · p. 3

  • Method

    The paper uses value and Bellman penalty functions to reformulate bilevel RL and develops a first-order policy-gradient algorithm, extending the approach to zero-sum game constraints.

    Source: §1 · p. 3

  • Results

    The paper establishes a provably convergent first-order algorithm for bilevel RL and a provably convergent algorithm for bilevel RL with a game constraint.

    Source: §1 · p. 3

  • Takeaways & Limitations

    An approximate solution to the penalty reformulation is also an effective solution to the original bilevel problem.

    Source: §1 · p. 3

    The differentiability analysis relies on assumptions including reward differentiability, policy uniqueness conditions, and MDP regularity.

    Source: §4.1 · p. 11, §4.2 · p. 12

Abstract

from arXiv · show

Bilevel optimization has been recently applied to many machine learning tasks. However, their applications have been restricted to the supervised learning setting, where static objective functions with benign structures are considered. But bilevel problems such as incentive design, inverse reinforcement learning (RL), and RL from human feedback (RLHF) are often modeled as dynamic objective functions that go beyond the simple static objective structures, which pose significant challenges of using existing bilevel solutions. To tackle this new class of bilevel problems, we introduce the first principled algorithmic framework for solving bilevel RL problems through the lens of penalty formulation. We provide theoretical studies of the problem landscape and its penalty-based (policy) gradient algorithms. We demonstrate the effectiveness of our algorithms via simulations in the Stackelberg Markov game, RL from human feedback and incentive design.

1 Introduction

Bilevel RL models a leader whose objective depends on a follower’s best response, but its nonconvex discounted-return objective lacks structures required by existing bilevel methods. The paper develops penalty-based first-order methods for this setting and its zero-sum game extension.

1 Introduction

  • The framework models a follower optimizing an MDP parameterized by the leader’s decision, with the leader optimizing subject to the follower’s best-response policy.The formulation covers applications including Stackelberg Markov games, reward learning, and RLHF.Source: §1 · p. 2
  • Existing implicit-gradient methods rely on strong convexity, while established penalty methods rely on uniform error bounds that discounted-return objectives do not satisfy.The introduction therefore identifies the penalty function, its relation to the original problem, and a convergent solution method as open needs for bilevel RL.Source: §1 · p. 2, §1.1 · p. 3

1.1 Our contributions

  • The paper extends penalty-based bilevel optimization with value and Bellman penalties, proves properties connecting the reformulation to the original problem, and establishes a provably convergent first-order policy-gradient algorithm.The value penalty uses the optimal value function; the Bellman penalty uses Bellman error.Source: §1 · p. 3
  • For a zero-sum Markov game with two lower-level RL agents, the paper derives a value penalty from the Nikaido-Isoda function and gives the first provably convergent algorithm for a game constraint.The authors describe this extension as a route toward more complicated lower-level problems.Source: §1 · p. 3

1.2 Related works

  • Prior penalty-based bilevel analyses use uniform Polyak-Łojasiewicz error bounds, whereas the paper notes that discounted-return objectives do not satisfy those bounds to its knowledge.The related-work discussion also contrasts the paper’s fully first-order approach with PARL, which uses implicit gradients and second-order derivatives.Source: §1.2 · pp. 4–5

2 Problem Formulations

This section formulates bilevel RL as a leader-follower problem in which the follower’s MDP can depend on the leader’s decision. It then shows how game-theoretic, human-feedback, and reward-shaping problems fit that structure.

2.1 Bilevel reinforcement learning formulation

  • A parameterized MDP Mτ(x) uses leader-controlled parameter x to specify the reward and transition distribution, with τh providing policy regularization.The policy is parameterized by y, and τ=0 gives an unregularized MDP.Source: §2.1 · pp. 5–6
  • The generic problem constrains the upper-level decision by a lower-level optimality condition, with compact convex decision sets and a positive initial-state distribution.The follower’s policy class is parameterized by y.Source: §2.1 · p. 6
  • In bilevel RL, a leader’s objective is constrained by the follower’s optimal policy for an MDP determined by the leader’s decision.The section introduces a generic formulation and then presents applications.Source: §2 · p. 5, §2.1 · p. 6

2.2 Applications of bilevel reinforcement learning

  • A Stackelberg Markov game fits the framework by treating the follower’s best response to the leader’s policy as the lower-level problem.The leader and follower act simultaneously at each time step, receive respective rewards, and affect state transitions jointly.Source: §2.2 · pp. 6–7
  • In deep RLHF, preference labels on pairs of agent-collected segments train a reward predictor while an RL agent learns from that reward.The preference probability is modeled with Bradley-Terry.Source: §2.2 · p. 7
  • Under a special tabular, single-step setting with KL regularization and γ=0, the lower-level optimality relation yields a DPO-like upper-level objective.The paper distinguishes its on-policy upper-level loss from DPO’s off-policy dataset loss.Source: §2.2 · p. 7
  • Reward shaping seeks a parameterized reward that supports more efficient learning while preserving the original task, and is formulated as a special case of bilevel RL.The motivating setting includes sparse reward signals.Source: §2.2 · pp. 7–8

3 Penalty Reformulation of Bilevel RL

The section reformulates bilevel RL as a single-level objective with a penalty for lower-level suboptimality. It develops value and Bellman penalties and establishes conditions under which their solutions relate to those of the original problem.

3 Penalty Reformulation of Bilevel RL

  • The reformulation minimizes Fλ(x, y) = f(x, y) + λp(x, y) over x and y, using p to penalize lower-level suboptimality.λ is the penalty constant.Source: §3 · p. 8

3.1 Value penalty and its landscape property

  • The value penalty measures the lower-level optimality gap, and the follower's value function is gradient-dominated under the stated policy-class condition.The cited result generalizes a gradient-dominance property from unregularized to regularized MDPs.Source: §3.1 · pp. 8–9
  • Assuming f is bounded on X × Y, for accuracy δ > 0, choosing λ ≥ Cδ^-1 makes an ε-minimum of the value-penalized problem an ελ-minimum of the relaxed problem, with ελ ≤ δ + λ^-1ε.Source: §3.1 · p. 8
  • Under direct policy parameterization and Lipschitz continuity of f in y, a sufficiently large value-penalty weight transfers local solutions to the relaxed bilevel problem.The stated threshold is λ ≥ LCuδ^-1, with relaxed-problem accuracy ελ ≤ δ.Source: §3.1 · p. 9

3.2 Bellman penalty and its landscape property

  • For τ > 0, the Bellman construction uses g(x, y) − v(x) as an optimality metric, and its minimum identifies the unique optimal policy of the regularized MDP.The Bellman penalty is introduced for tabular policies; its penalty function is τ-strongly convex in y.Source: §3.2 · pp. 9–10
  • With Lipschitz continuity of f in y and a sufficiently large penalty weight, local or global solutions of the Bellman-penalized problem are solutions of BMε with ελ ≤ δ.The stated threshold is λ ≥ Lτ^-1δ^-1.Source: §3.2 · p. 10

4 A Penalty-based Bilevel RL Algorithm

The section derives gradients for both penalty choices, then uses them in a first-order algorithm. Under stated regularity and oracle-accuracy conditions, it establishes smoothness, convergence, and an iteration-complexity bound.

4.1 Differentiability of the value penalty

  • Under Assumption 1, the value penalty is differentiable in x, with closed-form gradient results specified for the applications.The stated assumptions include differentiability of the reward and, in some settings, uniqueness of the optimal policy; τ > 0 is given as a sufficient condition for uniqueness.Source: §4.1 · p. 11

4.2 Differentiability of the Bellman penalty

  • Under Assumption 2, the Bellman penalty is differentiable in x with gradient ∇xp(x, y) = ∇xg(x, y) − ∇v(x).The cited regularity assumption includes τ > 0 and a condition on the MDP's induced Markov chain unless the discount factor is zero.Source: §4.2 · p. 12

4.3 A gradient-based algorithm and its convergence

  • PBRL optimizes the penalized objective using penalty gradients, with the lower-level RL problem solved approximately at each outer iteration.The approximate penalty gradient replaces the optimal follower policy with its approximate policy estimate.Source: §4.3 · p. 13
  • Under Assumptions 2 and 4, either penalty is Lipschitz-smooth; with Assumptions 2–5 and the stated step-size bound, PBRL satisfies the convergence theorem.The step size must satisfy α ≤ 1/(Lf + λLp).Source: §4.3 · p. 14
  • With policy mirror descent as the oracle, the stated iteration complexity is Õ(λε^-1).The general bound is O(λε^-1 com(εorac)); the mirror-descent oracle has complexity O(log(λ²/εorac)).Source: §4.3 · p. 14

5 Bilevel RL with Lower-level Zero-sum Games

The section extends penalty-based bilevel RL to a zero-sum Markov game, where a leader chooses a parameter and followers seek a Nash equilibrium. It uses the Nikaido-Isoda function to penalize equilibrium error and establishes finite-time convergence for a projected policy-gradient update.

5.1 Formulation

  • The formulation minimizes the leader’s objective over its parameter and a selected Nash equilibrium; incentive design is one motivating example.In that example, the designer changes agents’ rewards to maximize the designer’s reward net of cost.Source: §5.1 · pp. 15–16

5.2 Nikaido-Isoda function as a penalty

  • The Nikaido-Isoda function ψ is nonnegative, is at most 2ϵ for an ϵ-NE, and ψ≤ϵ implies an ϵ-NE.Under Assumption 6, ψ is also μ-gradient dominated in π, with μ=(1−γ) mins ρ(s).Source: §5.2 · pp. 16–17
  • Penalizing the objective with λψ recovers a solution of the relaxed bilevel problem within accuracy δ when λ≥δ^-1.This result holds under Assumption 6 and when f is L-Lipschitz-continuous in π.Source: §5.2 · pp. 16–17

5.3 A policy gradient based algorithm and its convergence analysis

  • At each iteration, estimating ∇ψ uses approximate optimal policies for each player against the other player’s current policy, obtainable with single-agent policy optimization.The section extends the earlier single-agent lower-level MDP penalty method to zero-sum Markov games.Source: §5 · p. 15, §5.3 · pp. 17–18
  • The projected update has iteration complexity O(λϵ^-1com(ϵorac)); with policy mirror descent, its complexity is ˜O(λϵ^-1).The convergence theorem assumes Assumptions 5–8; policy mirror descent gives com(ϵorac) = O(log(λ2/ϵorac)).Source: §5.3 · p. 19

6 Simulation

The simulations evaluate PBRL in Stackelberg games, RLHF, and incentive design, comparing its performance with independent gradients, DRLHF, and Meta-Gradient. Across these settings, the reported results favor PBRL on leader return, selected Atari games, and designer reward.

6 Simulation

  • The simulations test PBRL across a Stackelberg Markov game, Atari RLHF, and incentive design.The Atari preferences are assigned to the segment with the higher ground-truth reward, and true-reward A2C is used as an oracle baseline.Source: §6 · p. 19, §6.2 · p. 20

6.1 Stackelberg Markov game

  • Both PBRL penalties outperform independent policy gradients on leader return, with the value penalty reaching the highest value.The follower’s optimality gap diminishes to zero; the experiment covers 10 randomly generated MDPs.Source: §6.1 · p. 19, Figure 1 · p. 20

6.2 Deep reinforcement learning from human feedback

  • On BeamRider and MsPacman, PBRL achieves higher episode return than DRLHF and becomes comparable to the true-reward A2C oracle.On Seaquest, PBRL and DRLHF have similar asymptotic performance, while DRLHF is more unstable during training.Source: §6.2 · p. 20, Figure 2 · p. 20

6.3 Incentive design

  • In incentive design, PBRL and Meta-Gradient both exceed the no-incentive start line, and PBRL yields higher designer reward.Both methods find an approximate Nash equilibrium under their respective incentive rewards.Source: §6.3 · p. 21, Figure 3 · p. 21

7 Concluding Remarks

The conclusion presents PBRL as a penalty-based first-order algorithm for bilevel reinforcement learning. It reports penalty reformulation results, gradient analysis, convergence, and evaluations in three application settings.

7 Concluding Remarks

  • The paper proposes a penalty-based first-order algorithm for bilevel reinforcement learning.Its reported evaluations cover the Stackelberg Markov game, RLHF, and incentive design.Source: §7 · pp. 21–22
  • The authors identify penalty functions with landscape properties such that the induced reformulation admits solutions to the original bilevel RL problem.The conclusion also reports studying penalty differentiability and deriving closed-form gradients.Source: §7 · pp. 21–22
  • The paper proposes a convergent PBRL algorithm based on the penalty and gradient results.The algorithm is evaluated on the Stackelberg Markov game, RLHF, and incentive design.Source: §7 · p. 22

A Preliminary results

The preliminary result establishes uniqueness and local Lipschitz continuity of the optimal-policy mapping under stated conditions. The proof uses strong convexity of the policy objective and a variational-inequality solution mapping.

A Preliminary results

  • Under Assumption 2, τ > 0, and compact X, the optimal policy is locally Lipschitz continuous in the parameterized MDP setting.The lemma considers optimal policies in a convex policy class Π.Source: A Preliminary results · p. 28
  • The optimal policy is unique for the parameterized MDP on the convex policy class Π.The proof attributes uniqueness to τ-strong convexity of the policy objective.Source: A Preliminary results · p. 28
  • The proof characterizes the optimal policy through a parameterized variational inequality and applies a solution-mapping result to obtain local Lipschitz continuity.The cited solution-mapping result gives local τ^-1-Lipschitz continuity.Source: A Preliminary results · p. 28

B.1 Proof that Stackelberg Markov game is a bilevel RL problem

This appendix develops the proof framework for casting the Stackelberg Markov game as a bilevel RL problem, covering penalty guarantees, policy gradients, differentiability, and smoothness. It establishes regularity and penalty-smoothness bounds and gives examples of penalty-gradient estimators and inner-policy updates.

B.1 Proof that Stackelberg Markov game is a bilevel RL problem

  • Averaging rewards and transitions over the leader’s policy defines a follower-view parametric MDP, allowing the Stackelberg game to be rewritten as a bilevel problem.The construction uses the follower’s policy as the MDP policy and the leader’s policy as the parameter.Source: §B.1 · p. 29, Then we have V · p. 29

B.2–B.4 Proofs of Lemmas 1–3

  • The penalty proofs bound the follower residual by δ + λ^-1 ε and choose λ so that ελ ≤ δ, including a result for feasible points in a neighborhood.The neighborhood result follows from local optimality of the penalized solution.Source: §B.2 · p. 30, §B.4 · p. 32

B.3 Proof of Lemma 2

  • The regularized policy-gradient proof derives a gradient using the discounted visitation distribution and applies the performance difference lemma and convexity in the subsequent bound.The cited performance difference lemma is attributed to [Lan, 2023, Lemma 2] and [Zhan et al., 2023, Lemma 5].Source: §B.3 · pp. 30–31

B.5 Proof of Lemma 4

  • The follower’s regularized optimum is unique under τ-strong convexity, and the proof identifies the minimizing policy with the follower’s optimal-policy set.The proof also establishes a bound between regularized and unregularized values.Source: §B.5 · pp. 32–33

C.1 Proof of Lemma 6

  • Under Assumption 1, generalized Danskin’s theorem gives the gradient of the maximum follower value using any optimal follower policy.The proof invokes the theorem’s differentiability conditions for the maximized value function.Source: §C.1 · p. 34

C.2 Proof of Lemma 7

  • The leader-view MDP averages rewards and transitions over the follower’s policy and includes the follower’s regularization penalty; the proof uses the policy-gradient theorem to derive the leader gradient.The leader-view MDP does not regularize the leader’s policy.Source: §C.2 · p. 34

C.3–C.4 Proofs of Lemmas 8–9

  • The differentiability proof relies on uniqueness of the follower’s optimal policy and an irreducible Markov chain, while the gradient proof uses the log-trick and policy-gradient identities.These arguments establish differentiability of the relevant value-based quantities and derive their gradients.Source: §C.3 · p. 35, §C.4 · pp. 36–37

C.5 Sufficient conditions of the smoothness assumption

  • The smoothness conditions cover bounded policy derivatives, smooth regularization when τ > 0, smooth value dependence on x, and bounded rewards; direct and softmax parameterizations satisfy the policy conditions.The reward-smoothness condition is application-specific; RLHF and reward shaping can be checked from the value-gradient formula.Source: §C.5 · p. 37, Then it holds for any s that V · p. 37

Then it holds for any s that V

  • The proof establishes uniform-in-x Lipschitz continuity in y for π_y^{s0,x}(s), using a decomposition and the lemma’s assumptions to complete the result.The supplied passages describe the decomposition and its continuation but do not state the final bound.Source: Then it holds for any s that V · pp. 38–39

C.6.1 Smoothness of the value penalty

  • Under the two assumptions, y(x) is unique and τ−1CJ-Lipschitz continuous on X, so every y, y′ ∈ Y∗(x) induces the same optimal policy π∗.Source: §C.6.1 · p. 39

C.6.2 Smoothness of the Bellman penalty

  • The Bellman penalty is Lipschitz smooth with modulus L_g(2 + τ^-1C_J), yielding F_λ smoothness modulus L_f + λL_g(2 + τ^-1C_J).Here L_g = 2C_J + L_v(1 + τ^-1C_J) + L_h.Source: §C.6.2 · p. 40

C.7 Example gradient estimators of the penalty functions

  • The value- and Bellman-penalty gradient estimators can use approximate inner policies, with policy mirror descent or policy-gradient updates as options.The value-estimator error is bounded by the policy optimality gap, while the supplied examples give logarithmic inner-iteration complexities and mention REINFORCE and advantage actor-critic as policy-gradient methods.Source: §C.7 · p. 41

C.8 Proof of Theorem 1

Theorem 1’s proof derives a descent bound for the penalized objective and telescopes it to obtain the stated result; the OS case follows similarly without an oracle. The Lemma 12 proof also establishes uniqueness of the starred policies under its parameterization.

C.8 Proof of Theorem 1

  • The update yields a descent inequality that bounds the squared step by penalized-objective decrease plus gradient-estimation error.The proof uses smoothness of Fλ, the projected update’s optimality condition, and Young’s inequality.Source: §C.8 · p. 42
  • Telescoping the descent bound and applying the stated assumptions completes the first inequality of Theorem 1.The proof uses nonnegativity of the penalty to lower-bound Fλ; the OS result follows similarly with Lv-smoothness and no oracle.Source: §C.8 · pp. 42–43, §D.6 · p. 47

D.1 Proof of Lemma 12

  • Treating (π2, x) or (π1, x) as the parameter makes the corresponding starred policy unique; Assumption 6 and Lemma 17 then establish part (a).The proof states that part (b) follows using Lemma 2.Source: §D.1 · p. 43

D.2 Proof of Lemma 13

These proofs relate penalized local solutions to approximate feasibility, justify smoothness and policy regularity under stated assumptions, and bound the effort needed for accurate gradient estimation.

D.2 Proof of Lemma 13

  • The chosen λ ensures ψ(xλ, πλ) ≤ δ, and local optimality minimizes f + λψ over feasible points in the solution’s region.The proof uses this comparison to establish the result for feasible points in the region.Source: §D.2 · p. 44

D.3 Justification of the smoothness assumption in the two-player case

  • Under the lemma’s assumptions, the proof establishes uniform Lipschitz-smoothness constants for the two players’ value functions.The assumptions include bounds on rewards and smoothness conditions for the players’ policies.Source: §D.3 · pp. 44–45

D.4 Proof of Lemma 14

  • The starred policy mappings are Lipschitz-continuous in their parameters, yielding a Lipschitz bound for the combined mapping.The proof applies the result to each player’s policy and then uses Lemma 12(a).Source: §D.4 · p. 45

D.5 Gradient Estimator Accuracy

  • O(−log(ϵorac/λ^2)) inner iterations suffice to meet Assumption 7’s gradient-estimator accuracy requirement.The proof obtains this complexity by using policy mirror descent to approximate each player’s starred policy.Source: §D.5 · p. 45

D.6 Proof of Theorem 2

Theorem 2’s proof establishes descent for the penalized objective despite gradient-estimation error, then combines that bound with projection and telescoping arguments to complete the result.

D.6 Proof of Theorem 2

  • The proof derives a descent inequality bounding the squared update by penalized-objective decrease and a term for gradient-estimation error.The error term is scaled by αλ^2 and the squared difference between the exact and estimated gradients of ψ.Source: §D.6 · p. 46
  • Projection non-expansiveness and the descent inequality are combined to obtain the bound used in the convergence argument.The proof states that these inequalities together imply the next result.Source: §D.6 · pp. 46–47
  • Telescoping the bound, using nonnegativity of the penalty and Assumption 7, completes the theorem’s proof.The proof lower-bounds Fλ by the infimum of f over the feasible set.Source: §D.6 · p. 47

E Additional Experiment Details

The appendix reports implementation choices for experiments in the Stackelberg Markov game, deep RLHF, and incentive design, including model architecture and training schedules.

E.1 Stackelberg Markov game

  • For the Stackelberg Markov game, PBRL uses learning rate 0.1 and Monte Carlo trajectories of length 5 with batch size 16; value-penalty and Bellman-penalty settings use λ = 2, T = 1 and λ = 7, T = 10, respectively.The independent policy-gradient baseline uses the same trajectory length and batch size.Source: §E.1 · p. 47

E.2 Deep reinforcement learning from human feedback

  • For Atari RLHF, A2C policy and critic share a convolutional base, while the reward predictor uses a separate convolutional model and is trained synchronously with the policy.The reward predictor uses dropout with probability 0.5 between fully connected layers to prevent over-fitting.Source: §E.2 · pp. 47–48

E.3 Incentive design

  • For incentive design, PBRL uses λ = 4 and one inner policy-gradient iteration per outer iteration, with Monte Carlo trajectory length 5 and batch size 24.The meta-gradient method uses the same learning rate, trajectory length, batch size, and inner iteration number.Source: §E.3 · p. 48
Loading 2402.06886v3…