UBP2: Uncertainty-Balanced Preference Planning Mohamed Nabail, Leo Cheng∗ , Jingmin Wang∗ , Nicholas Rhinehart
arXiv:2606.19328v1 [cs.LG] 17 Jun 2026
Learning, Embodied Autonomy, and Forecasting (LEAF) Lab, University of Toronto
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. Project Website: https://mohamedgnabail.github.io/ubp2/
at=1
Figure 1 Illustration of planning in UBP2. The agent rolls out the dynamics model to generate imagined trajectories. The planner selects the trajectory maximizing expected return and epistemic uncertainty by solving Equation (1) and then executes at=1 . ∗ There authors contributed equally. Author order determined by coin flip.
1
1
Introduction
In reinforcement learning, agents must learn effective behavior from interaction, often without prior knowledge of the dynamics of the environment or access to a densely specified reliable reward signal. This challenge is particularly acute in unsafe environments, where interactions are costly or risky, forcing agents to identify favorable actions and infer dynamics using as few samples as possible. Meanwhile, designing accurate reward functions for complex tasks is difficult and often leads to misspecified objectives that fail to capture the true intention [1, 30]. Preference-based reward learning addresses this limitation by replacing manual reward design with pairwise trajectory comparisons provided by an oracle, such as a human [6, 36]. Instead of relying on explicit numerical rewards, the agent learns a reward model from preferences, which can then guide decision-making [17]. However, effectively leveraging such learned rewards requires reasoning over hypothetical trajectories rather than relying solely on real environment interactions. Model-based reinforcement learning offers a natural solution by learning a predictive dynamics model and allowing planning over imagined trajectories, thereby improving sample efficiency, particularly in the early stages of learning [22, 23]. In this paper, we propose an approach that learns environment dynamics and a reward model from preference feedback, and uses this learned reward as the only supervision to train a value function and policy. A modelbased planner is used for efficient exploration and action selection under a limited preference feedback budget. After the budget is exhausted, control transitions to the learned policy. In the context of preference-based reward learning, purely model-free schemes passively collect trajectories (often via random or heuristic exploration) and only afterward fit the reward model [12, 17, 18]. In contrast, our model-based approach plans under the learned reward, dynamics, and value ensembles to actively select trajectories by maximizing a trajectory-level acquisition function that combines predicted cumulative return with epistemic uncertainty in the reward, dynamics, and value estimates. This objective jointly handles exploration and exploitation during learning. This tight coupling of planning and reward learning has been found to empirically accelerate convergence by focusing data collection on areas of interest. We propose Uncertainty-Balanced Preference Planning (UBP2), which uses ensembles for reward, dynamics, and value models and plans trajectories using a unified score that combines expected return and uncertainty. Executing informative trajectories allows UBP2 to efficiently learn the dynamics model and to learn the reward from preferences. Uncertainty-guided planning is used only when preference feedback is available; once the feedback budget is exhausted, the agent switches to fast policy execution under the learned reward. Our experiments show that uncertainty-guided planning in the preference-based setting achieves earlier success than both non-optimistic model-based and model-free baselines in most tasks. We establish a sublinear regret bound for UBP2 in the number of environment interactions, under standard regularity assumptions. Our main contributions are: • We introduce an uncertainty-guided planning objective that combines predicted cumulative return with epistemic uncertainty over reward, dynamics, and value estimates, enabling optimistic exploration that outperforms planning based on return alone or uncertainty in a single model component. • We introduce an optimistic preference label query strategy that prioritizes trajectory segment pairs with high predicted reward and high reward-model epistemic uncertainty. • Under standard smoothness/RKHS assumptions and well-calibrated GP models, we establish finite- and infinite-horizon regret bounds for UBP2 that are sublinear in the number of N episodes of environment interactions, with explicit dependence on the maximum information gain of the learned dynamics and reward models. • We empirically demonstrate on Meta-World manipulation [38] tasks that UBP2 operating on proprioceptive observations, achieves on average improved sample-efficiency compared to state of the art model-free [18] [19] and non-optimistic model-based preference baselines.
2
2
Related Work
We now compare our method to areas of closely-related prior work. UBP2 is a preference-based reinforcement learning approach that relies on uncertainty-driven exploration.
2.1
Preference-based Reinforcement Learning
Although comparisons provide an intuitive form of feedback [6], achieving high sample-efficiency in human feedback is a challenge [12]. Furthermore, collecting environment interactions for online learning can be impractical, particularly in high-risk environments. To address feedback sample-efficiency, prior work explored smart querying strategies, such as selecting the most informative queries through disagreement sampling after collecting an initial batch of seed interactions [17]. Few-Shot PbRL [12] uses a pretrained reward model on similar tasks for meta-learning the reward of the target task. [21] uses a pre-collected dataset of transitions to encode environment dynamics in reward learning. MoP-RL [20] handles learning efficiency by relying on pre-collected human demonstrations and a pretrained dynamics model to improve sample-efficiency before reward learning begins. In contrast, our method, to the best of our knowledge, is the first model-based PbRL approach to learn the reward function entirely from scratch with no precollected agent interactions, no pretrained reward models, and no offline pretrained dynamics model. This setup forces the agent to learn purely through online environment interactions and preference feedback, providing a more realistic test of feedback sample efficiency.
2.2
Epistemic uncertainty in model-based RL
Prior work on Model-based RL estimate multiple sources of uncertainty for different intentions. MOPO estimates epistemic uncertainty in learned dynamics and penalizes the reward proportionally to this uncertainty, to discourage rollouts that leave the dataset support [37]. RUNE, however, uses epistemic reward uncertainty optimistically as an intrinsic exploration bonus to improve feedback efficiency in PbRL [18]. UNISafe treats epistemic uncertainty of the dynamics model as a proxy for Out Of Distribution (OOD) risk, so uncertainty is used to synthesize a safety filter [28]. [9] uses epistemic uncertainty in the value function as a regularizer of test-time planning, so candidate action sequences that appear high return but uncertain are de-prioritized, reducing extrapolation error during online finetuning. In contrast to using uncertainty in a single component of the world model, UBP2 plans with a unified trajectory score that combines predicted reward, terminal value, and epistemic uncertainty of the reward, dynamics, and value models, explicitly balancing exploration and exploitation in PbRL in a single planning objective.
3
Preliminaries
We consider a discrete-time Markov decision process (MDP) defined by the tuple (S, A, p, r, γ, ρ0 ), where s ∈ S with dimension ds and a ∈ A with dimension da denote the continuous state and action spaces, p : S × A 7→ S is the unknown transition distribution that captures process noise, r : S × A 7→ R is a generic reward function, ρ0 is the initial state distribution and γ ∈ (0, 1) is the discount factor. The transition dynamics p∗ (s′ | s, a) is assumed to be unknown. A parametric dynamics model pϕ (s′ | s, a) is learned from data collected through online interaction with the environment and used during planning. During model-based planning, the Model Predictive Control (MPC) planner generates the control actions using the models learned after the nth episode of interaction with the real environment, n ∈ {1, ..., N }. This implicitly defines a policy through the solution of the planning objective, of which its first action is executed in the real environment; we denote this induced policy by πnind . In addition, we learn an explicit parametric policy πφlearn , trained using Equation (6). The objective is to learn a control policy π : S 7→ A that maximizes the expected discounted return over H timesteps under real dynamics p∗ and the true reward r∗ denoted by η, despite the absence of direct true
3
reward observations. This desired policy, denoted π ∗ , is: ∗
⋆
π = arg max η(π, p ) = arg max Eπ,p∗ π
3.1
π
"H−1 X
# t ∗
γ r (st , at ) .
t=0
Preference-Based Reward Learning
In the PbRL setting, the agent does not observe samples from the ground-truth reward function r∗ (s, a).Instead, learning is guided by preference-based feedback in the form of pairwise comparisons between trajectory segments, where segments are uniformly sampled over the entire trajectory, rather than restricted to any particular sub-trajectory. Because preference queries are available only on a limited budget, their acquisition must be carefully optimized through informative exploration. Given trajectory segments τi and τj each of length L, a preference oracle indicates which segment is preferred. We model preferences using a learned reward function rθ (s, a). Reward model parameters, θ, are optimized by maximizing the log-likelihood of observed preference labels using the Bradley-Terry model, in which the probability that τi is preferred over τj is P (τi ≻ τj ) =
3.2
exp
exp
PL−1
PL−1
rθ (sit ,ait ) . PL−1 j j +exp t=0 rθ (st ,at )
t=0
i i t=0 rθ (st ,at )
Modeling True Reward in PbRL
In PbRL, preference data cannot uniquely identify r∗ , as the oracle only reveals comparisons, i.e. information about which segment has larger relative utility, not the actual reward scale [36]. This creates an equivalence class of rewards - a set of functions that are indistinguishable from the perspective of the preference oracle because they induce the exact same preference outcomes. Rather than claiming that PbRL recovers the literal, unobserved r∗ , we interpret r∗ as a latent utility function that generates preferences and contend that the learning procedure produces some proxy reward, denoted r(s, a), within the same equivalence class as r∗ , and hence is the subject of our study instead. We list further conditions on this equivalence class in Section 5, Assumption 5.2, which makes our theoretical analysis feasible. For more discussion of modeling the true reward in PbRL, please refer to Appendix D.1. Unless otherwise noted, henceforth the “reward” we refer to in all subsequent methodology and theory denotes the proxy reward r. The notation r∗ and the phrase “true reward” are reserved specifically when referring to the unobserved true reward.
4
UBP2: Uncertainty-Balanced Preference Planning
We now present UBP2, (Algorithm 1), our approach for efficient optimistic exploration in model-based reinforcement learning (MBRL). UBP2 uses ensembles of reward, dynamics, and value function models to guide exploration using epistemic uncertainty, enabling preference-based learning that is both sample-efficient and informative. Uncertainty Representation: The uncertainty in the dynamics and value function models is estimated through ensemble disagreement, following previous work [16, 37], which reflects total uncertainty that encompasses both epistemic and aleatoric effects. For the reward model, we adopt the Jensen–Rényi divergence (JRD) [25] to estimate epistemic uncertainty, following [28]. JRD explicitly separates epistemic and aleatoric uncertainty and has been shown to outperform total-uncertainty and density-based estimates in downstream tasks [28]. In our setting, we apply JRD for measuring epistemic uncertainty only to the reward model, and use disagreement based total uncertainty for the dynamics and value function ensembles. We found doing so to be empirically sufficient for our results. Our approach is based on TD-MPC2 [11], using MPC as the primary control mechanism during the preference-based reward learning phase to allow efficient exploration under uncertainty. Unlike TD-MPC2 our approach does not assume access to scalar rewards throughout training, rather learns the reward from preference feedback and transition to a learned policy initialized from the value function once the preference budget is exhausted and model uncertainty decreases.
4
4.1
Uncertainty-Guided Optimistic Planning
At each decision step t in the real environment during the nth episode, UBP2 optimizes the parameters (µmpc , σmpc ) of a time-dependent Gaussian distribution over action sequences by approximately solving Equation (1).
arg max Eat
µmpc , σmpc
" H−1 X
γ
t
r d µrn (ŝt , at ) + λr σepi (ŝt , at ) + λd σepi (ŝt , at ) n n
+γ
H
q µqn (ŝH , aH ) + λq σepi (ŝH , aH ) n
# ,
t=0
(1) 2 learn where {at }H−1 (· | ŝH ). t=0 ∼ N (µmpc , σmpc ) and aH ∼ π
(·)
(·)
Here, (µmpc , σmpc ) parameterize the planner’s time-dependent action distribution, while µn and σepin denote ensemble means and epistemic uncertainties of the learned reward, dynamics, and value models in the nth episodes of actual environment interaction (hence, these means are fixed for each planning procedure). During planning, the trajectory (ŝ1 , . . . , ŝH ) is generated by rolling out the learned dynamics model conditioned on sampled action sequences pϕ (ŝ′ | ŝ, a). Only the first action a0 of the optimized sequence is executed in the real environment, after which the planning process is repeated at the next timestep. For long horizon reasoning beyond the planning horizon, we learn a value function qψ (s, a) using temporal-difference learning under the learned reward using rθ (s, a). The coefficients λr , λd , and λq control the relative contribution of the reward, the dynamics, and the value uncertainty, respectively, and are automatically tuned online. We use the “autotuning” approach from [33] to optimize each coefficient using the current policy π and a Polyak-averaged target policy π̄ to define the loss: h i (.) (.) L(λ) = Es∼B log(λ) σepi (s, π(s))−σepi (s, π̄(s)) . (2)
4.2
Optimistic Preference Selection
Prior preference-based RL methods use query strategies such as uniform sampling, disagreement or entropybased uncertainty, and coverage-based diversity [17]. While these approaches encourage exploration, they are largely reward-agnostic and select queries from locally sampled candidate batches. In contrast, our labeling approach (Algorithm 2) differs in two ways: 1) Optimistic scoring: we combine predicted preference likelihood with reward uncertainty to prioritize informative, high-value queries that improve downstream planning performance. 2) Global selection: instead of sampling from local batches, we rank candidate trajectory pairs globally across the replay buffer, reducing the chance of missing highly informative comparisons.
4.3
UBP2 World Model Objective
Model objective. UBP2 is based on the TD-MPC2 world-model objective, with the following key differences: (i) rewards are learned from preference feedback, (ii) TD targets are computed using predicted rewards, and (iii) reward and dynamics models are represented as ensembles. Given a horizon segment H and a preference pair segment L, we optimize a loss per-member that combines consistency, reward, and value losses, (m)
(m)
L(m) = αc L(m) cons + αr Lpref + αv Lval . The ensemble members m of the dynamics and rewards models are independently trained as in [28]. To (m) optimize the dynamics model, we compute a rollout {ŝt } and minimize the consistency loss, 1 L(m) cons = H
H−1 X
(m)
∥ŝt+1 − st+1 ∥22 ,
t=0
5
(3)
Algorithm 1 UBP2 1: Input: Replay buffer B, preference buffer P 2: Params: model ensembles (pϕ , rθ , qψ ), policy
πφlearn 3: Hyperparameters: horizon H, total steps after N episodes N T , pref period Fpref , pref budget Npref 4: Initialize t ← 0, npref ← 0,s ← env.reset 5: while t ≤ N T do 6: if (npref < Npref ) then 7: a ← MPC(s; λr , λd , λq ) 8: else 9: a ← πφlearn (s) 10: end if 11: Step env: (s′ , r) ← env.step(a); 12: B ← B ∪ (s, a) ; s ← s′ , t ← t + 1 13: Update pϕ , rθ , qψ using Equations (3) to (5) 14: Update πφlearn using Equation (6) 15: Update {λr , λd , λq } using Equation (2) 16: if episode done (t==T) then 17: if Fpref and (npref < Npref ) then 18: P ← P ∪ (τ1 , τ2 , y) using Algorithm 2; 19: update npref 20: end if Figure 2 UBP2 uses ensembles of reward, dynamics, and 21: s ← env.reset() value models to perform optimistic planning over imagined 22: end if trajectories by combining predicted returns with epistemic 23: end while uncertainty. Preferences are then selected optimistically from uniformly sampled segments of executed trajectories. Both trajectories and preferences are used to jointly train the dynamics, reward, and value models.
For the reward model, we minimize the preference loss, h i (m) Lpref = −E(τi ,τj ,y)∼P y log Pθ(m) (τi ≻ τj )+(1−y) log 1−Pθ(m) (τi ≻ τj ) ,
(4)
where y is the true preference label provided by a preference oracle, which in our experiments is simulated using the true reward from the environment. The value function ensemble is trained using temporal-difference regression. For each member m, the TD targets tdt are computed using the predicted reward of the same member, (m)
tdt
(m)
= rθ
(st , at ) + γ q̄min (st+1 , π learn (st+1 )),
where q̄min denotes the target value function. The value loss is: (m)
Lval =
EQ H−1 i 1 X t X h (j) (m) 2 γ qψ (ŝt , at ) − tdt . HEQ t=0 j=1
(5)
We adopt the stochastic maximum-entropy policy objective from TD-MPC2. The policy πφ (a | s) is optimized to maximize value while encouraging exploration through entropy regularization over imagined rollouts. The policy objective is given in Equation (6). H−1 1 X t q L(φ) = − γ µ (ŝt , πφ (ŝt ))+αH H(πφ ) , (6) H t=0 where H(πφ ) denotes the policy entropy. Gradients of Lπ are taken with respect to the policy parameters φ only, with value estimates fixed during policy optimization.
6
5
Theoretical Results
First, we make standard continuity common in literature [7, 32, 35]. Assumption 5.1 (Continuous dynamics, continuous and bounded reward). The true dynamics p∗ , all policies π ∈ Π, and true reward r∗ are continuous. Furthermore, r∗ is bounded, i.e. r∗ : S × A 7→ [0, Rmax ], and the process noise, captured by p∗ (s′ |s, a), is i.i.d. Gaussian with variance σ 2 . After the nth episode of environment interaction, we learn a model of the true dynamics, with a mean µdn (s, a) d and uncertainty estimate σepi (s, a), and also a model of the proxy reward with a mean µrn (s, a) and uncertainty n r σepin (s, a). Like previous work, we model each dynamics coordinate p∗j , j ∈ {1, ..., ds } using a Gaussian process (GP) [34]. The scalar reward r is treated analogously. For clarity, we limit our discussion on the setting with Gaussian noise, though this is not a requirement for the final regret bound. Please refer to Appendix D.4 for settings with subgaussian process noise. Assumption 5.2 (RKHS regularity). The ground truth dynamics p∗ and the proxy reward r reside in their own respective Reproducing Kernel Hilbert Spaces (RKHS). Each dynamics coordinate p∗j lies in an RKHS Hkd induced by kernel kd with bounded norm ∥p∗j ∥Hkd ≤ Bd , and kd is uniformly bounded along the diagonal, i.e. d sup(s,a)∈S×A kd (s, a), (s, a) ≤ σmax . Furthermore, we assume r ∈ Hkr with bounded norm ∥r∥Hkr ≤ Br and r with kr uniformly bounded along the diagonal, i.e. sup(s,a)∈S×A kr (s, a), (s, a) ≤ σmax . Assumption 5.2, common in the Bayesian optimization and RL literature [7, 31, 34], allows the modeling and learning of the true dynamics and reward using GPs. This admits a closed form for the posterior mean and variance of the learned models, presented in Appendix D.1. Given dynamics p we define the mean action-value function Q̄ after the nth episode of environment interaction to be "∞ # X π t r Q̄p,n (s, a) := Eπ,p γ µn (st , at ) s0 = s, a0 = a , t=0
and denote our estimated action-value function as Q̂πn . Following a similar approach as Theorem 1 in [29], we have: Assumption 5.3 (Q-function sub-optimality). At the nth episode of environment interaction, the estimation error of Q is bounded by a constant ϵq,n , i.e., ∥Q̂πn − Q̄πp,n ∥∞ ≤ ϵq,n .
(7)
To quantify the epistemic uncertainty of Q, we have the following assumption: Assumption 5.4 (Epistemic Uncertainty of Q). For every episode n, the epistemic uncertainty σnq of the action–value function Q̄πn is proportional to the discounted cumulative epistemic uncertainty arising from the dynamics and reward models, i.e., "∞ # "∞ # X X q t d t r σn ∝ Eπ,p γ ∥σn (st , at )∥ + Eπ,p γ σn (st , at ) . t=0
t=0
This assumption reflects that uncertainty in the action–value function accumulates from epistemic uncertainty in both the dynamics and reward models along trajectories induced by policy π. A detailed theoretical justification and discussion are provided in Appendix D.3. For the purposes of theoretical analysis, we provide the following definition of the UBP planner objective that (·) is optimized by the planner at each planning procedure using fixed learned dynamics and reward models µn , (·) σepi : n
7
Definition 5.5 (Optimistic planning objective). Let π denote a candidate policy (i.e. the policy at the current planning iteration) and consider rollouts using the mean dynamics µdn with the same process noise as the true environment. We define the UCB planning objective over planner horizon H: " H−1 # X UCB ηγ,n (π, µdn ) := Eπ,µdn γ t µrn (ŝt , at ) + an σnr (ŝt , at ) + bn σnd (ŝt , at ) + γ H µqn (ŝt , at ) + cn σnq (ŝt , at ) , (8) t=0
where an , bn , cn are parameters that control the strength of reward, dynamics exploration, with their theoretical bounds provided in Lemma 5.7 and justified in Appendix D.3. They correspond to hyperparameters λr , λd , λq in Equation (1). To derive a bound, we define the cumulative infinite-horizon regret for a particular policy πn as Rγ,N =
N X n=1
rγ,n =
N X
ηγ (π ∗ , p∗ ) − ηγ (πn , p∗ ),
(9)
n=1
where
ηγ (π, p) = Eπ,p
"∞ X
# γ t r(st , at ) .
(10)
t=0
We make the following assumption on the planner: Assumption 5.6 (No planner sub-optimality). The policy πnind induced by the planning procedure through optimizing (8) is equal to the policy πn∗ that maximizes (8). That is, πnind = πn∗ . We refer to the policy as being induced by the planner because the planner does not train a parameterized policy. Rather, it repeatedly solves (8) and executes only the first action. Note that this is not the learned parametric policy πφlearn . For theoretical results that consider planner suboptimality, please refer to Appendix D.5. We now present the following lemma, which states: Lemma 5.7 (Optimism of the UCB planning objective). Let Assumptions 5.1, 5.2, 5.3, and 5.4 hold,pand let the learned GP dynamics and reward models be well-calibrated. Then there exist coefficients an ∈ O Γr,N , p p p bn ∈ O Γd,N and cn ∈ O max{ Γr,N , Γd,N } such that for all episodes n and all policies π, with probability at least 1 − δ: UCB ηγ (π, p∗ ) ≤ ηγ,n (π, µdn ) + γ H ϵq,n . (11) Γd,N is the state and action dimension-dependent maximum information gain [31, 34], defined in Appendix D.1, for a function in the RKHS induced by kd over N episodes of environment interaction and similarly for Γr,N . Lemma 5.7 states that the planning objective upper-bounds the true return uniformly over policies up to some Q estimation error, i.e., the agent plans using an optimistic estimate of the environment. We first establish an analogous result for a reduced objective in Appendix D.2. The proof of Lemma 5.7 builds on this reduced case and is provided in Appendix D.3. We now present the following theorem, which bounds the regret of our method: Theorem 5.8 (Infinite-horizon regret bound with dynamics, reward and value uncertainty bonuses). Let Assumptions 5.1, 5.2, 5.3, 5.4 and 5.6 hold, and let the learned GP dynamics and reward models be well-calibrated. The cumulative infinite-horizon regret after N episodes of real environment interaction satisfies, with probability 1 − δ: √ PN 3/2 3/2 Rγ,N ≤ O H 3 N Γd,N log N +Γr,N log N +eq,N , (12) where eq,N := n=1 ϵq,n . If we additionally assume that the growth rate of eq,N is sublinear, then Theorem 5.8 yields a sublinear regret rate governed by the maximum information gain of the dynamics and reward kernels, and parallels the infinitehorizon GP regret guarantees established for past work [34]. Our additional dependence on the decomposition into (Γd , Γr ) arise from separately controlling (1) the simulation error due to transition uncertainty and (2) reward estimation uncertainty across the planning horizon. As in previous work, we first establish an analogous 8
1.0 1e6
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
1.0 1e6
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
Episode Success (%)
Episode Success (%) 0.0
0.0
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
100 80 60 40 20 0
0.0
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
(c) Handle Press (1000)
(d) Coffee Button (1000)
(e) Faucet Open (2000)
100 80 60 40 20 0
100 80 60 40 20 0
100 80 60 40 20 0
100 80 60 40 20 0
0.0
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
(g) Door Unlock (2500)
0.0
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
MBP
0.0
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
(i) Drawer Open (5000)
(h) Sweep Into (5000)
UBP2
Episode Success (%)
Episode Success (%)
Episode Success (%) 0.0
(f) Door Open (2000)
0.2 0.4 0.6 0.8 Environment Steps
100 80 60 40 20 0
(b) Window Close (500)
(a) Door Close (500) 100 80 60 40 20 0
0.0
100 80 60 40 20 0
Episode Success (%)
0.2 0.4 0.6 0.8 Environment Steps
Episode Success (%)
0.0
100 80 60 40 20 0
Episode Success (%)
Episode Success (%)
Episode Success (%)
100 80 60 40 20 0
MRN
0.0
0.2 0.4 0.6 0.8 Environment Steps
1.0 1e6
(j) Hammer (10000)
Rune
Figure 3 Performance comparison across MetaWorld tasks. Curves show IQM performance with shaded standard error. Numbers between brackets indicate the maximum preference feedback budget.
sublinear result for a reduced finite-horizon objective in Appendix D.2 and then prove Theorem 5.8 in Appendix D.3. Note that our regret bound as written is implicit in the state and action dimension ds and da ; it depends on the information gain in order to present it in a kernel-agnostic manner. Please refer to Appendix D.3 for the explict regret bound for different kernels.
6
Experiments
We design our experiments to answer the following: Q1: Does uncertainty-guided optimistic planning improve sample efficiency in preference-based reinforcement learning compared to model-free and non-optimistic modelbased baselines? Q2: How do uncertainty-aware planning components (dynamics, reward, and value), together with optimistic preference selection, individually affect success-rate performance? Q3: Can UBP2 benefit from varying the planning horizon? Q4: How does UBP2 perform under constrained preference feedback budgets? Overall, we find that our approach is substantially more sample-efficient than strong model free approaches, and that all components of uncertainty are generally helpful to boosting sample efficiency and robustness. Table 1 Performance comparison on MetaWorld tasks over 1M environment steps. We report IQM success averaged over seeds and time, with standard error at the final evaluation step, along with average performance and average rank across tasks. Best results are shown in bold. Asterisks denote statistically significant improvements of UBP2 over the corresponding baseline under a Wilcoxon signed-rank test (p < 0.05), reported only when sufficient seeds are available. Method
Door Close
Window Close
Handle Press
Coffee Button
Faucet Open
Door Open
Door Unlock
Sweep Into
Drawer Open
Hammer
Avg ↑
Rank ↓
UBP2
97.8 (1.0)
93.1 (1.0)
93.4 (0.0)
86.2 (3.7)
89.5 (7.9)
70.7 (10.4)
85.8 (2.1)
36.4 (12.3)
77.7 (4.9)
29.9 (11.5)
76.1
1.20
MBP Rune [18] MRN [19]
95.6 (0.0)* 91.5 (10.0)* 94.2 (0.0)*
74.7 (13.5)* 85.2 (9.2)* 84.9 (10.2)*
90.8 (4.0)* 84.5 (8.1)* 81.5 (10.3)*
75.3 (10.7) 90.4 (1.0) 87.7 (11.0)
83.6 (5.0)* 83.1 (1.4)* 80.1 (0.0)*
61.5 (11.7) 41.7 (13.1)* 48.2 (11.4)
85.2 (10.0) 56.1 (13.1)* 65.6 (9.8)*
1.1 (2.0)* 35.5 (9.3) 28.3 (14.5)
43.7 (12.8)* 35.9 (12.8)* 38.1 (14.9)*
14.8 (9.1) 4.6 (0.9)* 7.5 (7.2)*
62.6 60.8 61.6
2.60 3.10 3.10
Experimental Setup. We evaluate UBP2 on 10 manipulation tasks of varying degrees of complexity from the MetaWorld benchmark [38] using proprioceptive observations and task-specific success metrics [24]. For our main results, we use feedback budgets commonly reported in the literature [17, 18, 19]. We compare UBP2 with the following baselines: (1) RUNE [18]: A model-free, preference-based reinforcement learning method that enhances exploration by augmenting the extrinsic reward with an intrinsic bonus based on disagreement among an ensemble of learned reward models , (2) MRN (Meta-Reward-Net) [19]: A model-free, preference-based reinforcement learning approach that employs bi-level optimization to learn the reward function based on the performance of the Q-function. (3) MBP (Model-Based PbRL): A non-optimistic variant of UBP2 that removes uncertainty-aware planning and optimistic preference selection. This baseline isolates the effect of model-based learning without optimism or uncertainty-driven exploration. This variant is conceptually similar to MoP-RL [20]. Evaluation Metric: We report Interquartile Mean (IQM) success rate over environment steps where 9
Table 2 Ablation study on MetaWorld tasks over 1M environment steps. Config indicates whether Reward, Dynamics, Value Uncertainty, and Optimistic Preference Selection are enabled (1) or disabled (0). We report IQM success averaged over seeds and time, with standard error at the final evaluation step. We also report overall mean performance and worst-case regret (maximum per-task drop relative to the best variant). Best results are in bold. Approach
Config
Door Close
Window Close
Handle Press
Coffee Button
Faucet Open
Door Open
Door Unlock
Sweep Into
Drawer Open
Hammer
Avg ↑
Regret ↓
UBP2
1111
97.8 (1.0)
93.1 (1.0)
93.4 (0.0)
86.2 (3.7)
89.5 (7.9)
70.7 (10.4)
85.8 (2.1)
36.4 (12.3)
77.7 (4.9)
29.9 (11.5)
76.1
13.8
0111 1001 1011 0101 1101 0011 1110
96.7 (2.0) 94.8 (0.0) 95.9 (0.0) 97.5 (0.0) 96.4 (20.0) 97.3 (8.0) 96.0 (0.0)
94.4 (0.0) 93.1 (0.0) 94.2 (0.0) 93.6 (0.0) 94.2 (0.0) 94.0 (0.0) 89.8 (0.0)
91.6 (0.0) 91.2 (0.0) 92.0 (0.0) 91.0 (20.0) 92.1 (20.0) 93.7 (0.0) 88.8 (0.0)
90.4 (9.7) 82.8 (11.7) 84.2 (8.7) 76.8 (9.7) 90.1 (3.2) 89.8 (8.0) 79.5 (12.4)
82.8 (4.0) 90.3 (0.0) 89.7 (2.4) 85.6 (20.0) 84.7 (6.0) 80.0 (23.3) 87.7 (0.0)
84.5 (2.4) 67.7 (14.0) 58.5 (24.5) 82.4 (16.0) 78.8 (0.0) 66.9 (2.0) 68.9 (8.0)
85.0 (0.0) 78.4 (0.0) 83.3 (0.0) 81.1 (0.0) 88.6 (0.0) 83.6 (4.0) 76.0 (6.3)
38.9 (16.4) 20.3 (13.2) 29.3 (13.2) 31.4 (22.3) 4.0 (15.6) 15.0 (15.7) 11.2 (9.7)
68.3 (4.0) 44.0 (19.6) 58.4 (19.6) 53.9 (18.9) 76.0 (4.0) 71.8 (2.0) 68.2 (16.0)
22.4 (19.7) 14.1 (18.1) 12.4 (10.7) 15.1 (19.4) 38.2 (22.6) 33.3 (19.7) 27.4 (19.6)
75.5 67.7 69.8 70.8 74.3 72.5 69.3
15.8 33.7 26.0 23.8 34.9 23.9 27.7
0000
95.6 (0.0)
74.7 (13.5)
90.8 (4.0)
75.3 (10.7)
83.6 (5.0)
61.5 (11.7)
85.2 (10.0)
1.1 (2.0)
43.7 (12.8)
14.8 (9.1)
62.6
37.8
MBP
faster gains indicate higher sample efficiency in preference-based learning. We report results using the IQM in our plots and tables, as a robust aggregate metric for reinforcement learning evaluation [2]. IQM reduces the influence of outlier runs and provides a more stable estimate compared to the mean, especially in the low-seed regime typical of computationally intensive RL experiments. In our experiments we run 5-15 seeds per task based on task complexity. Sample Efficiency in PbRL: Figure 3 shows that UBP2 reaches high success rates earlier in 9 out of the 10 tasks, demonstrating the benefit of optimistic planning in PbRL. Table 1 further supports these findings by reporting IQM success averaged across evaluation steps over the course of training. UBP2 achieves both the highest overall average performance and the best average rank across tasks. Component Ablations: We evaluate the contribution of each component in UBP2 through 5-seed ablation studies against the full model and MBP, where all optimistic components are disabled. The results in Table 2 show that UBP2 provides the best overall balance between performance and robustness. Although some ablations slightly improve individual tasks, these gains are inconsistent and often incur much higher regret elsewhere. Since regret reflects worst-case degradation across tasks of varying difficulty, these results highlight the importance of each optimistic component to achieve strong and consistent performance. UBP2 also achieves the highest average overall performance, indicating improved sample efficiency and robustness across tasks Optimistic Preference Selection We compare UBP2 against two commonly used preference-query heuristics discussed in prior work: disagreement-based and entropy-based selection. Figure 4(a) shows that optimistic preference selection consistently outperforms both alternatives in challenging manipulation tasks Varying Horizon We evaluate the effect of the planning horizon in Figure 4(b), a key degree of freedom in model-based planning enabled by UBP2. Unlike model free methods, UBP2 allows adapting the planning horizon to better match task structure. Because the optimal horizon varies across tasks, we tune it separately for each environment. Across the three representative tasks shown, increasing the horizon from 7 to 11 consistently improves performance. Although even longer horizons may provide additional gains, they come at a significantly higher computational cost during planning. Refer to Section C.3 for more experiments. Reduced Feedback We evaluate UBP2 under constrained preference feedback budget, 2 tasks receive a limited feedback budget. The results in Figure 4(c) indicate that UBP2 retains its effectiveness even when the available feedback is significantly reduced.
7
Conclusion
This work presents Uncertainty-Balanced Preference Planning (UBP2), a model-based preference-based reinforcement learning approach that actively guides exploration by optimistically planning over uncertainty in the reward, dynamics, and value functions. By planning under a unified optimistic objective that balances expected return with epistemic uncertainty, UBP2 explicitly balances exploitation and information acquisition, avoiding uninformative data collection.
10
0
20
40
75.2 (19.6) 77.7 (4.9)
60
Success (%)
80
100
(a) Preference Selection Method
0
11
0.2 FB Budget
7
77.2 (6.4)
70.7 (10.4) 63.1 (14.2)
73.1 (13.6) 85.8 (2.1) 82.5 (5.0) 92.8 (2.0) 86.2 (3.7) 77.5 (10.4) 20
40 60 Success (%)
(b) Planning Horizon
80
100
Door Close
42.0 (16.6)
70.7 (10.4)
2000 2000 2000 2500 2500 2500 1000 1000 1000
100 250 500
Drawer Open
45.5 (23.2)
15
Optimistic Door Open
5000 5000 5000
Disagreement
Coffee Door Button Unlock
Door Open Drawer Open
Entropy
2000 2000 2000
1000 2500 5000 0
0.5 FB Budget
Main FB Budget
92.5 (19.6) 97.2 (0.0) 97.8 (1.0) 40.1 (21.5)
20
40 60 Success (%)
78.7 (7.5) 77.7 (4.9) 80
100
(c) Reduced Feedback
Figure 4 Sensitivity analysis of UBP2 across three key hyperparameters on MetaWorld tasks. (a) Preference selection strategy. (b) Planning horizon. (c) Feedback budget. IQM success averaged over seeds and evaluation steps is reported in white, with standard error at the final evaluation step in brackets; feedback budget is reported in black on the left of each bar. In (a), entropy-based preference selection achieves the lowest average IQM of 0.1% (20) on Door Open.
8
Limitations
A limitation of our work is that the role of preference learning in the theoretical guarantees could be made more explicit. There exist past work that tightly characterizes preference-learning error via more stringent assumptions on the reward function [39]. Furthermore, the modeling of the dynamics and reward as wellcalibrated GPs in the theory may not coincide with practical implementation (i.e. deep ensembles). For more discussion regarding both these points, please refer to Section D.1.
9
Acknowledgment
This research was supported by an NSERC Discovery Grant, the NVIDIA Academic Grant Program, and computing resources from the Digital Research Alliance of Canada (alliancecan.ca).
References [1] David Abel, Will Dabney, Anna Harutyunyan, Mark K. Ho, Michael L. Littman, Doina Precup, and Satinder Singh. On the expressivity of markov reward, 2022. URL https://arxiv.org/abs/2111.00876. [2] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice, 2022. URL https://arxiv.org/abs/2108.13264. [3] Arthur Argenson and Gabriel Dulac-Arnold. Model-based offline planning, 2021. URL https://arxiv.org/abs/2008. 05556. [4] Arunkumar Byravan, Leonard Hasenclever, Piotr Trochim, Mehdi Mirza, Alessandro Davide Ialongo, Yuval Tassa, Jost Tobias Springenberg, Abbas Abdolmaleki, Nicolas Heess, Josh Merel, and Martin Riedmiller. Evaluating modelbased planning and planner amortization for continuous control, 2021. URL https://arxiv.org/abs/2110.03363. [5] Sayak Ray Chowdhury and Aditya Gopalan. On kernelized multi-armed bandits, 2017. URL https://arxiv.org/abs/ 1704.00445. [6] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017. [7] Sebastian Curi, Felix Berkenkamp, and Andreas Krause. Efficient model-based reinforcement learning through optimistic policy search and planning, 2020. URL https://arxiv.org/abs/2006.08684. [8] Pieter-Tjerk de Boer, Dirk P. Kroese, Shie Mannor, and Reuven Y. Rubinstein. A tutorial on the cross-entropy method. Annals of Operations Research, 2005. URL https://people.smp.uq.edu.au/DirkKroese/ps/aortut.pdf. [9] Yunhai Feng, Nicklas Hansen, Ziyan Xiong, Chandramouli Rajagopalan, and Xiaolong Wang. Finetuning offline world models in the real world, 2023. URL https://arxiv.org/abs/2310.16029.
11
[10] Saeed Ghadimi and Guanghui Lan. Stochastic first- and zeroth-order methods for nonconvex stochastic programming, 2013. URL https://arxiv.org/abs/1309.5549. [11] Nicklas Hansen, Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for continuous control. In The Twelfth International Conference on Learning Representations, 2024. URL https://openreview.net/forum?id= Oxh5CstDJU. [12] Joey Hejna and Dorsa Sadigh. Few-shot preference learning for human-in-the-loop rl, 2022. URL https://arxiv.org/ abs/2212.03363. [13] David Janz, David R. Burt, and Javier González. Bandit optimisation of functions in the matérn kernel rkhs, 2023. URL https://arxiv.org/abs/2001.10396. [14] Jihwan Jeong, Xiaoyu Wang, Jingmin Wang, Scott Sanner, and Pascal Poupart. Reflect-then-plan: Offline model-based planning through a doubly bayesian lens, 2025. URL https://arxiv.org/abs/2506.06261. [15] Sham Kakade, Akshay Krishnamurthy, Kendall Lowrey, Motoya Ohnishi, and Wen Sun. Information theoretic regret bounds for online nonlinear control, 2020. URL https://arxiv.org/abs/2006.12466. [16] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel : Model-based offline reinforcement learning, 2021. URL https://arxiv.org/abs/2005.05951. [17] Kimin Lee, Laura Smith, Anca Dragan, and Pieter Abbeel. B-pref: Benchmarking preference-based reinforcement learning. In Thirty-fifth Conference on Neural Information Processing Systems Datasets and Benchmarks Track (Round 1), 2021. URL https://openreview.net/forum?id=ps95-mkHF_. [18] Xinran Liang, Katherine Shu, Kimin Lee, and Pieter Abbeel. Reward uncertainty for exploration in preference-based reinforcement learning, 2022. URL https://arxiv.org/abs/2205.12401. [19] Runze Liu, Fengshuo Bai, Yali Du, and Yaodong Yang. Meta-reward-net: Implicitly differentiable reward learning for preference-based reinforcement learning. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL https://openreview.net/forum?id= OZKBReUF-wX. [20] Yi Liu, Gaurav Datta, Ellen Novoseller, and Daniel S Brown. Efficient preference-based reinforcement learning using learned dynamics models. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pages 2921–2928. IEEE, 2023. [21] Katherine Metcalf, Miguel Sarabia, and Barry-John Theobald. Rewards encoding environment dynamics improves preference-based reinforcement learning. arXiv preprint arXiv:2211.06527, 2022. [22] Thomas M. Moerland, Joost Broekens, Aske Plaat, and Catholijn M. Jonker. Model-based reinforcement learning: A survey, 2022. URL https://arxiv.org/abs/2006.16712. [23] Anusha Nagabandi, Gregory Kahn, Ronald S. Fearing, and Sergey Levine. Neural network dynamics for model-based deep reinforcement learning with model-free fine-tuning, 2017. URL https://arxiv.org/abs/1708.02596. [24] Maxime Oquab, Timothée Darcet, Théo Moutakanni, Huy Vo, Marc Szafraniec, Vasil Khalidov, Pierre Fernandez, Daniel Haziza, Francisco Massa, Alaaeldin El-Nouby, Mahmoud Assran, Nicolas Ballas, Wojciech Galuba, Russell Howes, Po-Yao Huang, Shang-Wen Li, Ishan Misra, Michael Rabbat, Vasu Sharma, Gabriel Synnaeve, Hu Xu, Hervé Jegou, Julien Mairal, Patrick Labatut, Armand Joulin, and Piotr Bojanowski. Dinov2: Learning robust visual features without supervision, 2024. URL https://arxiv.org/abs/2304.07193. [25] A. R’enyi. On measures of entropy and information. In Proceedings of the Fourth Berkeley Symposium on Mathematical Statistics and Probability, Volume 1: Contributions to the Theory of Statistics, volume 4, pages 547–562. University of California Press, 1961. [26] Jonas Rothfuss, Bhavya Sukhija, Tobias Birchler, Parnian Kassraie, and Andreas Krause. Hallucinated adversarial control for conservative offline policy evaluation, 2023. URL https://arxiv.org/abs/2303.01076. [27] Rajiv Sambharya, Georgina Hall, Brandon Amos, and Bartolomeo Stellato. End-to-end learning to warm-start for real-time quadratic optimization, 2022. URL https://arxiv.org/abs/2212.08260. [28] Junwon Seo, Kensuke Nakamura, and Andrea Bajcsy. Uncertainty-aware latent safety filters for avoiding out-ofdistribution failures, 2025. URL https://arxiv.org/abs/2505.00779. [29] Harshit Sikchi, Wenxuan Zhou, and David Held. Learning off-policy with online planning, 2021. URL https: //arxiv.org/abs/2008.10066.
12
[30] Satinder Singh, Richard L. Lewis, and Andrew G. Barto. Where do rewards come from? In Proceedings of the Annual Conference of the Cognitive Science Society, pages 2601–2606. Cognitive Science Society, 2009. [31] Niranjan Srinivas, Andreas Krause, Sham M. Kakade, and Matthias W. Seeger. Information-theoretic regret bounds for gaussian process optimization in the bandit setting. IEEE Transactions on Information Theory, 58(5):3250–3265, 2012. doi: 10.1109/TIT.2011.2182033. [32] Bhavya Sukhija, Lenart Treven, Cansu Sancaktar, Sebastian Blaes, Stelian Coros, and Andreas Krause. Optimistic active exploration of dynamical systems, 2023. URL https://arxiv.org/abs/2306.12371. [33] Bhavya Sukhija, Stelian Coros, Andreas Krause, Pieter Abbeel, and Carmelo Sferrazza. Maxinforl: Boosting exploration in reinforcement learning through information gain maximization, 2025. URL https://arxiv.org/abs/2412.12098. [34] Bhavya Sukhija, Lenart Treven, Carmelo Sferrazza, Florian Dörfler, Pieter Abbeel, and Andreas Krause. Sombrl: Scalable and optimistic model-based rl, 2025. URL https://arxiv.org/abs/2511.20066. [35] Scott Sussex, Anastasiia Makarova, and Andreas Krause. Model-based causal bayesian optimization, 2023. URL https://arxiv.org/abs/2211.10257. [36] Christian Wirth, Riad Akrour, Gerhard Neumann, and Johannes Fürnkranz. A survey of preference-based reinforcement learning methods. Journal of Machine Learning Research, 18(136):1–46, 2017. [37] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Model-based offline policy optimization, 2020. URL https://arxiv.org/abs/2005.13239. [38] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Avnish Narayan, Hayden Shively, Adithya Bellathur, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning, 2021. URL https://arxiv.org/abs/1910.10897. [39] Wenhao Zhan, Masatoshi Uehara, Wen Sun, and Jason D. Lee. Provable reward-agnostic preference-based reinforcement learning, 2024. URL https://arxiv.org/abs/2305.18505.
13
A
Algorithms
Algorithm 2 Optimistic Preference Pair Selection (m)
Input: replay buffer B, reward ensemble {rθ }E m=1 , pref. buffer P Hyperparameters: segment length L, pairs to add K Candidate set C ← ∅ for each batch do Sample candidate segment pairs (τ1 , τ2 ) (length L) from B for each (τ1 , τ2 ) do L−1 P (m) i i (m) Ri ← rθ (st , at ), i ∈ {1, 2}, m ∈ [E] t=0
(m)
(m)
p(m) ← softmax([R1 , R2 ])1 score ← µ(p(1:E) ) + σ(p(1:E) ) C ← C ∪ {(τ1 , τ2 , score)} end for end for Add top-K pairs in C by score to P UBP2 selects the most informative preference pairs for labeling by scoring each candidate segment pair using both predicted reward and reward-model epistemic uncertainty, then globally ranking and selecting the top-K pairs to add to the preference buffer P.
14
B
Implementation Details
B.1
Hyperparameters
UBP2 learns a reward model from preference feedback using short trajectory segments. Preferences are periodically collected and the reward model is updated in batches. Table 3 Preference Learning Hyperparameters Parameter number of seeds pref_segment_length reward_train_horizon num_pref_per_addition max_pref_feedback pref_update_freq reward_batch_size
Door Close 10 10 10 12 500 2500 50
Window Close 10 10 10 12 500 2500 50
Handle Press 10 10 10 12 1000 2500 50
Coffee Button 10 10 10 25 1000 2500 50
Faucet Open 15 10 10 12 2000 2500 50
Door Open 15 10 10 12 2000 2500 50
Door Unlock 10 10 10 12 2500 2500 50
Sweep Into 10 10 10 25 5000 2500 100
Drawer Open 10 10 10 25 5000 2500 100
Hammer 15 10 10 50 10000 2500 200
We use Model Predictive Control (MPC) with the Cross-Entropy Method (CEM) for planning. At each timestep, candidate trajectories are sampled and iteratively refined using elite selection Table 4 MPC (CEM) Planning Hyperparameters Parameter iterations number of samples num of elites num of policy trajectories horizon
B.2
Door Close 6 512 64 24 7
Window Close 6 512 64 24 7
Handle Press 6 512 64 24 7
Coffee Button 6 512 64 24 11
Faucet Open 6 512 64 24 7
Door Open 6 512 64 24 11
Door Unlock 6 512 64 24 11
Sweep Into 6 512 64 24 7
Drawer Open 6 512 64 24 7
Hammer 6 512 64 24 7
Compute Resources
All experiments were conducted in a shared high-performance computing cluster. We used GPU-accelerated nodes for training and evaluation. Specifically, we utilized A100 or H100 GPU partitions depending on availability. Table 5 summarizes the hardware used in our experiments. Table 5 Compute resources used in all experiments.
Resource GPU (A100) GPU (H100)
Specification 5GB MIG partition 10GB MIG partition (HBM3)
15
Quantity 1 1
C
Additional Experiments
C.1
Visual Results
We evaluate the generalization of preference-based reinforcement learning to high-dimensional visual observations and using DinoV2 encoding [24] of the observations. Experiments are conducted on 2 standard continuous-control benchmarks: Walker Walk, and Cheetah Run. On Walker Walk, both methods UBP2 and MBP achieve comparable performance and outperform model free baselines. In contrast,in Cheetah Run, model free baseline outperform model based approach and UBP2. However, UBP2 still outperforms MBP. These results suggest that optimism helps mitigate the additional uncertainty associated with high-dimensional visual observations in model based approaches.
Episode Reward
300 250 200 150 100 50 0
0.4
0.6
Environment Steps
0.8
(b) Visual: Cheetah Run (5000)
(a) Visual: Walker Walk (5000)
UBP2
0.2
MBP
MRN
Rune
Figure 5 Performance comparison across visual DMControl tasks. Curves show IQM performance with shaded standard error over 5 seeds. Numbers between brackets indicate the maximum preference feedback budget.
C.2
Conservative Evaluation
We additionally evaluate whether the uncertainty estimates learned by UBP2 can be used conservatively at test time by reversing the optimism term during evaluation, i.e., penalizing trajectories with high epistemic uncertainty instead of rewarding them. This experiment is motivated by the observation that optimism is beneficial during data collection and learning, where exploration is required, but may become undesirable during deployment if a risk-averse or safety-oriented policy is preferred. Importantly, optimism is not used during standard evaluation in UBP2; the learned policy is evaluated without any exploration bonus. Here, we instead explicitly introduce a conservative evaluation objective by multiplying the uncertainty bonus by −1 during planning. Results shows that conservative evaluation substantially degrades performance on most tasks We speculate that uncertainty in UBP2 primarily acts as a useful exploratory signal rather than a reliable proxy for risk, and that naively enforcing conservatism at test time can significantly reduce task performance. Table 6 Performance comparison across MetaWorld tasks for UBP2 with conservative evaluation. Entries report average IQM success with standard error at the final evaluation step. Best results are shown in bold Method
C.3
Coffee Button
Faucet Open
Sweep Into
Drawer Open
UBP2
86.2 (3.7)
89.5 (7.9)
36.4 (12.3)
77.7 (4.9)
UBP2 w conservative eval
51.0 (22.3)
85.9 (20.0)
6.4 (0.0)
30.8 (23.7)
Varying Horizon
In the ablation study, we compare planning horizons of 7 and 11 across tasks. While horizon 11 yields better performance on a subset of tasks reported in the experiments section, it does not consistently improve results on the tasks considered in this appendix. In contrast, a horizon of 7 provides more stable and higher overall performance for these tasks. Therefore, we adopt horizon 7 for the results reported in this section.
16
Table 7 Effect of planning horizon on UBP2 performance across a subset of MetaWorld tasks over 1M steps. Entries report IQM success with standard error at the final evaluation step. Horizon
Door Close
Window Close
Handle Press
Sweep Into
Drawer Open
7 11
97.8 (1.0) 96.3 (0.0)
93.1 (1.0) 91.6 (0.0)
93.4 (0.0) 86.1 (8.0)
36.4 (12.3) 15.0 (7.5)
77.7 (4.9) 71.5 (19.6)
D
Proofs
D.1
Definitions and Clarifications
We first give some definitions and characterizations of our learned models. Later, we clarify some of the choices we have made in regard to the theory. Closed forms of the mean and variance of GPs Let Kn ∈ RM ×M denote the kernel matrix over Dn , kn (s, a) ∈ RM denote the vector of kernel evaluations between (s, a) and all inputs in Dn , and s′j 1:M denote the noisy measurements of p∗j . For GP regression of the true dynamics with observation noise variance σ 2 , the posterior mean and variance for coordinate j have the following standard closed forms [34]: −1 ′j µdn,j (s, a) = kn (s, a)⊤ Kn + σ 2 I s1:M , 2 −1 d σepin,j (s, a) = k((s, a), (s, a)) − kn (s, a)⊤ Kn + σ 2 I kn (s, a).
(13) (14)
d We collect these coordinate-wise values into vector estimates µdn (s, a) = [µdn,j (s, a)]j≤ds and σepi (s, a) = n d r [σepin,j (s, a)]j≤ds . We similarly fit a GP model using (14) on the reward with posterior mean µn (s, a) and r standard deviation σepi (s, a) using the reward kernel kr . Furthermore, we say that these learned models are n well-calibrated [26, 34].
Definition of well-calibrated models To connect epistemic uncertainty to optimism, in section 5 we assumed the learned models are well-calibrated: their posterior standard deviations provide valid confidence radii around the true functions they are attempting to estimate. Concretely, we give the following definition. Definition D.1 (Well-calibrated learned dynamics and reward [26]). Simultaneously for all episodes n and all (s, a) ∈ S × A, with probability 1 − δ:
∀j ≤ ds :
r r(s, a) − µrn (s, a) ≤ βnr (δ) σepi (s, a), n
(15)
p∗j (s, a) − µdn (s, a)
(16)
d ≤ βnd (δ) σepi (s, a), n,j
r d r where ∈ Rp ≥0 arenon-decreasing sequences. Moreover, in the GP/RKHS setting, typically βn ∈ p βn(δ), βn (δ) d O Γr,n and βn ∈ O Γd,n [5].
We invoke a lemma in [26] to show that GP estimates of the dynamics and reward are well-calibrated models if Assumption 5.2 holds: Lemma D.2 (Well-calibrated confidence intervals for RKHS [26]). Let p∗ be in a RKHS with kernel kd and the d characteristics in Assumption 5.2. Suppose µdn and σepi are the posterior mean and variance of a GP with kernel n d d kd as defined in equation (14). Then, there exists βn (δ) for which the tuple (µdn , σepi , βnd (δ)) is a well-calibrated n model of p∗ . The same is true for the proxy reward function r(s, a). Proof. Lemma D.2 is directly stated in [26] for an RKHS of vector-valued, scalar-output functions. Hence, this can directly be applied to the reward model and element-wise for the dynamics model.
17
Definition of the maximum information gain For a kernel k and noise variance σ 2 , the maximum information gain is defined to be: 1 ΓM (k) := max log det I + σ −2 KA , (17) A⊂S×A, |A|≤M 2 where KA is the kernel matrix over the set A. Intuitively, ΓM (k) measures the maximum amount of information that M noisy function evaluations can reveal about a function in the RKHS class parameterized by k. This quantity is sublinear in M for many common kernels [34]. See Table 8 for a list of the asymptotic growth of ΓM for different kernels. Table 8 Maximum information gain bounds for common choice of kernels [34].
Kernel k(s, s′ )
ΓN
⊤ ′
Linear: s s ∥s − s′ ∥2 RBF: exp − 2l2 √ 2ν ∥s − s′ ∥ ν √2ν ∥s − s′ ∥ 1 Matérn: Bν Γ(ν) 2ν−1 l l
O d log N
O logd+1 N 2ν d O N 2ν+d log 2ν+d (N )
Further discussion on modeling the true reward r∗ in preference-based RL. As discussed in section 3.1 and the Limitations, in PbRL, the agent never observes the real scalar rewards, and is instead trained on preference labels, giving rise to the definition of the proxy reward function r(s, a) within the same equivalence class of r∗ (s, a). This “latent-reward" view is standard in the PbRL literature, where preferences provide information about an underlying reward signal and hence allows the learning of a reward model whose ordering of information matches the ground-truth ordering [6, 36]. We assume that the well-calibrated GP reward model effectively approximates the proxy reward function r(s, a), satisfying equation 15. This view has enabled us to effectively produce sublinear regret bounds for the problem setting presented in this paper. However, as we have acknowledged earlier, it may be possible to make the reward learning error stemming from PbRL more explicit. Namely, [39] establishes value approximation guarantees under a linear reward parameterization, where the preference learning error can be more clearly characterized and directly incorporated into the final bound. In contrast, our setting does not assume such a parametric structure on the reward. As a result, explicitly propagating preference-learning error into the regret bound would require introducing additional assumptions on the reward model and its convergence rate. Introducing assumptions like a linear reward model could incorporate reward convergence into Lemma D.4 and Lemma 5.7, but this would complicate the analysis while introducing harsher assumptions. Our current result isolates planning and provides guarantees under general reward uncertainty, avoiding strong reward model conditions. Clarification between implemented dynamics and reward estimators and theoretical assumptions As introduced in the Limitations, practically the dynamics and reward are implemented via deep ensemble networks, while in the theoretical analysis they are assumed to be GPs. We acknowledge that this may be a strict assumption, and that our actual implementation may not fully reflect the theory. However, we note that in practice, deep ensembles are ubiquitous in model-based RL to quantify epistemic uncertainty or model disagreement. Furthermore, similar approximations have been adopted in prior work, such as [34], providing a practical starting point for theoretically characterizing performance. Our intent is not to claim that the implemented ensemble estimator is an exact GP posterior. Our proof requires two properties: 1) a high-probability confidence relationship between the reward/dynamics model error and their uncertainty radii, and 2) a cumulative uncertainty bound on these uncertainty radii (in our case, via GP information gain bounds). GP posterior standard deviations provide these quantities in closed form, which is why they are used in Lemma 5.7 and Theorem 5.8. Hence, both our implementation and theory aim to capture the same underlying notion of posterior-predictive uncertainty, but using different model classes and approximations. In the implementation, ensemble disagreement and JRD are used as sample-based approximations to the spread of the posterior predictive distribution 18
of the reward and dynamics. We view the ensemble members as approximate samples from a posterior or posterior-like distribution over possible reward and dynamics models. Then, the ensemble mean/disagreement can be viewed as Monte Carlo estimates of the mean/variance of the approximate posterior predictive distribution, and the JRD is a divergence-based measure of disagreement among each ensemble member’s predictive distributions. Therefore, both our implementation and theory use predictive uncertainty as a measurement of confidence, but only the GP case yields an explicit, clean, closed form regret bound. We acknowledge, however, that finite-neural ensembles are not guaranteed to be as well-calibrated as GPs, potentially presenting challenges reconciling the derived regret bounds with the practical results.
D.2
Analysis of Discounted Finite Horizon Planning with UCB Objective
Finite-horizon objective and regret Prior to showing the theoretical justification for the full infinite-horizon regret bound and the full UCB objective, we first consider a reduced UCB objective. Define: ηnUCB (π, µdn ) := Eπ,µdn
" H−1 X
γ
t
µrn (ŝt , at ) + an σnr (ŝt , at ) + bn σnd (ŝt , at )
# .
(18)
t=0
This objective corresponds to a finite-horizon variant of the full UCB objective in (8), with an , bn defined as in Lemma 5.7. Since the objective is finite-horizon, we omit the terminal value estimation and its uncertainty bonus. In this section, we derive a finite-horizon regret bound for policies obtained by optimizing (18) and how that it is sublinear. This informs our infinite horizon regret bound for the full UCB planning objective. Analogous to equation (9), we introduce a performance metric for the finite-horizon setting. We define the cumulative finite-horizon regret for a particular policy πn as: RN =
N X
rn =
n=1
N X
η(π ∗ , p∗ ) − η(πn , p∗ ),
(19)
n=1
where η(· , ·) is the finite-horizon analogue of (10) η(π, p) = Eπ,p
"H−1 X
# γ t r(st , at ) .
(20)
t=0
Discounted Finite-Horizon Cumulative Uncertainty and Return We further define the following quantities. For real episode n and policy π, define (as in [34]): e r (π) := Eπ,µd Σ n n
h H−1 X
i r γ t σepi (ŝ , a ) , t t n
(21)
i r γ t σepi (st , at ) , n
(22)
t=0
Σrn (π) := Eπ,p∗
h H−1 X t=0
e dn (π) := Eπ,µd Σ n
h H−1 X
d γ t σepi (ŝt , at ) n,j
i ,
(23)
d γ t σepi (st , at ) n,j
i .
(24)
t=0
Σdn (π) := Eπ,p∗
h H−1 X t=0
These quantities represent the sum of the estimated uncertainties of the learned reward and dynamics across timesteps under policy rollouts in the true and mean dynamics. Definitions (21) to (24) will be useful later proofs, and will be generalized to the infinite horizon setting. Lemma D.3 (Upper Bound of Discounted Return). Consider the equations (20), (23), and p (24) with the true dynamics p∗ and mean dynamics µdn having the same process noise. Then, there is a λn ∈ O( Γd,N ) such that for any policy π: 19
|η(π, p∗ ) − η(π, µdn )| ≤ λn Σdn (π).
(25)
|η(π, p∗ ) − η(π, µdn )| ≤ λn Σ̃dn (π).
(26)
Proof. Lemma D.3 is simply a discounted variation of Lemma B.1. from [34]. It can be seen that adding the discount factor has no effect on the proof. Lemma D.4 (Optimism of the UCB planning objective). Let Assumptions 5.1, 5.2phold, and let the learned p GP dynamics and reward models be well-calibrated. Then there exist coefficients an ∈ O Γr,N and bn ∈ O Γd,N such that for all episodes n and all policies π, with probability at least 1 − δ: η(π, p∗ ) ≤ ηnUCB (π, µdn ).
(27)
Proof. From Lemma D.3, we have:
η(π, p∗ ) ≤ λn Σ̃dn (π) + η(π, µdn ) h H−1 i X = λn Σ̃dn (π) + Eπ,µdn γ t r(ŝt , at ) t=0
= λn Σ̃dn (π) + Eπ,µdn
h H−1 X
i h H−1 X i γ t r(ŝt , at ) − µrn (ŝt , at ) + Eπ,µdn γ t µrn (ŝt , at )
t=0
≤ λn Σ̃dn (π) + Eπ,µdn
h H−1 X
t=0
i
γ t r∗ (ŝt , at ) − µrn (ŝt , at ) +Eπ,µdn
t=0
|
h H−1 X
i γ t µrn (ŝt , at ) .
t=0
{z
}
(A)
We bound (A) using Definition D.1, arriving at the optimism lemma:
η(π, p∗ ) ≤ λn Σ̃dn (π) + βnr Σ̃rn (π) + Eπ,µdn
h H−1 X
γ t µrn (ŝt , at )
i
t=0
= ηnU CB (π, µdn ),
(by definition)
p p where an , bn correspond to the functions βnr , λn , respectively. Since λn ∈ O( Γd,N ), and βnr ∈ O( Γr,N ), we arrive at the theoretical bounds for coefficients an , bn presented in Lemma D.4. Now, we introduce the analogue of Assumption 5.6 for the reduced UCB objective. This will be used in the proof of Lemma D.6 below. Assumption D.5 (No planner sub-optimality for reduced UCB objective). The policy πnred induced by the planning procedure through optimizing the reduced UCB objective (18) is equal to the optimal policy πn∗,red which maximizes (18). That is, πnred = πn∗,red . Lemma D.6 (Single episode regret bound). Let Assumptions 5.1, 5.2, and D.5 hold, and assume that the learned dynamics and reward are well-calibrated. Consider the definition of rn in (19), and let an , bn be as defined in Lemma D.4 . Then, ∀n > 0 with probability at least 1 − δ, ′ d red 2 r red rn ≤ (b′2 n + 2bn )Σn (πn ) + (an + an )Σn (πn ),
where b′n = max{bn , λn } and πnred is the planner induced policy, assumed to maximize equation (18).
20
(28)
Proof. We start with the definition of rn presented in (19):
rn = η(π ∗ , p∗ ) − η(πn , p∗ ) ≤ ηnU CB (π ∗ , µdn ) − η(πnred , p∗ )
(Lemma D.4)
≤ ηnU CB (πn∗,red , µdn ) − η(πnred , p∗ ) = ηnU CB (πnred , µdn ) − η(πnred , p∗ ) = η(πnred , µdn ) + an Σ̃rn (πnred ) + bn Σ̃dn (πnred ) − η(πnred , p∗ )
(Optimization of (18)) (Assumption D.5)
≤ λn Σdn (πnred ) + an Σ̃rn (πnred ) + bn Σ̃dn (πnred )
(Lemma D.3)
Let b′n = max{bn , λn }. Then
rn ≤ 2b′n Σdn (πnred ) + b′n (Σ̃dn (πnred ) − Σdn (πnred )) + an Σrn (πnred ) + an (Σ̃rn (πnred ) − Σrn (πnred )) d red 2 r red ≤ (2b′n + b′2 n )Σn (πn ) + (an + an )Σn (πn ).
In the last inequality, we use the fact that σnr (·) and σnd (·) are bounded and positive. Hence, they qualify as valid “reward signals" which can be substituted into Lemma D.3, yielding Σ̃dn (πnred ) − Σdn (πnred ) ≤ b′n Σdn (πnred ) and Σ̃rn (πnred ) − Σrn (πnred ) ≤ an Σrn (πnred ) [34] Theorem D.7 (Regret bound with dynamics and reward uncertainty bonuses). Let Assumptions 5.1, 5.2, and 5.6 hold, and let the learned GP dynamics and reward models be well-calibrated. The cumulative regret after N episodes of real environment interaction satisfies, with probability 1 − δ: √ 3/2 3/2 RN ≤ O H 3 N Γd,N + Γr,N . (29) Proof. We begin with the definition of the cumulative regret RN : RN =
N X
rn
n=1
≤
N h i X d ind 2 r red (2b′n + b′2 )Σ (π ) + (a + a )Σ (π ) n n n n n n n n=1
≤ (2b′N + b′2 N)
N X
Σdn (πnred ) + (a2N + aN )
n=1
≤ (2b′N + b′2 N)
N X
N X
Σrn (πnred )
n=1
Ep∗
n=1
h H−1 X
d σepi (st , πnred (st )) n
i
+ (a2N + aN )
t=0
N X n=1
Ep∗
h H−1 X
r σepi (st , πnred (st )) n
i
t=0
v uN h H−1 i uX X √ ′ ′2 d red (s )) 2 Ep∗ σepi (s , π ≤ (2bN + bN ) N H t t t n n n=1
t=0
v uN h H−1 uX X √ i 2 r red (s )) 2 σepi (s , π + (aN + aN ) N H t Ep∗ t t n n n=1
≤ C1 (2b′N + b′2 N )H
(Cauchy-Schwartz + Jensen’s inequality)
t=0
p p N Γd,N H + C2 (a2N + aN )H N Γr,N H
(From Lemma 17 in [7])
Finally, notice that from Lemma D.4 and Lemma D.3 that aN ∝ Hβnr , bN ∝ λN , λN ∝ Hβnd , βnd ∈ O( p and βnr ∈ O( Γr,N ). Substituting, we arrive at (29).
21
p
Γd,N ),
D.3
Analysis of Discounted Infinite Horizon Planning with the Complete UCB Objective
We now utilize and generalize the ideas for the finite horizon regret proof in Appendix D.2 to derive our final infinite-horizon regret bound with the complete UCB objective (8). We start by extending the definitions of the finite-horizon cumulative uncertainty (21)-(24) to be infinitehorizon. Discounted Infinite-Horizon Cumulative Uncertainty For episode n and policy π, define
e rγ,n (π) := Eπ,µd Σ n Σrγ,n (π) := Eπ,p∗
t=0 ∞ hX
e dγ,n (π) := Eπ,µd Σ n Σdγ,n (π) := Eπ,p∗
∞ hX
t=0 ∞ hX
t=0 ∞ hX
i r γ t σepi (ŝt , at ) , n
(30)
i r γ t σepi (st , at ) , n
(31)
d γ t σepi (ŝt , at ) n,j
i
d γ t σepi (st , at ) n,j
i
,
(32)
.
(33)
t=0
By Assumption 5.4, there exists constant αd and αr such that for all episode n, ∞ ∞ hX i hX i d t r σnq (s, a) = αd E γ t ∥σepi (s , a )∥ + α E γ σ (s , a ) . t t r t t epi n n n=1
n=1
We give a brief justification for this assumption by proving a lemma bounding the epistemic uncertainty of Q. We assume that the underlying value function V is bounded by Vmax and Lipchitz with a Lipchitz constant LV . We start by the following lemma: Lemma D.8. Let Assumption 5.1, 5.2 hold. Given dynamics p∗ , µdn and rewards r, µrn . We have the following bound
|Q̄πµn (s, a) − Qπp∗ (s, a)| ≤ βnr Eπ,µdn
"∞ X
# γ t σnr (st , at )
+ γLVpπ∗ βnd Eµdn ,π
t=0
"∞ X
# t
γ ∥σnd (st , at )∥
.
(34)
t=0
Proof. Define the transition operators (Ppπ f )(s, a) := Es′ ∼p(·|s,a) f (s′ , π(s′ )) ,
(35)
(Pp g)(s, a) := Es′ ∼p(·|s,a) g(s′ ) .
(36)
Then we have, Q̄πµdn − Qπp∗ = (µrn − r) + γ Pµπdn Q̄πµdn − Ppπ∗ Qπp∗
= (µrn − r) + γPµπdn (Q̄πµdn − Qπp∗ ) + γ(Pµπdn − Ppπ∗ )Qπp∗ . Rearranging gives the resolvent form (I − γPµπdn )(Q̄πµdn − Qπp∗ ) = (µrn − r) + γ(Pµπdn − Ppπ∗ )Qπp∗ .
22
Since rewards are bounded, (I − γPµπd )−1 exists on bounded functions, use the Neumann series (I − γPµπd )−1 = n n P∞ π t t=0 (γPµd ) to get n
Q̄πµdn − Qπp∗ =
∞ X
(γPµπdn )t (µrn − r) + γ
t=0
∞ X
(γPµπdn )t (Pµπdn − Ppπ∗ )Qπp∗ .
t=0
Take absolute values pointwise and apply the triangle inequality: Q̄πµdn − Qπp∗ ≤
∞ X
γ t (Pµπdn )t (µrn − r) + γ
t=0
∞ X
γ t (Pµπdn )t (Pµπdn − Ppπ∗ )Qπp∗ .
t=0
Now interpret (Pµπd )t as expectation over a length-t rollout under (µdn , π) starting from (s, a): for any bounded n g, ((Pµπdn )t g)(s, a) = Eµdn ,π [g(st , at ) | s0 = s, a0 = a] . Applying this with g = µrn − r yields
((Pµπdn )t (µrn − r))(s, a) ≤ Eµdn ,π µrn (st , at ) − r(st , at ) s0 = s, a0 = a ≤ βnr Eπ,µdn [σnr (s,t , at )] For the second term, note that (Pµπdn − Ppπ∗ )Qπp∗ (s, a) = Eµdn Vpπ∗ (s′ ) − Ep∗ Vpπ∗ (s′ ) = (Pµdn − Pp∗ )Vpπ∗ (s, a), so ((Pµπdn )t (Pµπdn − Ppπ∗ )Qπp∗ )(s, a) ≤ Eµdn ,π (Pµdn − Pp∗ )Vpπ∗ (st , at ) s0 = s, a0 = a ≤ LVpπ∗ Eµdn ,π [W (pµdn (·|st , at ), pp∗ (·|st , at ))] = LVpπ∗ Eµdn ,π [|µdn (st , at ) − p∗ (st , at )|] ≤ LVpπ∗ βnd Eµdn ,π [∥σnd (st , at )∥]. W (·, ·) is the Wasserstein distance between two distributions. For two Gaussian distributions, this reduces to the distance between their means. The true Q-error |Qπp∗ − Q̄πµn | is unobservable. However, Lemma D.8 shows that this error can be upperbounded by a discounted accumulation of the posterior standard deviations of the dynamics and reward models. Therefore, we utilize this bound as a computable epistemic uncertainty proxy for the Q-function. Intuitively, because Q is a functional of the reward and dynamics, it’s within reason that the epistemic uncertainty of Q should be built from uncertainty the reward model and dynamics model. Note that although we explicitly reference Assumption 5.1, the Gaussian nature of the dynamics and reward in order to characterize their Wasserstein distance, this does not mean that the process noise also has to be Gaussian. Please refer to Appendix D.4 for further discussions on settings with sub-gaussian noise. We can give a proof of Lemma 5.7 following the previous proof of Lemma D.4 and also treat Q as the tail return estimator. Proof of Lemma 5.7. By Lemma B.7 in [34], ηγ (π, p∗ ) ≤ λn Σ̃dγ,n (π) + ηγ (π, µdn ) "∞ # "∞ # X X t d t = λn Eπ,µdn γ ∥σepin (st , at )∥ + Eπ,µdn γ r(st , at ) t=0
t=0
23
= λn Eπ,µdn
"H−1 X
# t
d γ ∥σepi (st , at )∥ n
+ Eπ,µdn
t=0
+ Eπ,µdn
"H−1 X
"H−1 X
# t
γ (r(st , at ) − µrn (st , at ))
t=0
# γ t µrn (st , at ) + γ H Eπ,µdn
"∞ X
t=0
# d γ t ∥σepi (st+H , at+H )∥ + r(st+H , at+H ) n
t=0
≤ λn Σ̃dγ,n (π) + βnr Σ̃rγ,n (π) + E
"H−1 X
# γ t µrn (st , at )
t=0
"∞ # "∞ # X X H t d H r t r γ ∥σepin (st+H , at+H )∥ + γ βn E γ ∥σepin (st+H , at+H )∥ + γ λn E t=0
t=0
"∞ # X H t r +γ E γ µn (st+H , at+H ) t=0
≤ λn Σ̃dγ,n (π) + βnr Σ̃rγ,n (π, µdn ) + E
"H−1 X
# γ t µrn (st , at )
t=0 H
+ γ (cn E
h
q σepi (sH , aH ) n
i
+ E Q̄πµn (sH , aH ) )
≤ λn Σ̃dγ,n (π) + βnr Σ̃rγ,n (π) + E
"H−1 X
# γ t µrn (st , at )
t=0
h i + γ E Q̂πµn (sH , aH ) + cn σnq (sH , aH ) + γ H ϵq,n H
UBP = ηγ,n (π, µdn ) + γ H ϵq,n , max{λ ,β r }
where cn = min{αrn,αdn} ∈ O max{
p
Γr,N ,
p
Γd,N } . □
Lemma D.9 (Single episode regret bound for infinite-horizon return). Let Assumptions 5.1, 5.2, and 5.6 hold, and assume that the learned dynamics and reward are well-calibrated. Consider the definition of rn in (9), and let an , bn , cn be as defined in Lemma D.4 . Then, ∀n > 0 with probability at least 1 − δ, d r ′ ′ ′ ′ ind 2 ′ ′ ind H rγ,n ≤ b′2 (37) n + bn cn + 2bn + cn Σγ,n (πn ) + 2an + an cn + 2an + cn Σγ,n (πn ) + 2γ ϵq,n where b′n = max{bn , λn }, c′n = cn max{αr , αd } and πnind is the planner induced policy, assumed to maximize equation (1). Proof. rn = ηγ (π ∗ , p∗ ) − ηγ (πn , p∗ ) UCB ∗ ≤ ηγ,n (π , µdn ) − ηγ (πnind , p∗ ) + γ H ϵq,n
(Lemma 5.7)
UCB ∗ ≤ ηγ,n (πn , µdn ) − ηγ (πnind , p∗ ) + γ H ϵq,n
(Optimization of (8))
UCB ind = ηγ,n (πn , µdn ) − ηγ (πnind , p∗ ) + γ H ϵq,n
(Assumption 5.6)
= η(πnind , µn ) + an Σ̃rn (πnind ) + bn Σ̃dn (πnind ) h ind i π q + γ H Eµdn Q̂µnn (sH , aH ) + cn σepi (sH , aH ) − ηγ (πnind , p∗ ) + γ H ϵq,n n = η(πnind , µn ) − ηγ (πnind , p∗ ) + an Σ̃rn (πnind ) + bn Σ̃dn (πnind ) "∞ # X q H H r + γ Eµdn γ µn (st+H , at+H ) − r(st+H , at+H ) + cn σepin (sH , aH ) + 2γ H ϵq,n t=0 ′2 ′ ≤ (bn + 2bn ) Σdγ,n (πnind ) + (a2n + an ) Σrγ,n (πnind )
(Lemma B.7 from [34])
24
" H
+ γ Eµdn
βnr
∞ X
r γ t σepi (st , at ) + c′n n
t=0
∞ X
γ
t
r d σepi (st+H , at+H ) + ∥σepi (st+H , at+H )∥ n n
# + 2γ H ϵq,n
t=0
′ d ind 2 r ind ≤ (b′2 n + 2bn ) Σγ,n (πn ) + (2an + 2an ) Σγ,n (πn ) # " ∞ X r d H ′ t + γ Eµdn cn γ σepin (st+H , at+H ) + ∥σepin (st+H , at+H )∥ + 2γ H ϵq,n t=0 ′2 ′ ′ ≤ (bn + bn cn + 2b′n + c′n ) Σdγ,n (πnind ) + (2a2n + an c′n + 2an + c′n ) Σrγ,n (πnind ) + 2γ H ϵq,n .
where c′n = cn max{αr , αd }.
This lemma yields our bound for the infinite-horizon regret. Proof of Theorem 5.8. First note that since an , bn are monotonic, we know that c′n is also monotonic.
RN =
N X
rn
n=1
≤
N X
d r ′ ′ ′ ind 2 ′ ′ ind H b′2 n + bn cn + 2bn + cn Σγ,n (πn ) + 2an + an cn + 2an + cn Σγ,n (πn ) + 2γ ϵq,n
n=1
≤
′ ′ ′ b′2 N + bN cN + 2bN + cN
N X
Σdγ,n (πnind ) +
2a2N + aN c′N + 2aN + c′N
n=1
:= BN
N X
Σdγ,n (πnind ) + AN
n=1
Since O max{
N X
p
N X
n=1
Σrγ,n (πnind ) + 2γ H
n=1
Γr,N ,
Σdγ,n (πnind ) ≤
p
√
N X
N X
Σrγ,n (πnind ) + 2γ H
N X
ϵq,n
n=1
ϵq,n .
n=1
p p Γd,N } = O( Γr,N + Γd,N ), we have AN , BN ∈ O(Γr,N + Γd,N ). Consider
v uN uX 2 Σdγ,n (πnind ) Nt n=1
n=1
v u !2 N ∞ X X √ u u d E γ t ∥σepi ∥ ≤ Nt n n=1
t=0
v !# uN " ∞ ! ∞ X X √ uX t d t t 2 ≤ N E γ γ ∥σepin (st , at )∥ n=1
s = s ≤
N 1−γ
t=0
t=0
v uN "∞ uX X t E γ t ∥σ d
#
2 epin (st , at )∥
n=1
t=0
d Cγ N ΓN log N log N C (σmax )2 N log N + 1−γ (1 − γ)2
(Proof of Theorem. 5.5 from [34])
with σmax as defined in Assumption 5.2 and for some constants C and Cγ . Following similar steps, we can PN prove an analogous inequality for Σrγ,n . By defining eq,N = n=1 ϵq,n , we obtain the bound in (12). □
25
Explicit infinite-horizon regret bound While Theorem 5.8 is written in terms of the GP information gain, since both reward and dynamics are modeled as functions of state and action (i.e., input is (s, a)), the input dimension in the kernel is ds + da . Therefore, once a kernel is chosen, the asymptotic growth of ΓN explicitly depends on the dimension. We demonstrate that for each of the linear, RBF, and Matérn kernels, our regret bound remains sublinear. Let dx = ds + da and M = N log N . Since both the reward model and each coordinate of the dynamics model √ 3/2 3/2 dx 3 are functions of the joint input x = (s, a) ∈ R , our Theorem 5.8 gives: Rγ,N ≤ H N Γd,M + Γr,M + eq,N . Algebraically substituting the information-gain bounds for each of the kernels yields the following kernel-specific, dimension-dependent variants of Theorem 5.8: • Linear kernel: ΓM = O(dx log M ) =⇒
√ Rγ,N = Õ H 3 N d3/2 x + eq,N
• RBF kernel: ΓM = O(logdx +1 M )
√ 3 Rγ,N = Õ H 3 N log 2 (dx +1) N + eq,N dx 2ν • Matérn kernel: ΓM = O M 2ν+dx log 2ν+dx M 3dx 1 =⇒ Rγ,N = Õ H 3 N 2 + 2(2ν+dx ) + eq,N =⇒
Therefore, it is seen that for linear and RBF, that our regret bound remains sublinear in N for fixed state and action dimension. For the Matérn with smoothness ν, for sublinear growth a sufficient condition is 3dx 1 2 + 2(2ν+dx ) < 1, or dx < ν. However, this may be too conservative. There exist work that show function optimization in the RKHS of a Matérn kernel obtain sublinear regret under broader conditions, such as [13].
D.4
Analysis for Settings With Sub-Gaussian Process Noise
We now discuss how our infinite horizon regret bound extends beyond Gaussian process noise. In Assumption 5.1, we assume that the process noise of the true dynamics p∗ is Gaussian. As our theoretical work is current presented, the Gaussianity of the noise is necessary in two locations: (i) to obtain high-probability confidence intervals for the estimated dynamics and reward, and (ii) the invocation of Lemma B.7 from [34], which is a simulation-lemma type bound for the value function under the true and discounted dynamics. Therefore, relaxing the Gaussian noise to subgaussian noise requires the addressing of both these points. The first use is not inherently Gaussian - we simply assumed so to simplify the analysis. The high-probability confidence relation required between the reward and dynamics model error and their posterior GP uncertainty can indeed be obtained under conditionally subgaussian noise. In particular, we refer to analysis done by [5]; in the paper, the authors operate within an agnostic RKHS setting and also utilize the standard closed form for the GP posterior mean and variance as presented in (14). It is explictly stated that the Gaussian prior assumption are only an aid to algorithm design, and that the true process noise may be a conditionally R-subgaussian martingale-difference sequence. They then provide a uniform confidence bound in their Theorem 2, stated below. Note that this bound is precisely the well-calibrated confidence bound our theory requires. p |µt−1 (x) − f (x)| ≤ B + R 2(γt−1 + 1 + ln(1/δ)) σt−1 (x) (Theorem 2, [5])
However, as Lemma B.7 from [34] relies on a Gaussian smoothing argument through their use of Lemma C.2 from [15], we need explicitly address settings with subgaussian noise. We will show that even under the subgaussian noise, the regret bound asymptotically remains sublinear. We assume that the true dynamics has additive subgaussian noise: st+1 = g ∗ (st , at ) + ξt , where ξt is a conditionally mean-zero σd -subgaussian martingale-difference sequence. The learned dynamics is ŝt+1 = µdn (ŝt , at ) + ξt . To avoid overloading the notation, we denote the true dynamics as g ∗ . However, in actuality, g ∗ = p∗ . We also make the following assumption:
26
Assumption D.10 (Lipschitzness of dynamics, policies, and reward). The true dynamics g ∗ , the reward r, and all π ∈ Π are Lg , Lr , and uniformly Lπ Lipschitz respectively and satisfy a discounted stability condition γLg (1 + Lπ ) < 1 An immediate corollary of the above assumption is the following: Corollary D.11 (Lipschitz-ness of value function). Let all conditions in Assumption D.10 hold. Then, the value function Vgπ (s) induced by g ∗ under any policy π ∈ Π is uniformly LV -Lipschitz with LV ≤
Lr (1 + Lπ ) . 1 − γLg (1 + Lπ )
(38)
Proof. Take any two starting states s0 and s̃0 and run the same deterministic policy from each state, yielding two actions, a0 and ã0 . Furthermore, consider the one-step dynamics which yield next states s1 and s̃1 , where s1 = g ∗ (s0 , a0 ) + ξ0 and s̃1 = g ∗ (s̃0 , ã0 ) + ξ0 . Let (s0 , a0 ) denote the concatenated vector of s0 and a0 , and let (s̃0 , ã0 ) denote the concatenated vector of s̃0 and ã0 . We have that: p ∥(s0 , a0 ) − (s̃0 , ã0 )∥ = ∥s0 − s̃0 ∥2 − ∥a0 − ã0 ∥2 ≤ ∥s0 − s̃0 ∥ + ∥a0 − ã0 ∥ ≤ (1 + Lπ )∥s0 − s̃0 ∥. In the last step, we use the Lipschitz-ness of the policy, or the fact that ∥a0 − ã0 ∥ ≤ Lπ ∥s0 − s̃0 ∥. Now, we consider the difference between s1 and s˜1 . Here, we cancel the process noise: s1 − s̃1 = g(s0 , a0 ) − g(s̃0 , ã0 ). Then, using the Lipschitz-ness of the dynamics we have the following: ∥s1 − s̃1 ∥ ≤ Lg ∥(s0 , a0 ) − (s̃0 , ã0 )∥ ≤ Lg (1 + Lπ )∥s0 − s̃0 ∥ = ρ∥s0 − s̃0 ∥ where we define ρ = Lg (1 + Lπ ) Again consider the one-step dynamics beginning at s1 and s̃1 after running the same deterministic policy from each state and taking actions a1 and ã1 . We obtain two more states s2 = g ∗ (s1 , a1 ) + ξ1 and s̃2 = g ∗ (s̃1 , ã1 ) + ξ1 It can easily be seen that ∥(s1 , a1 ) − (s̃1 , ã1 )∥ ≤ (1 + Lπ )|s1 − s̃1 ∥ and hence ∥s2 − s̃2 ∥ ≤ ρ2 ∥s0 − s̃0 ∥ via the same argument as above. Iterating this argument, for some timestep t, ∥(st , at ) − (s̃t , ãt )∥ ≤ (1 + Lπ )∥st − s̃t ∥. ∥st − s̃t ∥ ≤ ρt ∥s0 − s̃0 ∥ Now, we compare rewards. Using the Lipschitz-ness of r and the equations above, |r(st , at ) − r(s̃t , ãt )| ≤ Lr ∥(st , at ) − (s̃t , ãt )∥ ≤ Lr (1 + Lπ )∥st − s̃t ∥ ≤ Lr (1 + Lπ )ρt ∥s0 − s̃0 ∥. We then define the value functions Vgπ (s0 ) = Eξ
"∞ X
# t
γ r(st , at )
t=0
Vgπ (s̃0 ) = Eξ
"∞ X
# t
γ r(s̃t , ãt )
(39)
t=0
Then, comparing one coupled rollout we have |Vgπ (s0 ) − Vgπ (s̃0 )| ≤
∞ X
γ t |r(st , at ) − r(s̃t , ãt )|
t=0
≤
∞ X
γ t Lr (1 + Lπ )ρt ∥s0 − s̃0 ∥
t=0 ∞ X ≤ Lr (1 + Lπ ) (γρ)t t=0
27
! ∥s0 − s̃0 ∥.
Note that the infinte sum is only finite if γρ = γLg (1 + Lπ ) < 1, which is precisely the discounted stability P∞ 1 condition required in Assumption D.10. If this is the case, then, the result of the infinte sum is t=0 (γρ)t = 1−γρ . Therefore, Lr (1 + Lπ ) |V π (s0 ) − V π (s̃0 )| ≤ ∥s0 − s̃0 ∥. 1 − γLg (1 + Lπ ) So, for any policy π ∈ Π and any starting states, Vgπ is LV -Lipschitz with LV upper bounded by the coefficient above. Now, we introduce the following lemma, which is a simulation lemma similar to Lemma D.3 and Lemma B.7 of [34]. It will be shown that Lemma D.12 (Simulation lemma for settings with subgaussian noise). Let Assumptions 5.1, 5.2, D.10 hold, and let LV be as defined in Corollary D.11. Then, there is a λdn,SG = γLV βnd . such that for every policy π ∈ Π with probability at least 1 − δ, ηγ (π, p∗ ) − ηγ (π, µdn ) ≤ λdn,SG Σdγ,n (π) (40) e dγ,n (π), ηγ (π, p∗ ) − ηγ (π, µdn ) ≤ λdn,SG Σ
(41)
Proof. For simplicity, assume there is a deterministic policy π. The stochastic policy variant of the proof is the same, with an extra expectation over the action. Define two value functions, V∗ (s0 ) = Vgπ∗ (s0 ) and Vµ (s0 ) = Vµπd (s0 ), corresponding to true dynamics and learned n dynamics. Consider their Bellman expansions, which are V∗ (s0 ) = r(s, π(s0 )) + γEξ0 [V∗ (g ∗ (s0 , π(s0 )) + ξ0 )] and Vµ (s0 ) = r(s0 , π(s0 )) + γEξ0 [Vµ (µdn (s0 , π(s0 )) + ξ0 )]. Let ∆(s) = Vµ (s) − V∗ (s). Then, |∆(s0 )| = |γEξ0 Vµ (µdn (s0 , π(s0 )) + ξ0 ) − γEξ0 [V∗ (g ∗ (s0 , π(s0 )) + ξ0 )] | = |γEξ0 Vµ (µdn (s0 , π(s0 )) + ξ0 ) − V∗ (µdn (s0 , π(s0 )) + ξ0 ) + γEξ0 V∗ (µdn (s0 , π(s0 )) + ξ0 ) − V∗ (g ∗ (s0 , π(s0 )) + ξ0 ) | ≤ γEξ0 [|∆(sµ1 )|] + γEξ0 [|V∗ (sµ1 ) − V∗ (s∗1 )|] ≤ γEξ0 [|∆(sµ1 )|] + γEξ0 [LV ∥sµ1 − s∗1 ∥]
(Lipschitz-ness of V∗ ) µ d ∗ = γEξ0 [|∆(s1 )|] + γEξ0 [LV ∥µn (s0 , π(s0 )) − g (s0 , π(s0 ))∥] d ≤ Eξ0 [|∆(sµ1 )|] + γEξ0 [LV βnd ∥σepi (s0 , π(s0 ))∥] (Definition D.1) n µ d = Eξ0 [|∆(s1 )|] + γLV βnd ∥σepi (s0 , π(s0 ))∥ n where sµt+1 = µdn (st , π(st )) + ξt , and s∗t+1 = g ∗ (st , π(st )) + ξt . Now, we repeat the above process for the next state produced by the learned model sµ1 , yielding d |∆(sµ1 )| ≤ γEξ1 [|∆(sµ2 )|] + γLV βnd ∥σepi (sµ1 , π(sµ1 ))∥ n
Substituting into the previous inequality, we obtain d d |∆(s0 )| ≤ γLV βnd ∥σepi (s0 , π(s0 ))∥ + γ 2 LV Cd βnd Eξ0 [∥σepi (sµ1 , π(sµ1 ))∥] + γ 2 Eξ0 Eξ1 [|∆(sµ2 )|]. n n
We repeat this expansion of sµt forever. Notice that the last term vanishes as t → ∞ due to the boundedness of the reward. Therefore, "∞ # X d t d |∆(s0 )| = |Vµ (s0 ) − V∗ (s0 )| ≤ λn,SG Eξ0 γ ∥σepin (st , at )∥ | s0 t=0
where λdn,SG = γLV βnd . Using (32) and taking the expectation over the starting state s0 yields (41). We can reuse the above procedure but instead add and subtract γEξ0 [Vµ (g ∗ (s0 , π(s0 )) + ξ0 ] rather than γEξ0 [V∗ (µdn (s0 , π(s0 )) + ξ0 ]. This will result in terms involving ∆(s∗t ) after repeating the same expansion step, yielding (40).
28
Theorem D.13 (Infinite-horizon regret bound in settings with subgaussian noise). Let all conditions and assumptions stated in Theorem 5.8 hold, except that the process noise is a conditionally mean-zero σd -subgaussian martingale-difference sequence. Then, the cumulative infinite-horizon regret after N episodes of real environment interaction satisfies, with probability 1−δ, the same sublinear bound as presented in Theorem 5.8 up to constants/log factors: √ 3/2 3/2 Rγ,N ≤ O H 3 N Γd,N log N +Γr,N log N +eq,N where eq,N :=
PN
n=1 ϵq,n .
Proof. We can see that Lemma D.12 has the same form of Lemma B.7. in [34]. The subgaussian replacement only changes the coefficient of the simulation lemma. In the Gaussian proof, the coefficient λn = √ p γ (1+ dx )βn−1 Cmax 1−γ , while in the subgaussian proof, we instead use λdn,SG = γLV βnd . Since βnd = O( Γd,N ), σ the coefficient has the same information-gain order as the original coefficient. Therefore, the regret rate remains the same up to constants/log-factors.
D.5
Analysis of UCB Objective With Planner Suboptimality
Introducing Planner Sub-optimality. Recall Assumption 5.6, where we assumed that the policy πnind induced by the finite-horizon optimization of (8) after the nth episode of real-environment interaction coincides with the ideal policy πn∗ induced by exact maximization of (8). This assumption effectively treats the planner as an oracle that, under MPC-style re-planning, always selects the globally optimal first action according to the optimistic objective. In practice, however, planning is implemented via approximate optimization over some non-convex objective, hence the final solution action sequence may be suboptimal with respect to the planning objective [3, 14, 29]. Hence, it is prudent to also consider this imperfection in the theoretical analysis. We use the formulation introduced in [29], and introduce a planner sub-optimality parameter ϵp , which upper bounds the gap between the optimistic objective value attained by the ideal MPC-induced policy and that achieved by the implemented planner. Using this, we re-state Assumption 5.6: Assumption D.14 (Planner suboptimality in the optimistic objective). At the nth episode of real environment −1 interaction, with each episode containing T timesteps, there exist nonnegative numbers {ϵp,n,t }Tt=0 such that the policy induced by the UCB objective (8) satisfies: UCB ∗ UCB ind ηγ,n (πn , µdn ) − ηγ,n (πn , µdn ) ≤ ϵγp,n ,
ϵγp,n :=
T −1 X
γ t ϵp,n,t .
(42)
t=0 T
In particular, if ϵp,n,t ≤ ϵp for all t, then ϵγp,n ≤ 1−γ 1−γ ϵp . Because MPC re-solves the UCB objective at each timestep and executes only the first action, suboptimality can accumulate across re-planning steps; we therefore quantify a discounted, episode-level planner error ϵγp,n , where ϵp,n,t bounds the one-step optimization error at real environment step t in episode n. This is unlike [29], where the MPC setting is not considered. Ultimately, our definition yields a direct additive degradation in the per-episode regret bound, which we state below. Lemma D.15 (Single episode regret with planner suboptimality). Let Assumptions 5.1, 5.2, 5.3, 5.4, and D.14 hold (note: Assumption 5.6 no longer holds), and assume that the learned dynamics and reward are well-calibrated. Consider the definition of rn in (9), and let an , bn be as defined in Lemma D.4, Then, ∀n > 0 with probability at least 1 − δ, ′ ′ ′ ′ d ind 2 ′ ′ r ind H γ rn ≤ (b′2 n + bn cn + 2bn + cn ) Σγ,n (πn ) + (2an + an cn + 2an + cn ) Σγ,n (πn ) + 2γ ϵq,n + ϵp,n ,
where an , b′n , c′n are as defined in Lemma D.9 and ϵγp,n is as defined in Assumption D.14 Proof. We follow the proof of Lemma D.9:
29
(43)
rn = ηγ (π ∗ , p∗ ) − ηγ (πn , p∗ ) UCB ∗ ≤ ηγ,n (π , µdn ) − ηγ (πnind , p∗ ) + γ H ϵq,n UCB ∗ UCB ∗ UCB ,∗ UCB ind = ηγ,n (π , µdn ) − ηγ,n (πn , µdn ) + ηγ,n (πn , µdn ) − ηγ,n (πn , µdn ) | {z } | {z } (A)
(B)
UCB ind + ηγ,n (πn , µdn ) − ηγ (πnind , p∗ ) + γ H ϵq,n . | {z } (C)
UCB Since πn∗ ∈ arg maxπ ηγ,n (π, µdn ) by definition, (A) ≤ 0. Notice that term (B) is precisely the definition of the MPC planner suboptimality in Assumption D.14, so (B) ≤ ϵγp,n . Finally, notice that (C) is equivalent to the original regret bound presented in the proof of Lemma D.9. Putting the terms together, we obtain the regret bound in (43).
Theorem D.16 (Regret bound with planner suboptimality). Let Assumptions 5.1, 5.2, 5.3, 5.4, and D.14 hold (note: Assumption 5.6 no longer holds), and assume that the learned dynamics and reward are well-calibrated. Let ϵγp,n be as defined in Assumption D.14. Then, the cumulative regret after N real environment episodes satisfies, with probability 1 − δ: N √ X 3/2 3/2 Rγ,N ≤ O H 3 N Γd,N log N + Γr,N log N + eq,N + ϵγp,n ,
(44)
n=1
Moreover, if ϵγp,n ≤ ϵγp,n,t for all n, t, then: √ 1 − γT 3/2 3/2 Rγ,N ≤ O H 3 N Γd,N log N + Γr,N log N + eq,N + N · ϵp 1−γ
(45)
Proof. The single episode regret bound from Lemma D.15 can be substituted into the proof procedure for Theorem 5.8 with minimal changes. The planner suboptimality simply introduces an additive term which has no effect on the rest of the proof. Planner suboptimality in the finite horizon regret bound It is easily shown that the same additive terms corresponding to the planner suboptimality which appear in (44) and (45) appear in the finite horizon regret bound in Theorem D.7, contributing linearly. Discussion of planner suboptimality and regret Evidently, in order for the cumulative regret to remain sublinear, the cumulative planner error should be sublinear rather than linear √ (as it currently is in Theorem D.16 . If an PN √ assumption is made that ϵγp,n ∈ O(1/ n), then n=1 ϵγp,n ∈ O( N ) and the overall regret remains sublinear. The question then is whether such a decay assumption is warranted in practice. Under a fixed planning budget per replanning step (e.g., a constant number of trajectory samples or gradient steps), there is generally PN no reason to expect ϵp,n,t to vanish with n, in which case n=1 ϵγp,n is linear and can dominate the regret bound asymptotically. On the other hand, a decaying ϵγp,n becomes plausible if planning accuracy improves over time, either by allocating increasing compute to the planner or by reducing the effective optimization difficulty through warm-starting and amortization. Past work has demonstrated that in nonconvex stochastic optimization, many algorithmic guarantees improve monotonically with the number of iterations or gradient evaluations, yielding smaller optimization residuals as compute increases [10]. Similarly, sampling-based optimizers commonly used in MPC such as the cross-entropy method improve solution quality as the sampling budget increases [8]. Moreover, even without explicitly increasing the compute budget, which may not always be feasible, it has been shown that methodological innovations such as warm-starting and amortization has been explicitly studied as a mechanism to accelerate nonlinear trajectory optimization by predicting high-quality initial solutions [4, 27]. 30
√ Taken together, these considerations suggest that in practice, conditions for assuming ϵγp,n ∈ O(1/ n) may be a reasonable abstraction in certain planning regimes with the aforementioned computational and methodological optimizations. In particular, practical TD-MPC2-style systems often warm-start planning [11]. However, establishing a general, specific decay rate for these methods remain challenging and are an interesting direction for future theoretical work.
31