Model-based Offline RL via Robust Value-Aware Model Learning with Implicitly Differentiable Adaptive Weighting
Authors: Zhongjian Qiao
Jiafei Lyu
Boxiang Lyu
Yao Shu
Siyang Gao
Shuang Qiu
Venue: ICLR 2026 Poster (2026)
Keywords: Offline RL; Model-based RL
Primary area: reinforcement learning
TL;DR: We propose an model-based offline RL method called ROMI to address the limitations of RAMBO.

Abstract: Model-based offline reinforcement learning (RL) aims to enhance offline RL with a dynamics model that facilitates policy exploration. However, model exploitation could occur due to inevitable model errors, which degrades algorithm performance. Adversarial model learning offers a theoretical framework to mitigate model exploitation by solving a maximin formulation, and RAMBO provides a practical implementation with model gradient. However, we empirically observe that severe Q-value underestimation and gradient explosion can occur in RAMBO with only slight hyperparameter tuning, suggesting that it tends to be overly conservative and suffers from unstable model updates. To address these issues, we propose RObust value-aware Model learning via Implicitly differentiable adaptive weighting (ROMI). Instead of updating the dynamics model with model gradient, ROMI introduces a novel robust value-aware model learning approach. This approach requires the dynamics model to predict future states with values close to the minimum Q-value within a scale-adjustable state uncertainty set, enabling controllable conservatism and stable model updates. To further improve out-of-distribution (OOD) generalization during multi-step rollouts, we propose implicitly differentiable adaptive weighting, a bi-level optimization scheme that adaptively achieves dynamics- and value-aware model learning. Empirical results on D4RL and NeoRL datasets show that ROMI significantly outperforms RAMBO and achieves competitive or superior performance compared to state-of-the-art methods on datasets where RAMBO typically underperforms.

Published as a conference paper at ICLR 2026
MODEL-BASED OFFLINE RL VIA ROBUST VALUE- AWARE MODEL LEARNING WITH IMPLICITLY DIFFER- ENTIABLE ADAPTIVE WEIGHTING
Zhongjian Qiao', Jiafei Lyu", Boxiang Lyu', Yao Siyang Gao', Shuang Qiu CHiyUHK, Tennent, 'UChicago 'HKUST(GZ) zhongqiao2 cer.oo--mmmmccce.mm..ee www.y.eeuu...........oomooo..... hk, shuangiuEcityu edu.
ABSTRACT
Model-based offline reinforcement learning (RL) aims to enhance offline RL with a dynamies model that facilitates policy exploration. However, model exploitation could occur due to inevitable model errors, degrading algorithm performance. Ad- versarial model learning offers a theoretical framework to mitigate model exploita- tion by solving maximin formulation. Within such a paradigm, RAMBO (Rigter et al., 2022) has emerged as representative and most popular method that pro- vides practical implementation with model gradient. However, empirically reveal that severe Q-value underestimation and gradient explosion can occur in RAMBO with only slight hyperparameter tuning, suggesting that it tends be overly conservative and suffers from unstable model updates. To address these issues, we propose RObust value-aware Model learning with Implicitly differ- entiable adaptive weighting (ROMI). Instead of updating the dynamics model with model gradient, ROMI introduces a novel robust value-aware model learn- ing approach. This approach requires the dynamics model to predict future states with values close to the minimum Q-value within scale-adjustable state uncer- tainty set, enabling controllable conservatism and stable model updates. To fur- ther improve ditt--ddistriuttinn (OOD) generalization during multi-step rollouts, we propose implicitly differentiable adaptive weighting, bi-level optimization scheme that adaptively achieves dynamics- and value-aware model learning. Em- pirical results D4RL and NeoRI. datasets show that ROMI significantly out- performs RAMBO and achieves competitive or superior performance compared other state-of-the of-the-rr methods on datasets where RAMBO typically underper- forms. Code is available at https: /ttp::/gtttttt com/ z RROMI git.
INTRODUCTION
Model- offline RL (Yu al., 2020; Kidambi al., 2020; Rigter et al., 2022; Qiao et al.,
2025; 2024) aims to enhance offline RL (Levine et al., 2020; Prudencio et al., 2023) by learning an environmental dynamics model in which the policy could explore. Compared with model-free offline RI. (Kumar et al., 2020; Lyu et al., 2022; Jin al., 2021; Fujimoto & Gu, 2021; Wu et al., 2019: Qiao et al., 2026), model-based offline RL exhibits higher data efficiency and could generalize better to states not present in the dataset. However, model-base offline RL suffers from model exploitation (Kidambi et al., 2020), that is. the policy might explore the regions where the model cannot accurately predict the dynamics. Utilizing transitions from these inaceurate regions could degrade the algorithm's 's performance. Therefore, incorporating conservatism (a.k.a. pessimism) is necessary for model-based offline RL to mitigate model exploitation.
One mainstream approach is uncertainty estimation with various heuristics (Yu et al., 2020; Kidambi al., 2020; Sun et al., 2023; Kim & Oh, 2023). However, uncertainty estimation could be unreli- able for neural networks (Rigter et al., 2022; Yu et al., 2021; Lu et al., 2021). Adversarial model learning (Uchara & Sun, 2021) offers another theoretical framework mitigate model exploitation by solving maximin formulation, without uncertainty estimation. Among its practical implemen- tations (Rigter et al., 2022; Bhardwaj al., 2023; Yang et al., 2023), RAMBO (Rigter et al., 2022)
"Corresponding Author

---

Published as a conference paper at ICLR 2026
has emerged representative and widely adopted method, which (1) forces the dynamies model to reduce the value function in OOD regions by minimizing the adversarial loss with model gradient, (2) introduces a tradeoff coefficient À to the adversarial term to balance the adversarial training and maximum likelihood estimation (MLE). Although RAMBO provides a practical implementation for adversarial model learning, we observe that the tradeoff coefficient À is set to an extremely small value (c.g., 3e-4) in its implementation, minimizing the effect of the adversarial term. In our further investigation, we empirically find that a slightl larger yet still small À (e.g., 0.05, 0.1) could cause severe Q-value underestimation and even gradient explosion, resulting in training collapse. This indicates that RAMBO has difficulties in controlling the conservatism level and can be unstable during training.
To address these issues, we propose RObust value-aware Model learning with Implicitly differ- entiable adaptive weighting (ROMI). In contrast to RAMBO, ROMI discards the model gradient method that could be overly conservative and unstable, and introduces a novel robust value-aware model learning (Farahmand et al., 2017; Voelcker et al., 2022) approach to incorporate conservatism. Specifically, we require the dynamics model to predict future states with values close to the minimum Q-value within scale-adjus aajjusaajeteble state uncertainty set, and adjusting the scale of the uncertainty set is equivalent to controlling the degree of conservatism. This offers flexible and controllable conser- vatism adjustment and stabilizes model update. However, this value-aware learning scheme solely enforces value awareness and overlooks dynamics awareness that crucially benefits OOD general- ization. To close this gap, we further propose implicitly differentiable adaptive weighting, which achieves both dynamics and value awareness in a bi-level optimization framework. Concretely, the inner level incorporates an adaptive weighting network that re-weighs each training sample in an MLE loss, while the outer level updates the weighting network by minimizing the robust value-awan aware model loss with implicit differentiation. This hierarchical framework adaptively balances dynamics and value-aware model learning. Lastly, we empirically evaluate ROMI on various D4RL (Fu et al., 2020) and NeoRL (Qin et al., 2022) datasets. Our experimental results demonstrate that ROMI significantly surpasses RAMBO, matching surpassing oftee-of-theart (SOTA) methods such as MOBILE (Sun et al., 2023) and Count-MORL (Kim & Oh, 2023), especially on datasets where RAMBO typically underperforms.
2 BACKGROUND
RL. We consider a Markov Decision Process (MDP) (Puterman, 1990) that can be specified by a tu- plc AS T,, TT.TT. p. 7), where S and A are the state space and action space, respectively, T S Xx A + (SS) is the transition dynamics, and A--) is the probability simplex. r S x A + R is the reward function, P is the initial state distribution, and y E [0, 1) is the discount factor. The goal of RL is to obtain a policy Ta which maximizes the objective function: J(0)= = Ex, E 7 a:)l00 80 ~ p].
Model-based Offline RL. In offline RL, the agent only has access to pre-collected offline dataset: D ((s, dz Fa, 88+1 1 Model-based offline RL typically trains a dynamics model to augment the dataset. The policy could perform h-step rollouts within the dynamics model starting from states in D, and store the collected samples in D modell However, the dynamics model could deviate from the true dynamics. Therefore, conservatism is necessary prevent model exploitation.
Adversarial Model Learning for Model-based Offline RL. Uehara & Sun (2021) first proposes a theoretical framework that incorporates adversarial model learning into model-based offline RL to introduce conservatism. The offline RL problem is cast as a zero-sum game against an adversarial dynamics model within an uncertainty set, and thus the object is to find a policy from the policy class II by solving the maximin problem:
arg max min VT TEll TEME
s.t. Mg ( ( Dis(TMuB(()))) a), ((1s, a))l <t).
where Mg is the dynamics uncertainty set and Dis(.-) is the distributional discrepancy measure, TMLE represents the dataset dynamics learned with maximum likelihood estimation (MLE), i.c, TMLE = arg maxx (eelaog T(s' s, a)). While this offline RL framework offers strong theo- retical guarantees, its practical algorithm remains be explored. A widely-adopted implementation
2

---

Published as a conference paper at ICLR 2026
1000
hallchhefalh-medinn--xppet
300
0.0
tefhaltcheelll-mmedl-----s
A 0.1
Gradient: Norn
Q value
500
200
-1.0
100
-1.5
À - 0.1
-500
= 30 Se
1e-1
100
200
300
-2.
100
300
Step
Step
163
Figure 1: Left: Q-value estimation for different values of X. The slash-shaded bars indicate that the Q-value has diverged. Middle: Gradient norm of the adversarial loss recorded during training. Right: Q-value estimation recorded during training.
of this framework is RAMBO (Rigter al., 2022), which optimizes the following objective:
arg max min { AV (s(,,,,')) log [sssss,a)]:
f
where the first term is the adversarial loss that corresponds to minimizing the value in OOD regions, optimized via model gradient, and the second term relates minimizing the negative log-likelihood on the offline data, which is standard MLE. Here À is treated as a Lagrange multiplier but, for practicality, hand-tuned and fixed per task.
MOTIVATION EXAMPLE: THE LIMITATIONS OF RAMBO
In this section, we empirically examine and identify key limitations of RAMBO: the conservatism degree hard to control via X, and training can be unstable. Our key observation that the adversarial weighting coefficient À is set to 3e-4 in the original paper, which is an exceedingly small value for the weighting coefficient. This implics that the effect of the adversarial term should be minimized. We would like explore the effect of larger À on training.
We vary À across (0, 3e-4, Se-3, 5e-2, le-1}, and run RAMBO on D4RL datasets. We show the re- sults on halfcheet.e ah-medi amm-xpeet-v2 dataset in Figure 1. We find that setting À across (0, 3e-4,5e-3) does not show clear difference in Q-value estimation as illustrated in Figure Left, indicating that the adversarial training has little effect on conservatism with too small X. However, when increasing to larger but still small values within (5e-2, le-1), }. we observe that the Q-value diverges very negative value. This suggests that RAMBO highly sensitive to the choices of A. When À = 0.1, we even observe the gradient of the adversarial loss explodes, and the Q-value drops sharply, leading to training collapse, shown in Figure Middle, Right. This phenomenon implies that RAMBO tends to be overly conservative and unstable. More details and results on all 9 datasets are presented Appendix F.2.
We take deeper analysis of this phenomenon. On the one hand, À is the Lagrange multiplier and should be optimized with gradient ascent (Equation (7) Rigter al. (2022)). hard to precisely control the degree of conservatism by simply tuning X. On the other hand, model gradient would encourage the dynamics model to explore regions with a sharp decrease in value, which may lead to Q-value underestimation and gradient explosion. These factors force À to remain extremely small to minimize the impact of the adversarial term, thereby limiting RAMBO's potential. Therefore, a more practical implementation for Equation that allows flexible adjustment of the conservatism level and enhances training stability is highly desirable.
METHODOLOGY
Section 3 highlights two major limitations of RAMBO: (1) difficulties in controlling the degree of conservatism, resulting in severe value underestimation, i.e., over-conservatism: (2) gradient explo- sion, leading to unstable training. To address these issues, we develop our method ROMI within the framework of Equation 1. Our central contribution is novel treatment of adversarie model learn- ing, whose core components, robust value-aware model learning and implicitly differentiable adaptive weighting, are described in Section 4.1 and Section 4.2 respectively.
3

---

Published as a conference paper at ICLR 2026
4.1 ROBUST VALUE-AWARE MODEL LEARNING
Recall that the adversarial model learning paradigm in Equation seeks an adversarial dynam- ies model for any policy π via the inner constrained minimization problem minfeM; V with
Mg = ( (D [Dis(TMLL.e-ls,.. a), T(|ss, a))] S E). As stated previously in Equation 2, RAMBO explores an implementation of this inner constrained optimization by using a modified Lagrangian-multinn
approach, leading to min s') log T(slls, a)]} with hand- tuned multiplier A. This reformulation induces over-conservatism and unstable training, according to our analysis in Section 3.
To tackle such challenge, we propose an alternative implementation of Equation a robust value- aware model learning framework that incorporates conservatism, from a perspective of explicitly characterizing the one-step value estimation error induced by adversarial model learning. Letting Tm be the dynamics model be learned, we reformulate Equation 1 as follows:
π = arg max V TEll TT
S.t. V = min V
TEM;
We next focus on the condition V min EM T Equation 3, which motivates us to minimize a value-aware model error
E
defined in terms of the value difference as follows:
8T T = Pemppo())
V (8)
min
(i)(s)
Tm
TEMe
(4)
= z.- (egpp())) g (115,a)
V Pm
(5) min E (1. V
TEM;
We consider only pD(() here, since the rollout initial states drawn from the dataset D. Given that Me is only defined on D and a ~ (-|ss) contains OOD regions, directly minimizing E7 could lead to arbitrarily low Q-value estimates in OOD regions. RAMBO's over-conservatism arises for essentially the same reason. To resolve this issue, we propose minimize E as an alternative:
cn m == Y (7..e-pp())a-vle B Ba (1))))
(s)) min E. (.-((- (T))
TEMe
The blue parts highlight the difference between Ex and E where su(1(s) is the behavior policy generating D, and V= (abbreviated as V hereafter) represents the value function during training. Intuitively, we learn a mildly conservative value function the in-distribution (ID) regions, and the conservatism in OOD regions relies on generalization, which we analyse in Proposition 4.2.
The main challenge of minimizing Eq. lies in computing minfeM, B.r-T(ls..) (s'), (s'). since are not accessible to the dynamics uncertainty set Mr. which is crucial for controlling the level of conservatism. To address this, take the distance measure Dis to be the Wasserstein distance and define the Wasserstein dynamics uncertainty set in Definition 4.1, and exploit its dual form in Proposition 4,1, which provides new insights for handling this issue.
Definition 4.1 (Wasserstein Dynamics Uncertainty Set). We define the dynamics uncertainty set via the Wasserstein distance as
Me = { ED [w (TTLEE(-1s a), (-11s,))) [ $E).
E B(ss,gg)<cld(s s/)] is the Wasserstein distance between dynamics TALEE and T, I(..)) the set of joint distributions whose marginals are PALLE and T, and d(...) can be chosen as the Euclidean distance.
where W (TALE(-1s, a), f(-1s,0)) = inf (fea(fana(1saa)
Proposition 4.1 (Dual Reformulation). Let Me be the Wasserstein dynamics uncertainty set defined in Definition 4.1. Then we have
min E TEMs
EyTT a.0) V(s') Ee-Tarnttt a.))
min V(3) BEUe(o")
where Ug(s')= {8 d(s,, 8') E is the state uncertainty
4

---

Published as a conference paper at ICLR 2026
Proposition 4.1 states that, under the Wasserstein distance, the dynamics uncertainty set could be transformed into the state uncertainty Based on this reformulation, rewrite & T
T =Y Berpp())errrlll (. V(s)) E,
& af min BEUg(o")
(5)
To practically minimize Ex we further consider empirical version of E based on the sample 7m
from the offline data, i.e., (s, a, s') ) ~ D, in place of s pD(+), (((-ss), s' MLLE(-Is. a). since D was collected through the same sampling process. Thus, given that 8' is present in D, we can simply construct Us()") by perturbing 8 with noise of scale & and performing random sampling. We further parameterize the dynamics model F. by th, denoted by T. Eventually, Equation 5 is reformulated as a Robust Value-aware model Learning (RVL) loss:
(1)).(1)= = ((E [y (11s1a)))))) V (5')
min V(3))² (
(6)
where V is chosen the target value network for training stability, { are sampled from Ug(s'), and N is the sample size. We note that Equation 6 differs from the typical value-aware model loss (Farahmand et al., 2017; Voelcker al., 2022). Our approach contributes to incorpo- rating conservatism via a state uncertainty set into the loss for offline RL. Specifically, Equation requires the dynamics model to predict the future state whose value is close to the minimum value for states within the state uncertainty set. Moreover, the scale & of the uncertainty set precisely quantifies the level of conservatism according to Proposition 4. Thus, we could simply adjust the uncertainty set scale & to achieve controllable conservatism and enhance training stability.
Note that Equation only utilizes the dataset samples, and does not explicitly regularize the values for the OOD region that may be encountered during policy rollouts. To address this concern, we first analyze the impact of potential generalization errors on training. We define the maximum generalization eтror of Equation 6 during rollouts as
-B -
min F(3)]}- (lUt(())
61 =
max [s.0]EDmoce madell
(-1x.))
where Dmooeel denotes the dataset consisting of synthetic transition trajectories generated via pol- icy rollouts, within which OOD data can reside. We also define the maximum value discrepancy between true next states and those within the state uncertainty set Ug as
max
Ex-T((19.))
((s')
- ( ('U((()))
min
V(33]]-
(8.0) modell
Then the learned Q-function by minimizing Equation maintains a bounded value during policy rollouts, guaranteed by Proposition 4.2.
Proposition 4.2 (Bounded Q-value). The learned Q-function Q satisfies:
Queue (s, a) 7(1 +
< (((sa)) S Qs
761
V(s,a) Dmodel-
1 7
1-7"
where Qteue denotes the true Q-function.
Proposition 4.2 indicates that the learned Q-value would stay bounded el and e2 are bounded. e2 monotonically increasing w.r.t. ç, indicating that we can bound cg by adjusting We demonstrate how we minimize the impact of 61 in Section 4.2.
4.2 IMPLICITLY DIFFERNNIABLE ADAPTIVE WEIGHTING
We note that the generalization emror e could be large during multi-step rollouts, since LRVI. only focuses value awareness while ignoring dynamics awareness, which could be crucial for OOD generalization. To demonstrate this, let 8p be the rollout starting point, and ag ~ x-1[sp) be the policy output. The dynamics model needs to predict the next state 81. Since LRVL only regularizes the value of $1, the predicted 81 may deviate significantly from the true dynamics transition. In subsequent steps, the policy will be forced explore OOD regions, and the dynamics model will
5

---

Published as a conference paper at ICLR 2026
also need to predict transitions in these OOD regions. This may lead to compounding prediction errors and a large é1. Therefore, how can we achieve value awareness to ensure conservatism while improving dynamies awareness for better generalization?
Inspired by previous works (Shu et al., 2019; Jiang et al., 2022; Yuan et al., 2023), we propose to optimize Equation 6 with implicitly differentiable adaptive weighting, which integrates dynamics- and value- aware model learning in a bi-level optimization framework. Specifically, we introduce an weighhie-weighting network (s, a, s') to assign an individual weight to each transition (s, a, s') and optimize the following objective:
min CAVV((V((V))
(8)
s.t. ()( - arg min (V, v) ((,,,,))) [ a, 8') log (o(6//s,a))]]
where LWSL denotes the Weighted Supervised Learning (WSL) loss, compared to the standard su- pervised learning loss Lst. := =)),2 log (Tos (5'(s, (s'(s, a))] Specifically, we update the dynamics model through weighted supervised learning in the inner level to achieve dynamics awareness, and optimize wy in the outer level by minimizing LRVL with implicit differentiation to achieve value awareness.
Inner Level: Optimizing de to Improve Dynamics Awareness. In the inner level, we fix Wy and update the dynamics model by minimizing LWSL with typical gradient descent:
(5(++1) +- wy(t) - Bi(e) &e ((()),
where Bi(0) represents the 1-step learning rate of the dynamics model.
Outer Level: Optimizing y to Achieve Value Awareness. The outer level objective is formalized in Equation 6. Since have established the analytical relationship between t/ and through the inner level, we can implicitly calculate the gradient of LRVL (15()) w.r.t. v using the chain rule:
JRVI. V (1 (T(())T (((0))) (uns(Ct))()(1)) = h Vewaco (s,
(t)
V log (Pero (d'1s, a)) The derivation is deferred in Appendix A. Then we can update using automatic differentiation in Pytorch (Paszke al., 2019): ((+1) t vl0) S2(1) IRVL.
where h = - V LRVI. (1bltt)))))))))))) (v(())) (V
(8)
where (2(1) is the t-step learning rate of the weighting network. Through this bi-level optimiza- tion process, the weighting network learns to prioritize samples that contribute most to mini- mizing LRVL. and the dynamics model still learns to reconstruct the environmental dynamics, achieving both dynamics and value awareness. Assume (1) V is Lv-Lipschitz continuous, that is, (sz)) < Lv|ls1 -s2l| holds for Vsi, 52 E S; (2) the dynamics model prediction error dur- ing rollouts is bounded by a constant e: max(s. Egr,T, (3,ss,(()ss.) (-(ss.))ss' ('...a).s'~aaaaa((( (a.LBaL || - s'Il S e, which could be satisfied through bi-level optimization and short-branch rollouts (Janner et al., 2019). Then the generalization error el could be bounded according to Corollary 4.1.
Corollary 4.1 (Bounded e1). Under the assumptions of Lipschitz continuity and bounded dynamics model prediction error, the generalization error E1 could be bounded by
6i S Ly (c+ E),
where & is the uncertainty set scale. We also provide the convergence rate analysis for this bi-level optimization framework in Proposition 4.3.
Proposition 4.3 (Convergence Analysis), Assume LRVL and LWSL are L-Lipschitz smooth, and the gradient of LRVL, LSI and LWSL are bounded by p. Let wy be twice differentiable, with its gradient and Hessian bounded by and B, respectively. Denote T'' T== p + +B), and the total steps as K. For some C1, C2 > 0, we assume the learning rate of the inner and outer loop Br and B2 satisfies: e2 VK S 331,82 S min t. b. where VK S t.
C2
=2 1, + < 1. Then the outer
VK
loss LRVL, and inner loss Lwst, can achieve a convergence rate of 0
i.e.
VK
min oStSK-1
E
50
VK
min 0S15K-1
E. V
00
VK

---

Published as a conference paper at ICLR 2026
Table I: Normalized score comparison between different offline RL algorithms and ROMI on 12 D4RL MuloCo datasets. We abbreviate "random" as "r", "medium" as "m", "medium-replay" as "mr", "medium-exper"" as "me". Each method rum with 1M steps.
Task Name
CQL
IQL
MОPO
RAMBO
COUNT
MOBILE
ROMI (Ours)
hopper-r
7.9
7.6
27.8
23.9
30.0
28.5
31.7 + 0.3
halfcheetah-r
17.5
13.1
36.7
29.5
37.3
39.4
38.41:.11
walker2d-r
5.1
5.4
0.0
0.0
19.9
23.6
22.4 #:2.6
hopper-m
53.0
66.2
72.2
95.2
100.2
77.6
105.0:13.3
halfcheetah-m
47.0
47.9
64.0
76.2
75.9
74.6
76.2:1:0.6
walker2d-m
73.4
67.2
82.1
83.9
88.4
90.7
93.6:11.3
hopper-mr
88.7
94.0
90.5
77.2
97.6
85.7
102.0+-2.4
haffcheetah-mr
45.3
44.3
52.1
69.7
70.5
67.7
70.7++1.8
walker2d-mr
81.8
74.8
73.6
81.2
87.5
75.0
85.21:..6
hopper-me
105.9
92.7
84.5
99.7
109.8
95.6
110.5:+6.8
halfcheetah-me
75.6
86.3
92.3
93.9
100.0
96.1
104.5 +2.2
walker2d-me
107.9
110.7
93.8
73.7
110.4
103.2
113.3
Total
709.1
710.2
769.6
804.1
927.5
857.7
953.5
ROMI Algorithm. Eventually, combining the dynamics model learning in Equation 8 and policy learning π arg max F. yields the optimization problem for ROMI. The key distinction
from RAMBO lies in the training strategy of the dynamics model, as described in Section 4.1 and Section 4.2. The policy is learned with SAC (Haarnoja et al., 2018) by using the data from D and Dmodel More implementation details of ROMI are presented in Appendix D.2.
EXPERIMENTS
In this section, we empirically evaluate ROMI. We compare the performance between ROMI and other baselines on standard benchmarks in Section 5.1. We conduct an ablation study to examine the effect of adaptive weighting in Section 5.2, and we verify that ROMI offers flexible conservatism control and avoids gradient explosion across a wide range of & in Section 5.3.
5.1 BENCHMARK RESULTS
D4RL. We compare ROMI with several offline RL algorithms, including model-free methods CQL (Kumar et al., 2020) and IQL (Kostrikov et al., 2021); and model-based methods: MOPO (Yu et al., 2020), RAMBO (Rigter et al., 2022), Count-MORL (Kim & Oh, 2023), and MOBILE (Sun et al., 2023).
For the MuloCo (Todorov et al., 2012) domain, we consider three tasks (hal fcheetah, hopper, walker2d) with four types of datasets (random, medium, medium-repias medium-experts for each task. The datasets we use are of the "v2" version. Experiments are with 5 seeds. More evaluations Antmaze tasks are deferred to Appendix F.I due to space limitations.
In Table 1, summarize the evaluation results of ROMI and the compared baselines. We that ROMI outperforms RAMBO on 11 out of 12 datasets, especially on datasets where RAMBO underperforms, such as hopper-medi cem-repl ay and wa 2er2dermme wt--mmiuu-expppeeeee um-expert. ROMI achieves a total score of 953.5, which is 18.6% higher than that of RAMBO. Furthermore, even when compared other algorithms, including MOBILE and Count-MORL, ROMI demonstrates superior performance, achieving (one of) the best results on 11 out of 12 datasets. On the remaining datasets, ROMI also delivers performance comparable the best persoreefooming methods.
33 steps used in MOBILE paper, and only report the results 1M steps for fair comparison.

---

Published as a conference paper at ICLR 2026
Table 2: Normalized score comparison between different offline RL algorithms and our method ROMI on NeoRL tasks. Each method is run with 1M steps over 5 seeds.
Task Name
CQL
IQL
EDAC
MОPO
RAMBO
MOBILE
ROMI (Ours)
Hopper-L.
16.8
14.3
12.2
12.0
15.1
18.1
0.5
HalfCheetah-L.
35.0
34.1
31.3
34.2
30.1
38.9
35.8+33.9
Walker2d-L
42.4
40.1
43.1
15.6
24.3
39.0
36.4
Hopper-M
58.1
56.8
37.2
0.0
37.5
47.3
46.6 6.8
HalfCheetah-M
50.0
52.0
54.6
53.1
51.2
55.6
1..7
Walker2d-M
52.5
51.2
53.9
35.7
36.2
51.4
54.9:12.0
Hopper-H
65.6
62.4
56.7
17.9
48.7
65.0
65.9114.5
HalfCheetah-H
74.6
76.3
72.0
63.2
72.3
69.4
7.74+2.7
Walker2d-H
71.3
74.0
73.1
26.4
67.4
71.7
75.1++1.8
Total
466.3
461.2
434.1
258.1
382.8
456.4
472.2
adaptive weighting
adaptive weighting hopper-medum-replayy
halfcheetah-mmdiiu--eedaa
halfchhestah-meduu--rppaa
hopper-medium-replay
Prediction Eror
Prediction Error
0.8
Nomalzze Score
100
Nomalized
0.6
50
20
25
Step
Step
Step
Stop
1e6
(a)
(b)
(c)
(d)
Figure 2: Comparison the prediction error and performance and w/o adaptive weighting.
NeoRL. We also test our method on NeoRL (Qin al., 2022) benchmark, which is often used to evaluate model-based offline RL (Sun et al., 2023; Lin et al., 2024). We consider nine datasets, which involves three tasks (half cheetah, hopper, wa ker2d) and three types of data qualities (low, medi um, high). For each dataset, we select 1,000 trajectories uniformly for our experiments. We choose model-free methods (CQL, IQL, EDAC (An et al., 2021)) and model- base- methods (MOPO, RAMBO, MOBILE) compared basclines. We exclude Count-MORL since there no reference hyperparameter settings available the original paper or the NeoRL. paper.
We summarize the normalized methods outperform model-based baselines on most datasets, while ROMI surpasses all baselines on 6 out of 9 datasets, and also achieves the highest total score. Notably, ROMI outperforms RAMBO on all 9 datasets, verifying the effectiveness of ROMI.
comparison results in Table 2. We observe that model-free
5.2 STUDY THE EFFECT OF DYNAMICS AWARENESS
In this section, we conduct an ablation study to test the effect of dynamics awareness on the performance and OOD generalization. Specifically, we solely apply Equation for model updates and drop the adaptive weighting component in Section 4.2 to examine its effect on the performance and prediction emror during multi-step rollouts. We conduct experiments on wel.fchhetth-mmeiinnee.
and wwwpreer-mdiuu--eeeee
datasets.
We present the experimental results Figure 2, where panels (a) and (c) show the model predic- tion error under different rollout steps with and without adaptive weighting. while panels (b) and (d) display the corresponding learning curves. The results clearly demonstrate that adaptive weight- ing leads to improved performance and reduced prediction enror across various rollout steps. This indicates that the dynamics awareness is crucial for OOD generalization and achieving strong per- formance.
8

---

Published as a conference paper at ICLR 2026
{ 0.01
E- 0.1
{-1
{
half-me
half-me
half-me
half-me
Gradient Norm
0.07
Outer Gradient Norm
750
0.02
Value
500
0.06
50
0.01
250
0.05
0.00
1.0
0.0
1.0
Step
1e6
Step
Step
Step
1e6
(a)
(b)
(c)
(d)
Figure 3: Comparison for different & in terms of Q-value, normalized score, inner loss gradient norm, outer loss gradient norm.
5.3 STUDY ON THE EFFECT OF
In this section, we examine the effect of the key hyperparameter: the uncertainty set scale &, on training. Specifically, we conduct experiments on halfcheetah meii-hheeea--edium medium expert-t v2 dataset and vary the value of across wide range of values: {0.01, 0.1, 1.0, 10). A larger value of cor- responds to a higher degree of conservatism. We record four metrics during training: the estimated Q-value, normalized score, gradient norm of the inner loss, and gradient norm of the outer loss.
The experimental results are presented in Figure 3. Compared to the results of RAMBO in Section 3, observe that: (1) neither severe Q-value underestimation
training collapse occurs even for
& = 10, and the gradient norm stays small throughout training: (2) the Q-values are clearly distin- guishable for different &, with larger & corresponding to a lower Q-value. These verify that ROMI could provide a more controllable degree of conservatism while ensuring training stability.
6 RELATED WORK
Model-based Offline RL. The core issue of model-based offline RL how to incorporate conser- vatism into model learning or policy update to prevent model exploitation. One line of work is to adopt uncertainty estimation. MOPO (Yu et al., 2020) and MoRel. (Kidambi et al., 2020) use the uncertainty of model prediction to penalize the reward function to achieve pessimistic value esti- mation; SUMO (Qiao et al., 2025) utilizes KNN-based entropy estimator for uncertainty quantifi- cation; Count-MORL (Kim & Oh, 2023) utilizes the count estimates of state-action pairs to measure the uncertainty; MOBILE (Sun et al., 2023) proposes
-edel-Bellman inconsistency to improve the uncertainty quantification. COMBO (Yu et al., 2021) and RAMBO (Rigter et al., 2022) do not rely on uncertainty estimation and minimize the value estimation on model-generated samples.
Robust RI. Similar RAMBO, ROMI is related to Robust RL. Robust RL optimizes the policy in a Robust MDP (Bagnell et al., 2001; Cohen & Hutter, 2020; Nilim & El Ghaoui, 2005; Wiesemann al., 2013; Rigter et al., 2021), that is, finding the policy with the best worst-case performance over a set of possible MDPs. Typically, the uncertainty set of MDPs is specified a priori. In contrast, our problem setting differs in that the uncertainty set of MDPs assumed be unknown. We need to leamn dynamics model to simulate the worst-case MDP.
Bi-level Optimization in RL. Bi-level optimization is classic framework for optimizing hierarchi- cal learning objectives. There are several studies leveraging bi-level optimization to solve different RL tasks. CAIL (Zhang et al., 2021) re- re-weights demonstrations with different optimality in imita- tion learning; Meta-Rewant Sett-Reeardd-eer (Liu et al., 2022) incorporates the performance of the Q function into reward learning for preference-based RL; VACO (Jiang et al., 2025) incorporates bi-level opti- mization in offline RL to balance policy improvement and behavior cloning; TEMPO (Yuan et al., 2023) is closest to our work, which leverages both task-specific and semantic information for effec- tive world model learning. But our method is different since focus on offline RL and have to consider conservatism in the outer level. There are also recent studies (Li et al., 2024; Chen et al., 2025) incorporating bi-level optimization into Large Language Model -oriened RI.

---

Published as a conference paper at ICLR 2026
7 CONCLUSION
In this paper, we first identify the limitations of RAMBO. We empirically find that RAMBO tends be overly conservative and incurs instability during training. To address these issues, we propose our method ROMI, which adopts novel robust value-aware model learning scheme by requiring the dynamics model predict future states with values close the minimum value in the state uncertainty set. To further improve OOD generalization ability, we introduce a bi-level optimization framework called implicitly differentiable adaptive weighting. which achieves dynamics and value awareness. Compared to RAMBO, ROMI provides more controllable conservatism and stabilizes training. Empirical Results on D4RL and NcoRL datasets demonstrate the effectiveness of ROMI.
Limitations. One limitation of ROMI is that additional computational cost for bi-level optimiza- tion required compared to the original RAMBO. Another limitation is that the conservatism de- gree governed by has to be specified before training, unlike methods that support runtime adjust- ment (Ghosh et al., 2022; Hong et al., 2022; Swazinna et al., 2022; Wang et al., 2023). Addressing these limitations presents interesting direction for future work.
ACKNOWLEDGEMENTS
This research supported in part by the Hong Kong Research Grants Council (GRF 11217925, GRF 16209124) and the National Science Foundation of China (Grant 72371214). The authors would also like to thank the anonymous reviewers for their valuable comments on our manuscript.
REFERENCES
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-based offline re- inforcement learning with diversified q-ensemble. Advances in neural information processing systems, 34:7436-7444 7447, 2021.
J Andrew Bagnell, Andrew Ng, and Jeff G Schneider. Solving uncertain markov decision pro- cesses. 2001.
Atilim Gunes Baydin, Robert Cornish, David Martinez Rubio, Mark Schmidt, and Frank Wood. Online learning rate adaptation with hypergradient descent. arXiv preprint arXiv:170.004782, 2017.
Mohak Bhardwaj. Tengyang Xie, Byron Boots, Nan Jiang, and Ching-An Cheng. Adversarial model for offline reinforcement learning. Advances in Neural Information Processing Systems, 36: 1245- 1023.
Jérôme Bolte. Edouard Pauwels, and Samuel Vaiter. One-step differentiation of iterative algorithms. Advances in Neural Information Processing Systems, 36.77089-77103, 2023.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv:1606.01540. 2016.
Liang Chen, Xueting Han, Li Shen, Jing Bai, and Kam-Fai Wong. Beyond two-stage training: Cooperative sft and for Ilm reasoning. arXiv preprint arXiv:2509.06948, 2025.
Michael K Cohen and Marcus Hutter. Pessimism about unknown unknowns inspires conservatism. In Conference on learning theory, pp. 1344-1373. PMLR, 2020.
Amir-massoud Farahmand, Andre Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. In Artificial Intelligence and Statistics, pp. 1494. PMLR, 2017.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach to offline reinforcement learning. Advances in neural information processing systems, 34:20132-20145, 2021.
10

---

Published as a conference paper at ICLR 2026
Dibya Ghosh, Anurag Ajay, Pulkit Agrawal, and Sergey Levine. Offline rl policies should be trained to be adaptive. In International Conference on Machine Learning, Pp- 7513-7530. PMLR, 2022.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off Off-policc maximum entropy deep reinforcement learning with stochastic actor. In International confer- ence on machine learning, pp. 1861-1870. Pmir, 2018.
Joey Hong, Aviral Kumar, and Sergey Levine. Confidence-conditionee value functions for offline reinforcement learning. arXiv preprint arXiv:2212.0460) 2022.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in neural information processing systems, 32, 2019.
Xingyu Jiang, Hongkun Dou, Chengwei Fu, Bingquan Dai, Tianrun Xu, and Yue Deng. Boosting su- pervised dehazing methods via bi-level patch reweighting. European Conference on Computer Vision, pp. 57-73. Springer, 2022.
Xingyu Jiang, Ning Gao, Xiuhui Zhang, Hongkun Dou, and Yue Deng. Value-aligned behavior cloning for offline reinforcement learning via bi-level optimization. In The Thirteenth Interna- tional Conference on Learning Representations, 2025.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline ri? In International conference on machine learning, pp. 5084-5096. PMLR, 2021.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. Advances in neural information processing systems, 33: 21810-21823, 2020.
Byeongchan Kim and Min-hwan Oh. Model-based offline reinforcement learning with count-based conservatism. International Conference on Machine Learning, pp. 16728-16746. PMLR, 2023.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization, arXiv preprint arXiv:1412.6980, 2014.
Ilya Kostrikov, Ashvin Nair, and Sergey Levine. Offline reinforcement learning with implicit q- learning. arXiv preprint 1 2021.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. reinforcement learning. Advances neural information processing systems, 33:1 18:1179-110 11779-99991 1191, 2020.
Conservative q-learning for offline
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tuto- rial, review, and perspectives on open problems. arXiv preprint arXiv:2005.016443 2020.
Jiaxiang Li, Siliang Zeng, Hoi-To Wai, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Getting more juice out of the sft data: Reward learning from human demonstration improves sft for Ilm alignment. Advances in Neural Information Processing Systems, 37:124292-124311 2024.
Haoxin Lin, Yu- Yan Xu, Yihao Sun, Zhilong Zhang, Yi-Chen Li, Chengxing Jia, Junyin Ye, Jiaji Zhang, and Yang Yu. Any-step dynamics model improves future predictions for online and offline reinforcement learning. In The Thirteenth International Conference Learning Representations.
Haoxin Lin, Yu- Yan Xu. Yihao Sun, Zhilong Zhang, Yi-Chen Li, Chengxing Jia, Junyin Ye, Jiaji Zhang, and Yang Yu. Any-step dynamics model improves future predictions for online and offline reinforcement learning. arXiv preprint arXiv:2405. 17031, 2024.
Risheng Liu, Pan Mu, Xiaoming Yuan, Shangzhi Zeng, and Jin Zhang. A generic first-order al- gorithmic framework for bi-level programming beyond lower-level singleton. In International conference on machine learning, pp. 2305-63155 PMLR, 2020.
Runze Liu, Fengshuo Bai, Yali Du, and Yaodong Yang. Meta-reward-nett Implicitly differentiable reward learning for preference-basee reinforcement learning. Advances in Neural Information Processing Systems, 35:22270-22284, 2022.
11

---

Published as a conference paper at ICLR 2026
Cong Lu, Philip Ball, Jack Polder- Michael A Osborne, and Stephen J Roberts. Revisiting design choices in offline model-base reinforcement learning. arXiv preprint arXiv:2110.04135, 2021.
Jiafei Lyu, Xiaoteng Ma, Xiu Li, and Zongqing Lu. Mildly conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 35:1711-1724, 2022.
Yurii Nesterov et al. Lectures on convex optimization, volume 137. Springer, 2018.
Arnab Nilim and Laurent El Ghaoui. Robust control of markov decision processes with uncertain transition matrices. Operations Research, 53(5):780-798, 2005.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, high- performance deep learning library. Advances in neural information processing systems, 32, 2019.
Rafael Figueiredo Prudencio, Marcos ROA Maximo, and Esther Luna Colombini. A survey on offline reinforcement learning: Taxonomy, review, and open problems. IEEE Transactions on Neural Networks and Learning Systems, 35(3)110237-10257 2023.
Martin L Puterman. Markov decision processes. Handbooks in operations research and management science, 2:331-434, 1990.
Zhongjian Qiao, Jiafei Lyu, and Xiu Li. Mind the model, not the agent: The primacy bias in model- based rl. In ECA/ 2024, pp. 1824-18311 IOS Press, 2024.
Zhongjian Qiao, Jiafei Lyu, Kechen Jiao, Qi Liu, and Xiu Li. Sumo: Search-based uncertainty es- timation for model-based offline reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 39, pp. 20033-20041, 2025.
Zhongjaan Qiao, Rui Yang. Jiafei Lyu, Xiu Li, Zhongxiang Dai, Zhuoran Yang, Siyang Gao, and Shuang Qiu. Dual-robust cross-domain offline reinforcement learning against dynamics shifts. In The Fourteenth International Conference on Learning Representations, 2026.
Rong-Jun Qin, Xingyuan Zhang, Songyi Gao, Xiong-Hui Chen, Zewen Li, Weinan Zhang, and Yang Yu. Neorl: A near real-world benchmark for offline reinforcement learning. Advances in Neural Information Processing Systems, 22244322277755 2022.
Mare Rigter, Bruno Lacerda, and Nick Hawes. Minimax regret optimisation for robust planning in uncertain markov decision processes. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 11930-11938, 2021.
Mare Rigter, Bruno Lacerda, and Nick Hawes. Rambo-rl: Robust adversarial model-based offline reinforcement learning. Advances in neural information processing systems, 35:16082-16097, 2022.
Jun Shu, Qi Xie, Lixuan Yi, Qian Zhao, Sanping Zhou, Zongben Xu, and Deyu Meng. Meta- weight-net: Learning explicit mapping for sample weighting. Advances in neural information processing systems, 32, 2019.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-bellman inconsistency for model-based offline reinforcement learning. In International Conference on Machine Learning, pp. 33177-33194, PMLR, 2023.
Phillip Swazinna, Steffen Udluft, and Thomas Runkler. User-interactive offline reinforcement learn- ing. arXiv preprint X205 10629, 2022.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference intelligent robots and systems, pp. 5026-5033. IEEE, 2012.
Masatoshi Uehara and Wen Sun. Pessimistic model-based offline reinforcement learning under partial coverage. arXiv preprint arXiv:2107.06226, 2021.
12

---

Published as a conference paper at ICLR 2026
Claas Voeleker, Victor Liao, Animesh Garg, and Amir-massout Farahmand. Value gradient weighted model-based reinforcement learning. arXiv preprint arXiv:2204.01664 2022.
Shenzhi Wang, Qisen Yang, Jiawei Gao, Matthieu Lin, Hao Chen, Liwei Wu, Ning Jia, Shiji Song, and Gao Huang. Train once, get family: State-adaptive balances for offline-to-online reinforce- ment learning. Advances in Neural Information Processing Systems, 36:47081
2023.
Xinyi Wang, Hieu Pham, Paul Michel, Antonios Anastasopoulos, Jaime Carbonell, and Graham Neubig. Optimizing data usage via differentiable rewards. In International Conference on Ma- chine Learning, pp. 9983-9995. PMLR, 2020.
Wolfram Wiesemann, Daniel Kuhn, and Benç Rustem. Robust markov decision processes. Mathe- matics of Operations Research, 38(1):153-183, 2013.
Yifan Wu, George Tucker, and Ofir Nachum. Behavior regularized offline reinforcement learning. arXiv preprint arXiv:1911.11300 2019.
Junming Yang, Xingguo Chen, Shengyuan Wang, and Bolei Zhang. Model-based offline policy optimization with adversarial network. arXiv preprint arXiv:2309.02157, 2023.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Mopo: Kodell-ased offiin policy optimization. Advances in Neural Information Processing Systems, 33:14129- 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. Combo: Conservative offline model-based policy optimization. Advances in neural information processing systems, 34:28954 2296777,2 2021.
Huining Yuan, Hongkun Dou, Xingyu Jiang, and Yue Deng. Task-aware world model learning with meta weighting via bi-level optimization. Advances in Neural Information Processing Systems, 36:54167-54186, 2023.
Songyuan Zhang, Zhangjie Cao, Dorsa Sadigh, and Yanan Sui. Confidence-aware imitation learn- ing from demonstrations with varying optimality. Advances in Neural Information Processing Systems, 34: 13340-12355, 2021.
13

---

Published as a conference paper at ICLR 2026
A DERIVATION OF THE OUTER LOSS GRADIENT
In Section 4.2, we calculate the gradient of the outer loss at step t by:
(e) JRVL.
V VRLRVI (yptrv (epp(ct)(((1))
= h (s,,,,
where
= - V & ( Vp log (Taco(s'ls (8'(1,a))
(10)
We provide the complete derivation in this part. We denote the initial value of vp and y as ei(() and (0), and make the following stipulations for the outer loop and inner loop updates: (1) at each update step i, the inner loop is updated first, followed by the outer loop; (2) at each update step only one step of gradient descent is performed. Under above conditions, at step I, the inner loop first takes (t) and y() as imput and performs one-step update to obtain ((++1), Then the outer loop takes (i(+++) as input, which a function of p(t), and performs implicit update to obtain ((++), Then we can derive the gradient as follows:
(e) IRVI.
V.LRVI (1bli (1blt+ (v(o))
CRVL (w(O))) (
VeLLLL (w V. (u Bi(o) (.wsL((D0,)
Bi(0) V V CCRRT, ( V. V. (VoLwsn(1bl), vl11))
-Bi0 CRLL ( V log (s)o( V0
The fourth equality holds under the Markov assumption, under which V Ve (x(*)) = 0 since ((())
is determined by (((-)) and the effect of y()) is ignored. This technique is widely used in the field of bi-level optimization (Bolte al., 2023; Liu et al., 2020) and meta-learning (Baydin al., 2017; Wang et al., 2020). Then we finish our derivation.
B MISSING PROOFS
In this section, we provide detailed proofs for our theoretical results in the main text.
B. PROOF OF PROPOSITION 4.1
We first introduce the following lemma in preparation for proving Proposition 4.1.
Lemma B.1. Let S be a measure space and P be a probability measure on S. In addition, let
S be any measure function and c S + Rzo be a cost function. Then for any scalar
A 2 0, the following equality holds:
(Be))(((3) AW(P,P)) ESNr inf (J($) Ac(s,8))
inf (NP(S)
BNS
where P(S) represents all probability measures on S, and W is the Wasserstein distance w.r.. the cost function c.
Proof. We prove this lemma by performing series transformations on the left-hand side (LHS) and proving its equivalence to the right-hand side (RHS).
According to the definition Wasserstein distance, the LHS can be written as
LHS = inf PEP(S)
[x
+X inf E(s,s),7 [c(s, PET(P.P)
14

---

Published as a conference paper at ICLR 2026
The optimization of P and the inner optimization over y [(P, P) could be combined into single optimization over all couplings y whose first marginal is P, and the second marginal, P, could be arbitrary in P(S). Then we have
LHS
inf (Bg)e(((8) (B) ((0,,) 6)])
7Er(P,P)
(11)
inf (,,,) [f($) + Ac(s, 8)]
PET(P,P)
where the second equality holds by the linearity of expectation.
By the disintegration theorem for measures, any coupling E T(P, P) could be represented as the product of its first marginal P and a stochastic kernel K (dd]ss) + P(S):
(ds, d8) = K (d))s))((dss
(12)
This implies that optimizing over all couplings E I(P, P) is equivalent to optimizing over all possible stochastic kernels K. We substitute Equation 12 into Equation 11:
LHS inf / (S(8) + Ac(s, 8)] ((d8)s) K
inf R Egn [Ban (f()) + Ac(s, 8)]]
K
We change the position of the infimum operator and the inner expectation:
LHS = Eg-P
inf [s~K(1))) [f()) + rc(s, 8)]
(13)
(()))))(())
We then solve the inner minimization problem for fixed S:
inf K ((-))))((S)
((>)((1)) (f) + 7c(s, 8)]
8). The problem is to find a probability measure ((+ls) that minimizes Let (8) & f()) + 7c(s, (,,, the expectation of (5). is obvious that this minimum achieved by concentrating the entire probability mass point 8 where g5) attains its infimum.
Let 8* = arg infses (8). The optimal measure is Dirac measure dz. centered on 8*, Therefore, inf [ [/($) + 7c(8, 8)] {({(s))((S)
= Bind, [S)]
(14)
f(8²) + (c(s,
inf [f($) + 7c(s, 8)] RES
Finally, substituting Equation 14 back into Equation 13, we obtain the RHS of the lemma
LHS Exur inf (f($) + vc(s, BES
RHS.
This concludes the proof.
Now we are ready to give our formal proof for Proposition 4.1. We restate follows. Proposition B.1 (Proposition 4.1). Let Me be the Wasserstein dynamics uncertainty set. Then
have
min E TEMe
ByNT s.a) V())) = E Ban.Baan (..))
min
V(8)
(15)
(EUe((")
Proof. We start with the LHS and prove its equivalence to the RHS.
The LHS is a constrained optimization problem that can be solved using the Lagrange multiplier method. We define the Lagrange function as
L(T, >) = Ey D(8)) (WCT, PAA..F) - -6).
15

---

Published as a conference paper at ICLR 2026
The LHS is equivalent to solving the dual problem
LHS = sup inf L(T, X)-
(16)
X
For a fixed A, we solve the inner minimization problem inf L(T A) as
inf C(T, A) inf (E3)))) D(8)= + AW(T, TMLEE) - Xe.
(17)
According to Lemma B. we have
inf (Ex-F(1n.)) (55) + AW(T, TuLE)) - inf (V(8)) + )c(6
(18)
T
Substituting Equation 18 into Equation 17 and Equation 16, we obtain the dual reformulation of LHS
LHS
sup
E T inf ( X20
XE
(19)
The RHS in Equation 19 is exactly the Lagrange dual reformulation of the RHS in Equation 15. This implies Equation 15 holds, which concludes the proof.
B.2 PROOF OF PROPOSITION 4.2
Proof. The standard Bellman operator is defined
TQ(s, = r + 7B, - Q(s') aa).
For Bellman updates using Dmodll the Bellman operator could be expressed as Tmodel Q a) = (1a a's
For any (8.a) E D model, we have
min Eg Eg Q('', S E BanBaaa) (eU Q(3',a)) +
I By B (s' a') + 61.
Therefore, we have
Q < a))
V(s a) E Dmodel
Let @* denote the valuuu at iteration & by applying Tmodel. and Q² denote the k-iteration Q-value by applying T. Q" represents the initial Q-value. After one iteration, we have
1 1-7
@'(( S Q'
761.
1 1-
Suppose at iteration & we have
(() 1- nk
761.
(20)
1-7
Then for iteration & + 1, we have
(s) Tmodeo (x, (8, 7
In the meantime, we have
1 1-gh (Q((s, @t (*((,))++
1
1 1-9 y 761
1- 5 (*(5,0) 1 761
=r 7IE yf " (
1 1- g 2c1- QQ++1 +
16

---

Published as a conference paper at ICLR 2026
Therefore, have
(s) S Q + 1 1
761-
Therefore, Equation also holds for iteration + 1, and thus holds for all E Z=. If we set + 0o, then Q and Q would converge to the fixed points, and we have
((s, a) S ()) 1-7
761
(21)
In the meantime, also have
min E ( Q a') g T (E) Q(8',a') -1
E B Q(s' (a + (2)
Following a similar derivation process as above, we obtain
(s, a) Qtrue (s, a) 1-y
(22)
Combining the results of Equation 21 and Equation 22, we conclude the proof.
B.3 PROOF OF COROLLARY 4.1
Proof. The proof is straightforward.
((3))
V(3))
61 (.a)EDmoded
E Bg (+Ia.)), Paa (1),))
min s'E(((s))
[((3) V(s') V(s')
V(3')
max (s.d)EDmodel
E B (-19a)))) Tx
min
s'el((e)
5 mnax (s.a)EPuodde
Eg (-l.o) -
[[(3) - F(6)|+ V(6') - min
V(3')
< Lv (e + E)
This concludes the proof.
B.4 PROOF OF PROPOSITION 4.3
Definition B.1 (Lipschitz Smoothness). A function f(x) R" > R is L-Lipschitz smooth if the following inequality holds
( S L L| x2ll Ve1, e R".
(23)
We then have the following Lemmas:
Lemma B.2 (Nesterov al. (2018)). If function f(x) _8 - IR is L-Lipschitz smooth, then we have
(((:)) x1) 21)1= S 1|2- x111 Yei, 32 E R".
(24)
Lemma B.3. Assume the outer loss LRVL L-Lipsehitz smooth, and the gradient of LRVL. LsL are bounded by Let wy be twice differentiable, with its gradient and Hessian bounded by 5 and B, respectively, Then Luvi. is (L82 + -)-Liccchitz smooth w.rl. v.
Proof. Recall that the gradient of the outer loss Luvi(((()) w.r.t is
VoLRVL (2 ((())))))))))))))
v Cr = h Powpen(s,c a, 8'),
(25)
where h is defined in Equation 10. Taking the gradient of v in both sides of Equation 25, we have
Lyv (1())))))))))0000))))))))) Veh V Velewal a, s') +h V2 Viune (8, a, 8')
17

---

Published as a conference paper at ICLR 2026
The norm of the first term I could be bounded follows:
||/l|
T v log (Turos (611o.a)
v (V
vv (V log (s11s,))) (ompeo(s,a,s)) Vo log (Foco
log (Toco(s')s,,)) V vommvllra,ss) T V log (T
S 81((1 Lp²0².
The norm of the second term I2 could be bounded as below:
12ll = |/ V wos a')l
= B1(1) V CRVL CRV (((ltt))TVV..... Vo (5'1s,)) (s,no)
S Bp²
Combining the above results of || and |/21/1 we could bound the norm of L. (vS(((6(-))) as follows:
V CRVL (2b(+t)))
|11 + I2ll
< |illl + |12/|
(1()) (Bi Lo² B)
S p² (L8² + B).
According to Lagrange mean value theorem, for yı 2/2, we have
| (v(x)(( - V.LLxI s L' ||2 12l1.
where L' (L52 + B). This concludes the proof.
Then we prove our proposition 4.3. We restate it as follows.
Proposition B.2 (Proposition 4.3). Assume LRVL and LWSL are L-Lipschitz smooth, and the gra- dient of Luvi, LSt. and Lwst. bounded by P. Let wy be twice differentiable, with its gradient and Hessian bounded by and B, respectively. Denote Tí' = p + B), and the total steps as K. For some C1, C2 > we assume the learning rate of the inner and outer loop B1 and 32 satisfies: I c2 - 331,82 S min K. T. } where ca S &mnn a. T. b. /.V VK
1, K < 1. Then the outer
VK
loss LRVL and inner loss LWSL can achieve a convergence rate o
(
min E osesk-1
(26)
VK
(vs00)))
min OSASK-1
(27)
K
Proof. We first prove Equation 26. Consider the difference in the outer loss between the t-th and the 1-th iteration:
LIVI( (2(0+1))) (2((02)))(()+)))) LRVL (2i(t+1) (2(00))
LRVI (1ilt())))))))))))) (((+1))) LRVI. C (16(tt1)(y(tttt)))) + LRVI (plt)()))tt())(ttt))) CRVI, (1))) (v(()))
18

---

Published as a conference paper at ICLR 2026
According to Lemma B.2 and the assumption that LRVL is Lipschitz smooth, could bound the first term 1 by
4 w ( (((t1))
2
+ 2 [120 wlt s (ulttt)
For term 13. since (x(++)) (v(++1)) and 1)(++) (x((+1)) have the relation
2b(++2) (x(++1) wil+ (x(tt)) BI(+t) LWst. www.. (zb(t+))),
we could bound I by
Is V CRVL (1
( - 20((+2)
$ p: (1u411) V Lwst (ws((v)))
S S1(+t))p
Similarly, 14 is bounded by
11 = 2 (1(e+1)) V (ws
5 2 32 "1(41)00
Therefore, have
11 1/ + p1 + Pi
(28)
For term T2, regard (uvi((/(V)) function w.r.t v. By using Lemma B.2 and Lemma B.3, we have
I2 S V.LIVI (2 -0)) + 1' | -xtt
2
L'B2(0)
-211 V ( (/
2
I'PL.00 (
(2(0) 2
(29)
The first equality holds due to
((++)) (()) 82(1) (1
Combining the results of Equation 28 and Equation 29, we have
LuVI L (v((++2) (v((tt)))
(v(())()))))
L'32(0)
L 02 R0=n))2
(2(1)
2
2
Adjusting the order of the inequality and summing both sides from t to K we have
I/
)
B2(1)
2
K-1
L32 189
s CRVI. (v)(4)
LRVI.W (w(Ki)) +
Σ
31 +
2
K-1
< (L
+
2
19

---

Published as a conference paper at ICLR 2026
Therefore, we have
min OStSK-1
E ||| V
LLVL (2b(tt))
(B20) (B2(1)
1K-1
L183(1))
V. (u
2
1k-10 (32(1)
2
2C )) + E/= (2.3(1+1))p²
+LS3 (((())))²))
(2 -L(32(0)
K-1
1
(00 (z)0)) +
Σ
11(441) LB10TD))
$ K-1
EE= 82(1)
$
K minogiss-- R2(1)
2CHVL (v(0))) + K
0stsk-1
max
(2++
(v000)
+ maxosIsK 1INOSISK-11ICC)))2222 (1( T.)
K minosess- K- 12(1)
minosesk-1 32(0)
(2+((++L)
C2VK K
coVK
1 VK
The third inequality holds from KK1 B2(0) S -K-1 (202()) - (0)) given 82(0) S #.
The proof of Equation 27 is similar to Equation 26. For completeness, we also present the proof procedure. Consider the difference in the inner loss between the t-th and + 1-th iterations:
CWSI(vi(tH (vb((+1),,((t))) (((+1)) Cwst( (v)11, :(t)))
= = ,(((t1)) - LWsl (1)(et))) ,x(1) + ( ,10) (10)) (w:l
For term J, according to Lemma B.2, we have
L [(ttt) [ -000l -1(0)
T S V VeLwsL( (1)(t11), ((0))t)))0))(((())))))) ((00)T(((ttt) - ((1) + 2
L
= - V CLLws((bbtt ,(00)' V. (v + 2
s 320pp² + L o2
2 22
For term I2, we have
1/ S V Lw (10) (10)))))))000000)))))))))) ((0) + _/2 | 2//10/ /²
-11c)) V v. v. (4610) ( LWSs. (1b(t)) ult)) +
W. (16(1), "000))
2
- L800 ) (V w((t)>( /|?
2
Therefore, we have
LWSL (w)(()) ,2(0))))) - LwsL(2b(0) (211))
< p2 + 22 8200002 - ( (8100)
I.89(0)
)
|| V Lwst (1610) _0)0))
2
20

---

Published as a conference paper at ICLR 2026
Rearranging the inequality and summing both sides from t = to K 1, we have
K-1 1
(1))) - LS(()
V wwst (20(0) //0))l (
2
t=
k
- LWSi. (1ss..(//0))) (1)(0),(10))
LWsl (rb(K), _((K))
,
+
Σ
B2 + 2 8 22000
1=0
L Σ B2 +
S Lwsi (16(0), (10))
Therefore, we have
CWSL (wst(bb0) (2b(tt) b000)|1
min 0S1ST-1
E
V
KK1
LP Bi(0)
|| 1. (v:10) (()))
2
Σk=1
(1(0) 2
LPf
2LWsL (ei10), ,2(0)) + EKED 233(c) + L. L.3201002
Ek=1 (281(()) -L83(E)
2LWsL (0))
IK-1 -K- 2B2(0) + 20)
5
2 -
Bi(0)
2CWsL( (v)10), (210)
KK 0S1gK
max
(2
K mmnogisk-L 31(1)
2Lwsi( (00)
(2 +
K
C2
0
K
where the third inequality holds since Σk B1(0) S 2k=1(2210) IB(())) given Bi(1) S T. This concludes the proof.
C BENCHMARK SETTINGS
In this part, we provide details of the benchmarks and datasets we use for evaluating our method in our paper. The offline datasets are taken from the D4RL (Fu al., 2020) and NeoRL (Qin al., 2022) benchmarks, two popular benchmarks designed for evaluating offline RL algorithms.
C.I DARL
For the D4RL benchmark, we mainly evaluate our method on two kinds of datasets: MuJoCo datasets and Antmaze datasets. We first introduce the two datasets.
MuloCo datasets are collected through interactions with continuous control tasks in Gym (Brock- man et al., 2016) simulated by MuJoCo (Todorov et al., 2012). The tasks we use are half halfcheetah, hopper and walker2d, as illustrated in Figure 4. For each task, we use the four types of datasets: (1) Random: data collected with random policy. (2) Medium: IM samples collected by an early-stopped SAC policy.
(3) Medium-Repla: 1M samples from the replay buffer of the agent trained up to the performance of medium-leve agent. (4) Medium- Expert: 50-50 split of medium-level data and expert-level data. The MuJoCo dataset version we in work is "-v2". The metric use for evaluating the agent's performance MuJoCo
21

---

Published as a conference paper at ICLR 2026
Figure 4: D4RL. MuloCo tasks. Left: halfcheetah, Middle: hopper, Right: walker2d.
F
Figure 5: D4RL Antmaze tasks. Left: umaze, Middle: medium, Right: large.
tasks is Normalized Score (NS). NS computed follows.
Jx Jrandom
NS
x 100%
(30)
Jexpert Jraudom
where Jx is the performance of the evaluated policy, Jrandom is the performance of the random policy, Jexpert is the performance of the expert policy.
Antmaze datasets robot is required to reach a goal location. Antmaze tasks are more challenging than MuJoCo tasks for offline RL algorithms due their sparse reward setting. There are three maze layouts contained in Antmaze tasks: umaze, medium, large, as shown in Figure 5. The datasets are collected in three flavors: (i) the robot needs reach specified goal from a fixed start point (antmaze waaizze-umzz---eee (00), (ii) the robot is required reach random goal from random start point (the diverse datasets). (iii) the robot is commanded to reach specific locations from a different set of specific start locations (the play datasets). In our work, intmaze-mediunt diverse,
collected in Antmaze tasks. In Antmaze tasks, 8-DOF "Ant" quadraped
the six Antmaze datasets:
antmaze -arizzeuummar umaze, antmaze -eedi um-play, antmaze- antmaze-Lazge-dd- large- antmaze -aa-ge-arraaaaa The dataset version we use is "-v0". The metric we use for evaluating the policy performance on Antmaze tasks is the success rate of task completion.
setaazz-uazzee-ivve
C.2 NFORL
NeoRL is benchmark that aims to simulate real-world scenarios by collecting datasets with more conservative policy, which more aligns with the real-world data-collection strategy. NeoRL datasets contain narrow and limited data, making it challenging for offline RL methods. There are several tasks on the NeoRL benchmark. Following previous studies (Sun et al., 2023; Lin et al.), we Hopper- Walker2d-v3) For each task, we three types of datasets collected with policies of different qualities (low, medi um, high). NeoRL provides different numbers of trajectories (100, 1000, 10000) training data for each task. We select 1000 trajectories for our experiments. The evaluation metric the Normalized Score, the same metric used for MuJoCo tasks in D4RL..
focus on three tasks (Ha @Cleffhhhhhhhhe
22

---

Published as a conference paper at ICLR 2026
D IMPLEMENTATION DETAILS
In this part, we present the details of baseline implementations, ROMI implementations, and hyper- parameter setup.
D.1 BASELINE IMPLEMENTATION
We elaborate the baseline implementation details on D4RL and NeoRL benchmarks, respectively.
For D4RL datasets, our baselines include CQL (Kumar et al., 2020), IQL (Kostrikov et al., 2021), MOPO (Yu et al., 2020), RAMBO (Rigter et al., 2022), Count-MORL (Kim & Oh, 2023), MO- BILE (Sun et al., 2023). For Count-MORL, we retrain the algorithm using the official codebase" provided in the original paper. We note that there are three count estimation methods (LC, AVG, UC) chosen for Count-MORI, and we choose UC since the original paper reports that UC exhibits the best average performance. For other baselines including CQL, IQL, MOPO, RAMBO, and MO- BILE, we run the code in the codebase OfflineRLKiit which provides reliable implementations for different model-free and model-based offline RL. algorithms.
For NeoRL datasets, our baselines include CQL, IQL, EDAC (An et al., 2021). MOPO, RAMBO, and MOBILE. Since the codebase OfflineRL-Kit has provided the implementations for these base- lines, we directly use it as our baseline codebase.
For Antmaze tasks, we run RAMBO according to the hyperparameter setting the original pa- per, since MOPO, COMBO, and MOBILE do not provide the detailed hyperparameter setting for Antmaze tasks, we directly copy the evaluation results of these algorithms as reported in (Lin al., 2024).
worth noting that the default training steps OfflineRL-Kit are 3M. To consistent with the general training settings and ensure a fair comparison, we uniformly set the training steps for all algorithms to IM.
D.2 IMPLEMENTATION OF ROMI
We then provide the implementation details of ROMI from the following aspects:
Codebase. Since ROMI is an improved version of RAMBO, we choose OfflineRL-Kit as our code- base, and implement ROMI on top of the implementation of RAMBO.
Dynamics Model Training, Following previous works (Rigter et al., 2022; Sun et al., 2023; Kim & Oh. 2023), we represent the dynamics model as an ensemble of neural networks that output Gaussian distribution over the next state given the current state and action:
(s'ls, = (s a). (s,
(31)
Note that we do not predict the reward r, since we assume the reward function r is known as in (Sun et al., 2023), although r can be considered as part of the model if unknown. Similar to RAMBO, we perform an initial maximum likelihood model training. The difference lies that we fix the training epochs to 50, rather than stopping the training based on the held-out loss, since we find this could be quite cime-consuming. After the initial training, ROMI updates all the ensemble members in bi-level optimization framework. For each model rollout, we randomly select ensemble member to simulate the rollout.
weeitive Network Training. We represent the alaptiive-eighttee network as sim- ple Multi-Layer Perceptron (MLP) with 2 hidden layers. The input vector is a concatenation of (s, a, s'), and the output is scalar representing the allocated weight for each training sample. To map the network's 's output to the interval [a, b), we apply the tanh function and then perform an affine transformation (scaling and shifting) regulate the final range:
weight ht = (out(o)))
a
+ b
(32)
2
2
"https://github... "htpp://gitub..co/oh-bbbii http://iiwiiccccm//aaaaaaaoaoo-- hhtps://gthuh..om/yiaaaaaa h240raaaaaaaa.0000000000eeelelee 124/OfffineRI-Kit..o..
23

---

Published as a conference paper at ICLR 2026
In experiment, choose 0.5 and 2, that is, the assigned weight is constrained in the range [0.5, 2]. Since the performance of ROMI is already satisfactory, we do not further tune the values of and b.
The policy optimization method we adopt is SAC, and most hyperparameters follow its standard implementations. For each training step, sample batch of 256 samples where 50% them are from the offline dataset D, and another 50% are from the model buffer D model- We use automatic entropy tuning, where the entropy target set to the standard heuristic dimm (AA).
Policy Optimization.
We present the base hyperparameter configurations for ROMI in Table 3 and the pseudo-code for ROMI in Algorithm 1. The key difference between ROMI and RAMBO lies in the training strategy for the dynamics model, as highlighted by the blue parts in the pseudo-code.
Algorithm ROMI
Offline dataset D, number of epochs N, maximum rollout length H, dynamics model pretraining steps model. policy pretraining steps T policy. learning rate X, real data ratio f
Require:
2: Initialization: model data buffer D model + Ø, ensemble dynamics model T. {T 1. weighting network Wy, initial policy тa. other components of SAC (Haarnoja et al., 2018) algo- rithm
3: // Pre-train the dynamics model
4: for step in to T model do
5: Compute loss: LSL - (s,,, (s')ss))
6: Perform gradient descent: va + 1h AV (vb
7: end for
8: // (Optional) Pre-train the policy
9: for step in to pollicy do
10: Compute loss: LBC (DD [log To (s|a)
11: Perform gradient descent: 0 - 0
12: end for
13: for epoch from 1 to N do
14: Sample an initial state so from
15: for l in 0 H do
16:
Take an action an ~ mo(-1sn) Randomly pick an ensemble member T from (To J.M and sample next state 8841
17:
T(- sh, an)
18:
Add the model generated transition (sh, ah, Th: 8h++1 into Dindel
Sample batch of transitions from D U Dmodel with ratio f optimize To with SAC: 0 <- 0 AVVLSAC(()
20:
21:
end for
22:
// Inner level update
Compute loss Lwst. (,b, 100=--Enp -E (u(s,,,ss) log (To(o)ss,a))
23:
24:
Perform gradient descent: + vb - AV LWsL
25:
I/ Outer level update
26:
Compute h according to Equation 10
27:
Compute gradient ynv1. = h V.w 8')
28:
Perform gradient descent: v +- u X JRVl.
end for
D.3 HYPERPARAMETER TUNING OF ROMI
There are three main hyperparameters for ROMI: the uncertainty set scale &. the uncertainty set sampling number N, and the rollout length H. We show how tune those hyperparameters below.
Uncertainty set scale &. For D4RL and NeoRL tasks, we search & in the range of (0.01, 1.0, and choose the value with the best performance.
24

---

Published as a conference paper at ICLR 2026
Table 3: Base hyperparameter setup for ROMI shared across all runs.
Hyperparameter
Value
Dynamies Model
Ensemble size
7
Architecture
(input, 200, 200, 200, 200, output) 0.5
Real data ratio (f)
Model learning batch size
256
Model learning rate
3 x 10
Model pretraining epochs
50
Weighting network
Architecture
(input, 256, 256, 256, 1)
Output range
[0.5, 2]
learning rate
1 x 10 4
Batch size
256
Policy optimization
Require pretraining
False
Critic learning rate
3 x 10-4
Actor learning rate Optimizer
1 x 10-4
Adam (Kingma & Ba, 2014)
Discount factor
0.99
Soft update parameter
5 x 10-3
Target entropy
-dim(A)
Batch size
256
Sampling number N. controls the approximation error between the sampled minimum Q-value and the true minimum Q-value within the uncertainty set. A larger N generally leads to smaller error incurs higher computational costs. Furthermore, larger & often necessitates larger N. In experiments, select N from the set {10, 20, 30) for all the tasks. To conserve computational resources, the performance difference between two candidate values of is not significant, we choose the smaller one.
Rollout length H. We perform short-horizon rollouts similar to previous works (Sun et al., 2023; Rigter et al., 2022), and tune H across 52,5} for Gym and NeoRL tasks. For Antmaze tasks, vary the parameters on H E (1,3).
The selected hyperparameters for each task are listed in Table 4.
E COMPUTATIONAL COST
E.1 COMPUTE INFRASTRUCTURE
We list our hardware specifications as follows:
GPU: NVIDIA RTX A6000 (x4)
CPU: AMD Ryzen Threadripper PRO 5975WX
We also list our software specifications as follows:
Python: 3.8.18
Pytorch: 1.1 1.1.00ccu13
Tensorflow: 2.9.1
Numpy: 1.24.4
Gym: 0.22.0
MuJoCo: 2.0
D4RL: 1.1
25

---

Published as a conference paper at ICLR 2026
Table 4: Hyperparameters of ROMI.
Domain Name
Task Name
N
H
halfcheetah-random
0.01
10
5
hopper-random walker2d-random
0.1
10
5
0.01
10
5
halfcheetah-meddum
0.1
10
2
hopper-medium
1.0
10
2
D4RL MuJoCo
walker2d-medium
0.01
10
5
malfcherahh-meduu--eplaa
0.1
10
5
hopper-mediim--eplayy
1.0
10
5
walker2d-mediun--rplay
0.01
10
5
halfcheeraa--meduu--pppee
1.0
10
2
hopper-medine -expertexpeertt
paaper-medium-exppee
1.0
10
2
malker2d-mmdiumg medium-expertt expert
0.01
10
5
caimaze-umaze
0.1
10
1
calmazz-umaze-divee
1.0
10
1
D4RL Antmaze
antmaze-medium-ppay antmaze-medium-ddives
0.1
10
3
1.0
10
antmaze-large-play
1.0
10
antmaze-large-diverse
1.0
10
1
HalfCheetah-L
0.1
10
2
Hopper-L
0.1
10
2
Walker2d-L
1.0
10
2
NeoRL
HalfCheetah-M
0.01
10
5
Hopper-M
0.1
10
2
Walker2d-M
0.01
10
2
HalfCheetah- -H
1.0
10
2
Hopper-H
0.1
10
2
Walker2d-H
0.01
10
26

---

Published as a conference paper at ICLR 2026
E.2 COMPUTING RESOURCE CONSUMPTION
We compare the computational resource consumption of ROMI and RAMBO. including model pa- rameter size, GPU memory usage during training, and runtime, as shown in Table 5. The results indicate that while ROMI and RAMBO show no significant difference in model parameter size or GPU memory usage, ROMI requires more runtime. This is primarily due the additional training overhead of the weighting network.
Table 5: Comparison of ROMI and RAMBO on computing resource consumption
Size of Parameters (MB)
GPU Memory (GB)
Runtime (s/epoch)
Task Name
ROMI
RAMBO
ROMI
RAMBO
ROMI
RAMBO
hopper-random
7.93
7.66
3.30
3.28
28.20
23.76
halfeheetah-randoon
8.19
7.91
3.33
3.32
29.41
22.24
walker2d-random
8.19
7.91
3.33
3.32
30.06
23.89
hopper-medium
7.93
7.66
3.30
3.28
28.43
23.31
halfcheetah-mmeiumm
8.19
7.91
3.33
3.32
29.91
22.63
Balker2ddmmddum
8.19
7.91
3.33
3.32
29.67
23.15
hopper- mooper-mmdium-rellayy
7.93
7.66
3,30
3.28
28.46
22.27
serccheetah-medium---
8.19
7.91
3.33
3.32
29.13
23.52
8.19
7.91
3.33
3.32
29.64
23.11
hopper-mediung expert
7.93
7.66
3.30
3.28
28.33
22.93
8.19
7.91
3.33
3.32
29.21
23.47
felker2d-mediunte expert
8.19
7.91
3.33
3.32
30.02
23.09
F MORE EXPERIMENTAL RESULTS
In this part, we provide more experimental results missing from the main text.
F.1 D4RL ANTMAZE RESULTS
We compare ROMI with several model-based offline RL methods on the challenging D4RL Antmaze tasks; results presented in Table 6. Performance varies substantially across methods: MOPO fails to achieve meaningful performance on all six tasks; COMBO performs well on umaze and umaze-di verse but fails on others: RAMBO achieves nonzero scores on four of the six tasks, but its total score remains low. MOBILE performs strongly overall, achieving nonzero scores on all tasks with a total of 173.4. BILE and the highest among all compared methods. Notably, ROMI clearly outperforms MOBILE on umaze diverse and medi un-diverse and exceeds RAMBO on five the six tasks. The only exception is large -play, where both methods score zero. These results demonstrate ROMI's effectiveness in addressing challenging Antmaze domains.
contrast, ROMI achieves total score of 186.5- higher than MO-
F,2 MORE RESULTS ON MOTIVATION EXAMPLE
We conduct experiments on nine D4RL datasets to examine the limitations of RAMBO. In Sec- tion 3, we presented results on the wtwhheethh-medimmeeeeeeeee constraints. Here, we provide results across all nine datasets, including gradient norms, normalized scores, and Q-value estimates.
dataset due to space
Figure 6 shows the gradient norm of the adversarial loss during RAMBO training for X = 0.05 and = 0.1. We omit curves for smaller À such as
x 10-4 since we find the gra- dient norm stays small for those X. as the gradient norm remains consistently small in those cases. As shown, gradient explosion occurs on two tasks Sopper-mediume -n-eemert-vv) with À 0.05. For > 0.1, this issue becomes more fre- quent, occurring in 7 out of 9 tasks. These results demonstrate that RAMBO suffers from training instability unless À is set extremely small value.
and
27

---

Published as a conference paper at ICLR 2026
Table Successful Success rate comparison on Antmaze tasks. Results are averaged over 5 seeds. The top two scores for each task are bold.
Task Name
MОPO
COMBO
RAMBO
MOBILE
ROMI (Ours)
antmaze-umaze
0.0
80.3
42.6 + 11.9
77.0
70.3+ 11.4
antmaze-umazz-divens
0.0
57.3
12.9 +5.6
20.4
32.8 + 13.0
antmaze-medium-plae
0.0
0.0
18.7 + 9.2
64.6
51.3 + 19.7
nntmaze-medium-diieee
0.0
0.0
20.4
1.6
30.1 + 10.0
untmaze-large-play
0.0
0.0
0.0 0.0
2.6
0.0 + 0.0
antmaze-large-dee
0.0
0.0
0.0 0.0
7.2
6,4 + 4.2
Total
0.0
137.6
94.6
173.4
186.5
A=
A=
halfcheetah-medium
halfchedtah-medduu-rrppaay
30d hophiicchheaalmmmduue-epprrr
0.005
150
Gradient Norm
Gradient Norm
0.004
Gradient Norm
200
100
0.003
100
50
0.002
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Step
Step
1e6
Step
1e6
hopper-mediim
hopper-medium-reepla
hopper-medium-expert
125
250
Gradient' Norm
200
Gradient Norm
100
Gradient Norm
150
150
75
100
100
50
50
0
0.5
1.0
0.0
2.5
5.0
Step
1e6
Step
1e5
Step
1e5
walker2d-medium
walker22-medium--eplay
walker2dd-eeduum-exper
250
0.020
Gradlent Norm
200
Gradient Norm
Gradient Norm
1000
150
0.015
100
0.010
500
50
0.005
0
0
0.0
1.0
Step
1e6
Step
1e6
Step
1e6
Figure 6: Gradient norm curves on 9 D4RL datasets for RAMBO when À = 0.05 and X = 0.1
Figure 7 and Figure 8 present the normalized score and Q-value estimation on 9 D4RL datasets for RAMBO across different A. A clear trend can be observed: setting À to very small values within (0, 3e 4, 5e 3} has almost larger value such as 0.05 and 0. leads to diverged Q-value and performance collapse, indicating that RAMBO is highly sensitive to À and prone excessive conservatism.
impact performance and Q-value estimation. In contrast,
28

---

Published as a conference paper at ICLR 2026
À ==0 0
A= 3e- 4
5e -
A = 5e 2
A= le
haifcheetah-mediim
80
Normalized Score
Normalzzed Score
Normalized Score
60
75
40
50
20
20
25
hopper-medium
hopper-medium-replay
hopper-medilm--eppel
100
80
100
Normalized Score
Normalized Score
Normalized Score
75
80
75
50
25
20
25
0
walker2d-medium
walkkezz--eddium-relay
walkkr22d-medium-exper
80
80
Normalized Score
Normalized Score
Normalized Score
60
60
60
40
40
20
20
20
0
Figure Normalized score comparison on 9 D4RL datasets across different À for RAMBO.
29

---

Published as a conference paper at ICLR 2026
À == 0
A= 3e-4
5e -3
5= 2
=
halfcheetah-mediun
halfcheetah-meddiu--rrlaaay
roopalicchetthh-mediu----
500
500
Q Value
Value
250
Q Value
500
-250
-500
-500
-500
hopper-medium
hopper-mediummreplay
hopper-medium-exper
200
200
200
Q Value
O Value
Value
-2000
-200
-200
-400
-400
4400
walker2d-medium
walker22-meedium-replay
walker2d-mmedium-exper
200
400
200
200
Value
Value
Value
-200
-200
-200
-400
-400
4400
Figure Q-value estimation on 9 D4RL datasets across different À for RAMBO.
30

---

Published as a conference paper at ICLR 2026
8 0.01
E- 0.1
{-1
{
walker2d-me
walker2d-me
walker2d-me
walker2d-me
400
100
0.08
Outer Gradient Norm
VVllue
300
lener Gradient
0.07
0.15
200
0.06
0.10
100
0.05
0.00
1.0
Step
Step (b)
Step
Step
(a)
(c)
(d)
Figure 9: Comparison for different & in terms of Q-value, normalized score, inner loss gradient
norm, outer loss gradient norm.
Table 7: Performance comparison for different N on D4RL datasets.
Datasets
N 10
N 20
N 30
half-mr
70.7:4.88
71.8.14.2
72.2::5.1
hopper-mr
102:2.4
100.0
103.3+2.8
walker2d-mr
85.242..
87.4 4.6
81.5 + 3.8
half-me
108.5:4.2
110.445.3
111.013.4
hopper-me
111.5+33.8
110.0+1.44
113.9 4.7
walker2d-me
113.3+1.9
115.6
111.7+005
Average
98.5
99.2
99.1
F.3 MORE PARAMETER STUDY
In this section, we further study the sensitivity of ROMI the choices of the uncertainty set scale and sampling number N.
We present the parameter study results of on the halfcheetah senieeetth-mediure expert-v2 dataset in Section 5.3. Here we present more results on bolkkerrd-medduu-eexxrrer vary & across {0.01, 0.1, 1.0, 10} and compare the normalized score, Q-value estimate and the gra- dient norm of the inner and outer loss. The results are shown in Figure 9. We can see that no severe Q-value underestimation or gradient explosion occurs across different Ç, verifying the superiority of ROMI over RAMBO.
dataset. We
We also study the sensitivity of N. We sweep across {10, 20, for each task, and present the results 6 D4RL datasets in Table 7. We find that setting N 20 and 30 does not show clear performance improvement compared to N 10. For the consideration of computational cost, we fix N 10 for all the experiments without further tuning.
F.4 LEARNING CURVES
We provide the learning curves of RAMBO, MOBILE, and ROMI on 12 D4RL datasets in Figure 10 and Figure 11. The training steps are 1M for all the methods.
G LLM USAGE STATEMENT
In this work, we use the large language model (LLM) to polish the grammar our manuscript. However, we confirm that no LLMs are involved in the key steps of this work, including the con- ceptualization of the idea, the development of methodology, and theoretical proofs.
31

---

Published as a conference paper at ICLR 2026
RAMBO
ROMI
halfchhetta--anddom
hopper-random
walker2d-random
Normalzed Score
40
Normalized Score
30
MMMM
Normalized Score
20
30
10
20
10
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5 Step
1.0
Step
1e6
Step
1e6
halfcheetah--medium
hopper-medium
walker2d-medium
80
Normalzzed 2ed Score
Normalized Score
100
Normalized Score
60
75
40
50
50
20
25
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Step
1e6
Step
Step
1e6
haffchhetta--meiiu---elleg
hopper-medium-replla
walker2d maalkerdd-mddium-rrplaayy
100
100
Nomalized Score
60
Normalized Score
Jormalized Score
75
40
50
50
20
25
0
0.5
1.0
Step
1e6
Step
e
Step
1e6
nalichhetah-mmdduu--eeppe
hopper-medium-expert
walker2d-mediuu-eeeae
Normalized alizeddzd Score
100
alizalized Score
Normalzee alized Score
100
100
50
50
0
0.0
0.5 Step
1.0
0.0
Step
1e6
1e6
Step
1e6
Figure 10: Learning curves of RAMBO and ROMI on 12 D4RL datasets. Averaged over seeds.
32

---

Published as a conference paper at ICLR 2026
MOBILE hopper-random
halfcheetah-randomm
walker2d-random
40
30
20
Normalized Score
Normalized Score
Normalized Score
30
WW
10
20
10
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5 Step
1.0
Step
1e6
Step
1e6
halfcheetaa--medium
hopper-medium
walker2d-medium
80
Normaized Score
Normalized Score
100
Noddlizze Score
100
60
75
40
50
25
0.0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Step
1e6
Step
Step
1e6
haffchhetta--mediuu--elllle
moppur-mediumrreplay
walkkrzd--medium-reell
Normalized Score
Normalized Monnalizzddoooo Score
100
Normalized lized Score
100
60
75
75
40
50
20
25
25
0.0
1.0
0.5
Step
1e6
Step
1e6
Step
1e6
hallchhetah-mmddimm-xppeer
hopper-medium-exper
walker2d--mediu-eexpee
alized Score
Nlizzl Score
100
Nommalized Score
100
75
75
50
50
25
25
25
0.0
0.5 Step
1.0
0.0
Step
1e6
1e6
Step
1e6
Figure 11: Learning curves of MOBILE on 12 D4RL datasets. Averaged over seeds.
33


Reviews:
[{"review_id": "hSJ0sIjC6a", "reviewer": "Reviewer_7d1B", "summary": "This paper looks into an instability noticed in RAMBO with underestimated Q-values and proposes a new method called ROMI to fix this. Their method aims to prevent model exploitation by training the model to balance predicting a close-by state with the lowest value, as well as getting the dynamics correct.", "questions": "**Q1.** Where are the results from for each method for both the D4RL and Neo-RL benchmarks? If they are from the authors own implementations, how did they ensure they were implemented correctly and tuned fairly?\n\n**Q2.** How was Wasserstein distance chosen. Did the authors try or consider any other distance metrics?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Interesting finding in RAMBO, seemingly strong results, includes ablations, and an interesting method of balancing competing objectives that I have not come across in RL before.", "weaknesses": "**W1.** There are lots of moving parts. It seems significantly more complicated to implement and understand than RAMBO.\n\n**W2.** The authors say they compare all methods at 1M steps for fairness. However, since some baselines like MOBILE were tuned for 3M steps, it would be important to see if ROMI's performance holds up relative to MOBILE etc when trained for the full 3M steps.\n\n**W3.** Larger compute requirement.\n\n**W4.** They solve the problem of tuning lambda, but introduce more new hyperparameters: xi, H, N.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "z11EWYgtrp", "reviewer": "Reviewer_rQhm", "summary": "This paper introduces ROMI, a new method for model-based offline reinforcement learning (RL) that addresses instability and over-conservatism found in previous approaches like RAMBO. ROMI uses a robust value-aware model learning strategy and a bi-level optimization scheme to better control conservatism and improve generalization. Experiments show ROMI outperforms RAMBO and matches or exceeds other state-of-the-art methods on standard offline RL benchmarks.", "questions": "Am I correctly assuming that the parameter governing the size of the uncertainty set has to be set a priori?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Proposes a novel, robust approach that addresses instability and over-conservatism in model-based offline RL. The bi-level optimization scheme appears a promising solution for adaptive conservatism, improving both stability and performance\n- Demonstrate strong empirical results on MuJoCo datasets, outperforming or matching state-of-the-art baselines\n- Proposed framework appears to be less sensitive to hyperparameters and more generalizable across tasks, as shown e.g. in Fig. 3", "weaknesses": "- Even though robustness has clearly gone up, the approach may still require careful tuning of certain parameters - i.e. in Fig. 3 we see that while 0.01-1 are close together, setting the parameter to 10 yields very unstable results. Since all evaluations are in a relatively similar task domain (MuJoCo hopper, walker and halfcheetah), it is unclear whether the range of stable parameters is representable for other tasks as well, i.e. for other environments tuning the parameter may again become an issue.\n- More generally speaking, while the authors clearly demonstrate performance gains on a set of well established offline RL benchmarks, the distribution of these tasks is quite narrow - it would be interesting to see how the algorithm performs on a broader selection\n- Another downside I see is, if I am not mistaken, that the degree of uncertainty that governs the training needs to be set a priori and cannot be adjusted at runtime. Prior methods, such as [1,2,3,4] (could also help your related work section) have developed policies that can be adapted immediately when it is observed that the behavior is too far OOD or too conservative, which can be an advantage over having to go back and retrain. Of course, this point does not necessarily have to be addressed, the proposed method has other demonstrated merits which justify publication.\n- The method introduces additional complexity through the bi-level optimization, which may increase implementation difficulty and be a burden in practical application\n\n[1] Ghosh, D., Ajay, A., Agrawal, P., & Levine, S. Offline RL Policies Should be Trained to be Adaptive. Proceedings of the 39th International Conference on Machine Learning (ICML), 2022.\n\n[2] Hong, J., Kumar, A., & Levine, S. Confidence-Conditioned Value Functions for Offline Reinforcement Learning. International Conference on Learning Representations (ICLR), 2023.\n\n[3] Swazinna, P., Udluft, S., & Runkler, T. User Interactive Offline Reinforcement Learning. International Conference on Learning Representations (ICLR), 2023. \n\n[4] Zhang, Y., Liu, J., & Wang, Y. Train Once, Get a Family: State-Adaptive Balances for Offline-to-Online Reinforcement Learning. Advances in Neural Information Processing Systems (NeurIPS), 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QsbwfTw3j5", "reviewer": "Reviewer_VuTF", "summary": "The paper proposes an alternative approach to model-based offline reinforcement learning, building on the idea of value-aware model learning. A model is constructed that minimizes an adversarial value difference error as a weighing to a log likelihood reconstruction loss. The authors show the efficacy of their method on standard offline RL tasks, in addition to providing some theoretical insights into the method.", "questions": "How is the state distance metric defined? This seems to be a crucial implementation detail.\n\nHow are the uncertainty intervals in the empirical section computed? Are bolded results statistically significant? Why do aggregated results and baseline lack uncertainty intervals?\n\nCan the method work for higher dimensional problems as well? Pessimism over a state uncertainty set could conceivably be harder to tune for larger more high-dimensional problems, and all those tested here are rather low-dimensional.\n\nHow was hyperparameter tuning and model selection conducted? Was there a separate offline validation set, or where hyperparameters tuned by testing in the online environment.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "While the overall method is slightly complex, the paper presents the motivation and contribution of the work clearly. The proposed integration of conservative value-aware model learning into offline model-based RL is reasonable and shows some promising results. In addition, known issues about the generalization abilities of value-aware model learning are appropriately avoided with the bi-level optimization scheme.", "weaknesses": "I genuinely do not have any major issues with the paper, so the following are excessively nitpicky and mostly comments on the writing.\n\nI believe the writing of the paper could be strengthened slightly by defining the method less in contrast to RAMBO, as I believe the method conceptually stands on it's own.\n\nLine 175: I don't believe this is a proper example of an adversarial loss, as there is no min/max formulation with competing networks.\n\nIt would be polite to cite Farahmand for coining the term value-aware as soon as it appears. Since the Farahmand paper is cited, I don't believe this is an egregious error. Similarly, the method is somewhat conceptually similar to Voelcker et al., as a reweighted regular MLE loss is considered. Again, prior work is acknowledged, so this is not really an issue.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)