Any-step Dynamics Model Improves Future Predictions for Online and Offline Reinforcement Learning
Authors: Haoxin Lin
Yu-Yan Xu
Yihao Sun
Zhilong Zhang
Yi-Chen Li
Chengxing Jia
Junyin Ye
Jiaji Zhang
Yang Yu
Venue: ICLR 2025 Poster (2025)
Keywords: model-based reinforcement learning
any-step dynamics model
Primary area: reinforcement learning
TL;DR: We propose the Any-step Dynamics Model (ADM) to mitigate the compounding error by reducing bootstrapping prediction to direct prediction.

Abstract: Model-based methods in reinforcement learning offer a promising approach to enhance data efficiency by facilitating policy exploration within a dynamics model. However, accurately predicting sequential steps in the dynamics model remains a challenge due to the bootstrapping prediction, which attributes the next state to the prediction of the current state. This leads to accumulated errors during model roll-out. In this paper, we propose the Any-step Dynamics Model (ADM) to mitigate the compounding error by reducing bootstrapping prediction to direct prediction. ADM allows for the use of variable-length plans as inputs for predicting future states without frequent bootstrapping. We design two algorithms, ADMPO-ON and ADMPO-OFF, which apply ADM in online and offline model-based frameworks, respectively. In the online setting, ADMPO-ON demonstrates improved sample efficiency compared to previous state-of-the-art methods. In the offline setting, ADMPO-OFF not only demonstrates superior performance compared to recent state-of-the-art offline approaches but also offers better quantification of model uncertainty using only a single ADM.

Published as a conference paper at ICLR 2025
ANY-STEP DYNAMICS MODEL IMPROVES FUTURE PRE- DICTIONS FOR ONLINE AND OFFLINE REINFORCEMENT LEARNING
Haoxin Lin YunYan& Xu², Yihao Sun', Zhilong Zhang Yi-Chen Li
Chengxing Jia 12, Junyin Yel2, Jiaji Zhang', Yang Vn1
National Key Laboratory for Novel Software Technology, Nanjing University, Nanjing, China School of Artificial Intelligence, Nanjing University, Nanjing, China PPolixir Technologies, Nanjing, China
*Png Cheng Laboratory, Shenzhen, 518055, China
linhx01 1 amda. ju. edu. en, yuyan. xuepolixir. (sunyh, zhangzl, liye, jiacx, yejy, yuy B1 amda nju. leub.edd.com cn
ABSTRACT
Model-based methods in reinforcement learning offer a promising approach enhance data efficiency by facilitating policy exploration within a dynamics model. However, accurately predicting sequential steps in the dynamics model remains a challenge due to the bootstrapping prediction, which attributes the next state to the prediction of the current state. This leads to accumulated errors during model roll-out. In this paper, we propose the Any-step Dynamics Model (ADM) to mitigate the compounding error by reducing bootstrapping prediction to direct prediction. ADM allows for the use of variable-length plans inputs for pre- dicting future states without frequent bootstrapping. We design two algorithms, ADMPO-OX -N and ADMPO- OFF,, which apply ADM in online and offline model- based frameworks, respectively. In the online setting, ADMPO-ON demonstrates improved sample efficiency compared to previous state-of-the-art methods. In the offline setting, ADMPO-OFF not only demonstrates superior performance compared to recent state-of-the-art offline approaches but also offers better quan- tification of model uncertainty using only a single ADM. The code is available at https//ggtub..commLLMMAAAAALLLL
INTRODUCTION
Model-based Reinforcement Learning (MBRL) Luo et al.
2024b) has demonstrated empirical
success both online (Feinberg et al. 2018 Buckman et al. 2018
Chua et al. 2018 Luo et al.
2019 Janner et al. 2019 Lin et al. (023) and offline (Yu et al. 2020) Kidambi et al. 2020 Yu
ct al. 2021 Rigter et al. 2022 Sun et al. 2023 Luo et al. 2024a Chen et al. 2024| Zhang et al.
2024) settings. The essence of MBRL lies the dynamics model, where extensive explorations and evaluations of the agent can occur, thereby reducing the reliance on real-world samples. Embedded
leverage large Podate-To-Daaaa (UTD) ratio (Chen et al. 2021) to improve sample efficiency, while offline policy optimization can be completed using the model-augmented data beyond the dataset.
in the model-based framework, online policy optimization
Although some efforts aim to propose high-fidelity dynamics models, such as adversarial models (Chen et al. 2023 Bhardwaj al. (023), causal models (Zhu et al., 2022), and ensemble dynamics models (Chua et al. 20188] Janner et al. 2019: Yu et al. (020) adopted by the majority of MBRL algorithms, it is challenging to generate high-quality imaginary samples via long-horizon model roll-out. In dynamics model with the common form, the state-action pair at time step t, (st, at). used as input to predict the next state 8t+1. Thus, the bootstrapping prediction, which attributes the next state to the prediction of the current state, is inevitably employed to roll out states in the dynamics model. The deviation generated states increases with the roll-out length since the error
"Corresoonding Author

---

Published as a conference paper at ICLR 2025
accumulates gradually as the state transitions in imagination. If the policy is updated using unreliable samples with large compounding errors, it will be misled by biased policy gradients.
The impact of compounding error (Xu et al. 2021) on policy optimization restricts the utilization of the model, thereby hindering further improvements in the sample efficiency of online MBRL and the performance of offline MBRL. One potential way to deal with the issue of compounding error is to reduce bootstrapping prediction to direct prediction, considering the direct state transition after executing multi-step action sequence (Asadi et al. 2018 2019] Che et al.] 2018 Machado et al. 2023). Although state S is only dependent on state-action (8t, Gt) under the assumption of Markov property (Sutton& & Barto 2018), the prediction of st+1 can also leverage earlier infor- mation. Tracing back a prior k-step plan, i.e., $1+1-k followed by a sequence of k-step actions (0t+1-k; 0r+-2k, ... at), sufficient to constitute an attribution to predict 81+1-
However, a fixed-horizon multi-step dynamics model cannot allow for rolling out a sequence starting from a single state since the policy cannot make multi-step action decisions once. More importantly, it still relies on an ensemble to estimate model uncertainty in the offline setting (Yu et al.) [2020] Sun et al. 2023), which further increases the computational burden. To enhance multi-step direct prediction, we propose that a dynamics model should be capable of directly predicting future states across arbitrary horizon lengths. This would incorporate the flexible roll-out capacity of previous single-step dynamiis models, and the variability in horizon lengths could create conditions for quantification of the model uncertainty, making it possible to abandon the ensemble.
To handle the variable-length plans, introduce special Any-step Dynamics Model (ADM) that allows for the use of $t+1-k and (a++1-k, 4t+2-k,' ,ag) corresponding to any integer & within specified range inputs for predicting 8141. When the agent faces changes occurring in the trajectory distribution, the state predictions from different backtracking lengths will exhibit noticeable divergence. This feature naturally enables ADM to estimate model uncertainty without the ensemble. Replacing the ensemble dynamics model with ADM, we devise unique model roll-out method with random backtracking, which can be plugged into any existing MBRL algorithmic framework. In this paper, our main purpose is to demonstrate how the augmented data generated by ADM exhibits excellent effectiveness, both in improving future predictions and measuring the model uncertainty.
In general, our contributions are summarized as follows. (1) We present a generalized dynamics model called ADM to replace the dynamics model used in existing online and offline MBRL algorithms and demonstrate its superiority in reducing compounding errors. (2) We propose new online MBRL algorithm called ADMPO-ON based on ADM and show that it can outperform recent state-of-the-ar art online model-based algorithms in terms of sample efficiency while retaining competitive performance (3) We propose a new offline MBRL algorithm called ADMPO-OFH based ADM and show that it can effectively quantify the model uncertainty. achieving superior performance compared to recent state-of- tf-the-arr offline algorithms on D4RL (Fu
on MuJoCo (Todorov et al. 2012) benchmarks.
al. 2020) and NeoRL (Qin al. 2022) benchmarks.
2
PRELIMINARIES
2.1
MARKOV DECISION PROCESS AND REINFORCEMENT LEARNING
We consider a standard Markov Decision Process (MDP) specified by a tuple M M Po.7). 7). where S is the state space, A is the action space, T(s Fet1lss, az) is the dynamics function that calculates the conditioned distribution of $111 E S and re+1 E R given (st, a.), Pu is the initial state distribution, and is the discount factor. We p" to denote the on-policy distribution over states induced by the dynamics function 7 and the policy π. From a multi-step perspective, the attribution of state 8 and reward r++1 be traced back to the earlier k-step plan, 8 along
with the action sequence a--k+1:t = (a--++1, az-k-2, az)) in between. This relationship can be represented by the k-step dynamics model
1-1 T (1s--44 te Pe a-i))
T (s T441 8 at a+b+1:))
Σ
II ==
(4--2421 X
We use I (st- a- 4=1s444)))41 ) to denote the distribution over (st at a++++ conditioned on 88t1 induced by the dynamic function T and the policy æ.
2

---

Published as a conference paper at ICLR 2025
The optimization goal of Reinforcement Learning (RL) is to find policy i that maximizes the expected discounted return Ep- [_t=1' 2 Such a policy can be derived from the estimation
of the state-action value function Q" (8t, at) (1e44
[rt+1 + 7V* (84+1)), where V (s++1) 1 1891))))11) IQ" ($141, at+1)] ]] the state value function.
2.2 MODEL-BASED REINFORCEMENT LEARNING
MBRL. aims to find the optimal policy while transferring the agent's explorations and evaluations from the environment to the learned dynamics model. Given a dataset Denv collected via interaction in the real environment, the dynamics model I is typically trained to maximize the expected likelihood Eeer ((et,.r..e+111::1)---- [lo (8t+1, 1, 18,, 0+)]- The estimated dynamics model defines a surrogate MDP AM A, po: 7)- Then any RL algorithm be used recover the optimal policy with the augmented dataset Deny U D model, where D Duodel is the synthetic data rolled out in M.
The above-mentioned paradigm is adopted by model-based policy optimization (MBPO) (J Jaanner al., 2019) and much of its follow-up work (Lin al., [023] Li et al. 2022 Pan et al. 2020 Clavera et al. 2020) in the online setting. These works don't t need to consider the issue of model coverage, as the agent can explore online fill in the regions where the dynamics model uncertain.
However, in the offline setting, the limited dataset causes T to cover only a part of the state-action space. Therefore, MOPO Yu et al. (020)) and some of its subsequent offline MBRL algorithms
(Kidambi et al. 2020 Sun et al. 2023) use the ensemble-base based uncertainty as a penalty term in the
reward function, allowing the agent to sample within safe regions of Ā. This work can estimate the model uncertainty without ensemble models. A related approach is IVE (Filos al. 2022), which also proposes an ensemble-free uncertainty estimator through the divergence of values obtained by applying the surrogate Bellman operator the value function for different times.
METHOD
In this section, we propose a special Any-step Dynamics Model (ADM) to replace the mainstream ensemble dynamics models. ADM reduces bootstrapping prediction to direct prediction by backtrack- ing variable-length plans. Applying ADM to existing MBRL frameworks for policy optimization, introduce two algorithms, namely online ADMPO-ON and offline ADMPO-OFF.
3.1 ANY-STEP DYNAMICS MODEL
Currently, the prevalent dynamics models typically operate on single-step basis, with 8, and a, as inputs to predict S1-11 and re+1. In a broader context, dynamics models can also be multi-step (Asadi al. 2018 2019 Che et al. 2018), where inputs encompass 8; along with a k-step sequence of actions (a,, ae+1. ac+&-11 to predict s and Tr+kk. To introduce flexibility in the backtracking length of the model, we further extend the definition of the multi-step dynamics model to allow ie to be any positive integer within a specified range, as delineated in Definition 3.1
Definition 3.1 (Any-step Dynamics Model). Given the maximum backtracking length m, an any-step dynamics model T (s Tet. 88, a++++-1) 1) is the distribution of 8 E S and rt+ E IR conditioned on the k-step plan (st, (1:++&11) (St, ae 01+1, 1t) S x AA, where i can be any integer between [1, m..
To handle inputs with variable step sizes, utilize an RNN (Elman 1990) with a GRU (Cho al. 2014) cell to implement the any-step dynamics model, as depicted in the left part of Figure 1] Certainly, Transformer (Vaswani al., 2017) is also feasible choice, but we do not consider because the model structure is beyond the scope of this study. Since the input state consists of only one step, while the action may be sequence of multiple steps, we duplicate the state match the length of the action sequence, and then sequentially feed it into the RNN. The input (st, (+:+++1)) after being represented by the RNN, yields the hidden hţ, which is then fed into an MLP to obtain the mean and standard deviation of st+k and Tt+k. i.e., (Mi EeEx) and (ML k:
EfrE) Similar to
previous model-based methods (Janner et al.) [019 Pann et al.) 20200 Lin et al.) 2023), we model the distributions of 81+k and Ti+& as Gaussian distributions and predict them through sampling. We
call the Any-step Dynamics Model ADM and denote it to (8t+k, Tetalsa S (+)+ where
3

---

Published as a conference paper at ICLR 2025
4pt1
Q
4y
62
42
62
GRU
GRU
GRU
MLP
$
$2
3
Smn
1
4gom
Amystte
Dynamics Model (ADM)
ADM
ADM
ADM
ADM
PEOR
Se4k
m1
m
5m+1
random choice
Figure I: Illustration of any-step dynamics model (left) structured using RNN and its application for next-step prediction with random backtracking (right).
Algorithm Roll-out ADM: ADM-Roll(o, Tg. H, m, ($1, Q1, $2, 02, 1 8m- Am 8
Input: Learned ADM Ta with parameters 0, policy To with parameters é, roll-out length H, maximum backtracking length m, state-action sequence ($1, a1, $2, 02, ... Sm-1+ 4m-1 8m)
1: for = 0 to H - do
2: if' 0 then Sample am+r TT(+(Sm+r7)
3: else Sample am+r ~ T-(-)mm+7)
4: Randomly sample an integer is from [1, ml uniformly
5: i S k then Roll out via (sm+r+1; fm+r+1) ~ m am+ +1 -km+)7)))
6: else Roll out via (8m+ 1m) ~ T- Om kim+r)
7: end for
8: return (Sm, am: fm+1 Sm+1, 8m+ -, am fm 8m+H)
represents the neural parameters. With the real samples from the environment, To is trained to maximize the expected likelihood:
(
E
Don log To(SEtk, TetklNe (e
01)tk-1)
(2)
1=1
With To, the frequent bootstrapping during model roll-out can be reduced. Specifically, given the maximum backtracking length m, state-action sequence of length m, (81, (1, 82, '2, 8m,; am)) sampled from the data buffer to start the roll-out in Tg. To obtain the prediction $m+1, an integer between [1, mi is chosen uniformly at random as the backtracking length. If only one step is selected for backtracking, (Sm, am) will be fed into Ta to obtain the prediction result; if m- steps chosen, (82, (2mm) will be fed into To, and so forth. The right part of Figure illlssrrate the aforementioned process based on random backtracking. Similarly, subsequent state predictions can also backtrack at most m steps. After rolling out several steps, inevitably backtracks previously predicted states. For example, for the prediction 82m, it can only backtrack to one of (sm: 11+1,50 8m, 82m 1).
The backtracked state, one part of the attribution for the next state prediction, located several steps ahead, from the perspective of expectation. Thus, ADM reduces the actual bootstrapping count of rolled-out trajectory. The complete H H-step roll-out process in ADM is described in Algorithm [
Similar to existing MBRL algorithms, policy roll-out in ADM generate large number of fake samples for policy updates. We refer to the new dyna-style ADM-based policy optimization framework as ADMPO (ADM-based Policy Optimization). Any policy optimization algorithm be plugged into this framework. the subsequent subsections, we will introduce two new model-based algorithms, ADMPO-ON and ADMPO-OFF, for online and offline settings, respectively.
4

---

Published as a conference paper at ICLR 2025
3.2 ADMPO- ON: ADM FOR POLICY OPTIMIZATION IN ONLINE SETTING
In the online setting, the agent interacts with the real environment while simultaneously optimizing the policy. Like MBPO (Janner et al. 2019), ADMPO-ON can be divided into two alternating stages, namely updating the dynamics model with continuously collected samples and utilizing samples generated through model roll-outs additionally for policy optimization. ADMPO-ON replaces the ensemble dynamics model in MBPO framework with ADM. It trains ADM with the optimization objective shown in Equation (2) and generates large number of fake samples using the roll-out method depicted in Algorithm
The detailed pseudo-code is provided by Algorithm 2 in Appendix
During roll-outs, ADM randomly selects a backtracking length at each step and attributes the states to be predicted to variable-length plans. While backtracking steps, we view the sampling process ($.41, f(+1) ~ To(-15t- a(-x+1a) as ($1+1, fr+1) = μe (5t Qc-k+1a) + 7 with Mh-1 (0, [g($ dt-k+1:)), where d4 is the deterministic dynamics function and Eg is the standard deviation function used to construct the noise distribution with zero mean. In expectation,
the target value of Q(s,, ae) is estimated
E (3-- -t If (-1xe)
E (Be4s Poloa
where y(s++1, 11+1) f--1 ( (Q a)) Data generated via roll-outs in our ADM can be viewed as an implicit augmentation. The augmentation stems from two sources: (i) variation of the backtracking length while applying the learned ADM to predict the next state, and (ii) the noise introduced by the distribution (0, Eo(st ai 1=:))) at each backtracking length k. According to (Zheng et al., 2023), variations of state predictions Lipschitz condition of the Q network around regions where the model prediction is uncertain, thereby regulating the value-aware model error (Farahmand et al. 2017).
effectively implicitly regularize the local
3.3 ADMPO-OFF: ADM FOR POLICY OPTIMIZATION OFFLINE SETTING
In the offline setting, due to limitations of the behavior policy corresponding to the dataset, the learned ADM can only cover some regions of the state-action space. Beyond these safe regions lie the risky regions where the model is uncertain and unable to be fixed since online exploration is inaccessible to the agent. To prevent policy optimization collapse, exploitation
the learned model needs to be focused within the safe regions. Simultaneously. efforts should be made to explore beyond the boundaries of the safe regions to discover samples conducive to a better policy than the behavior policy. Achieving such balance between conservatism and generalization often requires measuring model uncertainty. Based on ADM, we will introduce a new uncertainty quantification method.
In our ADM, states predicted using different backtracking lengths exhibit discrepancies. Intuitively, these discrepancies are closely related to the data distribution. When the agent is in safe regions, the discrepancies are small. As the agent gradually moves toward risky regions, the discrepancies tend to increase. The difference among probabilistic predictions Th f- a 1 obtained with different backtracking k serve as a natural measure of model uncertainty, which can be quantified using variance (or standard deviation), as defined by Definition 3.2
Definition 3.2 AADM-Uncertiin) Quantifier). For any maximum backtracking length and the corresponding learned ADM Tg. the uncertainty of To at (st, a.) is quantified as
ADM (st, az) =E --pm--1(1aa
Var U-iiformmm))) (aaiii/foo-ls (-18: aiala kiia) [
(4)
1 (1E)) (1x))) (z)) k=1
=Ep =Bpm--((/aa)
for any 8t E S and ag A, where ΣS E0( ng م p a1-)+1a) for convenience, and 1 m k= ugg.
This uncertainty term corresponds to a combination of epistemic and alcatoric model uncertainty with similar form to the ensemble standard deviation (Lu et al. 2022 Lakshminarayanan et al. (017).
5

---

Published as a conference paper at ICLR 2025
However, the source of diversity has shifted from ensemble to variable backtracking lengths. Since estimating the approximation error via epistemic or aleatoric uncertainty has been applied in many works (Yu et al.) 2020; Bai et al., 2022 Sun et al., 20233 Lu et al. 2022), we assume that our ADM uncertainty 4 is an admissible error estimator (Yu et al. 2020). as described in Assumption 3.3
Assumption 3.3 (Admissible Error Estimator), Assume that there exists a positive b E R such that the following inequality (5) holds for any maximum backtracking length and any 8t E S,ae at A.
Drv (-|s;, ar), T(-ls;, at)) S b UADM (s,,a,);
(5)
where To.m is the overall conditioned distribution coming from
To (-
Σ
Ih a (1 (6)
m
01-641it-1
Under Assumption 3.3 and the E-uncertainty quantifier definition (see Appendix A for details)
proposed by PEVI (Jin et al. 2021). we present the following theorem, demonstrating that 4/ can serve as Eunceerainty quantifier to bound the Bellman error.
Theorem 3.4. 3 LAD is valid f-uncertainty quantifier, with B b "Tmax Specifically,
1-7
(=)())) -Tannaa))T((((,,, an) S 1--55B--DD UADM (s,, as),
where To is the proxy Bellman operator induced by ADM estimate the true Bellman operator T
Proof. See Appendix B
According to the suboptimality theorem (see Appendix A for details) presented by PEVI (Jin et al. 2021), the policy ă derived via pessimistic value iteration, which incorporates any E-uncertainty quantifier as a penalty term into the value iteration process (Sutton & Barto 2018), has a bounded optimality gap the optimal policy x². The optimality gap dominated by the Bellman error and the uncertainty quantification. Intuitively. the Bellman error is usually small in safe regions where the dynamics model has been trained with rich data and tends yield high consistency under different backtracking lengths, while large errors often appear in risky regions where data is scarce and the predictions via backtracking different lengths become inconsistent. The penalization prevents the policy from taking actions leading it to risky regions, thus the model cannot induce inaccurate value estimations on these actions. Thus, we can penalize the Bellman operator to obtain a pessimistic value estimation by
ADM ar) = (=Q -x) 3 -- ADM (st, az).
(8)
We expect the penalty term B 4AMM (st; 0t) as small as possible thereby constraining the optimality gap. While our Assumption 3.3 lacks theoretical guarantees and the tightness of the bound Theorem 3.4 is unclear, we have provided sufficient evidence in Section 4.3.3 that our uncertainty quantification effectively estimates the model error.
Overall, ADMPO-OFF is the offline version of ADMPO-ON, which introduces a penalized Bellman operator (8) into the policy optimization process of ADMPO-ON, following the algorithmic frame- work of MOPO (Yu ct al. 2020). The detailed pseudo-code is provided by Algorithm in Appendix
2
EXPERIMENTS
In this section, we conduct several experiments to answer: (1) Does ADM roll out samples with less compounding error than the ensemble dynamics model? (2) How well does ADMPO-ON perform in the online setting? (3) How well does ADMPO-OFF perform in the offline setting? Does ADM quantify the model uncertainty better than the ensemble dynamics model?

---

Published as a conference paper at ICLR 2025
hopper-medium
hopper-medium-repee
walker 2d-medium
walker?d--ediimm-rpla
100
1034
10
10"
Pred Enter
1024
100
207
Prodicton Emror
2030
1ote
Prediction
10'
2033
1000
10t4
10-4
101
10-
25
25
75
50
75
100
Length
ADM (ours)
Ensemble Dynamics Model
pootstrapping RNN Dynamics Model
Comparison among ADM, ensemble dynamics model, and bootstrapping RNN dynamics model, in terms of the growth curve of the compounding error (in log scale) as roll-out length increases, after offline learning. The overflow value is regarded as the maximum value of float 32.
Figure 2:
4.1 DYNAMICS MODEL EVALUATION
An essential metric for evaluating dynamics model quality is the compounding error, which increases with the roll-out length. We selected four D4RL (Fu et al.] 2020) datasets, hopper-medium--2. and
compare the compounding eπror between ADM and the commonly used ensemble dynamics model. To eliminate the influence of the RNN structure, we also compare the bootstrapping RNN dynamic model, which shares the same structure as ADM but makes the prediction 8,+1 using the historical state-action St, a, as input, where & is uniformly sampled from [I. m]. The dataset is divided into a training set and validation set, with the former used to train the dynamics model and the latter used to evaluate the compounding error. Figure 2 shows the growth curves the compounding eтror the roll-out length increases, with the y-axis in the log scale. The linear-scale version is shown in Appendix E. We observe that the curves of ADM remain close to zero, while the other two models exhibit exponential growth as the roll-out length exceeds certain threshold. This phenomenon suggests ADM can improve predictions for future states due to its any-step backtracking mechanism during model roll-outs. We compare the performance of offline RL training using ADM and the bootstrapping RNN model in Appendix E.4
baikker2d-mediumm-a
sequence ($--4+1, dt=k++
4.2 EVALUATION IN ONLINE SETTING
We evaluate ADMPO-ON on four difficult MuJoCo continuous control tasks (Todorov et al. 201 2),
including Hopper, Walker2d, Ant, and Humanoid. All the tasks adopt version v3 and follow the default settings. Five model-based methods and one model-free method are selected as our baselines. These include SAC (Haamoja al. 2018), which is the tfate-of-the-an model-free free RI. algorithm; STEVE (Buckman et a 2018). which incorporates an ensemble into the model- based value expansion; MBPO (Janner et al.[[20 2019), which updates the policy with mixture of real environmental samples and branched roll-out data; BMPO (Lai et al.) 2020), which builds upon MBPO and replaces the dynamics model with a bidirectional one; DDPPO (Li et al. 2022], which adopts a two-model-bass learning method to control the prediction error and the gradient error; and MACURA (Frauenknecht
ct al. 2024), which uses inherent model uncertainty to consider local accuracy to make roll-out.
Figure shows the learning curves of ADMPO-ON and six baselines, along with SAC's asymptotic performance. ADMPO-ON achieves competitive performance after fewer environmental steps than most of the baselines. Taking the most difficult Humanoid as an example, ADMPO-OP and MACURA have achieved 100% of SAC convergence performance (about 6000) after 150k steps, while DDPPO needs about 200k steps, and the other four methods can't get close to the blue dashed line even at step 300k. ADMPO-ON matches the performance of the previous state-of- the-art algorithm, MACURA, and dominates other baselines terms of learning efficiency on the Humanoid task. After training, ADMPO-OX can achieve final performance close to the asymptotic performance of SAC on all these four MuloCo tasks. These results demonstrate that ADMPÓ-ON has both high sample efficiency and competitive performance. Further study on why ADMPO-ON performs well in the online setting can be found in Appendix E.2
It would be (8t-k+1; Ot-k+1,"" St, 0t) from the dataset if were the beginning of the roll-out.
7

---

Published as a conference paper at ICLR 2025
Hopper
Walker2d
Ant
Humanoid
4090
9050
3000
5900
Sewward
8900
1998
3023
Time
Time
Time Steps
Time Steps
ADMPO-ON (ous)
MACURA
DOPPO
MBPO
STEVE
SAC
Converganoe
Figure 3: Online learning curves of ADMPO-ON (red) and other six baselines on four MuJoCo-v3 tasks. The blue dashed lines indicate the asymptotic performance of SAC for reference. The solid lines indicate the mean while the shaded areas indicate the standard error over five different seeds.
Table Normalized scores after offline learning on D4RI, MuloCo tasks, averaged over five seeds.
Task Name
BC
CQL
TD3
EDAC
MOPO
COMBO
RAMBO
CBOP
MOBILE
ADMPO-OFF (ours)
3.7
5.3
8.5
25.3
31.7
17.9
25.4
31,4
31.9
32.7=-0.2
salfcheetah-ruddom
22
28.4
39.5
32.8
39.3
45.4
walker2d-nandonn
1.3
5.4
1.6
16.6
7.4
7.0
0.0
17.8
17.9
22.2:10.2
hopper-mn mediun
51.1
61.9
59.3
101.6
62.8
97.2
87.0
102.6
106.6
107.4.0.6
halfchettah--nddiie
43.2
72.2:10.6
walker30-mediue
70.9
79.5
83.7
92.5
84.1
81.9
84.9
45.5
87.7
93. 13.2::111
hopper medium replay
16.6
$6.3
60.9
101.0
103.5
89.5
99.5
104.3
103.9
104.4:0.44
halfcheetuh-mediune replayy
37.6
45.3
44.6
61.3
55.1
68.7
66.4
71.7
6.
walker2d medium--ephle
20.3
76.8
$1.8
87.1
56.0
89.2
92.7
95.6-22.1
hopper- mxduun-eenn
53.9
96.9
98.0
110.7
81.6
111.]
88.2
111.6
112.6
112.7:10,3
halfcheetah medium expert
44.0
95.0
90.7
106.3
90.8
90.0
95.4
105.4
108.2
102.7.0..2
walkerdd medium- expert
90.1
109.]
110.1
114.7
103.3
56.7
117.2
115.2
111.9:1033
Average
36.5
61.6
58.2
76.0
70.3
66.8
67.7
79.3
80.0
$1.0
4.3
EVALUATION IN OFFLINE SETTING
4.3.1 D4RL BENCHMARK RESULTS
We compare ADMPO-OFF with four model-free methods: BC (behavioral cloning), which simply imitates the behavior policy of the dataset; CQL (Kumar et al. 2020), which equally penalized the Q values on the--fffth-ddsstiuuuttn state-action pairs; TD3+BC (Fujimoto & Gu (021). which simply incorporates a BC term into the policy optimization objective of TD3 (Fujimoto al., (018); and EDAC (An et al. 202T) which quantifies the Q uncertainty via ensemble; as well as five model- based methods: MOPO Yuu et al., 2020), which adds the uncertainty of the model prediction as a penalization term to the reward function; COMBO (Yu et al. (022I). which introduces the penalty function of CQL into the model-based framework: RAMBO (Rigter et al. 2022), which adversarially trains the dynamics model and the policy; CBOP (Jeong et al. (023), which adopts the variance of values under an ensemble of dynamics models to estimate the Q value conservatively under MVE (Feinberg et al. 2018) regime; and MOBILE (Sun et al. 2023). which proposes Model-Bellman
inconsistency to estimate the Bellman error.
Table reports the results twelve D4RL (Fu al. 2020) MuJoCo datasets (v2 version). The normalized score for each dataset is obtained via online evaluation after offline learning. The source of the reported performance is provided in Appendix [D.5] We observe that ADMPO-OFF outperforms the other nine baselines in most tasks and achieves the highest average score. Notably, ADMPO-OFF has significant performance advantage over MOPO. This directly demonstrates the effectiveness of
ADM, since ADMPO-OFF only replaces the dynamics model in MOPO with ADM.
4.3.2 NEORL BENCHMARK RESULTS
NeoRL QQnn et al. 2022). is offline RL benchmark that collects the data in manner more conservative and closer to real-world data-collection scenarios. We focus on nine datasets collected using policies of three different qualities (low, medium, and high) in three environments Hopper- v3, HalfCheetah- v3, and Walker2d-v3 respectively. In our evaluation, each dataset contains 1000 trajectories.
8

---

Published as a conference paper at ICLR 2025
Table 2: Normalized scores after offline learning NeoRL. tasks, averaged over five seeds.
Task Name
BC
CQL
TD3+BC
EDAC
MOPO
MOBILE
ADMPO-OFF (ours)
nearl- sarri-hopper--aee
15.1
16.0
15.8
18.3
6.2
17.4
22.3:10.1
neori-halle cerd-lalfccheeat-lor
29.1
38.2
30.0
31.3
40.1
54.7
52.81:122
walker2d low
40.2
55.9
neort-hopper--ne hopper-medium
51.3
64.5
70.3
44.9
1.0
51.1
51.5+:5.0
nearl hulfcheetah medium
49.0
54.6
52.3
54.9
62.3
77.8
6..311.7
walke--alker2d--- medeiium
48.7
57.3
58.5
$7.6
39.9
62.2
70.1::2.44
nood-hhoppe-higg woorf-halffhhetaaheee
43.1
76.6
75.3
52.5
11.5
87.8
87.6+t449
71.3
77.4
75.3
81.4
65,9
83.0
84.0-20.88
nearl walker2d high
72.6
75.3
69.6
75.5
18.0
74,9
12.2.1
Average
50.7
64.0
ADM (Correlation: 0.981
Ensemble (Correlation: 0.94)
ADM (Correlation: 0.98)
Ensemble (Correlation:
random astion
random action
Moddo Uhoorraie
heamespolayy dutaact
Hosel Eapoer mwreeettty
learedd pol
calvoet
dulaseet
Mode Erdr
Model Error
Model Error
Model Eroer
(a) hopper- mameediu--rppplay
(b) walker2d geedium-reppay-ve v2
Figure 4: Comparison between ADM and ensemble model in uncertainty quantification.
We compare our ADMPO-OFF with six baselines, including BC, CQL. TD3+BC, EDAC, MOPO, and MOBILE. Table 2 presents the normalized scores these methods. Due to the narrow and limited coverage of the NeoRL data, all the baselines experience decline in performance. In contrast, our ADMPO-OFF maintains a relatively high-level average score, still achieving superior performance in most tasks. This remarkable out-performance indicates the potential of our algorithm in more challenging real-world tasks.
4.3.3 UNCERTAINTY QUANTIFICATION
In our analysis, we sample a lot of state-action pairs in the learned ADM and the ensemble dynamics model respectively. These samples are obtained by model roll-out with three types of policy: random action selection, the learned policy after offline training, and the behavior policy of the dataset. Subsequently, we measure their model uncertainty and model error. The resulting scatter plots on two and
are illustrated in Figure We observe that our ADM provides better quantification for the model uncertainty. On the one hand, points sampled in ADM with greater model errors tend to exhibit greater quantified model uncertainty. The correlation coefficient of 0.98, observed across both tasks, surpasses that of the ensemble dynamics model. On the other hand, ADM can distinguish the samples from different policies better than the ensemble model. Samples generated from random actions deviate from the dataset distribution, whose uncertainty should be maximum in expectation. Conversely, when the learned policy is optimized within the safe regions covered by the dataset, model uncertainty is expected to be minimal. The experimental plots of ADM illustrate this phenomenon more clearly.
D4RL tasks,
medium-replay
RELATED WORK
This work is related to online and offline dyna-style MBRL Sutton 1990).
5.1 ONLINE MODEL-BASE BASED REINFORCEMENT LEARNING
In the online setting, MBRL algorithms aim to accelerate value estimation or policy optimization with model roll-out data. MVE (Feinberg al. 2018) enhances Q-value target estimation by allowing short-term imagination to a fixed depth using the dynamics model. STEVE (Buckman et al. 2018)
builds upon MVE by incorporating an ensemble into the value expansion to better estimate the Q

---

Published as a conference paper at ICLR 2025
value. SLBO (Luo et al. 2019) directly utilizes TRPO (Schulman et al. with synthetic data generated by rolling out to the end of trajectories in the dynamics model. MBPO (Janner et al. 2019) proposes a branched roll-out scheme truncate unreliable samples, thereby reducing the influence of compounding error (Xu et al. 2021), and employs SAC (Haarnoja et al. 2018) to update the policy with mixture of real-world data and model-generated data.
2013 to optimize the policy
Recent work improves MBRL performance mainly from two perspectives. One focuses on learning better dynamics model, such as bidirectional models (Lai et al.) 2020), adversarial models (Chen
et al.
2023
Bhardwaj et al. 2023), causal models (Zhu et al. 2022), multi-step models (Asadi et al.
2018
2019 Che al.) 2018), and energy-based models (Chen al., 2024). The other pursues
better
utilization of the learned model, enhancing the reliability of model-generated samples (Pan
al.
2020)
or applying model-based multi-step planning techniques (Chua et al. 2018 Clavera
et al.
2020
Karkus et al. 2019 Okada et al. 2017; Srinivas et al. 2018 Lin et al.
2023)
5.2
OFFLINE MODEL-BASED REINFORCEMENT LEARNING
Although some model- free RL algorithms (umaar et al. 2019 Fujimoto et al.) 2019 Fujimoto & Gui
2021 Kumar et al., 2020 An et al.) [021] Bai et al., 2022) have made significant contributions to offline RL research, MBRL algorithms appear to be more promising for the offline setting since they can utilize the dynamics model to extend the dataset and largely improve the data efficiency.
The core issue of offline MBRL lies how to effectively leverage the model. MOPO (Yu et al. 2020 and MOReL. (Kidambi et al. (020) add the uncertainty of the model prediction as a penalization term to the original reward function to achieve pessimistic value estimation. MOBILE (Sun al., 2023) improves the uncertainty quantification by introducing Model-Bellman inconsistency into the offline model-based framework. COMBO (Yu et al. 202T) applies CQL (Kumar et al. 2020) to force the Q value be small on model- ggeeeratd out-of-di -fi-ttiiiittiitto samples. RAMBO (Rigter al. 2022) achieves conservatism by adversarial model learning for value minimization while keeping fitting the transition function. CBOP (Jeong al. 2023) introduces adaptive weighting of short-horizon roll-out into MVE (Feinberg et al. 2018) technique and adopts the variance of values under an ensemble of dynamics models to estimate the value conservatively. MOREC (Luo al., 2024a) designs a neward-consistend dynamics model using an adversarial discriminator to let the model -ener-egeeeeted samples be more reliable.
6 CONCLUSION
In this work, we propose a new method for environment model learning and utilization, namely Any-step Dynamics Model (ADM). ADM is applicable in both online and offline MBRL frameworks, which yields two algorithms, ADMPO-ON and ADMPO-OFF, respectively. Several analyses and experiments show that ADM outperforms the ensemble dynamics model applied in previous MBRL approaches widely. Specifically, the compounding errors of multi-step roll-outs in ADM are much smaller than those in ensemble dynamics models. As result, ADMPO-ON achieves significantly higher sampling efficiency in the online setting compared to MBPO. Additionally, the uncertainty term computed using ADM is more consistent with the model error, enabling ADMPO-OFF to achieve performance improvements in the offline setting. The remaining issues include that RNNs may consume more resources during the training process, and the effectiveness of ADM may be reduced in highly stochastic environments. These could serve directions for future research. We believe ADM has powerful potential beyond the capabilities demonstrated in this paper. the future, we will explore the scalability of ADM in Man-Markovian visual RL scenarios, considering both online and offline settings.
ACKNOWLEDGMENTS
This work is supported by NSFC (62495093) and Jiangsu Science Foundation (BK20243039). The authors thank anonymous reviewers for their helpful discussions and suggestions for improving the article.
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song. Uncertainty-base based offline reinforcement learning with diversified q-ensemble. In Advances Neural Information Processing Systems 34 (NeurIPS 21), Virtual Event, 2021.
Kavosh Asadi, Evan Cater, Dipendra Misra, and Michael I.. Littman. Towards simple approach to multi-step model-based reinforcement learning. CoRR, abs/ 111s/1811/// 1118111..8 1.00188, 2018.
Kavosh Asadi, Dipendra Misra, Seungchan Kim, and Michael L. Littman. Combating the compounding- problem with a multi-step model. CoRR, abs/ 1tss/9955.3300, 2019.
Chenjia Bai, Lingxiao Wang. Zhuoran Yang, Zhi-Hong Deng, Animesh Garg, Peng Liu, and Zhaoran Wang. Pessimistic bootstrapping for uncertainty-drie driven offline reinforcement learning. In 10th International Conference Learning Representations (ICLR Virtual Event, 2022.
Mohak Bhardwaj, Tengyang Xie, Byron Boots, Nan Jiang, and Ching-An Cheng. Adversarial model for offline reinforcement learning. In Advances in Neural Information Processing Systems 36 (NeurIPS 23), New Orleans, LA, 2023.
Jacob Buckman, Danijar Hafner, George Tucker, Eugene Brevdo, and Honglak Lee. Sample- efficient reinforcement learning with stochastic ensemble value expansion. In Advances in Neural Information Processing Systems 31 (NeurIPS' 18), Montréal, Canada, 2018.
Tong Che, Yuchen Lu, George Tucker, Surya Bhupatiraju, Shane Gu, Sergey Levine, and Yoshua Bengio. Combining model-based and model-free RI. via multi-step control variates, httpp::/ openreview net. 2018. URL
Ruifeng Chen, Chengxing Jia, Zefang Huang, Tian-Shuo Liu, Xu-Hui Liu, and Yang Yu. Offline transition modeling via contrastive energy learning. In Proceedings of the 41st International Conference on Machine Learning (ICML 24), Vienna, Austria, 2024.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith W. Ross. Randomized ensembled double qelearning: Learning fast without model. In 9th International Conference Learning Representations (1), Virtual Event, 2021.
Xiong-Hui Chen, Yang Yu, Zhengmao Zhu, Zhihua Yu, Zhenjun Chen, Chenghe Wang, Yinan Wu, Rong-Jun Qin, Hongqiu Wu, Ruijin Ding, and Fangsheng Huang. Adversarial counterfactual envi- ronment model learning. In Advances in Neural Information Processing Systems 36 (NeurIPS'23), New Orleans, LA, 2023.
Kyunghyun Cho, Bart van Merrienboer, Çaglar Gülçehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP'14)) 144), Doha, Qatar, 2014.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems 31 (NeurIPS' 18), Montréal, Canada, 2018.
Ignasi Clavera, Yao Fu, and Pieter Abbeel. Model-augmented actor-critic: Backpropagating through paths. In 8th International Conference on Learning Representations (ICLR 20), Addis Ababa, Ethiopia, 2020.
Jeffrey Elman. Finding structure in time. Cognitive Science, 14(2):179-211, 1990.
Amir Massoud Farahmand, André Barreto, and Daniel Nikovski. Value-aware loss function for model-based reinforcement learning. Proceedings of the 20th International Conference on Artificial Intelligence and Statistics (AISTATS' 17),, Fort Lauderdale, FL, 2017.
Vladimir Feinberg, Alvin Wan, Ion Stoica, Michael I. Jordan, Joseph E. Gonzalez, and Sergey Levine. Model-based value estimation for efficient model-free -free reinforcement learning. CoRR, abs/180 001.00011, 2018.
11

---

Published as a conference paper at ICLR 2025
Angelos Filos, Eszter Vértes, Zita Marinho, Gregory Farquhar, Diana Borsa, Abram L. Friesen, Feryal M. P. Behbahani, Tom Schaul, André Barreto, and Simon Osindero. Model-value inconsistency as signal for epistemic uncertainty. In Proceedings of the 39th International Conference on Machine Learning (ICMI'22)) Baltimore, Maryland, 2022.
Bernd Frauenknecht, Artur Eisele, Devdutt Subhasish, Friedrich Solowjow, and Sebastian Trimpe. Trust the model where trusts itself model-based actor-critic with uncertainty-aware rollout adaption. In Proceedings of the 41st International Conference on Machine Learning (ICMI/24), 2024.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. CoRR, abs/2004.07219, 2020.
Scott Fujimoto and Shixiang Shane Gu. minimalist approach offline reinforcement learning. In Advances in Neural Information Processing Systems 34 (NeurIPS Virtual Event, 2021.
Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference on Machine Learning (ICML'I8), Stockholm, Sweden, 2018.
Scott Fujimoto, David Meger, and Doina Precup. Off-policy deep reinforcement learning without exploration. In Proceedings of the 36th International Conference on Machine Learning (ICML'19) Long Beach, California, 2019.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Proceedings of the 35th International Conference on Machine Learning (ICML' 18), Stockholm, Sweden, 2018.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems 32 (NeurIPS' 19), Vancouver, Canada, 2019.
Jihwan Jeong, Xiaoyu Wang, Michael Gimelfarb, Hyunwoo Kim, Baher Abdulhai, and Scott Sanner. Conservative bayesian model-based value expansion for offline policy optimization. In The 11th International Conference on Learning Representations (CCLR'33), Kigali, Rwanda, 2023.
Ying Jin, Zhuoran Yang, and Zhaoran Wang. Is pessimism provably efficient for offline RL. In Proceedings of the 38th International Conference on Machine Learning (ICML'21), Virtual Event, 2021.
Péter Karkus, Xiao Ma. David Hsu, Leslie Pack Kaelbling, Wee Sun Lee, and Tomás Lozano-Péee Pérez. Differentiable algorithm networks for composable robot learning. Robotics: Science and Systems XV (RSS') 19), Freiburg im Breisgau, Germany, 2019.
Rahul Kidambi, Aravind Rajeswaran, Pranceth Netrapalli, and Thorsten Joachims. Morel: Model- based offline reinforcement learning. In Advances in Neural Information Processing Systems 33 (NeurIPS'200 20), Virtual Event, 2020.
Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine. Stabilizing off-policy Q-learning via bootstrapping еттоr reduction. In Advances in Neural Information Processing Systems 32 (NeurIPS' 19), Vancouver, BC, 2019.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative Q-learning for offline reinforcement learning. In Advances in Neural Information Processing Systems 33 (NeurIPS'20) Virtual Event, 2020.
Hang Lai, Jian Shen, Weinan Zhang, and Yong Yu. Bidirectional model- policy optimization. In Proceedings of the 37th International Conference on Machine Learning (ICML'20), Virtual Conference, 2020.
Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in Neural Information Processing Systems 30 (NeurIPS' 17), Long Beach, CA, 2017.
12

---

Published as a conference paper at ICLR 2025
Chongchong Li, Yue Wang, Wei Chen, Yuting Liu, Zhi-Ming Ma, and Tie- Yan Liu. Gradient infor- mation matters in policy optimization by back prock-pppgggting through model. In 10th International Conference Learning Representations (ICLR'22) 22)), Virtual Event, 2022.
Haoxin Lin, Yihao Sun, Jiaji Zhang, and Yang Yu. Model-based reinforcement learning with multi-step plan value estimation. In Proceedings of the 26th European Conference Artificial Intelligence (ECA/'23) 23),, Kraków, Poland, 2023.
Cong Lu, Philip J. Ball, Jack Parker-Holder, Michael A. Oshorne, and Stephen J. Roberts. Revisiting design choices in offline model based reinforcement learning. In The 10th International Conference Learning Representations (ICIRR Virtual Event, 2022.
Fan-Ming Luo, Tian Xu, Xingchen Cao, and Yang Yu. Reward-consistent dynamics models are strongly generalizable for offline reinforcement learning. In The 12th International Conference on Learning Representations (ICLR'S 24), Vienna, Austria, 2024a.
Fan-Ming Luo, Tian Xu, Hang Lai, Xiong- Hui Chen, Weinan Zhang, and Yang Yu. survey on model-based reinforcement learning. SCIENCE CHINA Information Sciences, 2024b.
Yuping Luo, Huazhe Xu, Yuanzhi Li, Yuandong Tian, Trevor Darrell, and Tengyu Ma. Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. In 7th International Conference on Learning Representations (ICLR' 19), New Orleans, LA, 2019.
Marlos C. Machado, André Barreto, Doina Precup, and Michael Bowling. Temporal abstraction in reinforcement learning with the successor representation. Journal of Machine Learning Research, 24:80:1-80:69, 2023.
Masashi Okada, Luca Rigazio, and Takenobu Aoshima. Path integral networks: End-to-end differen- tiable optimal control. CoRR, 0bss//70..995/ 2017.
Feiyang Pan, Jia He, Dandan Tu, and Qing He. Trust the model when it is confident: Masked model-based actor-critic. In Advances in Neural Information Processing Systems 33 (NeurIPS 20), Virtual Event, 2020.
Rongjun 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 35 (NeurIPS'22) New Orleans, LA, 2022.
Mare Rigter, Bruno Lacerda, and Nick Hawes. RAMBO-RL. Robust adversarial model-based offline reinforcement learning. In Advances in Neural Information Processing Systems 35 (NeurIPS'22) 22), New Orleans, LA, 2022.
John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In Proceedings of the 32nd International Conference on Muchine Learning (ICML'IS)) 15), Lille, France, 2015.
Aravind Srinivas, Allan Jabri, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Universal planning networks: Leaniing generalizable representations for visuomotor control. In Proceedings the 35th International Conference on Machine Learning (1), Stockholm, Sweden, 2018
Yihao Sun. Offlinerl-kit: An elegant pytorch offline reinforcement learning library.
https:
2023.
Yihao Sun, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. Model-bellman incon- sistency for model-based offline reinforcement learning. In Proceedings of the 40th International Conference on Machine Learning (ICMI'23). Honolulu, Hawaii, 2023.
Richard S. Sutton. Integrated architectures for learning, planning, and reacting based on approxi- mating dynamic programming. In Proceedings of the 7th International Conference Machine Learning ( (CCLL'90), Austin, Texas, 1990.
Richard S. Sutton and Andrew G. Barto. Reinforcement learning: An introduction. MIT Press, 2018.
13

---

Published as a conference paper at ICLR 2025
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for modell-asee control. In IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS' 20), Vilamoura, Portugal, 2012.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30 (NeurIPS' 17), Long Beach, CA, 2017.
Tian Xu, Ziniu Li, and Yang Yu. Error bounds of imitating policies and environments for reinforce- ment learning. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y. Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. In Advances in Neural Information Processing Systems 33 (NeurIPS 20), Virtual Event, 2020.
Tianhe Yu, Aviral Kumar, Rafael Rafailov, Aravind Rajeswaran, Sergey Levine, and Chelsea Finn. COMBO: Conservative offline model-based policy optimization. In Advances in Neural Informa- tion Processing Systems 34 (NeurIPS'2I), Virtual Event, 2021.
Zhilong Zhang, Ruifeng Chen, Junyin Ye, Yihao Sun, Pengyuan Wang, Jingcheng Pang, Kaiyuan Li, Tianshuo Liu, Haoxin Lin, Yang Yu, and Zhi-Hua Zhou. WHALE: towards generalizable and scalable world models for embodied pocision-making.. CoRR, abs/2411.05619, 2024.
Ruijie Zheng, Xiyao Wang, Huazhe Xu, and Furong Huang. model ensemble necessary? Model- based RL via single model with lipschitz regularized value function. In The 11th International Conference on Learning Representations (ICLR'23) 23), Kigali, Rwanda, 2023.
Zheng-Mao Zhu, Xiong-Hor Hui Chen, Hong-Long Tian, Kun Zhang, and Yang Yu. Offline reinforcement learning with causal structured world models. CoRR, abs/2206. 01474, 2022.
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL INTRODUCTION TO PESSIMISTIC VALUE ITERATION (PEVI)
Pessimistic Value Iteration (PEVI) (Jin ct al. 2021) is a meta-algorithm for offline RL settings.
constructs an estimated Bellman operator 7 based on the given dataset Deny to approximate the true Bellman operator T that satisfies
V*V++1(5h, 9x) (sn (10%)))))) [r++1 + (*+1($)+1)),
(9)
where h is the step index less than the horizon H. Then the state-action value function is updated with
(n an) < Vn+1 ah) (s an)
(10)
for each (sh, ax), where A₁ is the penalty function that guarantees the conservatism of the learned policy. Especially, A, should be f-uncertainty quantifier follows.
Definition A.1 (f-Uncertainty Quantifier (proposed by (Jin et al. 2021))). The set of penalization (An)e forms E-uncertainty quantifier if
{ Vi+1 (Sh, an) - V an) S As (sh, as)
(11)
holds with probability at least - & for all (sh: an) E S x A.
The following theorem characterizes the suboptimality of PEVI.
Theorem A.2 (Suboptimality PEVI (proposed by (Jin et al. 2021))) Suppose {Am] in PEVI set of E-uncertainty quantifier. Then the derived policy satisfies
VT (si) 2 $2Σ Bpr (An(on,an))
(12)
with probability least for all starting S. Here E is the expectation with respect to the trajectory induced by the optimal policy xt in the underlying MDP given the function An-
Proof. See PEVI Jin al. 2021) for detailed proof.
B ADM UNCERTAINTY ANALYSIS
Theorem B.1. B BADD is a valid E-uncertainty quantifier, with 3 = b2 Specifically,
Q= az) - -(= an) S. 3. ADM (st, at),
(13)
where fo is the proxy Bellman operator induced by ADM estimate the true Bellman operator T
Proof. First, define 0(5141,P--1) fett + (Ba (Q(5+11, and expand these two Bellman operator to 7s' Q(t,
1
=B (st (sc-m+tit m jarm-tt(ae
(f ( f(+1)l
1=1
p (8-mm11, 41-
5t Ai-met1
R 8ot1 f2
4p-m+1:t-1
Σ pk- (1--- Σ
To(-Ise-kat (4)
k=l
St-k41 Me-b+1i-1
Pr-
Σ
$111.fu11
(14)
15

---

Published as a conference paper at ICLR 2025
and
T=Q(s):
(15)
Σ
1141
Then,
obtain
T- Q(s,)
-T Q(s,,
Σ
( f++1)
Σ
T( f anl (a-a(( [((5+11))]]
8111-f241
(16)
Mmmax
Σ
(044, az) T ( f211s,,,
Mmaxx Drv (To (-18,, at), T(.)st,4))
1-7 ?
Vmmax ADM (sc.0)))
Thus, let B b Fmaax
1-y we can say that 3 8-24 LADM is valid f-uncertainty quantifier, defined by
Definition A.1
C IMPLEMENTATION DETAILS
C.I ADMPO-ON
Our ADMPO-OX algorithm follows the framework of MBPO (Janner et al. 2019) as shown in Algorithm 2
The only difference between ADMPO-ON and MBPO lies in the way the dynamics model is trained and utilized, as indicated by the parts highlighted in blue in the pseudo-code.
Algorithm 2 ADMPO-ON
Input: Initial ADM Ta and policy Tg. roll-out length H, maximum backtracking length m. real data buffer Denv. model data buffer Dmodel warm-up size U, interaction epochs N, steps per epoch E
1: Explore for U environmental steps and add data to Deuv
2: for N epochs do
3: Train ADM T, on Deny by maximizing Equation
4: for to do
5: Sample action a, according to (-(s))
6: Perform a, in the environment and add the real sample (5t, des Te+1, S141) D Denv
7:
for M model roll-outs do
8:
Sample initial m-step state-action sequence (Sts4m-1 4s:i+m-2) from Denv
9:
Roll out f steps in 7. via ADM-Roll(Ts, Tos H, m, (9ts4m-1: 4m)4mm2)) and add the model roll-out data to Diodell
10:
end for
for G policy updates do
12:
Update current policy Th using samples from Dany U D model 13: end for
14:
end for
end for
16

---

Published as a conference paper at ICLR 2025
C.2 ADMPO-OFF
Our ADMPO-OFF algorithm follows the framework of MOPO Yuu al. 2020p, shown in Algorithm 3 The only difference between ADMPO-OFF and MOPO lies in the way the dynamics model is trained and utilized, as indicated by the parts highlighted in blue in the pseudo-code.
Algorithm 3 ADMPO-OFF
Input: Pre-collected dataset Denv, initial ADM T3 and policy Fas roll-out length H, maximum backtracking length m, model data buffer Doddel iterations N, penalty coefficient /
1: Train ADM T. Peny by maximizing Equation
2: for N iterations do
3: for M model roll-outs do
4: Sample initial m-step state-action sequence (, from Denv
5: Roll out H steps in To via ADM-Roll(To, Ta. H, m, (Sei+m-1 (i:i-mm2))
6: Penalize the reward via f BUADM (s, a) for each rolled-out step
7: Add the penalized model roll-out data to moodel
8:
end for
9: for G policy updates do
Update current policy ne using samples from Deny U D model 11: end for
end for
C.3 POLICY OPTIMIZATION
The policy optimization method used our ADMPO-ON and ADMPO-OFF SAC (Haarnoja et al., 2018), following MBPO (Janerr et al.) 2019) and MOPO (Yu et al.) 2020). The hyper-parameters
about SAC follow its standard implementation, as listed Table
Table Hyper-parameters of Policy Optimization ADMPO-ON and ADMPO-OFE.
Hyper-parameter
Value
Description
No
2
the number of crities.
actor network
FC(256.256)
fully connected (FC) layers with ReLU activations. fully connected (FC) layers with ReLU activations. target network smoothing coefficient. discount factor.
critic network
FC(256.256)
5 10-3
0.99
Taatoor
1 10
learning rate of actor. learning rate critic.
Iranie
10
optimizer
Adam
optimizers of the actor and critics. batch size for each update.
batch size
256
D EXPERIMENTAL DETAILS
D.1 RESOURCE REQUIREMENTS
All experiments can be completed with just one NVIDIA GeForee RTX 2080 Ti or any other type of GPU with larger graphic memory. There are no additional resource requirements. The time of execution for each task is about 24 hours.
D.2 ADMPO-ON SETTINGS
The experimental settings of our ADMPO-ON in Section 4.2 are listed in Table
17

---

Published as a conference paper at ICLR 2025
Table 4: Hyper-parameter settings of ADMPO-ON results presented in Figure + 6 denotes thresholded linear increasing schedule, i.e. the length of model roll-outs at step t is calculated by [(1) = min mmax
y
yy x), x
b-a
environment
Hopper
Walker2d
Ant
Humanoid
steps
50k
200k
300k
Update-To-Date ratio
20
maximum backtracking length m
I 1-+15 over
T- over
1-5 over
I-10 over
model roll-out schedule
0-+00k
0~+100k
10k 100k
10k-+100k
target entropy
-1
-3
-4
-8
D.3
ADMPO-OFF SETTINGS
The experimental settings of our ADMPO-OFF Section 4.3 are listed in Table 5
Table 5: Hyper-parameter settings of ADMPO-OFF results presented in Section 4.3
Domain Name
Task Name
H
e
hopper-random
5
halfchectah-randone
10
2.5
walker2d-random
2
50
2.5
hopper-medium halfeheetah-medde walker2d -mediumm
medium
1
2.5
D4RL MuloCo
10
5
hopper--nedium-rella
0.1
medium-replay
2.5
walker2d-medium- mediim--rellla
0.1
hopper-medium expert
20
20
50
10
walker2d-medium medimm- -xxpert
neorl- neorl-hopper-low
5
20
5
neorl neorl-halfcheetal halfcheetah-lov
10
neorl- walker2d-
2.5
NeoRL MuloCo
neorl- neorl-hopper- seorl-halfcheetah-nee neorl- maalker2d-mediune
20
50
20
10
5
neorl- neorl-hopper -hihh neorl-hall weat--alffhccah-hiii
20
50
10
50
neorl- worl-wallee22--- high
10
2.5
D.4 PARAMETER TUNING OF ADMPO-OFF
There three important hyper-paramette in ADMPO-OFF: the maximum backtracking length m. the roll-out length H, and the penalty coefficient B. Below, we will introduce how to tune these three hyper-parameters.
m: The first seper-parameter to tune is m, as it can be adjusted based on the validation error after the training of the dynamics model. Experiments show that a moderately large m can effectively ensure algorithm performance, and increasing does not lead to significant performance degradation. Considering that an excessively large m would consume excessive computational resources, initial value for between and 10 is recommended. Then, observe whether the validation error of the dynamics model after training meets the task requirements. If not, slightly decrease m, but ensure it does not go below 2.
18

---

Published as a conference paper at ICLR 2025
H: The second typer-parameter to tune is H. After training the dynamics model, H directly affects the compounding error generated during roll-outs using the model. A larger H value produce more diverse data, which benefits policy learning. It is recommended set the initial value of H to 5 and then gradually increase until the compounding error after H-step roll-outs approaches the acceptable threshold.
3: Tuning 33 relies the results of Q-value estimation. According the findings of the previous work (Lu et al.) 2022). in general, larger values require larger 3 values for support. An appropriate miiial value be selected based on the size of H, typically ranging from one-tenth to one-half of H. Then, adjust 3 based on the estimation bias of the Q-value. If the Q-value is overestimated, increase 3; otherwise, decrease B. Finally, identify the critical B value at which the Q-value is no longer overestimated.
D.5 SOURCE OF BASELINES' RESULTS
For the evaluation on D4RL (Fu et al. 2020 benchmarks, the results of the compared baselines come from Iwo sources:
Retraining on D4RL datasets of v2 version with OfflineRL-Kit (un 2023). for the algo- rithms whose original papers only report the performance on the v0 version, such as CQL (Kumar et al. 2020), MOPO (Yu al. 2020)
Including the scores in their papers, for the algorithms whose original papers report the performance on the v2 version, such as TD3+BC (Fujimoto & Gu [022]. EDAC (An et al.
202T RAMBO (Rigeer et al. 2022). CBOP (Jcong al. (023). and MOBILE (Sun et al.
2023
or who does not provide source codes, such as COMBO Yu et al. 2021
For the evaluation on NeoRL (Qin et al. 2022) benchmarks, we report the scores of BC, CQL, and MOPO from the original paper of NeoRL and retrain TD3+BC and EDAC with OfflineRL-Kit (Sun) 2023
E
ADDITIONAL EXPERIMENTS
E.1 SUPPLEMENTARY RESULTS OF DYNAMICS MODEL EVALUATION
hopper-medium
hopper-medium-relay
walker2d-medium
wallkerzd--edduu--repaaa
Prodiction 6
Prediction ~
25
50
75
100
25
50
75
100
25
50
75
100
25
50 Roll-out Langth
100
Roll-out Longth
Roli-out Length
Roll-out Length
ADM (ours)
Ensemble Dynamics Model
Bootstrapping RNN Dynamics Model
Figure 5: Comparison among ADM, ensemble dynamics model, and bootstrapping RNN dynamics model, in terms of the growth curve of the compounding enror as roll-out length increases, after offline learning.
In Section 4.1 we present the compounding envor curves of ADM, the bootstrapping RNN dynamics model, and the ensemble dynamics model as the roll-out length increases. The y-axis uses log scale, which highlights ADM's s advantage in future predictions distinctly. However, the log scale also makes difficult to depict the precise values ADM's s compounding error. Here, we also provide the linear-scale version of Figure 2 shown in Figure 5 In the hopper task, ADM's prediction error remains close to zero even after rolling out to 100 steps. In the more complex walker2d task, ADM's prediction error does not exceed when the roll-out length is less than 50 steps.
Additionally, we set m to 1. 5, 10, 15, and 20 respectively to observe the compounding error curves of ADM after offline training, shown in Figure 6 Overall, increases, the growth rate of
19

---

Published as a conference paper at ICLR 2025
hopper-medium
hosper-medlum-rella
walker2d-medlum
wellerrddmeedlum-rpplaa
0.20
0.20
h/! M/
0.15
Prediction Errar
0.15
2.0
2.0
Pedicten
Prediction
1.5
Predizzinn
1.5
0.05
0.05
0.5
0.5
0.00
25
50
100
0.00
23
50
75
100
0.0
25
50 75
100
0.0
25
50
75
100
Length
Length
m
m 15
m
Figure 6: Comparison among ADM with set to 5, 10, 15, 20, in terms of the growth curve of the compounding error as roll-out length increases, after offline learning.
the compounding eтror tends to slow down. This is because a larger m implies fewer expected bootstrapping iterations and thus fewer error accumulation steps. Specifically, when m 1, the growth rate of the compounding error is particularly rapid. This also indirectly demonstrates the effectiveness of ADM with an appropriate m.
E.2 STUDY ON WIYY ADMPO-ON PERFORMS WELL IN ONLINE SETTING
Value-aware model error (Farahmand et al. 2017) is dependable metric for measuring the learning quality of the dynamics model and the suboptimality of the MBRL algorithm. We conduct a study verify how well ADMPO-ON regulates the value-aware model error. Without loss of rigor, only choose MBPO for comparison since most other model-based methods follow the same way of learning and utilizing the ensemble dynamies model. Figure shows the results of the most difficult Humanoid task. The learned ADM in ADMPO-ON and the ensemble dynamics model in MBPO achieve similar mean squared errors, indicating their similar fitting abilities. However, ADMPO-O1 provides greater model roll-out standard deviation over diverse state prediction, forcing the agent to explore more uncertain areas. Therefore, since the variation of state prediction helps smoothen the Q target, the Q network in ADMPO-ON has significantly smaller Lipschitz constant, and then the value-aware model eтror, which measures the suboptimality of MBRL becomes smaller. This phenomenon explains why ADMPO- ONNOONNNNTT performs significantly better than MBPO in Figure For details of the metrics used in this experiment, refer to (Zheng et al. 2023).
1.0
ABIMP--DIN Zeano
N00s0 Mez Scrret
Eare
MISPO
Mate Ralter
10000)
200030
30000
(0000e
300000
100036
100e00
200000
520000
10000
200000
3E0000
Time Stepe
Time Stppe
Time Supa
Tine Stgge
(8)
Figure 7: Comparison between ADMPO-ON and MBPO on Humanoid, in terms of (a) model mean squared error, (b) model roll-out standard deviation over diverse predictions, (c) estimated Lipschitz constant (Zheng et al. 2023) of Q, and (d) value- model error (Farahmand et al.) 2017). Results are averaged over five seeds.
E.3 D4RL ANTMAZE RESULTS
We compare ADMPO-OFF with several offline MBRL algorithms on the D4RL (Fu et al.) 2020) AntMaze tasks, which test the capability of algorithms to find an effective goal-oriented policy. Table 6 reports the corresponding results. ADMPO-OFF still outperforms other offline MBRL baselines in most AntMaze tasks and achieves the highest average score. Notably, the policy learned by MOPO lacks any ability reach the goal, scoring zero on all AntMaze tasks. However, by simply replacing the dynamics model with ADM, ADMPO-OFF can learn policy with some capability to reach the goal, demonstrating potential in addressing such challenging tasks.
20

---

Published as a conference paper at ICLR 2025
Table 6: Normalized scores after offline learning on D4RL AntMaze tasks, averaged five seeds.
Task Name
MOPO
COMBO
RAMBO
MOBILE
ADMPO-OFF (ours)
antmaze-umas
0.0
80.3
25.0
77.0
88.4:11.2
antmaze-umaze-divee
0.0
57.3
0.0
20.4
81.711..6
antmaaze-medium-play aitmaze-meeduu--ddeessee
0.0
0.0
16.4
64.6
23.94-6.3
0.0
0.0
23.2
1.6
24.11:5.7
antmaze-large-play
0.0
8.3+4..1
antmaze-large-divers
0.0
0.0
2.4
7.2
0.0+0.00
Average
0.0
22.9
16.8
28.9
37.7
E.4 STUDY ON DYNAMICS MODEL DESIGN
To predict the next state, ADM uniformly chooses a k from [1, ml.. then backtracks the state from steps earlier along with the following k-step sequence of actions input. The prediction error of the next state is accumulated based on the deviation of the state from k steps earlier. The only one-step input state is duplicated match the length of the action sequence. To verify the effectiveness of ADM's design for policy optimization, we introduce two ablation versions of ADMPO-OFF:
BootRNN replaces ADM with bootstrapping RNN dynamics model, which uses the preceding k-sstep sequence of state-action pairs as input when predicting the next state, where k is sampled uniformly from [1, ml]. The prediction error of the next state is accumulated based on the deviation the current state. All other settings remain consistent with ADMPO-OFF.
ADMPO-OFF-zero takes state and the following any-step actions as imput to the dynamics model, the same ADMPO-OFF However, the state is only fed into the RNN during its first cell access, and subsequently, zero tensors are concatenated with the actions match the dimensions. All other settings remain consistent with ADMPO-OFF.
We compare the performance of these dynamics model designs nine D4RL Fu et al. (020 MuJoCo tasks. Table reportss the corresponding results.
Table 7: Normalized scores corresponding to different dynamics model designs.
Task Name
BootRNN
ADMPO-OFF-zero
ADMPO-OFF
hopper-medium
89.8+1.8
106.0+0.66
107.4+00.6
hopper-medium--rellay
87.2+2.77
101.2+0.1
104.4410.4
hopper-medium-expert
108.6+1.0
112.9+1.5
112.7+0.3
walker2d-medium
1..
94.1110.2
93.2
walkkerad-mcdium-rellay
74.6+009
92.6+
95.6+2.11
walker2d-meddum-experr
95.5+2.3
113.210.2
114.9-10.3
halfcheetah-mmedium
52.811.2
73.710.11
72.210.6
halfcheetah-mmcdiun--llga
58.2+1.66
66.3+L0.6
67.6+
halfcheetah-mecdiumeeprrrr
94.0+1..2
101.110.3
103.7-10.2
BootRNN experiences significant policy performance degradation compared to ADMPO-OFE, since the bootstrapping RNN dynamics model accumulates the prediction error each step based on the deviation from the last step, leading to larger compounding error. This has already been validated in the dynamics model evaluation experiment in Section 4.1l COMPO-OFF-ZEE significant performance differences compared to ADMPO-OFF as their inputs are identical, differing only in the way they forwarded. This difference has minimal impact on the overall performance.
does not exhibit
21

---

Published as a conference paper at ICLR 2025
E.5 STUDY ON UNCERTAINTY QUANTIFIER
ADMPO-OP CIMIPOOFFFOIO uses the standard deviation of predictions as a measure of model uncertainty, which actually a combination of epistemic and aleatoric model uncertainty. Other common uncertainty choices include maximum aleatoric uncertainty and maximum pairwise difference. Denoting the mean and standard deviation of predictions when backtracking steps as 16 and S.y. respectively, the descriptions for these uncertainties are follows.
Max aleatoric Yuu et al. 2020): max | F. which corresponds to the maximum aleatoric error.
Max Pairwise Difference (Kidambi et al., 2020): max., 11 p/ll2. which corresponds
the pairwise maximum difference of the predictions from different k.
Prediction Standard Deviation (Lu et al. 2022 #/ E ((E)) (xt)))
( - m7 k= m²))² 2, which corresponds combination of epistemic and aleatoric model uncertainty.
We compare the performance of these uncertainty choices on nine D4RL. (Fu et al. 2020) MuJoCo tasks. Table reports the corresponding results.
Table 8: Normalized scores corresponding to different uncertainty quantifier choices.
Task Name
Max Aleatoric
Max Pairwise Difference
Prediction Std (ADMPO-OFF)
hopper-medium
55.6=0.9
105.6:10.2
107.4:10.6
hupper-medium--eplay
103.3:10.1
101.8
104.4:10.4
hupper-medium-expert
113.1:10.2
111.7=1.2
112.7=0.3
walker2ddmmdium
87.7=0.5
79.71:1.8
93.2:11.1
walker24-mmediu-rreee
replay
93.4=0.3
96.3110.9
95.6112,1
walker2d-mediuum--ee expert
111.6100.4
110.410.7
114.9110.3
halfcheetah-mediin hallchectub-mmddiun--enpp
71.5==033
7..4+3.4
72.2:10.6
69.7-1.3
59.412.8
67.6
hallichecth--mddium--rpprrt
102.8+22.3
96.9+-299
103.7-+02
The prediction standard deviation, which similar to the ensemble standard deviation in the ensemble dynamics model, performs the best. Similarly, it is found by (Lu et al. (022) that ensemble standard deviation is the best while using the ensemble dynamics model. The key to ADM that provides a method for diversifying state predictions that is different from ensemble methods, and any form of uncertainty is feasible.
E.6 STUDY ON PREDICTION CHOICE
While rolling out a sequence using ADM, we choose the backtracking length & uniformly from 21, , then predict the next state by feeding the state from & steps ago and the following k-step sequence of actions into the ADM. The random sampling of the backtracking length can be viewed as an implicit augmentation. The variations of state predictions can effectively implicitly regularize the local Lipschitz. condition of the Q network around regions where the model prediction uncertain, thereby regulating the value-aware model error (Farahmand al. 2017), according to (Teng et al. 2023). Other intuitive choices to predict the next state include:
Priority Sampling:
Determining the sampling priority based on the fitting losses of different backtracking lengths. For instance, the sampling probability corresponding to i is given by
E 4,, where eg is the fitting loss of the k-step prediction.
Max Backtracking: Always choosing the maximum backtracking length m to predict the next state.
Average Prediction: Averaging the predictions corresponding to different backtracking lengths from (1,2,... m).
22

---

Published as a conference paper at ICLR 2025
We compare the performance of these prediction choices nine D4RL tasks. Table reports the corresponding results.
Fu et al. 2020 MuJoCo
Table 9: Normalized scores corresponding to different prediction choices.
Task Name
Priority Sampling
Max Backtracking
Average Prediction Unitorm (ADMPO-OFF)
hopper-medium
106.4:0.88
98.810,9
56.1=244
107.4:10.6
nopper-medium-eepe
105.3:10.8
95.6+0.8
51.6=199
104.4+0.4
hopper-nedium-exper
113.613.77
110.3:1.6
52.6=1.4
112.740,3
walker2d- malker2d-medium
89.0=1.5
84.5#1.1
85.9=188
93.2:1111
walker2d-meet saake22dmmdiuu--pppay
96.2=11.0
91.1+0.66
87.0=16
95.642.1
walker2d-medium-oo medium-cxpert
115.5:10.7
107.7
88.6=100
114.9:0.3
halfcheetah medium
71.4=:3.2
56.9+0.2
58.9=133
72.2++0.6
halfcheetah
wewichhtta-mmiium--pppee medium-reppaa
61.7=2.3
49.8:1,99
46.8=1.
67.6=1.4
hulfcheetab-mediin mediimm-exxeett
103.2:1144
67.0 1.4
63.2=009
103.7::..2
We find that there is no significant performance difference between the priority sampling and the uniform sampling, since both of them introduce diversity in state prediction. Both maximum backtracking and average prediction experience significant performance degradation.
E.7 STUDY ON m
7, and 10, respectively, to study its impact on performance of ADMPO-OFT In addition, since ADM is improvement for fixed multi-step dynamics model (Asadi et al. 2018 2019: Che et al. 2018), we further set the learning and roll-out steps of ADM to fixed lengths to highlight the importance of the any-step design. Table reportss the ablation results four D4RL (Fu et al. 2020) MuJoCo datasets (v2 version). The network structure is kept unchanged.
set m to 1,
Table 10: Normalized scores corresponding to different maximum backtracking lengths.
Task Nams
fised Dellep
Bined Salop
Exed 4step
fined Seteep
m
m
i
m
m
hepper-nraiim
6.7=114
9.++1.
66616
88--13
53++1.
252450
10681407
107.4066
106 1+0.
106.+0077
hopper medium replay
22
25.3112
26.1115
223119
23.311.
318111
992100
104410.4
1023103
1058113
walkerdd medium
6,3218
$7.15
$5116
53118
55113
936116
8666167
932.1.1
90.5111
81.2122
walkerld medium-mppyy
174116
245:22
219:22
22662.
1988227
66.3207
837206
95.112
94.3208
914210
When setting m to using a fixed multi-step model, performance severely degrades. This is because, in these cases, state predictions lack diversity and cannot estimate model uncertainty, which critical offline setting. Starting from and gradually increasing m to leads to continuous performance improvement. Larger values of no longer bring further performance advantages since predicting across too many steps poses challenge to the expressive capacity of neural networks. This experiment demonstrates the necessity of the any-step design and the proper selection of m.
E.8 COMPUTING RESOURCE EVALUATION
In Table IT we compare the computing resource consumption of ADMPO-OFF and MOPO, including the size of model parameters. GPU memory usage during training, and runtime.
The model parameter size of ADMPO-OFF is smaller than that of MOPO, primarily because MOPO uses an ensemble for its dynamics model. The GPU memory usage of ADMPO-OFF primarily depends on the value of m, and it is generally higher than that of MOPO. The runtime is roughly the same for both.
LIMITATIONS
In general, our work's limitations are summarized as follows.
Similar to most previous MBRL algorithms focused on locomotion tasks, the dynamics modeling approach and uncertainty estimation used by ADMPO are likely not suitable for highly stochastic environments, they are prone being influenced by randomness, leading
23

---

Published as a conference paper at ICLR 2025
Table 11: Comparison of ADMPO-OFF and MOPO in terms of computing resource consumption.
Size of Parameters (MB)
GPU Memory (GB)
Runtime (s/epoch)
Task Name
ADMPO-OFF
MOPO
ADMPO-OFF
MOPO
ADMPO-OFF
MOPO
hopper-random halfcheetah-random
3.77
7.66
7.85
3,46
10.20
7.55
3,83
7.91
3.98
3.54
8.59
8.37
walker2d randdmm
3.83
7.91
3,98
3.54
9.78
8.19
hopper-medium halfcheetah-menn meemmdu
hopper- medium
3.77
7.66
7.85
3.46
9.94
8. 3,
7.91
3,98
6,65
walker24 2d walker24-medium
3.83
7.91
9.06
3.54
10.61
8.54
hopper-mediune medium-repioe
3.77
7.66
7.85
3,46
9.51
8.16
halfcheetah- medium-replay
3,83
7.91
3.98
3,54
7.35
8.27
malkerrd-mediiu--eee replay
3.83
7.91
9.06
3.54
7.01
8.15
medium-experr halfcheetah- wawasa-mddiuddpeeenne
3.77
7.66
3.38
3,46
8.61
8.01
3.83
7.91
3.98
3.54
8.49
8.07
walker2d meadd--eeiiieeeeeeee
3.83
7.91
4.32
3.54
7,01
8.09
to poor performance. Improving ADMPO to address the challenges of highly stochastic environments is an important issue for future research. We will consider this as part of our future work.
The computing resources used by ADMPO are mainly influenced by the maximum back- tracking length m. When takes particularly large values, ADMPO requires significantly more GPU memory compared to the previous MOPO algorithm.
24


Reviews:
[{"review_id": "E3crZY0qYv", "reviewer": "Reviewer_pnqF", "summary": "This paper introduces the Any-Step Dynamics Model (ADM), an RNN-based dynamics model with a variable prediction horizon. ADM enhances the fidelity of dynamics modeling and provides uncertainty estimates by leveraging variability across different prediction horizons.", "questions": "[Q1] ADM currently utilizes 2 to 5 backtracking steps. What would happen with a larger number? Would increased steps result in model divergence due to compounding errors? (See Weakness 1.)\n\n[Q2] Could the multi-step rollout compounding error be mitigated? For example, including the ground-truth state $ s_{t+k} $ instead of only $s_t$ in this work for each roll-out step in an autoregressive manner, as done in models like the Trajectory Transformer, may improve accuracy. (See Weakness 1.)", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "[S1] The approach to uncertainty estimation is novel, contrasting traditional ensemble methods by using variance across horizons within a single model instead of variance across the prediction of multiple models.\n\n[S2] ADM demonstrates strong performance in both online and offline experiments, with uncertainty measures that correlate well with ground-truth errors.", "weaknesses": "[W1] ADM uses multiple-step rollouts, initializing from a state several steps back in time to predict the trajectory, thus potentially disregarding states closer to the current time step. This may increase compounding errors in rollouts compared to ensembles, which operate with single-step rollouts based on the current state.\n\n[W2] The reference formatting in the main text appears incorrect; minor revisions are recommended.", "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": "HCw30w7fZi", "reviewer": "Reviewer_eBkb", "summary": "The authors propose Anytime Dynamics Models, which are trained to predict k-steps-ahead states and rewards given a past state and a sequence of actions. ADM are designed to reduce bootstrapping in autoregressive predictions, and alleviate accumulating errors over long horizons. ADM can be directly utilized for model-based policy optimization, and lead to improved sample efficiency in standard mujoco tasks. This work also proposes an estimate for the mixture of aleatoric and epistemic uncertainty by computing the disagreement of the model's predictions for different number of steps. Uncertainty estimates can then replace those obtained through model ensembling in the MOPO framework, again showing improvements in performance on D4RL and NeoRL tasks.", "questions": "Three questions are raised as part of the weaknesses. I would additionally ask the following:\n- Figure 2 shows that error do indeed accumulate less drastically in ADMs. However, the scale of the plots is not meaningful. An error of 10^34 is probably not good for planning, but the error incurred by ADMs might as well be too high. This is hard to see due to the y-scale of the plots. Furthermore, it's not clear what the scale of the states is, and a relative measure of error might be more informative. Can the authors provide a more readable plot? Moreover, Which policy is used during the rollouts?\n- The idea of temporally consistency in any-step predictions as a proxy for epistemic uncertainty has been explored in the past. Among other works, [3] uses uncertainty of the value, which prevents from mixing aleatoric and epistemic uncertainty. These previous works should be acknowledged.\n- ADMPO tunes hyperparameters on each task. How are the hyperparameters for the baselines selected?\n- There are inconsistencies in the spelling of the word rollout/roll-out.\n- Equation 1: for the sake of clarity, can the authors specify over which spaces/sets the sum is computed?\n- line 204: what about these algorithms is \"foundational\"?\n\n**References:**\n\n[1] Hafner et al., Mastering Diverse Domains through World Models, arXiv 2023\n\n[2] Hanses et al., TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024\n\n[3] Filos et al., Model-Value Inconsistency as a Signal for Epistemic Uncertainty, ICML 2022", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The presentation is clear and the paper is well written.\n- The empirical evaluation is broad, and the proposed methods largely outperform the selected baselines.", "weaknesses": "- The main motivation of this work is that of alleviating autoregressive error accumulation. Nevertheless, ADMs rely on a recurrent architecture that implicitly bootstraps on previous hidden states. It is not entirely clear to me why this would not also result in error accumulation. Empirical performance provides some support to the method's effectiveness, but are there any principled reasons why ADMs suffer less from error accumulation? In my opinion, this one of the main issues with the current version of this work. The authors report a Theorem (3.4), but it appears to be a general result. How does it relate to the backtracking nature of the proposed method?\n- The second issue is in the limited comparison to recurrent forward models. Recurrent world models are at the core of SOTA model-based algorithms [1,2]. The only empirical comparison to recurrent models is, as far as I understand, in Figure 2, which is not easy to interpret. The RL evaluation does not include any recurrent architecture. This raises the following question: is the improvement in RL performance due to any-time dynamic sampling, or to the choice of architecture? For instance, how would the method compare in online RL to a RSSM [1] of similar size, using standard autoregressive sampling?\n- A third main issue is an insufficient discussion of the limitations of this method. A fundamental issue which is not mentioned in Section 6 is that the uncertainty estimation method proposed appears to mix aleatoric and epistemic uncertainty. This is in general suboptimal: how would if perform in an environment in which the optimal policy traverses stochastic parts of the MDP?", "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": "J9EiMSzz32", "reviewer": "Reviewer_NWuZ", "summary": "The authors propose a dynamics model that is trained to predict a future state and reward tuple given various length state-action sequences. Notably the input to the model is not is not $(s_{t-3}, a_{t-3}, s_{t-2},a_{t-2},s_{t-1},a_{t-1})$ as one would expect, but instead they are using the same state, i.e. $(s_{t-3}, a_{t-3}, s_{t-3},a_{t-2},s_{t-3},a_{t-1})$.\nThis somehow leads to a signficiantly better performance and improved results when used for model-based online and Offline RL.\nFor Offline RL, the authors also use the ensemble disagreement for different input-lenghts as a measure of model uncertainty.", "questions": "* Why aren't all $m$ known states used as input for the model?\n * Why is a random $k$ used to train the policy?\n * Fundamentally, this paper is missing a good explanation for why its method works better. Additional ablations would help elucidate this. For example, inputting the whole known state sequence in to the GRU instead of repeating the first state. Or only inputting the state into the first GRU-call should have the same information content.\n\nOverall I believe this paper has potential to be very useful, but it requires additional work to improve the presentation and explanation of its method before publication.\n\n\nAfter rebuttal:\n\nThe authors have provided helpful experiments that confirm their changes to the model to be advantageous.\nI still have concerns particularly about the lack of hyper-parameter optimization for baselines.\n\nOverall I am raising my score (3->5)", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "* The results look quite promising, performing well in Online RL and Offline RL benchmarks", "weaknesses": "* Clarity: The paper is overall hard to follow with poor writing and inconsistent notation. For example, sometimes $\\hat{s}$ is used to refer to states predicted by the dynamics model, sometimes $s$ is used to refer to predicted states. $s$ is also used to refer to ground-truth states. In addition to various typos this makes the paper hard to follow.\n * Soundness: The paper does not give a convincing reason for why their proposed method works better, nor does it include ablations or experiments that could elucidate why it performs well. There also isn't any theoretical analysis that might be helpful to understand it better.\n\nSome important implementation choices are not justified either. To list the most major issues:\n * As written in L152, and shown in Figure 1, the state $s_t$ is repeated as input for all GRU/ model steps. This is reasoned to prevent error accumulation of using model predictions as input. However the paper assumes access $m$ step sequences which are then used to predict the next $H$ steps. Why aren't all $m$ steps used as model inputs?\n * Using a random input length during training is reasonable to learn a model that performs well for different input lenghts. However, when using the model to train a policy it would seem more natural to average across lenghts to obtain a better prediction, or choose the longest length. Why is a random length used?\n\n* Experiments: Hyper-parameter choice is not discussed and only 5 seeds are used. For most baselines, the results were copied from the original papers. As hyper-parameter choice and implementation details are important for RL methods, this is not sufficient, especially when applying the algorithms to a new benchmark (NeoRL).\n\n\nMinor issues:\n * Theorem 3.5 follows rather straightforwardly from Assumption 3.3 and does not provide any new insights. It is also not specific to the ADM and could apply to any model-based RL method.\n * parantheses should be placed around the author name for citations, not just around the year number.\n * What the authors refer to as \"bootstrapping prediction\" issues seems more like error accumulation to me.  Using the term bootrstrap can be a bit confusing in an RL context, but this is a minor point.\n* Algorithm 2 and 1 should probably combined and both be in the main text, as they are important to understand the approach.\n* L241: \"explore beyond the boundaries of the risky regions\" should be \"beyond the boundaries of the safe regions\"?\n* L352: should this be $\\hat{s}_{t+m}$? If not, there can be no compounding error.\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": "oNephcEE7M", "reviewer": "Reviewer_i58Z", "summary": "Most convention rollout approaches in MBRL make prediction of the future state based on the current state alone. The main drawback is on the potential error accumulation as the prediction horizon grow. This work proposes Any-step Dynamics Model (ADM)  to mitigate this issue by making prediction based on previous state and actions information. This is achieved by using RNN structure to capture the dependence relationship among steps. Based on ADM, both online and offline compatible algorithms, i.e., ADMPO-ON and ADMPO-OFF, are introduced. Extensive experiments on MuJoCo and D4RL tasks show that the proposed method achieve better performance comparing with baseline methods. Meanwhile, it shows that ADM also help with the uncertainty estimation, where previous works normally use the ensemble method. The ablation studies on the rollout length $m$ is also validated in the experiment.", "questions": "- This work mentions about MBPO (Janner et al. 2019), where they also proposed a truncated rollout to address the error accumulation over long horizon. Many works on adaptive rollout length should be addressed and compared [Q1,Q2]. \n- line 046, what is deviation error? Line 68, what do you mean by trajectory distribution? \n- What is the computation complexity of the proposed method?\n\n\n\n[Q1] Frauenknecht, Bernd, et al. \"Trust the Model Where It Trusts Itself-Model-Based Actor-Critic with Uncertainty-Aware Rollout Adaption.\" ICML.\n\n[Q2] Zhang, Weinan, et al. \"Model-based multi-agent policy optimization with adaptive opponent-wise rollouts.\" arXiv preprint arXiv:2105.03363 (2021).", "limitations": "", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Novel method on addressing bootstrapping errors during rollout. The proposed ADM incorporate an RNN structure to capture the dependencies among previous actions and state information, which can be helpful to avoid the prediction error induced by making predictions based on prediction. The ADM model is versatile and can be used for many offline and online policy optimization MBRL algorithms.\n- The experiments on both online and offline environments. This work considers both online and offline experiments on verifying the performance. Meanwhile the studies on the uncertainty estimation also demonstrate that the proposed ADM can be a good alternative of ensemble method for estimation.", "weaknesses": "- The motivation of the proposed ADM rollout should be clarified further. In the proposed rollout in ADM (Algorithm 1), when predicting the future state such as $s_{m+1}$, $k=2$, it essentially is based on the information of $s_{m-1},a_{m-1:m}$. When $s_{m}$ is available in the given sequence, it seems not clear why the proposed ADM can perform good. Moreover, when the state transition and policies are stochastic, the proposed rollout essentially requires to consider the stochasticity in $s_m$ and $a_m$ when making prediction for $s_{m-1}. What is the intuition of using this rollout for the sake of the performance. \n- The proposed rollout algorithms also make prediction based on the previous prediction, e.g., when predict $s_{m+2}$, it may require $s_{m+1}$ (which is obtained in the previous steps).  It is thus very important to clarify why this method can effectively mitigate the bootstrapping errors comparing with the conventional approach.\n- The backtracking length $m$ is a newly introduced parameter, and the ablation studies show that it affects the performance a lot. How to choose such a parameter in practice for specific task? Will it work for all different tasks?", "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)