UBP2: Uncertainty-Balanced Preference Planning for Efficient Preference-based Reinforcement Learning

Abstract: Preference-based RL provides an approach to learning reward models from pairwise comparisons of behaviors, bypassing the need for explicit reward design. However, existing methods typically rely on passive data collection and suffer from poor sample efficiency, especially during the early stages of learning. We introduce a model-based approach that actively directs exploration by jointly reasoning over uncertainties in the reward, dynamics, and value functions. Our method, Uncertainty-Balanced Preference Planning (UBP2), uses ensembles of reward, dynamics, and value function models to evaluate candidate trajectories according to a unified score that combines expected reward, terminal value, and epistemic uncertainty. Planning under this objective yields an explicit tradeoff between exploitation and information acquisition without requiring ad hoc exploration heuristics. Under standard regularity assumptions, we establish sublinear regret guarantees for both finite-horizon and infinite-horizon settings. Empirically, experiments on the Meta-World benchmark show UBP2 achieves substantially higher sample efficiency than model-free preference-based methods and non-optimistic model-based baselines.

• The paper introduces UBP2, a model-based RL algorithm that integrates uncertainty from reward, dynamics, and value networks to efficiently guide preference-based learning.
• It employs an ensemble-based world model with MPC and an uncertainty-aware acquisition function that balances exploitation and exploration, validated by strong empirical performance on MetaWorld tasks.
• Empirical evaluations show that UBP2 attains higher success rates and robustness under limited preference feedback, while providing theoretical sublinear regret guarantees.

Uncertainty-Balanced Preference Planning for Efficient Preference-based Reinforcement Learning

Introduction

UBP2 (Uncertainty-Balanced Preference Planning) introduces a model-based paradigm for preference-based reinforcement learning (PbRL) that leverages joint epistemic uncertainty across reward, dynamics, and value function estimators to guide exploration and accelerate feedback-efficient learning. Unlike contemporary PbRL methods, which are typically model-free and passively collect preference data, UBP2 structurally couples preference acquisition, reward learning, and control through an ensemble-based world model and optimistic planning objective. This coordination enables explicit tradeoff between exploitation and information seeking, yielding both practical performance gains and theoretically justified regret guarantees.

Method: UBP2 Algorithm and Objective

UBP2 operates via iterative model-based planning and uncertainty-calibrated preference querying. The agent comprises ensembles of reward, dynamics, and value networks. It performs planning using a Model Predictive Control (MPC) framework: at each step, action sequences are sampled and evaluated using a unified acquisition function that balances expected cumulative proxy reward, predicted terminal value, and epistemic uncertainty in both the reward and transition models. The acquisition function is

Figure 1: Illustration of planning in UBP2. The agent rolls out the dynamics model to generate trajectories and selects the first action of the most informative planned sequence based on an objective integrating return and uncertainty.

$$\mathbb E_{\mathbf{a}_t}\left[ \sum_{t=0}^{H-1} \gamma^t \left( \mu_n^r(\hat{s}_t, a_t) + \lambda_r \sigma^r_{\text{epi}}(\hat{s}_t, a_t) + \lambda_d \sigma^d_{\text{epi}}(\hat{s}_t, a_t) \right) + \gamma^H \left( \mu_n^q(\hat{s}_H, a_H) + \lambda_q \sigma^q_{\text{epi}}(\hat{s}_H, a_H) \right) \right]$$

where $\mu_n^{(\cdot)}$ are ensemble means and $\sigma_{\text{epi}}^{(\cdot)}$ credible epistemic uncertainties. The weights $\lambda_{r,d,q}$ are autotuned online via a Polyak-averaged target policy loss to optimize efficiency of uncertainty-driven exploration. Planning and reward model training are tightly linked by using optimistic query selection: preference pairs are selected globally to maximize not only the predicted preference likelihood but also epistemic uncertainty—prioritizing feedback most valuable for future planning.

The overall architecture and flow of UBP2 is summarized in Figure 2.

Figure 2: UBP2's control loop: ensembles model reward, dynamics, and value; MPC is used for uncertainty-aware trajectory selection; preference queries target segments with high reward and high model uncertainty; all components are optimized jointly using queried preferences.

Theoretical Guarantees and Regret Bounds

UBP2's planning rule can be interpreted as a form of dynamic Upper Confidence Bound (UCB) optimization, where optimism is integrated over both the reward and transition uncertainty. Under standard assumptions (bounded, smooth dynamics and rewards in RKHS, well-calibrated GP models, and bounded $Q$-function estimation error), UBP2 admits a sublinear regret bound in the number of episodes:

$$R_{\gamma,N} \leq \mathcal{O} \left( H^3 \sqrt{N} (\Gamma_{d,N\log N}^{3/2} + \Gamma_{r,N\log N}^{3/2}) + e_{q,N} \right)$$

where $\Gamma_{d,N}$ and $\Gamma_{r,N}$ are the maximum information gain for dynamics and reward kernels, and $e_{q,N}$ is the total $Q$-estimation error. The kernel-agnostic formulation allows immediate instantiation under linear, RBF, or Matérn priors. Unlike previous PbRL bounds, this result directly incorporates preference-induced reward ambiguity and the compounding effect of epistemic transition uncertainty, providing a joint exploration-theoretic guarantee for model-based and preference learning settings.

Empirical Evaluation

UBP2 was extensively evaluated on 10 MetaWorld manipulation tasks with proprioceptive observations, using standard Interquartile Mean (IQM) success rate metrics for sample efficiency. Baselines include state-of-the-art model-free methods (RUNE, MRN) and a model-based non-optimistic control (MBP—UBP2 with all optimism components disabled).

Results demonstrate that UBP2 achieves both earlier and higher final performance compared to all baselines on 9/10 tasks. On average, UBP2 yields the best aggregate rank and absolute IQM, with pronounced gains for feedback-efficient preference budgets (see Figure 3). Ablation studies confirm that each component of uncertainty (reward, dynamics, value) and optimistic preference querying is necessary for overall regret minimization and robustness; removal of any one deteriorates the worst-case regret.

Figure 3: UBP2 outperforms model-free PbRL (RUNE, MRN) and non-optimistic model-based (MBP), achieving faster success, higher plateau, and best average rank across diverse MetaWorld tasks.

Sensitivity studies illustrate:

• Optimistic preference selection outperforms both disagreement and entropy-based querying by a large margin, particularly on sparse-reward or high-variance tasks.
• Planning horizon has nontrivial effects: longer horizons yield higher reward in some tasks, but increase computational overhead and can degrade in unstable regimes.
• Effectiveness under constrained feedback: UBP2 remains competitive even with severely reduced preference feedback budgets, highlighting the benefits of uncertainty-focused data acquisition.

Figure 4: Sensitivity of UBP2 on MetaWorld Door Open task: optimism-based preference selection (left), planning horizon (center), and feedback budget (right) sweep. Entropy-based querying collapses performance.

Generalization to high-dimensional observation (visual) settings was also validated on DeepMind Control tasks with DinoV2 encodings, where UBP2 maintained or improved over MBP, though state-of-the-art model-free approaches were occasionally more robust under severe visual noise.

Figure 5: UBP2's IQM performance on visual DMControl Walker and Cheetah tasks; optimism mitigates visual-model uncertainty, outperforming non-optimistic MBRL.

Practical and Theoretical Implications

UBP2 provides a reference framework for optimal integration of reward- and dynamics-model uncertainty in preference-based online learning. The method's unified optimistic planning objective removes the need for heterogeneous or ad hoc exploration schemes; instead, exploration-exploitation is automatically tuned by ensemble epistemic statistics and informed querying. Theoretical sublinear regret is established under minimal additional assumptions, setting a new baseline for sample efficiency in online PbRL.

Practically, UBP2 enables interactive RL agents with high feedback-efficiency, robust to reward mis-specification and sparse preference cues—a critical desideratum for scalable real-world RL in safety-critical domains. The framework suggests that future scalable PbRL will be built on joint, uncertainty-prioritized preference acquisition, together with adaptive, model-based control stacks that tightly couple feedback and world modeling.

Future Directions

Immediate extensions include explicit characterization and incorporation of preference noise-induced reward estimation error into regret analysis, as in recent work on reward-agnostic PbRL. Another crucial line is the improvement of uncertainty calibration for deep neural ensembles, since theoretical results presently rely on GP well-calibration. Further, amortization and warm-start strategies for planning may reduce planner sub-optimality and contribute to stronger learning guarantees.

Conclusion

UBP2 advances state of the art in model-based preference-based RL by structurally integrating reward, dynamics, and value epistemic uncertainty into all aspects of exploration, planning, and preference data collection. The resulting algorithm achieves strong empirical performance and sample efficiency, and theoretical regret bounds justify its optimism-driven planning. This work highlights that tight coupling between uncertainty modeling and data acquisition is essential for robust, scalable RL with human-in-the-loop or proxy reward feedback.