Any-step Dynamics Model Improves Future Predictions for Online and Offline Reinforcement Learning

## Abstract

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 A ny-step D ynamics M odel (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.

## 1 Introduction

Model-based Reinforcement Learning (MBRL) [] has demonstrated empirical success in both online [, , , , , ] and offline [, , , , , ] settings. The essence of MBRL lies in the dynamics model, where extensive explorations and evaluations of the agent can occur, thereby reducing the reliance on real-world samples. Embedded in the model-based framework, online policy optimization can leverage a large Update-To-Data (UTD) ratio [] to improve sample efficiency, while offline policy optimization can be completed using the model augmented data beyond the dataset.

Although some efforts aim to propose high-fidelity dynamics models, such as adversarial models [, ], causal models [], and ensemble dynamics models [, , ] adopted by the majority of MBRL algorithms, it is challenging to generate high-quality imaginary samples via long-horizon model roll-out. In a dynamics model with the common form, the state-action pair at time step $t$, $(s_{t},a_{t})$, is used as input to predict the next state $s_{t+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 error of generated states increases with the roll-out length since the error accumulates gradually as the state transitions in imagination. If updated on the unreliable samples with a large compounding error, the policy will be misled by biased policy gradients.

In the online setting, the impact of compounding error [] on policy optimization restricts the utilization of the model, thereby hindering further improvement in sample efficiency. In the offline setting, compounding error affects the accuracy of current model uncertainty estimation based on ensemble. For instance, using the behavior policy corresponding to the dataset for a long-horizon roll-out in the model still leads to a great compounding error, and the accumulated deviations of different learners cannot be similar. In this case, the divergence of the ensemble will inevitably be large, which is inconsistent with the situation that the actual trajectory is within the region covered by the dataset. Therefore, it is essential to reduce compounding error in both online and offline settings.

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 a multi-step action sequence [, , , ]. Although state $s_{t+1}$ is only dependent on state-action $(s_{t},a_{t})$ under the assumption of Markov property [], the prediction of $s_{t+1}$ can actually leverage earlier information. Tracing back a prior $k$-step plan, $s_{t+1-k}$ and the intermediate $k$-step actions $(a_{t+1-k},a_{t+2-k},\cdots,a_{t})$ are sufficient to constitute an attribution to predict $s_{t+1}$.

To handle the variable-length plans, we introduce a special Any-step Dynamics Model (ADM) that allows for the use of $s_{t+1-k}$ and $(a_{t+1-k},a_{t+2-k},\cdots,a_{t})$ corresponding to any integer $k$ within a specified range as inputs for predicting $s_{t+1}$. 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 ensemble. Replacing the ensemble dynamics model with ADM, we devise a unique model roll-out method with random backtracking, which can be plugged into any existing MBRL algorithmic frameworks. 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 the compounding error. (2) We propose a new online MBRL algorithm called ADMPO-ON based on ADM and show that it can outperform recent state-of-the-art online model-based algorithms in terms of sample efficiency while retaining competitive performance on MuJoCo [] benchmarks. (3) We propose a new offline MBRL algorithm called ADMPO-OFF based on ADM and show that it can effectively quantify the model uncertainty, achieving superior performance compared to recent state-of-the-art offline algorithms on D4RL [] and NeoRL [] benchmarks.

## 2 Preliminaries

### 2.1 Markov Decision Process and Reinforcement Learning

We consider a standard Markov Decision Process (MDP) specified by a tuple $\mathcal{M}=(\mathcal{S},\mathcal{A},T,\rho_{0},\gamma)$, where $\mathcal{S}$ is the state space, $\mathcal{A}$ is the action space, $T(s_{t+1},r_{t+1}|s_{t},a_{t})$ is the dynamics function that calculates the conditioned distribution of $s_{t+1}\in\mathcal{S}$ and $r_{t+1}\in\mathbb{R}$ given $(s_{t},a_{t})$, $\rho_{0}$ is the initial state distribution, and $\gamma$ is the discount factor. We use $\rho^{\pi}$ to denote the on-policy distribution over states induced by the dynamic function $T$ and the policy $\pi$. From a multi-step perspective, the attribution of state $s_{t+1}$ and reward $r_{t+1}$ can be traced back to the earlier $k$-step plan, $s_{t-k+1}$ along with the action sequence $a_{t-k+1:t}=(a_{t-k+1},a_{t-k+2},\cdots,a_{t})$ in between. This relationship can be represented by the $k$-step dynamics model

$$ $T^{k}(s_{t+1},r_{t+1}|s_{t-k+1},a_{t-k+1:t})=\sum_{(s_{t-k+2:t},r_{t-k+2:t})} \prod_{i=0}^{k-1}T(s_{t-i+1},r_{t-i+1}|s_{t-i},a_{t-i}).$ (1) $$

We use $\Gamma^{k}_{\pi}(s_{t-k+1:t},a_{t-k+1:t}|s_{t+1})$ to denote the distribution over $(s_{t-k+1:t},a_{t-k+1:t})$ conditioned on $s_{t+1}$ induced by the dynamic function $T$ and the policy $\pi$.

The optimization goal of Reinforcement Learning (RL) is to find a policy $\pi$ that maximized the expected discounted return $\mathbb{E}_{\rho^{\pi}}\left[\sum_{t=1}^{\infty}\gamma^{t-1}r_{t}\right]$. Such a policy can be derived from the estimation of the state-action value function
$Q^{\pi}(s_{t},a_{t})=\mathbb{E}_{(s_{t+1},r_{t+1})\sim T(\cdot|s_{t},a_{t})}
\left[r_{t+1}+\gamma V^{\pi}(s_{t+1})\right]$, where $V^{\pi}(s_{t+1})=\mathbb{E}_{a_{t+1}\sim\pi(\cdot|s_{t+1})}\left[Q^{\pi}(s_{t+
1},a_{t+1})\right]$ is 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 $\mathcal{D}_{\mathrm{env}}$ collected via interaction in the real environment, the dynamics model $\hat{T}$ is typically trained to maximize the expected likelihood $\mathbb{E}_{(s_{t},a_{t},r_{t+1},s_{t+1})\sim\mathcal{D}_{\mathrm{env}}}[\log
\hat{T}(s_{t+1},r_{t+1}|s_{t},a_{t})]$. The estimated dynamics model defines a surrogate MDP $\hat{\mathcal{M}}=(\mathcal{S},\mathcal{A},\hat{T},\rho_{0},\gamma)$. Then any RL algorithm can be used to recover the optimal policy with the augmented dataset $\mathcal{D}_{\mathrm{env}}\cup\mathcal{D}_{\mathrm{model}}$, where $\mathcal{D}_{\mathrm{model}}$ is the synthetic data rolled out in $\hat{\mathcal{M}}$.

The above-mentioned paradigm is adopted by model-based policy optimization (MBPO) [] and much of its follow-up work [, , , ] in the online setting. These works don’t need to consider the issue of model coverage, as the agent can explore online to fill in the regions where the dynamics model is uncertain. However, in the offline setting, the limited dataset causes $\hat{T}$ to cover only a part of the state-action space. Once the agent encounters out-of-distribution samples during roll-out in $\hat{\mathcal{M}}$, the learning process can collapse. Therefore, MOPO [] and some of its subsequent offline MBRL algorithms [, ] incorporate a penalty term in the reward function to measure the model uncertainty, allowing the agent to sample within safe regions of $\hat{T}$.

## 3 Method

In this section, we propose a special Any-step Dynamics Model (ADM) to replace the mainstream ensemble dynamics models. Applying ADM to existing MBRL frameworks for policy optimization, we introduce two algorithms, namely online ADMPO-ON and offline ADMPO-OFF. ADM improves future state predictions since it reduces bootstrapping prediction to direct prediction by backtracking variable-length plans. Consequently, ADMPO-ON can improve the sample efficiency in the online setting, while ADMPO-OFF can accurately estimate the model uncertainty in the offline setting.

### 3.1 Any-step Dynamics Model

Currently, the prevalent dynamics models typically operate on a single-step basis, with $s_{t}$ and $a_{t}$ as inputs to predict $s_{t+1}$ and $r_{t+1}$. In a broader context, dynamics models can also be multi-step [, , ], where inputs encompass $s_{t}$ along with a $k$-step sequence of actions $(a_{t},a_{t+1},\cdots,a_{t+k-1})$ to predict $s_{t+k}$ and $r_{t+k}$. To introduce flexibility in the backtracking length of the model, we further extend the definition of the multi-step dynamics model to allow $k$ 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 $\hat{T}(s_{t+k},r_{t+k}|s_{t},a_{t:t+k-1})$ is the distribution of $s_{t+k}\in\mathcal{S}$ and $r_{t+k}\in\mathbb{R}$ conditioned on the $k$-step plan $(s_{t},a_{t:t+k-1})=(s_{t},a_{t},a_{t+1},\cdots,a_{t+k-1})\in\mathcal{S}\times
\mathcal{A}^{k}$, where $k$ can be any integer between $[1,m]$.

Figure: Figure 1: Illustration of any-step dynamics model (left) structured using RNN and its application for next-step prediction with random backtracking (right).
Refer to caption: x1.png

To handle inputs with variable step sizes, we utilize an RNN [] with a GRU [] cell to implement the any-step dynamics model, as depicted in the left part of Figure 1. Certainly, Transformer [] is also a feasible choice, but we do not consider it 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 a sequence of multiple steps, we duplicate the state to match the length of the action sequence, and then sequentially feed it into the RNN. The input $(s_{t},a_{t:t+k-1})$, after being represented by the RNN, yields the hidden $h^{k}_{t}$, which is then fed into an MLP to obtain the mean and standard deviation of $s_{t+k}$ and $r_{t+k}$, i.e., $(\boldsymbol{\mu}^{s}_{t+k},\boldsymbol{\Sigma}^{s}_{t+k})$ and $(\boldsymbol{\mu}^{r}_{t+k},\boldsymbol{\Sigma}^{r}_{t+k})$. Similar to previous model-based methods [, , ], we model the distributions of $s_{t+k}$ and $r_{t+k}$ as Gaussian distributions and predict them through sampling. We call the Any-step Dynamics Model as ADM and denote it as $\hat{T}_{\theta}(s_{t+k},r_{t+k}|s_{t},a_{t:t+k-1})$, where $\theta$ represents the neural parameters. With the real samples from the environment, $\hat{T}_{\theta}$ is trained to maximize the expected likelihood:

$$ $J_{T}(\theta)=\frac{1}{m}\sum_{k=1}^{m}\mathbb{E}_{(s_{t},a_{t:t+k-1},r_{t+k}, s_{t+k})\sim\mathcal{D}_{\textrm{env}}}\left[\log\hat{T}_{\theta}(s_{t+k},r_{t +k}|s_{t},a_{t:t+k-1})\right].$ (2) $$

With $\hat{T}_{\theta}$, the frequent bootstrapping during model roll-out can be reduced. Specifically, given the maximum backtracking length $m$, a state-action sequence of length $m$, $(s_{1},a_{1},s_{2},a_{2},\cdots,s_{m},a_{m})$, is sampled from the data buffer to start the roll-out in $\hat{T}_{\theta}$. For predicting $s_{m+1}$, an integer between $[1,m]$ is chosen uniformly at random as the backtracking length. If only one step is selected for backtracking, $(s_{m},a_{m})$ will be fed into $\hat{T}_{\theta}$ to obtain the prediction result; if $m-1$ steps are chosen, $(s_{2},a_{2:m})$ will be fed into $\hat{T}_{\theta}$, and so forth. The right part of Figure 1 illustrates the aforementioned process based on random backtracking. Further prediction for $s_{m+2}$ should backtrack to at most $(s_{2},a_{2:m+1})$, as $\hat{T}_{\theta}$ is trained with a maximum sequence length of $m$ steps. Similarly, subsequent state predictions can also backtrack at most $m$ steps. The backtracked state, one part of the attribution for next state prediction, is located several steps ahead in expectation. Thus, ADM reduces the actual bootstrapping count of a rolled-out trajectory. The complete $H$-step roll-out process in ADM is described in Algorithm 1.

Figure: Algorithm 1 Roll-out in ADM: ADM-Roll($\hat{T}_{\theta}$, $\pi_{\phi}$, $H$, $m$, $(s_{1},a_{1},s_{2},a_{2},\cdots,s_{m-1},a_{m-1},s_{m})$)

Similar to existing MBRL algorithms, policy roll-out in ADM can generate a 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 can be plugged into this framework. In the subsequent subsections, we will introduce two foundational algorithms, AMRPO-ON and ADMPO-OFF, for online and offline settings, respectively.

###### Definition 3.1 (Any-step Dynamics Model) .

### 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 [], 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 a large number of fake samples using the roll-out method depicted in Algorithm 1. The detailed pseudo-code is provided by Algorithm 2 in Appendix C.1.

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 $k$ steps, we view the sampling process $(\hat{s}_{t+1},\hat{r}_{t+1})\sim\hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t})$ as $(\hat{s}_{t+1},\hat{r}_{t+1})=\mu_{\theta}(s_{t-k+1},a_{t-k+1:t})+\eta_{t+1}$ with $\eta_{t+1}\sim\mathcal{N}(0,\Sigma_{\theta}(s_{t-k+1},a_{t-k+1:t}))$, where $\mu_{\theta}$ is the deterministic dynamics function and $\Sigma_{\theta}$ is the standard deviation function used to construct the noise distribution with zero mean. In expectation, the target value of $Q(s_{t},a_{t})$ is estimated as

$$ $\mathbb{E}_{(s_{t-m+1:t-1},a_{t-m+1:t-1})\sim\Gamma^{m-1}_{\pi}(\cdot|s_{t})} \left[\frac{1}{m}\sum_{k=1}^{m}\mathbb{E}_{(\hat{s}_{t+1},\hat{r}_{t+1})\sim \hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t})}\left[y(\hat{s}_{t+1},\hat{r}_{t +1})\right]\right],$ (3) $$

where $y(\hat{s}_{t+1},\hat{r}_{t+1})=\hat{r}_{t+1}+\gamma\mathbb{E}_{a\sim\pi(\cdot|
\hat{s}_{t+1})}\left[Q(\hat{s}_{t+1},a)\right]$. 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 $\mathcal{N}(0,\Sigma_{\theta}(s_{t-k+1},a_{t-k+1:t}))$ at each backtracking-length $k$. According to [], variations of state predictions can effectively implicitly regularize the local Lipschitz condition of the Q network around regions where the model prediction is uncertain, thereby regulating the value-aware model error [].

### 3.3 ADMPO-OFF: ADM for Policy Optimization in 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 of the learned model needs to be focused within the safe regions. Simultaneously, efforts should be made to explore beyond the boundaries of the risky regions to discover samples conducive to a better policy than the behavior policy. Achieving such a balance between conservatism and generalization often requires measuring model uncertainty. Based on ADM, next 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 explores towards risky regions, the discrepancies tend to increase. The difference among probabilistic predictions $\hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t})$ 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 (ADM-Uncertainty Quantifier) .

For any maximum backtracking length $m$ and the corresponding learned ADM $\hat{T}_{\theta}$, the uncertainty of $\hat{T}_{\theta}$ at $(s_{t},a_{t})$ is quantified as

$$ $\displaystyle\mathcal{U}^{\mathrm{ADM}}(s_{t},a_{t})=$ $\displaystyle\mathbb{E}_{\Gamma^{m-1}_{\pi}(\cdot|s_{t})}\left[\left\|\mathrm{ Var}_{k\sim\mathrm{Uniform}(m),\hat{s}_{t+1}\sim\hat{T}_{\theta}(\cdot|s_{t-k+ 1},a_{t-k+1:t})}\left[\hat{s}_{t+1}\right]\right\|_{1}\right]$ (4) $\displaystyle=$ $\displaystyle\mathbb{E}_{\Gamma^{m-1}_{\pi}(\cdot|s_{t})}\left[\left\|\frac{1} {m}\sum_{k=1}^{m}\left((\Sigma_{\theta}^{k})^{2}+(\mu_{\theta}^{k})^{2}\right) -(\bar{\mu})^{2}\right\|_{1}\right]$ $$

for any $s_{t}\in\mathcal{S}$ and $a_{t}\in\mathcal{A}$, where $\Sigma_{\theta}^{k}=\Sigma_{\theta}(s_{t-k+1},a_{t-k+1:t})$, $\mu_{\theta}^{k}=\mu_{\theta}(s_{t-k+1},a_{t-k+1:t})$ for convenience, and $\bar{\mu}=\frac{1}{m}\sum_{k=1}^{m}\mu_{\theta}^{k}$.

This uncertainty term corresponds to a combination of epistemic and aleatoric model uncertainty with a similar form to the ensemble standard deviation [, ]. 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 [, , , ], we assume that our ADM uncertainty (4) is an admissible error estimator [], as described in Assumption 3.3.

###### Assumption 3.3 (Admissible Error Estimator) .

Assume that there exists a positive $b\in\mathbb{R}^{+}$ such that the following inequality (5) holds for any maximum backtracking length $m$ and any $s_{t}\in\mathcal{S}$, $a_{t}\in\mathcal{A}$.

$$ $D_{\mathrm{TV}}(\bar{T}_{\theta,m}(\cdot|s_{t},a_{t}),T(\cdot|s_{t},a_{t})) \leq b\cdot\mathcal{U}^{\mathrm{ADM}}(s_{t},a_{t}),$ (5) $$

where $\bar{T}_{\theta,m}$ is the overall conditioned distribution coming from

$$ $\bar{T}_{\theta,m}(\cdot|s_{t},a_{t})=\frac{1}{m}\sum_{k=1}^{m}\left[\sum_{ \begin{subarray}{c}s_{t-k+1}\\ a_{t-k+1:t-1}\end{subarray}}\Gamma^{k-1}_{\pi}(s_{t-k+1},a_{t-k+1:t-1}|s_{t}) \hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t}))\right].$ (6) $$

Under Assumption 3.3 and the $\xi$-uncertainty quantifier definition (see Appendix A for details) proposed by PEVI [], we present the following theorem, demonstrating that $\mathcal{U}^{\mathrm{ADM}}$ can serve as a $\xi$-uncertainty quantifier to bound the Bellman error.

###### Theorem 3.4 .

$\beta\cdot\mathcal{U}^{\mathrm{ADM}}$ is a valid $\xi$-uncertainty quantifier, with $\beta=b\frac{\gamma r_{\mathrm{max}}}{1-\gamma}$. Specifically,

$$ $\left|\hat{\mathcal{T}}^{\pi}Q(s_{t},a_{t})-\mathcal{T}^{\pi}Q(s_{t},a_{t}) \right|\leq\beta\cdot\mathcal{U}^{\mathrm{ADM}}(s_{t},a_{t}),$ (7) $$

where $\hat{\mathcal{T}}^{\pi}$ is the proxy Bellman operator induced by ADM to estimate the true Bellman operator $\mathcal{T}^{\pi}$.

###### Proof.

See Appendix B.
∎

According to the suboptimality theorem (see Appendix A for details) presented by PEVI [], the policy $\hat{\pi}$ derived via pessimistic value iteration, which incorporates any $\xi$-uncertainty quantifier as a penalty term into the value iteration process [], has a bounded optimality gap to the optimal policy $\pi^{*}$. The optimality gap is 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 to 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, otherwise the model will induce inaccurate value estimations on these actions. Thus, we can penalize the Bellman operator to obtain a pessimistic value estimation by

$$ $\hat{\mathcal{T}}^{\mathrm{ADM}}Q(s_{t},a_{t}):=\hat{\mathcal{T}}^{\pi}Q(s_{t} ,a_{t})-\beta\cdot\mathcal{U}^{\mathrm{ADM}}(s_{t},a_{t}).$ (8) $$

We expect the penalty term $\beta\cdot\mathcal{U}^{\mathrm{ADM}}(s_{t},a_{t})$ to be as small as possible thereby constraining the optimality gap. While our Assumption 3.3 lacks theoretical guarantees and the tightness of the bound in 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 framework of MOPO []. The detailed pseudo-code is provided by Algorithm 3 in Appendix C.2.

###### Definition 3.2 (ADM-Uncertainty Quantifier) .

###### Assumption 3.3 (Admissible Error Estimator) .

###### Theorem 3.4 .

###### Proof.

## 4 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?

### 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 [] datasets, hopper-medium-v2, hopper-medium-replay-v2, walker2d-medium-v2, and walker2d-medium-replay-v2, to compare the compounding error 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 predicts $s_{t+m+1}$ using the historical state-action sequence $(s_{t},a_{t},\cdots,s_{t+m},a_{t+m})$ as input. Figure 2 shows the growth curves of the compounding error as the roll-out length increases. 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 a certain threshold. This phenomenon suggests ADM can improve predictions for future states due to its any-step backtracking mechanism during model roll-outs.

Figure: Figure 2: Comparison among ADM, ensemble dynamics model, and bootstrapping RNN dynamics model, in terms of the growth curve of the compounding error as roll-out length increases, after offline learning. The overflow value is regarded as the maximum value of float32.
Refer to caption: x2.png

Figure: Figure 3: Online learning curves of ADMPO-ON (red) and other five 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.
Refer to caption: x3.png

### 4.2 Evaluation in Online Setting

We evaluate ADMPO-ON on four difficult MuJoCo continuous control tasks [], including Hopper, Walker2d, Ant, and Humanoid. All the tasks adopt version v3 and follow the default settings. Four model-based methods and one model-free method are selected as our baselines. These include SAC [], whcih is the state-of-the-art model-free RL algorithm; STEVE [], which incorporates an ensemble into the model-based value expansion; MBPO [], which updates the policy with a mixture of real environmental samples and branched roll-out data; BMPO [], which builds upon MBPO and replaces the dynamics model with a bidirectional one; and DDPPO [], which adopts a two-model-based learning method to control the prediction error and the gradient error.

Figure 3 shows learning curves of ADMPO-ON and other five baselines, along with SAC’s asymptotic performance. ADMPO-ON achieves competitive performance after fewer environmental steps than the baselines. Taking the most difficult Humanoid as an example, ADMPO-ON has 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 performs about 1.33x faster than DDPPO and dominates other baselines in terms of learning efficiency on the Humanoid task. After training, ADMPO-ON can achieve a final performance close to the asymptotic performance of SAC on all these four MuJoCo tasks. These results demonstrate that ADMPO-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.1.

### 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 [], which equally penalized the Q values on out-of-the-distribution state-action pairs; TD3+BC [], which simply incorporates a BC term into the policy optimization objective of TD3 []; and EDAC [], which quantifies the Q uncertainty via ensemble; as well as five model-based methods: MOPO [], which adds the uncertainty of the model prediction as a penalization term to the reward function; COMBO [], which introduces the penalty function of CQL into the model-based framework; RAMBO [], which adversarially trains the dynamics model and the policy; CBOP [], which adopts the variance of values under an ensemble of dynamics models to estimate the Q value conservatively under MVE [] regime; and MOBILE [], which proposes Model-Bellman inconsistency to estimate the Bellman error.

Table 1 reports the results on twelve D4RL [] MuJoCo datasets (v2 version). The normalized score for each dataset is obtained via online evaluation after offline learning. The source of the reported performance in provided in Appendix D.4. We observe that ADMPO-OFF outperforms the other nine baselines in most tasks and achieves the highest average score.

**Table 1: Normalized scores after offline learning on D4RL MuJoCo tasks, averaged over five seeds.**
| Task Name | BC | CQL | TD3+BC | EDAC | MOPO | COMBO | RAMBO | CBOP | MOBILE | ADMPO-OFF (ours) |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| hopper-random | 3.7 | 5.3 | 8.5 | 25.3 | 31.7 | 17.9 | 25.4 | 31.4 | 31.9 | 32.7$\pm$0.2 |
| halfcheetah-random | 2.2 | 31.3 | 11.0 | 28.4 | 38.5 | 38.8 | 39.5 | 32.8 | 39.3 | 45.4$\pm$2.8 |
| walker2d-random | 1.3 | 5.4 | 1.6 | 16.6 | 7.4 | 7.0 | 0.0 | 17.8 | 17.9 | 22.2$\pm$0.2 |
| hopper-medium | 54.1 | 61.9 | 59.3 | 101.6 | 62.8 | 97.2 | 87.0 | 102.6 | 106.6 | 107.4$\pm$0.6 |
| halfcheetah-medium | 43.2 | 46.9 | 48.3 | 65.9 | 73.0 | 54.2 | 77.9 | 74.3 | 74.6 | 72.2$\pm$0.6 |
| walker2d-medium | 70.9 | 79.5 | 83.7 | 92.5 | 84.1 | 81.9 | 84.9 | 95.5 | 87.7 | 93.2$\pm$1.1 |
| hopper-medium-replay | 16.6 | 86.3 | 60.9 | 101.0 | 103.5 | 89.5 | 99.5 | 104.3 | 103.9 | 104.4$\pm$0.4 |
| halfcheetah-medium-replay | 37.6 | 45.3 | 44.6 | 61.3 | 72.1 | 55.1 | 68.7 | 66.4 | 71.7 | 67.6$\pm$3.4 |
| walker2d-medium-replay | 20.3 | 76.8 | 81.8 | 87.1 | 85.6 | 56.0 | 89.2 | 92.7 | 89.9 | 95.6$\pm$2.1 |
| hopper-medium-expert | 53.9 | 96.9 | 98.0 | 110.7 | 81.6 | 111.1 | 88.2 | 111.6 | 112.6 | 112.7$\pm$0.3 |
| halfcheetah-medium-expert | 44.0 | 95.0 | 90.7 | 106.3 | 90.8 | 90.0 | 95.4 | 105.4 | 108.2 | 103.7$\pm$0.2 |
| walker2d-medium-expert | 90.1 | 109.1 | 110.1 | 114.7 | 112.9 | 103.3 | 56.7 | 117.2 | 115.2 | 114.9$\pm$0.3 |
| Average | 36.5 | 61.6 | 58.2 | 76.0 | 70.3 | 66.8 | 67.7 | 79.3 | 80.0 | 81.0 |

#### 4.3.2 NeoRL Benchmark Results

NeoRL [], is an offline RL benchmark that collects the data in a 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.

We compare our ADMPO-OFF with six baselines, including BC, CQL, TD3+BC, EDAC, MOPO, and MOBILE. Table 2 presents the normalized scores of these methods. Due to the narrow and limited coverage of the NeoRL data, all the baselines experience a 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.

**Table 2: Normalized scores after offline learning on NeoRL tasks, averaged over five seeds.**
| Task Name | BC | CQL | TD3+BC | EDAC | MOPO | MOBILE | ADMPO-OFF (ours) |
| --- | --- | --- | --- | --- | --- | --- | --- |
| neorl-hopper-low | 15.1 | 16.0 | 15.8 | 18.3 | 6.2 | 17.4 | 22.3$\pm$0.1 |
| neorl-halfcheetah-low | 29.1 | 38.2 | 30.0 | 31.3 | 40.1 | 54.7 | 52.8$\pm$1.2 |
| neorl-walker2d-low | 28.5 | 44.7 | 43.0 | 40.2 | 11.6 | 37.6 | 55.9$\pm$3.8 |
| neorl-hopper-medium | 51.3 | 64.5 | 70.3 | 44.9 | 1.0 | 51.1 | 51.5$\pm$5.0 |
| neorl-halfcheetah-medium | 49.0 | 54.6 | 52.3 | 54.9 | 62.3 | 77.8 | 69.3$\pm$1.7 |
| neorl-walker2d-medium | 48.7 | 57.3 | 58.5 | 57.6 | 39.9 | 62.2 | 70.1$\pm$2.4 |
| neorl-hopper-high | 43.1 | 76.6 | 75.3 | 52.5 | 11.5 | 87.8 | 87.6$\pm$4.9 |
| neorl-halfcheetah-high | 71.3 | 77.4 | 75.3 | 81.4 | 65.9 | 83.0 | 84.0$\pm$0.8 |
| neorl-walker2d-high | 72.6 | 75.3 | 69.6 | 75.5 | 18.0 | 74.9 | 82.2$\pm$1.9 |
| Average | 45.4 | 56.1 | 54.5 | 50.7 | 28.5 | 60.7 | 64.0 |

#### 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 D4RL tasks, hopper-medium-replay-v2 and walker2d-medium-replay-v2, are illustrated in Figure 4. We observe that our ADM provides a 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 policy 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.

Figure: (a) hopper-medium-replay-v2
Refer to caption: x4.png

## 5 Related Work

This work is related to online and offline dyna-style MBRL [].

### 5.1 Online Model-based Reinforcement Learning

In the online setting, MBRL algorithms aim to accelerate value estimation or policy optimization with model roll-out data. MVE [] enhances Q-value target estimation by allowing short-term imagination to a fixed depth using the dynamics model. STEVE [] builds upon MVE by incorporating an ensemble into the value expansion to better estimate the Q value. SLBO [] directly utilizes TRPO [] to optimize the policy with synthetic data generated by rolling out to the end of trajectories in the dynamics model. MBPO [] proposes a branched roll-out scheme to truncate unreliable samples, thereby reducing the influence of compounding error [], and employs SAC [] to update the policy with a mixture of real-world data and model-generated data.

Recent work improves MBRL performance mainly from two perspectives. One focuses on learning a better dynamics model, such as bidirectional models [], adversarial models [, ], causal models [], and multi-step models [, , ]. The other pursues a better utilization of the learned model, enhancing the reliability of model-generated samples [] or applying model-based multi-step planning techniques [, , , , , ].

### 5.2 Offline Model-based Reinforcement Learning

Although some model-free RL algorithms [, , , , , ] 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 in how to effectively leverage the model. MOPO [] and MOReL [] add the uncertainty of the model prediction as a penalization term to the original reward function to achieve a pessimistic value estimation. MOBILE [] improves the uncertainty quantification by introducing Model-Bellman inconsistency into the offline model-based framework. COMBO [] applies CQL [] to force the Q value to be small on model-generated out-of-distribution samples. RAMBO [] achieves conservatism by adversarial model learning for value minimization while keeping fitting the transition function. CBOP [] introduces adaptive weighting of short-horizon roll-out into MVE [] technique and adopts the variance of values under an ensemble of dynamics models to estimate the Q value conservatively. MOREC [] designs a reward-consistent dynamics model using an adversarial discriminator to let the model-generated 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 and yields two algorithms, ADMPO-ON and ADMPO-OFF, respectively. Several analysis and experiments show that ADM outperforms the ensemble dynamics model applied in previous MBRL approaches widely. The only problem is that RNN may consume more resources during the training process. We believe ADM has powerful potential beyond the capabilities demonstrated in this paper. In the future, we will explore the scalability of ADM in non-Markovian visual RL scenarios, considering both online and offline settings.

## Appendix A Additional Introduction to Pessimistic Value Iteration (PEVI)

Pessimistic Value Iteration (PEVI) [] is a meta-algorithm for offline RL settings. It constructs an estimated Bellman operator $\hat{\mathcal{T}}^{\pi}$ based on the given dataset $\mathcal{D}_{\mathrm{env}}$ to approximate the true Bellman operator $\mathcal{T}^{\pi}$ that satisfies

$$ $\mathcal{T}^{\pi}V_{h+1}(s_{h},a_{h})=\mathbb{E}_{(s_{h+1},r_{h+1})\sim T( \cdot|s_{h},a_{h})}\left[r_{h+1}+V_{h+1}(s_{h+1})\right],$ (9) $$

where $h$ is the step index less than the horizon $\mathcal{H}$. Then the state-action value function is updated with

$$ $Q_{h}(s_{h},a_{h})\leftarrow\hat{\mathcal{T}}^{\pi}V_{h+1}(s_{h},a_{h})- \Lambda_{h}(s_{h},a_{h})$ (10) $$

for each $(s_{h},a_{h})$, where $\Lambda_{h}$ is the penalty function that guarantees the conservatism of the learned policy. Especially, $\Lambda_{h}$ should be a $\xi$-uncertainty quantifier as follows.

###### Definition A.1 ( $\xi$ -Uncertainty Quantifier (proposed by [ 23 ] )) .

The set of penalization $\{\Lambda_{h}\}_{h\in[\mathcal{H}]}$ forms a $\xi$-uncertainty quantifier if

$$ $\left|\hat{\mathcal{T}}^{\pi}V_{h+1}(s_{h},a_{h})-\mathcal{T}^{\pi}V_{h+1}(s_{ h},a_{h})\right|\leq\Lambda_{h}(s_{h},a_{h})$ (11) $$

holds with probability at least $1-\xi$ for all $(s_{h},a_{h})\in\mathcal{S}\times\mathcal{A}$.

The following theorem characterizes the suboptimality of PEVI.

###### Theorem A.2 (Suboptimality of PEVI (proposed by [ 23 ] )) .

Suppose $\{\Lambda_{h}\}_{h=1}^{\mathcal{H}}$ in PEVI is a set of $\xi$-uncertainty quantifier. Then the derived policy $\hat{\pi}$ satisfies

$$ $\left|V_{1}^{\pi^{*}}(s_{1})-V_{1}^{\hat{\pi}}(s_{1})\right|\leq 2\sum_{h=1}^{ \mathcal{H}}\mathbb{E}_{\rho^{\pi^{*}}}\left[\Lambda_{h}(s_{h},a_{h})\right]$ (12) $$

with probability at least $1-\xi$ for all starting $s_{1}\in\mathcal{S}$. Here $\mathbb{E}_{\rho^{\pi^{*}}}$ is with respect to the trajectory induced by the optimal policy $\pi^{*}$ in the underlying MDP given the fixed function $\Lambda_{h}$.

###### Proof.

See PEVI [] for detailed proof.
∎

###### Definition A.1 ( ξ 𝜉 \xi italic_ξ -Uncertainty Quantifier (proposed by [ 23 ] )) .

###### Theorem A.2 (Suboptimality of PEVI (proposed by [ 23 ] )) .

###### Proof.

## Appendix B Theoretical Results

###### Theorem B.1 .

$\beta\cdot\mathcal{U}^{\mathrm{ADM}}$ is a valid $\xi$-uncertainty quantifier, with $\beta=b\frac{\gamma r_{\mathrm{max}}}{1-\gamma}$. Specifically,

$$ $\left|\hat{\mathcal{T}}^{\pi}Q(s_{t},a_{t})-\mathcal{T}^{\pi}Q(s_{t},a_{t}) \right|\leq\beta\cdot\mathcal{U}^{\mathrm{ADM}}(s_{t},a_{t}),$ (13) $$

where $\hat{\mathcal{T}}^{\pi}$ is the proxy Bellman operator induced by ADM to estimate the true Bellman operator $\mathcal{T}^{\pi}$.

###### Proof.

First, we define $y(\hat{s}_{t+1},\hat{r}_{t+1})=\hat{r}_{t+1}+\gamma\mathbb{E}_{a\sim\pi(\cdot|
\hat{s}_{t+1})}\left[Q(\hat{s}_{t+1},a)\right]$ and expand these two Bellman operator to

$$ $\displaystyle\hat{\mathcal{T}}^{\pi}Q(s_{t},a_{t})$ (14) $\displaystyle=$ $\displaystyle\mathbb{E}_{(s_{t-m+1:t-1},a_{t-m+1:t-1})\sim\Gamma^{m-1}_{\pi}( \cdot|s_{t})}\left[\frac{1}{m}\sum_{k=1}^{m}\mathbb{E}_{(\hat{s}_{t+1},\hat{r} _{t+1})\sim\hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t})}\left[y(\hat{s}_{t+1} ,\hat{r}_{t+1})\right]\right]$ $\displaystyle=$ $\displaystyle\sum_{\begin{subarray}{c}s_{t-m+1}\\ a_{t-m+1:t-1}\end{subarray}}\Gamma^{m-1}_{\pi}(s_{t-m+1},a_{t-m+1:t-1}|s_{t}) \left[\frac{1}{m}\sum_{k=1}^{m}\sum_{\begin{subarray}{c}\hat{s}_{t+1}\\ \hat{r}_{t+1}\end{subarray}}\hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t})y( \hat{s}_{t+1},\hat{r}_{t+1})\right]$ $\displaystyle=$ $\displaystyle\frac{1}{m}\sum_{k=1}^{m}\left[\sum_{\begin{subarray}{c}s_{t-k+1} \\ a_{t-k+1:t-1}\end{subarray}}\Gamma^{k-1}_{\pi}(s_{t-k+1},a_{t-k+1:t-1}|s_{t}) \sum_{\begin{subarray}{c}\hat{s}_{t+1}\\ \hat{r}_{t+1}\end{subarray}}\hat{T}_{\theta}(\cdot|s_{t-k+1},a_{t-k+1:t}))y( \hat{s}_{t+1},\hat{r}_{t+1})\right]$ $\displaystyle=$ $\displaystyle\sum_{\hat{s}_{t+1},\hat{r}_{t+1}}\bar{T}_{\theta,m}(\hat{s}_{t+1 },\hat{r}_{t+1}|s_{t},a_{t})y(\hat{s}_{t+1},\hat{r}_{t+1}),$ $$

and

$$ $\displaystyle\mathcal{T}^{\pi}Q(s_{t},a_{t})$ (15) $\displaystyle=$ $\displaystyle\mathbb{E}_{\hat{s}_{t+1},\hat{r}_{t+1}\sim T(\cdot|s_{t},a_{t})} \left[y(\hat{s}_{t+1},\hat{r}_{t+1})\right]$ $\displaystyle=$ $\displaystyle\sum_{\hat{s}_{t+1},\hat{r}_{t+1}}T(\hat{s}_{t+1},\hat{r}_{t+1}|s _{t},a_{t})y(\hat{s}_{t+1},\hat{r}_{t+1}).$ $$

Then, we can obtain

$$ $\displaystyle\left|\hat{\mathcal{T}}^{\pi}Q(s_{t},a_{t})-\mathcal{T}^{\pi}Q(s_ {t},a_{t})\right|$ (16) $\displaystyle=$ $\displaystyle\sum_{\hat{s}_{t+1},\hat{r}_{t+1}}\left|\bar{T}_{\theta,m}(\hat{s }_{t+1},\hat{r}_{t+1}|s_{t},a_{t})-T(\hat{s}_{t+1},\hat{r}_{t+1}|s_{t},a_{t}) \right|\cdot\left|y(\hat{s}_{t+1},\hat{r}_{t+1})\right|$ $\displaystyle=$ $\displaystyle\gamma\sum_{\hat{s}_{t+1},\hat{r}_{t+1}}\left|\bar{T}_{\theta,m}( \hat{s}_{t+1},\hat{r}_{t+1}|s_{t},a_{t})-T(\hat{s}_{t+1},\hat{r}_{t+1}|s_{t},a _{t})\right|\cdot\left|\mathbb{E}_{a\sim\pi(\cdot|\hat{s}_{t+1})}\left[Q(\hat{ s}_{t+1},a)\right]\right|$ $\displaystyle\leq$ $\displaystyle\dfrac{\gamma r_{\mathrm{max}}}{1-\gamma}\sum_{\hat{s}_{t+1},\hat {r}_{t+1}}\left|\bar{T}_{\theta,m}(\hat{s}_{t+1},\hat{r}_{t+1}|s_{t},a_{t})-T( \hat{s}_{t+1},\hat{r}_{t+1}|s_{t},a_{t})\right|$ $\displaystyle=$ $\displaystyle\dfrac{\gamma r_{\mathrm{max}}}{1-\gamma}D_{\mathrm{TV}}(\bar{T}_ {\theta,m}(\cdot|s_{t},a_{t}),T(\cdot|s_{t},a_{t}))$ $\displaystyle\leq$ $\displaystyle b\dfrac{\gamma r_{\mathrm{max}}}{1-\gamma}\mathcal{U}^{\mathrm{ ADM}}(s_{t},a_{t}).$ $$

Thus, let $\beta=b\frac{\gamma r_{\mathrm{max}}}{1-\gamma}$, we can say that $\beta\cdot\mathcal{U}^{\mathrm{ADM}}$ is a valid $\xi$-uncertainty quantifier, as defined by Definition A.1.
∎

###### Theorem B.1 .

###### Proof.

## Appendix C Implementation Details

### C.1 ADMPO-ON

Our ADMPO-ON algorithm follows the framework of MBPO [], 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 blue-highlighted parts in the pseudo-code.

Figure: Algorithm 2 ADMPO-ON

### C.2 ADMPO-OFF

Our ADMPO-OFF algorithm follows the framework of MOPO [], as shown in Algorithm 3. The only difference between ADMPO-OFF and MOPO lies lies in the way the dynamics model is trained and utilized, as indicated by the blue-highlighted parts in the pseudo-code.

Figure: Algorithm 3 ADMPO-OFF

### C.3 Policy Optimization

The policy optimization method used in our ADMPO-ON and ADMPO-OFF is SAC [], following MBPO [] and MOPO []. The hyper-parameters about SAC follow its standard implementation, as listed in Table 3.

**Table 3: Hyper-parameters of Policy Optimization in ADMPO-ON and ADMPO-OFF.**
| Hyper-parameter | Value | Description |
| --- | --- | --- |
| $N_{Q}$ | 2 | the number of critics. |
| actor network | FC(256,256) | fully connected (FC) layers with ReLU activations. |
| critic network | FC(256,256) | fully connected (FC) layers with ReLU activations. |
| $\tau$ | $5\times 10^{-3}$ | target network smoothing coefficient. |
| $\gamma$ | 0.99 | discount factor. |
| $lr_{\textrm{actor}}$ | $1\times 10^{-4}$ | learning rate of actor. |
| $lr_{\textrm{critic}}$ | $3\times 10^{-4}$ | learning rate of critic. |
| optimizer | Adam | optimizers of the actor and critics. |
| batch size | 256 | batch size for each update. |

## Appendix D Experimental Details

### D.1 Resource Requirements

All experiments can be completed with just one NVIDIA GeForce 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 4.

**Table 4: Hyper-parameter settings of ADMPO-ON results presented in Figure 3. $x\rightarrow y$ over $a\rightarrow b$ denotes a thresholded linear increasing schedule, i.e. the length of model rollouts at step $t$ is calculated by $f(t)=\min\left(\max\left(x+\frac{t-a}{b-a}\cdot(y-x),x\right),y\right)$.**
| environment | Hopper | Walker2d | Ant | Humanoid |
| --- | --- | --- | --- | --- |
| steps | 50k | 200k | 300k |  |
| Update-To-Date ratio | 20 |  |  |  |
| maximum backtracking length $\boldsymbol{m}$ | 5 | 2 |  |  |
| model rollout schedule | 1$\rightarrow$15 over | 1$\rightarrow$10 over | 1$\rightarrow$5 over | 1$\rightarrow$10 over |
| 0$\rightarrow$50k | 0$\rightarrow$100k | 10k$\rightarrow$100k | 10k$\rightarrow$100k |  |
| target entropy | -1 | -3 | -4 | -8 |

### D.3 ADMPO-OFF Settings

The experimental settings of our ADMPO-OFF in 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 | $\boldsymbol{m}$ | $\boldsymbol{H}$ | $\boldsymbol{\beta}$ |
| --- | --- | --- | --- | --- |
| D4RL MuJoCo | hopper-random | 5 | 50 | 5 |
| halfcheetah-random | 2 | 10 | 2.5 |  |
| walker2d-random | 2 | 50 | 2.5 |  |
| hopper-medium | 5 | 10 | 1 |  |
| halfcheetah-medium | 2 | 5 | 2.5 |  |
| walker2d-medium | 5 | 10 | 5 |  |
| hopper-medium-replay | 5 | 5 | 0.1 |  |
| halfcheetah-medium-replay | 2 | 5 | 2.5 |  |
| walker2d-medium-replay | 5 | 5 | 0.1 |  |
| hopper-medium-expert | 2 | 20 | 20 |  |
| halfcheetah-medium-expert | 2 | 50 | 10 |  |
| walker2d-medium-expert | 3 | 2 | 6 |  |
| NeoRL MuJoCo | neorl-hopper-low | 5 | 20 | 5 |
| neorl-halfcheetah-low | 2 | 20 | 10 |  |
| neorl-walker2d-low | 5 | 10 | 2.5 |  |
| neorl-hopper-medium | 5 | 20 | 50 |  |
| neorl-halfcheetah-medium | 2 | 5 | 20 |  |
| neorl-Walker2d-medium | 5 | 10 | 5 |  |
| neorl-hopper-high | 5 | 20 | 50 |  |
| neorl-halfcheetah-high | 2 | 10 | 50 |  |
| neorl-walker2d-high | 5 | 10 | 2.5 |  |

### D.4 Source of Baselines’ Results

For the evaluation on D4RL [] benchmarks, the results of the compared baselines come from two sources:

- •
Retraining on D4RL datasets of v2 version with OfflineRL-Kit [], for the algorithms whose original papers only report the performance on the v0 version, such as CQL [], MOPO [].
- •
Including the scores in their papers, for the algorithms whose original papers report the performance on the v2 version, such as TD3+BC [], EDAC [], RAMBO [], CBOP [], and MOBILE [], or who does not provide source codes, such as COMBO [].

For the evaluation on NeoRL [] 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 [].

## Appendix E Additional Experiments

### E.1 Study on Why ADMPO-ON Performs Well in Online Setting

Value-aware model error [] is a dependable metric for measuring the learning quality of the dynamics model and the suboptimality of the MBRL algorithm. We conduct a study to verify how well ADMPO-ON regulates the value-aware model error. Without loss of rigor, we only choose MBPO for comparison since most other model-based methods follow the same way of learning and utilizing the ensemble dynamics model. Figure 5 shows the results on 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-ON 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 a significantly smaller Lipschitz constant, and afterwards the value-aware model error, which measures the suboptimality of MBRL becomes smaller. This phenomenon explains why ADMPO-ON performs significantly better then MBPO in Figure 3. For details of the metrics used in this experiment, refer to [].

Figure: Figure 5: 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 [] of Q, and (d) value-aware model error []. Results are averaged over five seeds.
Refer to caption: x6.png

### E.2 Study on Maximum Backtracking Length

Figure: Figure 6: Illustration of ADMPO-OFF’s performance under different maximum backtracking length.
Refer to caption: x7.png

Maximum backtracking length $m$ is an important hyper-parameter in our algorithms. Figure 5 shows the influence of $m$ on the performance of hopper-medium-v2 and walker2d-medium-v2 tasks, respectively. We observe that the performance of ADMPO-OFF is not very sensitive to the hyper-parameter $m$.

## Peer Reviews

confidence: 4
contribution: 3
initial_rating: 6
presentation: 3
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.)

[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.)
soundness: 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.

[S2] ADM demonstrates strong performance in both online and offline experiments, with uncertainty measures that correlate well with ground-truth errors.
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.
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.

[W2] The reference formatting in the main text appears incorrect; minor revisions are recommended.

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 3
questions: Three questions are raised as part of the weaknesses. I would additionally ask the following:
- 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?
- 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.
- ADMPO tunes hyperparameters on each task. How are the hyperparameters for the baselines selected?
- There are inconsistencies in the spelling of the word rollout/roll-out.
- Equation 1: for the sake of clarity, can the authors specify over which spaces/sets the sum is computed?
- line 204: what about these algorithms is "foundational"?

**References:**

[1] Hafner et al., Mastering Diverse Domains through World Models, arXiv 2023

[2] Hanses et al., TD-MPC2: Scalable, Robust World Models for Continuous Control, ICLR 2024

[3] Filos et al., Model-Value Inconsistency as a Signal for Epistemic Uncertainty, ICML 2022
soundness: 2
strengths: - The presentation is clear and the paper is well written.
- The empirical evaluation is broad, and the proposed methods largely outperform the selected baselines.
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.
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?
- 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?
- 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?

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: * Why aren't all $m$ known states used as input for the model?
 * Why is a random $k$ used to train the policy?
 * 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.

Overall 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.


After rebuttal:

The authors have provided helpful experiments that confirm their changes to the model to be advantageous.
I still have concerns particularly about the lack of hyper-parameter optimization for baselines.

Overall I am raising my score (3->5)
soundness: 2
strengths: * The results look quite promising, performing well in Online RL and Offline RL benchmarks
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})$.
This somehow leads to a signficiantly better performance and improved results when used for model-based online and Offline RL.
For Offline RL, the authors also use the ensemble disagreement for different input-lenghts as a measure of model uncertainty.
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.
 * 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.

Some important implementation choices are not justified either. To list the most major issues:
 * 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?
 * 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?

* 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).


Minor issues:
 * 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.
 * parantheses should be placed around the author name for citations, not just around the year number.
 * 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.
* Algorithm 2 and 1 should probably combined and both be in the main text, as they are important to understand the approach.
* L241: "explore beyond the boundaries of the risky regions" should be "beyond the boundaries of the safe regions"?
* L352: should this be $\hat{s}_{t+m}$? If not, there can be no compounding error.
 *

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 2
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]. 
- line 046, what is deviation error? Line 68, what do you mean by trajectory distribution? 
- What is the computation complexity of the proposed method?



[Q1] Frauenknecht, Bernd, et al. "Trust the Model Where It Trusts Itself-Model-Based Actor-Critic with Uncertainty-Aware Rollout Adaption." ICML.

[Q2] Zhang, Weinan, et al. "Model-based multi-agent policy optimization with adaptive opponent-wise rollouts." arXiv preprint arXiv:2105.03363 (2021).
soundness: 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.
- 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.
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.
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. 
- 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.
- 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?

Meta-review: This paper proposes the Any-step Dynamics Model (ADM), which is a model-based method that learns to predict transitions multiple steps into the future to avoid accumulating errors due to bootstrapping. They provide an online and an offline variant of their method, with their efficacy demonstrated by empirical evaluations. Additionally, the authors provide a mechanism for uncertainty estimation using ADM in the offline setting.

The paper is clear and well-written, and the empirical evidence supports the claims.

The main criticism from reviewers seems to revolve around hyper-parameter optimization for the baselines, in particular in the NeoRL benchmark, as the baselines were likely not optimized for that benchmark. Given that most of the evaluations were done in more standard (and well-understood) benchmarks, I believe it is reasonable to use the hyper-parameters provided by the original algorithm implementations. It would be good if the authors could make this point more clearly for readers.

Although there are some concerns on how it can handle stochastic environments, many works do introduce novel ideas in systems with limited stochasticity, so I do not see this as a reason to reject.

Thus, I am recommending acceptance of this work, as it does introduce a novel algorithm with reasonable empirical evidence for its effectiveness, which may be of general interest to the RL community.

Decision: Accept (Poster)