PARL: A Unified Framework for Policy Alignment in Reinforcement Learning from Human Feedback

TL;DR

This work tackles policy alignment in reinforcement learning by introducing PARL, a bilevel optimization framework that explicitly couples reward design with the data generated by the learned policy. It generalizes RLHF and addresses distribution shifts by formulating the upper-level objective as dependent on the lower-level policy, and it provides the A-PARL algorithm with a provable O(1/T) convergence rate. Empirically, A-PARL delivers near-oracle performance with significantly improved sample efficiency on large-scale robotics benchmarks like the DM Control Suite and MetaWorld. Overall, the approach offers a rigorous, practical path toward robust human-feedback–driven policy alignment in RL.

Abstract

We present a novel unified bilevel optimization-based framework, \textsf{PARL}, formulated to address the recently highlighted critical issue of policy alignment in reinforcement learning using utility or preference-based feedback. We identify a major gap within current algorithmic designs for solving policy alignment due to a lack of precise characterization of the dependence of the alignment objective on the data generated by policy trajectories. This shortfall contributes to the sub-optimal performance observed in contemporary algorithms. Our framework addressed these concerns by explicitly parameterizing the distribution of the upper alignment objective (reward design) by the lower optimal variable (optimal policy for the designed reward). Interestingly, from an optimization perspective, our formulation leads to a new class of stochastic bilevel problems where the stochasticity at the upper objective depends upon the lower-level variable. {True to our best knowledge, this work presents the first formulation of the RLHF as a bilevel optimization problem which generalizes the existing RLHF formulations and addresses the existing distribution shift issues in RLHF formulations.} To demonstrate the efficacy of our formulation in resolving alignment issues in RL, we devised an algorithm named \textsf{A-PARL} to solve PARL problem, establishing sample complexity bounds of order $\mathcal{O}(1/T)$. Our empirical results substantiate that the proposed \textsf{PARL} can address the alignment concerns in RL by showing significant improvements (up to 63\% in terms of required samples) for policy alignment in large-scale environments of the Deepmind control suite and Meta world tasks.

Figures (5)

• Figure 1: (a) This figure shows the proposed PARL framework for policy alignment in reinforcement learning. The standard RL is at the lower level (LL), and the alignment objective is at the upper level (UL). (b) This figure shows the performance gap of the SOTA approach due to policy misalignment. The blue curve should be as close as possible to the red dotted line of oracle.
• Figure 2: In this figure, we compare the performance of our algorithm A-PARL against SOTA baselines Pebble lee2021pebble, PEBBLE+SURF park2022surf and Oracle (true reward) for Walker (DMSuite dm_suite), DoorOpen and ButtonPress (MetaWorld metaworld) w.r.t ground truth return (averaged over 5 seeds). It clearly demonstrates the superiority of our algorithm over existing baselines in terms of episodic return, where A-PARL achieves near-oracle performance in a much faster time. This highlights the importance of our bilevel framework which considers the dependence (missing from existing literature) of distribution on the lower-level policy parameter during training.
• Figure 3: A visualization of learned behavior for the baseline Pebble (top row) and proposed (in the bottom row) (with policy at Env-Step $0.5 \times 10^6$). We note that the proposed A-PARL algorithm has been able to learn the aligned behavior of opening the door in the generated trajectory (top-right) whereas PEBBLE gets stuck depicting our algorithm's efficiency in alignment .
• Figure 4: This figure describes the implementation flowchart of the iterative process of policy alignment in reinforcement learning. We start with some initial reward $r_0$, learn an optimal policy $\pi_0$ for that particular reward function at instant $t=0$, and utility evaluates the policy to generate an updated reward function $r_1$. Then at the next iterate $t=1$, we learn $\pi_1$ and so on.
• Figure 5: Figure demonstrates the visual representation of the environments considered in the experimental setup

Theorems & Definitions (15)

• Remark 1: Contrast with Standard Stochastic Bilevel Optimization
• Remark 2
• Remark 3: Gradient derivations for RLHF problem in Section \ref{['RLHF']}
• Lemma 1
• Lemma 2
• Theorem 1
• Lemma 3: Value function related upper bounds
• Lemma 4
• proof
• proof