MAD-TD: Model-Augmented Data stabilizes High Update Ratio RL

## Abstract

Abstract Building deep reinforcement learning (RL) agents that find a good policy with few samples has proven notoriously challenging.
To achieve sample efficiency, recent work has explored updating neural networks with large numbers of gradient steps for every new sample.
While such high update-to-data (UTD) ratios have shown strong empirical performance, they also introduce instability to the training process.
Previous approaches need to rely on periodic neural network parameter resets to address this instability, but restarting the training process is infeasible in many real-world applications and requires tuning the resetting interval.
In this paper, we focus on one of the core difficulties of stable training with limited samples: the inability of learned value functions to generalize to unobserved on-policy actions.
We mitigate this issue directly by augmenting the off-policy RL training process with a small amount of data generated from a learned world model.
Our method, Model-Augmented Data for Temporal Difference learning (MAD-TD) uses small amounts of generated data to stabilize high UTD training and achieve competitive performance on the most challenging tasks in the DeepMind control suite.
Our experiments further highlight the importance of employing a good model to generate data, MAD-TD’s ability to combat value overestimation, and its practical stability gains for continued learning.

## 1 Introduction

Instead of solely relying on data gathered by a target policy, *off-policy* reinforcement learning (RL) aims to leverage experience gathered by past policies (Sutton & Barto, ) to fit a value function for the target policy.
Ideally, training over many iterations should help extract important information from past data.
However, recent work has shown that simply increasing the number of gradient update steps, the *replay ratio* or *update-to-data (UTD) ratio*, can lead to highly unstable learning (Nikishin et al., ; D’Oro et al., ; Hussing et al., ; Nauman et al., ).

Prior work has stabilized the learning by using double Q minimization to reduce overestimation (Fujimoto et al., ), training ensemble methods to improve value estimation (Chen et al., ; Hiraoka et al., ), introducing architectural regularization (Hussing et al., ; Nauman et al., ), or resetting networks periodically throughout the learning process (D’Oro et al., ; Schwarzer et al., ; Nauman et al., ).
However, while useful, pessimistic underestimation and architectural regularization are insufficient by themselves to combat the problem (Hussing et al., ), and so most methods resort to either network resets or ensembles.
Critic ensembles can be expensive to train, and resetting has several important limitations: in real systems, re-executing a random policy can be expensive or unsafe, the resetting interval needs to be carefully tuned (Hussing et al., ), and when storing a full reset buffer is infeasible, resetting loses important information.

We narrow in on a key issue contributing to unstable training: *wrong value function estimation on unobserved on-policy actions* (Thrun & Schwartz, ; Tsitsiklis & Van Roy, ).
Off-policy RL uses the values of states sampled under old policies with actions from the target policy to update the value function.
However, these state-action pairs themselves are not in the replay buffer and hence their value estimate is not directly improved by training.
Consequently, a learned function which achieves low error on seen data is not guaranteed to generalize well to actions that *differ* from past actions.
This problem is related to overfitting (Li et al., ) and contributes to overestimation (Thrun & Schwartz, ; Hasselt, ; Fujimoto et al., ).
However, overfitting assumes that train and test set are drawn from the same distribution, while we focus on the distribution shift between data collection and target policy.
Previous work has investigated the difficulty of off-policy learning due to this shift (Maei et al., ; Sutton et al., ; Hasselt, ; Fujimoto et al., ), yet there are no tractable mitigation strategies that work well in the high UTD regime with deep RL.

To corroborate our hypothesis that generalization to unobserved actions is a major obstacle for training at high UTDs, we examine the behavior of value functions on on-policy transitions.
Our experiments reveal that value functions generalize significantly worse to unobserved on-policy action transitions than to validation data from the same distribution as the training set.
Building on this, we propose to improve on-policy value estimation by using *model-generated on-policy data*.

Previous investigations into model-based deep RL have focused on learning values fully in model roll-outs (Buckman et al., ; Janner et al., ; Hafner et al., ; Ghugare et al., ) and the associated challenges (Zhao et al., ; Hansen et al., ).
In contrast, we show that mixing a small amount of model-generated on-policy data with real off-policy replay data is sufficient to achieve stable learning in the high UTD regime.
Our method, Model-Augmented Data for Temporal Difference learning(MAD-TD), mitigates the generalization issues of the value function in the hardest tasks of the DeepMind control (DMC) benchmark (Tunyasuvunakool et al., ) and achieves strong and stable high UTD learning without resetting or redundant ensemble learning.

The main contributions of this work are twofold:
First, we empirically show the existence of misgeneralization from off-policy value estimation to on-policy predictions. We connect this issue to the challenge of stable learning with high update ratios and highlight how increasing the update ratio increases Q function overestimation.
Second, we provide a new method called MAD-TD that improves the value function accuracy on unobserved on-policy actions with model-generated data and stabilizes training at high update ratios. This method proves to have equivalent performance to or outperform previous strong baselines.

## 2 Mathematical background

We consider a standard RL setting, the discounted infinite-horizon MDP $(\mathcal{X},\mathcal{A},\mathcal{P},r,\rho,\gamma)$ with state space $\mathcal{X}$, action space $\mathcal{A}$, a transition kernel $\mathcal{P}:\mathcal{X}\times\mathcal{A}\rightarrow\mathcal{M}(\mathcal{X})$, a reward function $r:\mathcal{X}\times\mathcal{A}\rightarrow\mathbb{R}$, starting state distribution $\rho\in\mathcal{M}(\mathcal{X})$ and a discount factor $\gamma\in[0,1)$ (Puterman, ; Sutton & Barto, ).
For a space $Y$ we use $\mathcal{M}(Y)$ to denote the set of probability measures over the space.
Our goal is to learn a policy $\pi:\mathcal{X}\rightarrow\mathcal{M}(\mathcal{A})$  that maximizes the discounted sum of future rewards

$$ $\pi^{*}\in\operatorname*{arg\,max}_{\pi\in\Pi}\sum_{t=0}^{\infty}\mathbb{E}_{ \mathcal{P}^{\pi}}[\gamma^{t}r(x_{t},a_{t})|x_{0}\sim\rho]\enspace,$ $$

where actions are sampled according to the policy and new states according to the transition kernel.

### 2.1 Off-policy value function learning

As an intermediate objective, many algorithms attempt to simplify the direct policy optimization problem by first learning a policy value function $Q^{\pi}$, which is defined via a recursive equation

$$ $Q^{\pi}(x,a)=r(x,a)+\gamma\mathbb{E}_{x^{\prime}\sim\mathcal{P}(\cdot|x,a),a^{ \prime}\sim\pi(\cdot|x^{\prime})}\left[Q^{\pi}(x^{\prime},a^{\prime})\right]\enspace.$ $$

The policy can then be incrementally improved by picking
$\pi_{k+1}(x)\in\operatorname*{arg\,max}_{a\in\mathcal{A}}Q^{\pi_{k}}(x,a)$
at every time step $k$.
In practice, $Q^{\pi}$ and $\pi$ are often parameterized as neural networks and learned from data.
To increase the sample efficiency of the algorithm, it is common to store all collected interaction data independent of the collection policy in a replay buffer $\mathcal{D}=\{(x_{t},a_{t},r_{t},x_{t+1})_{t=0}^{T}\}$.
As the Q-value only depends on the policy via the policy evaluation at the next state, it is possible to estimate Q-values from past interaction data by minimizing the fitted Q-learning objective

$$ $\displaystyle\mathcal{L}\left(\hat{Q}\middle|{\color[rgb]{ 0,0.49609375,0.640625}\mathcal{D}},\pi\right)=\frac{1}{|\mathcal{D}|}\sum_{t=0 }^{T}\left|\hat{Q}({\color[rgb]{0,0.49609375,0.640625}x_{t},a_{t}})-\left[{ \color[rgb]{0,0.49609375,0.640625}r_{t}}+\gamma\hat{Q}\left({\color[rgb]{ 0,0.49609375,0.640625}x_{t+1}},{\color[rgb]{0.86328125,0.2734375,0.19921875}a^ {\prime}}\right)\right]_{\mathrm{sg}}\right|^{2}\quad\mathrm{with}~{}{\color[ rgb]{0.86328125,0.2734375,0.19921875}a^{\prime}}\sim\pi(\cdot|{\color[rgb]{ 0,0.49609375,0.640625}x_{t+1}})\enspace.$ (1) $$

Here $[\cdot]_{\mathrm{sg}}$ denotes the stop gradient operation introduced to avoid the double sampling bias and all data contained in the replay buffer is colored blue.
However, the Q value at the next state ${\color[rgb]{0,0.49609375,0.640625}x_{t+1}}$ is evaluated with an action ${\color[rgb]{0.86328125,0.2734375,0.19921875}a^{\prime}}$ that is *not* guaranteed to be in the replay memory, as the target policy can be different from the policy used to gather the sample.
This means that we require the Q value to generalize to potentially unseen actions.
We provide a visualization of this issue in Figure 1.

## 3 Investigating the root cause of unstable Q learning

Figure: Figure 1: A visualization of the core issue we investigate. Even if a replay buffer contains good coverage for two policies ($\pi_{\mathrm{old}}$ and $\pi_{\mathrm{new}}$) starting from $\rho=x_{0}$, this does not guarantee that it contains a transition for executing an action under the new policy on a state visited under the old. However, this state-action pair’s value estimate is used to update the value of state $x_{0}$ via Equation 1, without being grounded in an observed transition.
Refer to caption: x1.png

Minimizing Equation 1 finds the policy Q function over a replay buffer with sufficient coverage of all states and actions that this policy visits.
However, in most continuous control RL algorithms (Lillicrap et al., ; Haarnoja et al., ; Fujimoto et al., ), this update is interleaved with policy update steps .
The data in $\mathcal{D}$ then necessarily becomes *off-policy* as training progresses.

This means that the number of actor and critic optimization steps needs to be balanced with gathering new data.
Obtaining new on-policy data is vital to continually improve policy performance (Ostrovski et al., ), but performing more update steps before gathering new data ensures that the existing data has been used effectively to improve the policy.
The *replay ratio* (Fedus et al., ) or *update-to-data (UTD) ratio* (Nikishin et al., ), which governs the number of gradient steps per environment step, is therefore a vital hyperparameter.

Naively training with high UTD ratios can lead to collapse in off-policy deep RL (Nikishin et al., ).
We conjecture that one of the major causes of the instability of high UTD off-policy learning are wrong Q values on *unobserved actions*.
This is a well-known problem for off-policy TD learning (Baird, ; Tsitsiklis & Van Roy, ; Sutton et al., ; Ghosh & Bellemare, ).
To differentiate the problem from *overfitting* to the training distribution, we use the term *misgeneralization* to highlight the importance of the distribution shift in causing the issue.
Our experiments in Subsection 3.2 show that generalization to on-policy actions is more difficult than generalization to a validation dataset that follows the training distribution, and that higher UTDs exacerbate the issue.

### 3.1 Action distribution shift can cause off-policy Q value divergence

To highlight the role that on-policy actions play in stabilizing Q value learning, we show an analysis the stability of Q learning with linear features.
The core ideas follow Sutton et al. () and are also explored by Tsitsiklis & Van Roy (); Sutton ().
We assume that the Q function is parameterized with fixed features and weights as $Q(x,a)=\phi(x,a)^{\top}\theta$.
Let $X$ and $A$ be the sizes of the state and action space respectively. Let $P\in\mathbb{R}^{X\cdot A\times X}$ be the matrix of transition probabilities from state-action pairs to states.
A policy can then be expressed as a mapping $\Pi\in\mathbb{R}^{X\times X\cdot A}$ from states to the likelihood of taking each action.
$R\in\mathbb{R}^{X\cdot A}$ is the vector of rewards.
$D^{\pi}\in\mathbb{R}^{X\cdot A\times X\cdot A}$ is a matrix where the main diagonal contains the discounted state-action occupancies of $P^{\pi}$ starting from $\rho$.
If we assume access to a mixed replay buffer $\mathcal{D}=\bigcup\{D^{\pi_{1}},\dots,D^{\pi_{n}}\}$ gathered with different policies, the Q learning loss for a target policy $\Pi$ can be written as

$$ $\displaystyle L(\theta)=$ $\displaystyle\sum_{i=1}^{n}\left[D^{\pi_{i}}\left(\Phi^{\top}\theta-[R+\gamma P \Pi\Phi^{\top}\theta]_{\mathrm{sg}}\right))^{2}\right]\enspace.$ (2) $$

The stability of learning with this loss can be analyzed using the gradient flow

$$ $\displaystyle\dot{\theta}$ $\displaystyle=-2\Phi\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi\right)\Phi^{ \top}\theta+2\Phi\sum_{i=1}^{n}D^{\pi_{i}}R\enspace.$ (3) $$

This gradient flow is guaranteed to to be stable around a fixed point $\theta^{*}$ if the key matrix $\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi\right)$ is positive definite (Sutton, ).
Details and a proof of the following statement are provided in Appendix C.
We can decompose the key matrix and see that the positive definiteness depends on the difference in policy between the replay buffer and the target policy

$$ $\displaystyle\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi\right)={\color[rgb]{ 0,0.49609375,0.640625}\underbrace{\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi _{i}\right)}_{\text{positive definite}}}+\gamma{\color[rgb]{ 0.86328125,0.2734375,0.19921875}\underbrace{\sum_{i=1}D^{\pi_{i}}P(\Pi_{i}-\Pi )}_{\text{no guarantees}}}\enspace.$ (4) $$

In general, we can provide no guarantees for the second term outside of the on-policy case ($\Pi_{i}=\Pi)$ where it becomes 0 0.
The stability depends on the difference between the target policy and the data-collection policies.
If the target policy takes actions which are not well covered under the training policies, the remainder can be non positive definite.
This also matches the intuition that learning fails if we simply do not have sufficient evidence for the Q function of unobserved actions.

When using features, the eigenvalue conditions on the key matrix are only sufficient, not necessary, as the features can allow for sufficient generalization between observed and unobserved state-action pairs.
In deep RL, the features $\phi$ are updated alongside with the weights, making it hard to provide definitive mathematical statements on stability.
With good function approximation, we could hope that the learned value function generalizes correctly to unseen actions.
In the next section we investigate this for a non-trivial task from the DMC suite and highlight that, while the value function does not diverge irrecoverably, good generalization is not guaranteed either.

### 3.2 Empirical Q value estimation with off-policy data

Figure: Figure 2: Left: the train, validation, and on-policy validation error of the Q function at UTD 1. Right: the Q values and return curves of TD3 agents across different UTD 1, 8, and 16.
Refer to caption: x2.png

In environments with large state-action space, ensuring coverage is difficult.
To investigate whether learning is stable nonetheless, we train a model-free TD3 agent on the *dog walk* environment (Tunyasuvunakool et al., ).
The architecture is presented in Subsection 4.1, and is regularized to prevent catastrophic divergence (Hussing et al., ; Nauman et al., ) and uses clipped double Q learning (Fujimoto et al., ).
This means it uses the most common techniques which are designed to prevent misgeneralization and overestimation.

While training a TD3 agent (Fujimoto et al., ), we save transitions in a validation buffer with a 5% probability.
At regular intervals we compute the critic loss on this validation set.
In addition, we reset our simulator to each validation state and sample an action from the target policy.
We then simulate the ground truth on-policy transition and compute the loss over these.
This allows us to test how well our value function generalizes to target policy state-action pairs (as depicted in Figure 1).

The results are presented in Figure 2 and show a gap both between the train and validation sets, as well as the validation and the on-policy sets.
While we use the on-policy state-actions to update the Q value, these estimates are not actually consistent with the environment.
Furthermore, the Q value overestimation grows with increasing UTDs.
This phenomenon was previously discussed in the context of over-training on limited data (Hussing et al., ) .

The experiments show that the problem outlined in Subsection 3.1 is not merely a mathematical curiosity, but that Q value generalization to out-of-replay-distribution actions is difficult in practice, and becomes more difficult with increasing update ratios.
Even though full divergence is not observed as new data is continually added to the replay buffer, it takes a long time for the effects of severe early overestimation to dissipate.

### 3.3 Previous attempts to combat misgeneralization and overestimation

Prior strategies that deal with misgeneralization can be grouped into three major directions: architectural regularization to prevent divergence of the value function, pessimism or ensemble learning to combat overestimation, and networks resets to restart learning.
While all of these interventions help to some degree, they each either do not solve the problem in full or cause additional issues.

Architectural regularization   Architecture changes (Hussing et al., ; Nauman et al., ; ; Lyle et al., ) and auxiliary feature learning losses (Schwarzer et al., ; Zhao et al., ; Ni et al., ; Voelcker et al., ) are largely reliable interventions, and have shown to provide improvements without much drawbacks in prior work.
However, as Hussing et al. () and our experiment presented in Subsection 3.2 highlight, by themselves they can mitigate catastrophic overestimation and divergence, but do not guarantee proper generalization.

Pessimism and ensembles   To combat overestimation directly, the most prominent approach in continuous action spaces is Clipped Double Q Learning (Fujimoto et al., ).
Here, a Q value estimate is obtained from two independent estimates $\hat{Q}_{1}$ and $\hat{Q}_{2}$.
If the error of the two critic estimators is assumed to be independent noise on the true critic estimate
then using the minimum over both estimates is guaranteed to underestimate the true critic value in expectation.
However, in complex settings this assumption on the the error of the critic estimates may not hold.

Ensembles (Lan et al., ; Chen et al., ; Hiraoka et al., ; Farebrother et al., ) or online tuning of the rate of pessimism (Moskovitz et al., ) have been proposed to obtain tighter lower bounds on the Q value.
However, these strategies can be expensive as redundant models or hyperparameter tuning are needed.
As a simpler strategy, recent works have also employed clipping to obtain an upper bound of the Q function to prevent divergence (Fujimoto et al., ).

Resetting   Finally, network resets been shown to mitigate training problems (Nikishin et al., ; D’Oro et al., ; Schwarzer et al., ; Nauman et al., ) in high UTD regimes.
However, in cases where the agent fails to explore any useful parts of the state space within the reset interval, restarting the learning process will not improve performance (Hussing et al., ).
This makes tuning the resetting interval both important and potentially difficult and no tuning recipes have been presented.
Resetting is also a potentially hazardous strategy in real-world applications, where re-executing a random policy might be costly or infeasible due to safety constraints.
Finally, it heavily relies on the assumption that all past interaction data can be kept in the replay buffer.

All of these strategies are somewhat able to alleviate the problem of out-of-distribution value estimation, yet none of them directly address the issue at the root.
In the next chapter, we present an alternate approach that aims to directly regularize the action value estimates under the target policy.

## 4 Mitigation via model-generated synthetic data

As value functions misgeneralize due to lack of sufficient on-policy data, we propose to obtain synthetic data from a learned model instead.
However, model-based RL can also cause problems such as compounding world model errors and optimistic exploitation of errors in the learned model.
By using both real and model-generated data, we can trade-off these issues: on-policy data improves the value function and limits the impact of off-policy distribution shifts, while using only a limited number of model-generated samples prevents model errors from deteriorating the value estimates.

Our approach builds on the TD3 algorithm (Fujimoto et al., ) and uses an update ratio of 8 by default.
Our critic is updated with both model-based and real data following the DYNA framework (Sutton, ).
More precisely, we replace a small fraction $\alpha$ of samples $\{x,a,r,x^{\prime}\}$ in each batch with samples from a learned model $\hat{p}$ starting from the same state $\{x,\pi(x),\hat{r},\hat{x}^{\prime}\}$ with $\hat{r},\hat{x}^{\prime}\sim\hat{p}(\cdot|x,\pi(x))$. In our experiments, $\alpha$ is set to merely $5\%$.
We term this approach Model-Augmented Data for Temporal Difference learning (MAD-TD).

Model vs Q function generalization   We expect that a learned models will yield better generalization than the Q function for two reasons.
First, the policy is updated each step to find an action that maximizes the value function.
This means we are effectively conducting an adversarial search for overestimated values.
The model’s reward and state estimation error on the other hand are independent of this process.
We test the adversarial robustness of our model-augmented value functions in Subsection 5.3.
Second, our experiment shows that value functions primarily diverge at the beginning of training.
In these cases, coverage is low and on-policy state-action pairs are often not available.
Obtaining a slightly wrong, yet converging value estimate can then be more useful than a diverging one.
Even as more data is gathered, new policies might not revisit old states with a high likelihood.
Therefore even as training continues we expect the model data to provide some benefit.

### 4.1 Design choices and training setup

Our model is based on the successful TD-MPC2 model (Hansen et al., ) combined with the deterministic actor-critic algorithm TD3 (Fujimoto et al., ).
We aim to reduce the complexity of TD-MPC2 to the minimal necessary components to achieve strong learning in the DM Control suite, and thus forgo added exploration noise, SAC, ensembled critics, and longer model rollout for training or policy search. We outline several design choices here and refer to Appendix D for more detail.

Encoder:    Like TD-MPC2, we parameterize the state with a learned encoder $\phi:\mathcal{X}\rightarrow\mathcal{Z}$ with a SimNorm nonlinearity (Lavoie et al., ).
This transformation groups a latent vector into groups of $k$ entries and applies a softmax transformation over each group.
This bounds the norm of the features, which has been shown aid with stable training (Hussing et al., ; Nauman et al., ).

Critic representation and loss:   We use the HL-Gauss transformation to represent the Q function (Farebrother et al., ). The critic loss is the cross-entropy between the estimated Q function’s categorical representation and the bootstrapped TD estimate.
To stabilize learning, we initialize the critic network towards predicting 0 0 for all states.
We find that this stabilizes the initial update steps in which almost no data is available.

Model loss:   The world model predicts the next state latent representation and the observed reward from a given encoded state $\phi(x)$ and action $a$. The loss has three terms: the cross-entropy loss over the SimNorm representation of the encoded next state, the MSE between the reward predictions, and the cross-entropy between the next state critic estimate and the predicted state’s critic estimate.
This final term replaces the MuZero loss in TD-MPC2 with a simplified variant based on the IterVAML loss (Farahmand, ).
We provide the exact mathematical equations for the loss in Appendix D.

Training:   We train the architecture by interleaving one environment step with one round of updates with a varying number of gradient steps governed by the UTD parameter.
For each update step, a new mini-batch is sampled independently from a replay buffer of previously collected experience.
We found that varying the number of update steps only for the critic and actor while keeping the update ratio for the model and encoder updates at $1$ leads to significantly more stable learning.

Run-time policy improvement with MPC:   Following the approach outlined by Hansen et al. (), the learned model can also be used at planning time to obtain a better policy.
Using the model for MPC at planning time exploits the same benefit of models as the critic learning improvement: we obtain a model-corrected estimate of the value function and choose our policy accordingly.
As we only train our model for one step, we also conduct the MPC rollout for one step into the future.

## 5 Experimental evaluation

Figure: Figure 3: Return curves for the dog tasks with differing UTD values. Across all tested environments, the return increases or remains stable with increasing UTD when training with MAD-TD. Without model data, the performance decreases when increasing the UTD. MPC is turned off in these runs to cleanly evaluate the impact of model data on critic learning.
Refer to caption: x4.png

We conduct all of our experiments on the DeepMind Control suite (Tunyasuvunakool et al., ). Following Nauman et al. ()’s recommendations we focus our main comparisons and ablations on the two hardest settings, the *humanoid* and *dog* environments (which we will refer to as the *hard suite*).
Implementation details can be found in Appendix D.
Unless stated otherwise we evaluate MAD-TD with a UTD of 8 and use the same hyperparameters across all tasks.

Note that even though we refer to training MAD-TD without using model data for the critic as “model-free”, the algorithm still benefits from the model through feature learning which has proven to be a strong regularization technique in high UTD settings (Schwarzer et al., ).
All main result curves are aggregated across $10$ seeds per task and we plot mean and bootstrapped confidence intervals for the mean at the 95% certainty interval.
For aggregated plots, we use the library provided by Agarwal et al. ().
Results on more environments are presented in Appendix E.

### 5.1 Impact of using model-generated data

We first repeat the experiment presented in Subsection 3.2 and show the results in Figure 5.
Using model-based data closes the gap between on-policy and validation loss.
We also observe that the initial Q overestimation disappears, which is consistent across all hard environments (see Subsection E.1).
This provides evidence that we are indeed able to overcome the unseen action challenge.

Performance with and without model data at varying UTD ratios:

Figure: Figure 4: Performance comparison on the hard tasks for MAD-TD, BRO, and TD-MPC, with varying number of steps and action repeat settings. MAD-TD is on par with all baselines, has higher mean and IQM when trained for 2 million time steps and action repeat 2, and strongly outperforms TD-MPC2 and BRO at 1 million time steps with action repeat 2.
Refer to caption: x5.png

In Figure 3 we present the impact of using model-based data across different UTD ratios.
Humanoid results are found in Subsection E.2.
As is directly evident, across the dog tasks, we observe stagnating or deteriorating performance when increasing the update ratio, consistent with reports in prior work.

Figure: Figure 5: Average loss values with and without the model generated data (compare Figure 2) at UTD 1.
Refer to caption: x6.png

However, when using a small fixed amount of model generated data, this trend is reversed across all tested environments, with performance improving or at least remaining consistent.
We find that with model-based data, training is stable across a range of UTDs, even beyond those tested in recent high UTD work (Nauman et al., ).
We also note that we observe only limited benefits from increasing the UTD ratio when properly mitigating *misgeneralization*, except for the highly challenging dog run task, all other tasks perform on par with each other.

Comparison with baselines:
As our method combines model-free and model-based updates, we compare our method against both TD-MPC2 (Hansen et al., ), a strong model-based baseline, and BroNet (Nauman et al., ), a recent algorithm proposed for high UTD learning.
Since Nauman et al. () and Hansen et al. () trained with differing numbers of action repeats, and we found that the performance does not cleanly translate between these regimes, we present our method both with an action repeat value of 1 and 2.
Some hyperparameters are adapted to the AR=1 setting (compare Table 2).
The results are presented in aggregate in Figure 4, with per environment curves show in Subsection E.4.
We find that our method performs on par or above previous methods, and strikingly it is able to achieve higher returns faster than both TD-MPC2 and BRO.

Figure: Figure 6: Resetting evaluation of MAD-TD and BRO. Lighter color denotes performance with reset, and darker without. While MAD-TD’s performance only increases slightly when adding resetting, BRO is unable to achieve strong performance in and setting without resetting.
Refer to caption: x7.png

### 5.2 Performance and stability impact of resetting

Resetting comparison:   To investigate whether our technique benefits from more stable training, we set up a comparison in which we test the effects of resetting on our method. Figure 6 presents aggregate results comparing our approach and BRO, both with and without resetting.
Across all tasks we find that resetting barely improves MAD-TDs performance with the tested hyperparameters and update steps.
Benefits can only be observed on some seeds and can most likely be attributed to restarting the exploration process (Hussing et al., ).
However, the BRO algorithm is not able to achieve reliable performance without resets.
Overall, these results highlight that our model substantially improves the problems related to incorrect generalization of the value function, and that these are likely a major cause of the failure of high UTD learning in the DMC tasks.
Conjectured problems like the primacy bias effect (Nikishin et al., ) need to be carefully investigated as we do not find evidence that a primacy bias impacts MAD-TD’s performance in the DMC environments.
Our work of course does not preclude the existence of phenomena such as loss of stability in different environments, architectures, or training setups. More discussion on this can be found in Appendix B.

Figure: Figure 7: Average regret ($\downarrow$), mean over the hard suite. Lower regret corresponds to faster, more stable training. MAD-TD outperforms BRO.
Refer to caption: x8.png

Continued training:   To highlight the pitfalls of resets, we employ a common RL theory metric the per timestep average regret

$$ $\overline{\text{Reg}}(T)=\frac{1}{T}\sum\nolimits_{t=0}^{T-1}\left(\mathcal{R} ^{*}-\mathcal{R}_{t}\right)\enspace,$ $$

where $\mathcal{R}_{t}$ denotes the cumulative return in episode $t$ and $\mathcal{R}^{*}$ the optimal return. We use the maximum return any of the algorithms achieved $\hat{\mathcal{R}}^{*}$ as a lower bound on the optimal return $\mathcal{R}^{*}$. Regret quantifies how much better the algorithm could have performed throughout training. In other words, in situations where continued learning is crucial, such as many safety critical applications, regret might be a better measure of performance. It captures not only how good the final policy is, but also how well the algorithm adapts over time, and minimizes mistakes. We present a comparison of MAD-TD and the resetting-based BRO in Figure 7 using an action repeat of $1$. The results show, even though both algorithms are close in their final return, their training behavior differs vastly. MAD-TD has lower regret showcasing its strength in continued deployment.

### 5.3 Further experiments and ablations

Figure: Figure 8: Return curves for the dog tasks when using on-policy, random and no model-generated data. When generating model-based data with random actions, performance of MAD-TD drops close to the model-free baseline, highlighting the importance of *on-policy* actions.
Refer to caption: x9.png

To further test our approach, we present two additional experiments on the *hard suite*: changing the action selection for the model data generation, and reducing the model performance.
In addition, we investigate the impact of using model based data on the smoothness of the learned value function.

Off-policy action selection in the model:   To verify that the improvement in performance is due to the off-policy correction provided by the model, we repeat the *hard suite* experiments with a UTD of $8$ and $5\%$ model data, but we chose actions randomly from a uniform distribution across the action space.
The results are presented in Figure 8.
They highlight that random state-action pairs do not provide the necessary correction and the performance deteriorates to that of the model-free baseline.

##### Smaller model networks:

To study the effect of the modelling capacity on our method, we ablate the size of the latent model by reducing the network size across the hard suite.

Figure: Figure 9: Performance evaluation when reducing the model size of the latent model in MAD-TD. The performance predictably drops with decreasing hidden layer size, however only strongly decreasing the model size below 64 reduces the performance below that of the model-free ablation.
Refer to caption: x10.png

The results are presented in Figure 9.
We see that reducing the network size has an immediate and monotonic impact on the performance of our approach, suggesting that the model learning accuracy and prediction capacity is indeed vital for our approach to function well.
However, even with small models of 64 hidden units, we still see some benefits from training with the model predicted data.

Figure: Figure 10: Magnitude of the difference between $Q(x,\pi(x))$ and $Q(x,\tilde{a})$, where $\tilde{a}$ is an adversarial perturbation of $\pi(x)$. Across the training why see higher perturbation magnitude for the runs without model correction data.
Refer to caption: x11.png

Perturbation robustness of the model-corrected value function:   To motivate our method, we conjectured that one of the problems of training in actor-critic learning is that the actor is conducting a quasi adversarial search for overestimated values on the learned critic.(^1^11*Quasi* because the actor is not constrained to find an action close to the replay buffer sample.)
To substantiate this claim and provide additional insight into the benefits of our approach, we used the iterated projected gradient method Madry et al. () to estimate the smoothness of the learned value functions across training on the humanoid environments at a UTD of 1 with and without model data.

Results are presented in Figure 10.
Across the humanoid tasks we find that not using any model data leads to value functions with higher oscillations, either across the whole training run in humanoid_run, or in the middle of training like in stand and walk.

## 6 Conclusion

Our experiments allow us to conclude that wrong generalization of the value functions to unseen, on-policy actions is indeed a major challenge that prevents stable off-policy RL, both in theory and in practice.
Model-Augmented Data for Temporal Difference learning (MAD-TD) is able to leverage the learning abilities of latent self-prediction models to provide small, yet crucial amounts of on-policy transitions which help stabilize learning across the hardest DeepMind Control suite tasks.
With a relatively simple model architecture and learning algorithm, this method proves to be on par with, or even outperform other strong approaches, and does not rely on mechanisms such as value function ensembles or resetting which were previously conjectured to be necessary for stable learning in high UTD regimes.

Our work opens up exciting avenues for future work.
The issue of poor generalization in off-policy learning can likely be tackled with other approaches such as diffusion models (Lu et al., ) or better pretrained foundation models, and our presented experiments provide an important baseline for such work.
Furthermore, while we have purposefully kept our approach as simple as possible to validate our hypothesis, many ideas from the model-based RL community such as uncertainty quantification (Chua et al., ; Talvitie et al., ), multi-step corrections (Buckman et al., ; Hafner et al., ), or policy gradient estimation (Amos et al., ) can be combined with our approach.
Our insight that surprisingly little data is necessary to achieve strong correction can likely be leveraged in these other approaches as well to trade-off model errors and value function errors more carefully.
Finally, while we chose the data to roll out in our models at random, our insights can likely be combined with ideas from the area of DYNA search control (Pan et al., ; ) to select datapoints on which the correction has the most impact.

#### Acknowledgments

Omitted for submission

## Appendix A Limitations

The core limitation of our methodology relies in the assumption that a sufficiently strong environment model can indeed be learned online.
While a proof of feasibility exists for many interesting RL benchmarks in the forms of the Dreamer (Hafner et al., ) and TD-MPC2 (Hansen et al., ) lines of work among many others, for a completely novel environment a practitioner will still have to test if current model learning schemes are sufficient to achieve strong control policies.

Furthermore, we can only generate data from the states visited under a past policy.
There is still a difference between the state distribution of the replay buffer, and the target policy stationary distribution.
While this difference does not seem to lead to catastrophic failures in the DMC benchmarks, the distribution shift might be more problematic in other environments.

Finally, we observe an interesting failure cases of our idea: in some simple environments we surprisingly observe worse performance with our network architectures compared to the BRO baseline.
This issue is likely due to the fact that the TD-MPC2 architecture is tuned for learning in complex high-dimensional problems, which leaves it potentially over-parameterized on simple tasks.

While our work shows that reduced learning capacity due to plasticity does not seem to be the major contributor to learning problems in benchmarks like DMC, that does not exclude the possibility that related issues appear nonetheless after accounting for the off-policy value estimation problem.
We did not test increasing the reset ratio even further as other prior work has done, as we already observed no benefits from increasing the replay ratio from 8 to 16 in most of our experiments and performed on par or beyond previous baselines.
Issues in reinforcement learning are often entangled in a complex way, e.g. a failure in exploration can lead to stagnant data in the replay buffer which prevents a critic from further improving its estimates, leading to worse exploration and so on.

## Appendix B Extended related work

Beyond mitigating value function overestimation and unstable learning (see Subsection 3.3), other works have approached the difficulty of off-policy learning and high update ratios from other perspectives.
Here, we survey further related papers which do not provide direct background for this work, but are nonethless relevant as either alternative approaches or possible enhancements.

##### Other ways of incorporating off-policy data

Having access to more diverse data has been shown to be beneficial for reinforcement learning, when this data is carefully used to mitigate the problems resulting from off-policy training.
Ball et al. () show that a large offline replay buffer can be used to improve training by sampling online training batches both from online data and offline data, and labelling the offline transitions with a reward of 0.
Agarwal et al. () and Tirumala et al. () also highlight that previously collected replay buffers can be used to improve training performance on agents.
In this work, we focus on the online setting where we do not have access to a replay buffer of previously collected transitions.
These ideas however can easily be combined by e.g. training a model from an available larger offline data buffer.

Another approach to obtain additional data is the diffusion-based method proposed by (Lu et al., ).
In this work, the replay buffer data is augmented with additional samples obtained from a diffusion model that is trained on the replay buffer.
This proves to be a powerful regularization for both offline and online learning.
Due to the lack of open source code, we were unable to compare our method directly to this work.
They also do not provide results for the hardest DMC tasks.

In the online off-policy regime, Fujimoto et al. () recently proposed TD7, which incorporates similar architectural choices to MAD-TD.
They use a self-predictive encoder to learn good state representations, but concatenate them with the state and action representation provided by the environment to limit loss of information.
This design choice proved to be beneficial but would require learning a observation-space next-state prediction, which is difficult in practice, especially in high dimensional environments.
To address the policy distribution shift, TD7 does not update the actor at every timestep but instead collects several full trajectories with a fixed policy and then conducts update steps afterwards.
However, this interval still needs to be balanced as a hyperparameter.
TD7 was not evaluated on DMC, which is why we do not present a comparison.

##### Model-based reinforcement learning

As surveying model-based reinforcement learning is a rather sizeable tasks, we refer readers to the survey by (Moerland et al., ) for reference.
Decision-aware latent models such as the one Hansen et al. () and we use have been studied specifically in several different variants.
Silver et al. () proposes a latent model that is trained with TD learning, which provides the basis for the Schrittwieser et al. () algorithm.
The addition of a latent self-prediction loss was first proposed by Li et al. () to stabilize learning problems with the TD learning loss.
This interplay was further studied by Ni et al. () and Voelcker et al. () in recent works.

From a theoretical angle, decision-aware losses similar to those used in MuZero where first studied by Farahmand et al. () and Farahmand ().
Grimm et al. () and Grimm et al. () further study the loss landscape and minimizers of such losses, while Kastner et al. () studied the extension of the loss to distributional settings.

While previous works have called the stability of the VAML loss into question (Lovatto et al., ; Voelcker et al., ), we find that it is stable and performant when combined with the HL-Gauss representation Farebrother et al. () and an auxiliary BYOL style loss Grill et al. (); Li et al. ().
Compared to MuZero it is also significantly easier to implement.

A more thorough overview on the topic of decision-aware learning can be found by Wei et al. ().

##### Loss of plasticity

A phenomenon that was originally reported in continual learning is that tendency for neural network based agents to lose their ability to learn over time. This phenomenon has also been investigated in the realms of RL (Igl et al., ), as RL can effectively be thought of as a type of continual learning problem. Sometimes the phenomenon is referred to as plasticity loss (Lyle et al., ; Abbas et al., ).
As highlighted before, we do not find strong evidence for the primacy bias or loss of plasticity during our experiments on the DMC suite.

However, that does not imply that the phenomenon does not exist. In fact, we believe that resolving stability issues such as those presented in our paper will help us to better isolate other nuanced issues such as plasticity loss more clearly.
Previous studies have identified and combated plasticity loss using feature rank maximization (Kumar et al., ), regularization (Lyle et al., ), additional neural network copies (Nikishin et al., ), minimizing dormant neurons (Sokar et al., ; Xu et al., ), various neural network architecture changes (Lee et al., ), slow and fast network updates (Lee et al., ) or weight clipping (Elsayed et al., ).

It is unclear how many improvements obtained by these changes can be explained by divergence effects (Hussing et al., ) or stability issues such as those established in our work as there seems to be a non-zero overlap in techniques that combat either.
Nauman et al. () have argued that many RL training problems can be difficult to disentangle from the plasticity loss phenomenon.
An interesting direction of future work is to test for plasticity loss with well regularized off-policy value function learning, for instance by combining our method with separate solutions established for plasticity loss such as those from Lyle et al. ().

It is also not unlikely that the training dynamics of the state-based dense-reward tasks on the DMC suite are more benign than those found in Atari games.
Many works on plasticity loss have studied sparse image-based control tasks with pure Q learning approaches, such as DQN on the Atari benchmark (Sokar et al., ; Lee et al., ).
The problem may be more prevalent when replay buffers cannot be maintained in full and the RL setting becomes a true continual learning problem.

##### Other stability perspectives

Our work studies the stability of losses during training.
We highlight that forgoing resetting decreases regret as the executed policies are more stable in the sense that they are not reset at regular intervals.
We also highlight that model-generated data can somewhat improve the stability of policies against adversarial attacks.
However, there are other notions of *stability* that should be considered relevant and that are orthogonal to our work.
Here we will give a non extensive overview into the different directions that exist as a starting point for the reader. For instance from a theoretical perspective, stability can be formulated as differential privacy (Vietri et al., ) or algorithmic replicability to obtain identical policies (Eaton et al., ). From a theoretical as well as practical perspective, issues such as robustness to adversarsial attacks (Nilim & Ghaoui, ; Iyengar, ; Wiesemann et al., ; Pinto et al., ). Finally, from an empirical perspective robustness to hyperparameters (Ceron et al., ; Patterson et al., ) and attempts at variance reduction to get more reliable solutions (Anschel et al., ; Kuang et al., ) can be considered notions of stability.

## Appendix C Mathematical derivations

While the proof by Sutton () which we use as a basis discusses the stationary distribution of the Markov chain $P^{\pi}$, we define our loss in terms of a discounted state-action occupancy.
We therefore briefly prove an auxiliary result to extend the analysis to the case of discounted state occupancy probabilites.
Note that when we talk about positive-definiteness, we use a definition which applies to potentially non-symmetric matrices, and merely requires that $u^{\top}Xu\>0$ for all vectors $u$.

###### Proposition 1 .

Let $P$ be a stochastic matrix. Define the discounted state occupancy distribution $\mu$ of $P$ for some starting state distribution $\rho$ and some discount factor $\gamma\in[0,1)$ as

$$ $\mu^{\top}=(1-\gamma)\sum_{n=0}^{\infty}\gamma^{n}\rho^{\top}P^{n}.$ $$

Let $D$ be a diagonal matrix whose entries correspond to the discounted state occupancy distribution.
Then the matrix $D(I-\gamma P)$ is positive definite.

###### Proof.

First, note that

$$ $(1-\gamma)\rho^{\top}+\gamma\mu^{\top}P=\mu^{\top}$ $$

by the definition of $\mu$ and the properties of the infinite sum.
Therefore,

$$ $\mu^{\top}P=\frac{1}{\gamma}\left(\mu-(1-\gamma)\rho\right)\enspace.$ $$

Sutton () asserts that a matrix $A$ is positive definite iff $A+A^{\top}$ is positive definite.
Furthermore, if the diagonal entries of a symmetric matrix are positive and its off-diagonal entries are negative, then it suffices to show that the row and column sums of matrix are positive.

For

$$ $D(I-\gamma P)+(I-\gamma P^{\top})D^{\top}$ $$

the off-diagonal terms are clearly non positive as D is diagonal.
On the main diagonal, we have $2(\mu_{i}-\gamma p(i|i)\mu_{i})$ which is positive as $p(\mu_{i}|\mu_{i})\leq 1$.
It now suffices to show that the row and column sums of $D(I-\gamma P)$ are positive.
For the row sum, we can make use of the fact that $P$ is a stochastic matrix, so

$$ $D(I-\gamma P)\mathbf{1}=D(\mathbf{1}-\gamma\mathbf{1})\geq 1\enspace.$ $$

For the column sum, we make use of the fact that $\mathbf{1}D=\mu$.
Then

$$ $\mu(I-\gamma P)=\mu-\gamma\frac{1}{\gamma}\left(\mu-(1-\gamma)\rho\right)=(1- \gamma)\rho\geq 1\enspace.$ $$

As $\rho$ is a probability vector the final inequality holds for all $\gamma\in[0,1)$.

All conditions presented by Sutton () hold, and therefore we have $D(I-\gamma P)$ is positive definite.
∎

To derive the gradient flow stability conditions in Subsection 3.1, we first restate the loss function

$$ $\displaystyle L(\theta)=$ $\displaystyle\sum_{i=1}^{n}\left[D^{\pi_{i}}\left(\Phi^{\top}\theta-[R+\gamma P ^{\pi}\Phi^{\top}\theta]_{\mathrm{sg}}\right))^{2}\right].$ (5) $$

The stability of learning with this loss can be analyzed using the gradient flow (Sutton et al., ). To derive the gradient flow, we compute the gradient of the loss function with regard to the parameters $\theta$.
As the loss has a relatively simple quadratic form and the derivative is a linear transformation, it decomposes nicely as

$$ $\displaystyle\nabla_{\theta}L(\theta)$ $\displaystyle=2\Phi\sum_{i=1}^{n}D^{\pi_{i}}\left(\Phi^{\top}\theta-R-\gamma P \Pi\Phi^{\top}\theta\right)$ (6) $\displaystyle=2\Phi\sum_{i=1}^{n}D^{\pi_{i}}\left(\left(I-\gamma P\Pi\right) \Phi^{\top}\theta-R\right)$ (7) $\displaystyle=2\Phi\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P^{\pi}\right)\Phi^ {\top}\theta-2\Phi\sum_{i=1}^{n}D^{\pi_{i}}R\enspace.$ (8) $$

Using the equation for the gradient flow $\dot{\theta}=-\frac{\eta}{2}\nabla_{\theta}L(\theta)$ with learning rate $\frac{\eta}{2}$, we obtain

$$ $\displaystyle\dot{\theta}$ $\displaystyle=-\eta\Phi\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P^{\pi}\right) \Phi^{\top}\theta+\eta\Phi\sum_{i=1}^{n}D^{\pi_{i}}R\quad,$ (9) $$

This gradient flow is guaranteed to be stables (meaning it will not diverge around the stationary point $\theta^{*}$) if the key matrix $\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi\right)$ is positive definite (Sutton, ).

We can decompose our key matrix into the on-policy key matrix and a remainder easily

$$ $\displaystyle\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi\right)$ (10) $\displaystyle=$ $\displaystyle\sum_{i=1}^{n}D^{\pi_{i}}\left(I-\gamma P\Pi+\gamma P\Pi_{i}- \gamma P\Pi_{i}\right)$ (11) $\displaystyle=$ $\displaystyle{\color[rgb]{0,0.49609375,0.640625}\sum_{i=1}^{n}D^{\pi_{i}}\left (I-\gamma P\Pi_{i}\right)}+\gamma{\color[rgb]{0.86328125,0.2734375,0.19921875} \sum_{i=1}D^{\pi_{i}}P(\Pi_{i}-\Pi)}\enspace.$ (12) $$

The first group of summands are all positive definite, following 1.
As the sum of positive definite matrices is positive definite, the claim stands.

However, the second group has no such guarantees.
This highlights the role that the target policy action selection plays in the stability of Q learning.

###### Proposition 1 .

###### Proof.

## Appendix D Implementation

**Table 1: Network architecture for MAD-TD.**
| Encoder $\Phi$ | Dense Layer | Mish | in_size=$|\mathcal{X}|$, out_size=$512$ |
| --- | --- | --- | --- |
|  | Dense Layer | Simnorm(8) | out_size=$512$ |
|  | Dense Layer | Mish | in_size=512 + $|\mathcal{A}|$, out_size=$512$ |
| Latent Model $F$ | Dense Layer | Mish | out_size=$512$ |
|  | Dense Layer | Simnorm(8) | out_size=$512$ |
|  | Dense Layer | Mish | in_size=512 + $|\mathcal{A}|$, out_size=$512$ |
| Q head $\hat{Q}$ | Dense Layer | Mish | out_size=$512$ |
|  | Dense Layer | – | out_size=$1$ |
|  | Dense Layer | Mish | in_size=512, out_size=$512$ |
| Actor $\hat{\pi}$ | Dense Layer | Mish | out_size=$512$ |
|  | Dense Layer | tanh | out_size=$|\mathcal{A}|$ |

**Table 2: Hyperparameters. We adapted three parameters to the action repeat = 1 setting, as the magnitude of the reward changes.**
| Parameter |  |
| --- | --- |
| Initial steps (Random policy) | $5000$ |
| Batch size | $512$ |
| RL learning rate | $0.0003$ |
| Model learning rate | $0.0003$ |
| Encoder learning rate | $0.0001$ |
| Soft update $\tau$ | $0.995$ |
| Discount factor $\gamma$ | $0.99$ |
| Model forward prediction steps | $1$ ($4$ for AR=1) |
| Gradient clipping | $10.0$ |
| HL-Gauss vmin | $-150\cdot\mathrm{AR}$ |

Our experiments are implemented in the jax library to allow for easy parallelization of multiple experiments across seeds.
All networks follow the standard architecture from Hansen et al. () with two changes: instead of using an ensemble of critics, we opt for a single double critic pair.
We also do not use a stochastic policy, instead simply using a deterministic network with a tanh activation as used in Lillicrap et al. (); Fujimoto et al. ().
Full hyperparameters are presented in Table 2 and the architecture can be found in Table 1.
We use mish activation functions (Misra, ) and the adam optimizer to train our models (Kingma & Ba, ).

##### Loss functions:

As we use the HL-Gauss representation (Farebrother et al., ) for the critic, the loss is the cross-entropy between the estimated Q function’s categorical representation $Q_{\mathrm{rep}}$ and the bootstrapped TD estimate,

$$ $\mathcal{L}_{Q}=\sum_{i=1}^{m}{\mathrm{TD}(\hat{Q}_{\mathrm{rep}})}_{i}\log \hat{Q}_{\mathrm{rep}_{i}}\enspace,$ $$

where the indices $i$ denote the positions of the categorical vector representation used by HL-Gauss.
This is the same loss that is used for the two-hot encoding in Hansen et al. (), the only difference is the target encoding function.
For more details, see Farebrother et al. ().

We use a latent encoder $\phi:\mathcal{X}\times\mathcal{A}\rightarrow\mathcal{Z}$ that maps into the simnorm space, the space of n k-dimensional simplicies (Lavoie et al., ).
Writing $\hat{p}$ for the learned world model and $\hat{r},\hat{x}^{\prime}\sim\hat{p}(|x,a)$ for reward and next latent-state samples, the loss for our model and encoder is

$$ $\displaystyle\mathcal{L}_{\mathrm{model}}(x,a,r,x^{\prime})$ $\displaystyle=\mathcal{L}_{\mathrm{rew}}(x,a,r)+\mathcal{L}_{\mathrm{forward}} (x,a,x^{\prime})+\mathcal{L}_{Q}(x,a,r,x^{\prime})$ $\displaystyle\mathcal{L}_{\mathrm{rew}}(x,a,r)$ $\displaystyle=\left(r-\hat{r}\right)^{2}$ $\displaystyle\mathcal{L}_{\mathrm{forward}}$ $\displaystyle=-\sum_{i=1}^{n\cdot k}\phi(x^{\prime})_{i}\log\hat{x}^{\prime}_{ i}\enspace,$ $$

where the index $i$ is again element-wise across the simplex representation used for the latent state.
Note that we propagate the critic learning gradients into the encoder only for the real data and not the model generated one to prevent instability.

##### Baseline results

We took available results from Nauman et al. () and Hansen et al. () for all plots where possible, and used the official implementation of BRO to rerun the experiments without resetting and with differing action repeats.
Other hyperparameters were left as-is.

## Appendix E Further results

### E.1 Q value overestimation

Figure: Figure 11: Return curves and Q values with differing UTD values.
Refer to caption: x12.png

We plot the return curves and corresponding Q estimates for different UTD values and with and without model-generated data on the hard suite.
The results are presented in Figure 11.
As we see, across all tasks the model free variant strongly overestimates the Q values, especially in the beginning.

### E.2 Humanoid results

Figure: Figure 12: Return curves for the humanoid tasks when using on-policy (blue), random (green) and no model-generated data (orange). The observed performance impacts are comparable to the dog case.
Refer to caption: x13.png

For several experiments, we only showed the dog results from the main suite to avoid cluttering the main body of the paper.
The corresponding humanoid results are presented in Figure 11 and Figure 12, corresponding to Figure 3 and Figure 8 respectively.
As the plots highlight, the main insights transfer across the hard tasks.

### E.3 Different quantities of model data

We evaluate using more model data to update our value functions and provide the results in Figure 13.
We observe that the majority of gain is obtained when using limited amount of model data, and larger amounts only provided limited gains in some humanoid runs.

Figure: Figure 13: Return curves on the hard suite. We see that using substantially more data than 5% does not improve performance in a statistically significant way.
Refer to caption: x14.png

### E.4 MAD-TD, BRO, TD-MPC2 per env on the hard suite

Figure: Figure 14: Return curves with action repeat set to 2.
Refer to caption: x15.png

Figure: Figure 15: Return curves with action repeat set to 1.
Refer to caption: x16.png

We present the return curves for our method and the baselines per environment on the hard suite.
Figure 14 shows the results with action repeat 2 and Figure 15 for an action repeat value of 1.
Perhaps surprisingly, the results of the algorithms are not fully consistent across this regime.
Partially, this can be explained by the fact that our method and TD-MPC2 were first developed in the regime of action repeat 2, while BRO was only evaluated in the action repeat = 1 setting.
This suggests that the performance of each method depends in a non-trivial fashion on hyperparameter tuning.
Nonetheless, across both action repeat setting MAD-TD outperforms BRO without resetting consistently and only under-performs any previous algorithm significantly on the dog trot task in the action repeat 1 setting.

We conjecture that the remaining gap in performance seems to be most likely attributable to exploration and optimism.
While we focus on learning accurate value functions, Bro contains several components which are specifically designed to improve exploration.
Investigating the interplay and tension between exploration and accurate value function fitting is an important direction for future work.

Finally, Bro and TD-MPC2 are explicitly evaluated without their exploration bonuses in separate evaluation rollouts.
We however do not conduct such as separate evaluation as we do not add any additional exploration noise to our training.
When plotting training performance, the gap between MAD-TD and Bro further closes, suggesting an important trade-off between test time performance and training returns.

### E.5 Results across further DMC environments

We conducted more experiments on all DMC environments which were shown to benefit from the interventions in prior work (D’Oro et al., ; Nauman et al., ).

Figure: Figure 16: Return curves evaluating the impact of model-based data for critic learning and MPC. Overall, MPC and model-based critic learning both stabilize the learning process, as conjectured.
Refer to caption: x17.png

Figure: Figure 17: Return curves for the impact of resetting on MAD-TD with and without MPC. Without MPC, resetting can still improve the performance, but with MPC, we see no significant benefits from resetting across environments except pendulum. The hopper results highlight the importance (and danger) of the reset interval, as seemingly the reset algorithm is not able to recover “in time” to improve performance.
Refer to caption: x18.png

Figure: Figure 18: Comparison of MAD-TD and TD-MPC2 across more environments of the DMC suite. We observe gains compared to TD-MPC2 in the hard tasks, especially in terms of early learning performance, while TD-MPC2 has advantages on the pendulum_swingup and acrobot_swingup tasks. These seem to be exploration and stability issues for which the longer model rollouts of TD-MPC2 seem to help.
Refer to caption: x19.png

## Peer Reviews

confidence: 3
contribution: 2
initial_rating: 3
presentation: 3
questions: 1. Have you explored different values for the percentage of model generated data in MAD-TD?
2. How does the performance of MAD-TD vary with different levels of world model accuracy? Have you conducted any ablation studies on less accurate models to understand how sensitive MAD-TD is to the world model quality?
soundness: 2
strengths: 1. The paper tackles an interesting and critical problem setting in sample-efficient reinforcement learning (RL), which has significant implications for many real-world applications where data collection is costly or time-consuming.
2. The authors provide both empirical evidence and theoretical analysis to demonstrate the importance of addressing incorrect Q-value learning in off-policy RL with limited samples, making a compelling case for their proposed solution.
3. The proposed method, MAD-TD, demonstrates competitive performance compared to existing baselines on challenging DeepMind Control (DMC) tasks, showcasing its potential as a viable solution for improving sample efficiency in RL.
summary: This paper addresses the challenge of unstable training in off-policy reinforcement learning (RL) methods when the update-to-data ratio is high. The authors identify the root cause of this instability as the difficulty in learning accurate value functions from limited data. To mitigate this issue, they propose a novel approach called Model-Augmented Data for Temporal Difference learning (MAD-TD). MAD-TD leverages model-generated data to improve the accuracy of value functions on unobserved on-policy actions, thereby stabilizing training even at high update ratios. They empirically show that MAD-TD achieves competitive performance on tasks from the DeepMind control suite.
weaknesses: 1. Limited Baseline Comparison: The paper only compares the proposed method, MAD-TD, against two baselines (BRO and TD-MPC) and their variants, which may not be sufficient to demonstrate its performance comprehensively. A more extensive comparison with other state-of-the-art methods would strengthen the paper's claims.
2. Lack of Ablation Study: Although the authors outline critical design choices in Section 4.1, they do not conduct an ablation study to investigate the importance of these choices. This omission makes it difficult to understand the individual contributions of each design element to the overall performance of MAD-TD.

---

confidence: 2
contribution: 3
initial_rating: 8
presentation: 3
questions: 1. In the un-numbered equation on Page 3 (Please number equations!) a certain matrix is partitioned into a positive definite and a non-positive definite term. I was not quite sure what to make of the discussion below this. It of course, follows that if \gamma is shown sufficiently small, or that the first term strongly dominates the second. I would appreciate more discussion of this.

2. It wasn't clear to me what sort of guarantees we expect with the authors approach - it seemed to be more of a "it seemed to work most of the time"....except when it didn't. I didn't quite understand the insight about why in certain circumstances the approach didn't work. Is that related to the assumption that a sufficiently high fidelity model cannot always be learned? Some more insight would be useful.
soundness: 3
strengths: I found this an interesting paper and the authors approach seems to generate useful results. The paper is well written and the presentation allows the reader to understand the rationale behind most of the work.

The paper combines some mathematical and intuitive insight into the stability problem. It then uses this insight as motivation for their new approach MAD-TD which seems to show some success based on the experimental results. 

I thought the authors gave quite a balanced presentation of the strengths, but also possible weaknesses of their work, which is commendable.
summary: In this paper the authors consider the reasons for RL to become unstable in UTD situations and they suggest a fix for this which they call MAD-TD. Basically, the data is augmented by that generated from a model to enable stabilisation. The authors then illustrate the success of their results using experiments.
weaknesses: The main limitation of the paper is precisely that which the authors point out themeselves i.e. that the assumption that a sufficiently high fidelity of the model can be learned online is valid. This is necessary as the "augmented data" is generated from this model. However, the authors have been quite up-front with this and, despite this short-coming, their results seem to show success with their MAD-TD approach. 

Another shortfall of the techniques is that of course, the main "proof" of the results is via experiments, although some mathematical insight is given. In other words, the authors do not prove that their MAD-TD approach is able to prevent instability, they observe, through their experiments that it seems to do well.

---

confidence: 4
contribution: 3
initial_rating: 5
presentation: 3
questions: 1. Referring to Weakness 1.: What would the results be on the other environments or benchmarks, like Meta-World (especially environments like Stick-pull, coffee-push/pull, assembly, and others) or MyoSuit?
2. Is there any advantage to using AR=2 rather than AR=1? Could you extend Figure 5 with AR=1 and 2 million time steps?
soundness: 3
strengths: 1. Authors propose a method that works effectively with high UTD without resets.
2. The topic is relevant.
3. The paper is well-written.
summary: The authors address the instability problem in off-policy reinforcement learning with the Replay Ratio. Building on previous results showing that neural networks trained on a specific set of experiences with the UTD lead to various instabilities in training, they show that the current solutions mitigate some of the problems to a greater or lesser extent, e.g., feature normalization mitigates the overestimation problem. Still, it does not allow the model to generalize to unseen actions. Therefore, they propose using a world model to augment a small part of the experiences from the replay buffer with its help.
weaknesses: 1. The results look pretty good, but then there is a question about robustness. Authors sometimes refer to [1], but the same paper points out that different effects can be obtained in RL in different environments. I agree with the theses in the manuscript, but I think it is valuable to validate these strong results on other benchmarks so that scientists in the future will know to what extent this is a general solution. Please see question 1 for more details.
2. Minor: Equations are without numeration.
3. Lack of an analysis of the impact of the percentage of augmented samples on performance. Could you provide plots showing performance across a range of alpha values (e.g., 1%, 5%, 10%, 25%, 50%) on a subset of representative tasks? This would give a clearer picture of the method's sensitivity to this hyperparameter.


[1] Nauman, M., Bortkiewicz, M., Miłoś, P., Trzcinski, T., Ostaszewski, M., & Cygan, M. Overestimation, Overfitting, and Plasticity in Actor-Critic: the Bitter Lesson of Reinforcement Learning. In Forty-first International Conference on Machine Learning

---

confidence: 3
contribution: 2
initial_rating: 5
presentation: 2
questions: How will the model data proportion affect the result? I suggest the author provide some preliminary results or insights regarding the impact of model data proportion and conduct ablation studies, if necessary, to show the trade-off on this key parameter in your method. Besides, is it possible to develop a self-adaptive mechanism to control the model data proportion during the training to enhance the final performance? To my understanding, the model error could be large at the beginning of the training process. Therefore, using fewer model-generated data at first and gradually adding the proportion as the world model converges might be helpful.
soundness: 3
strengths: 1. The authors provide thorough theoretical analysis and experimental results, demonstrating the effectiveness of MAD-TD in stabilizing the training process and achieving strong performance.
2. The paper is well-organized and the key idea is clearly delivered. Open-source code is also provided for reproducibility.
summary: This paper mainly investigates the issue of overestimation of unseen on-policy actions and instability caused by a high update-to-data (UTD) ratio in off-policy reinforcement learning via theoretical analysis and experiments. To address this issue, the authors introduce a new method named Model-Augmented Data for Temporal Difference Learning (MAD-TD), which combines model-generated synthetic data with real data to enhance and stabilize the off-policy RL training process. The experiments conducted on the DeepMind Control benchmark demonstrate that MAD-TD outperforms other baselines and leads stable learning even in high UTD settings.
weaknesses: 1. Using an incorrect format (ICLR 2024) for submission.
2. Typos (Also, please add indexes to each equation to make it easy to be referred to):
- In the first equation about the definition of $L(\theta)$, chapter 3.1, there is an extra ")".
- In the last equation of chapter 3.1, $n$ is missing on the top of the $\Sigma$ colored in red.
3. There's not much originality in the key idea of simply combining the model-generated data with the real data to augment the training process, which has already been used in previous works [1] and other model-based reinforcement learning algorithms [2-3]. As the distribution shift problem has been widely researched in offline RL, it's not surprising that a similar problem will occur in off-policy RL. Although the authors provide theoretical analysis to show the instability brought by target policy action selection, more profound results are expected, such as how will the model error affect the performance and when to trust model-generated data in off-policy RL.
4. The algorithm is mainly based on TD-MPC2, while several components are changed. While the proposed MAD-TD is compared with the original TD-MPC2, the influence of these changes on performance has not been clearly explained or demonstrated.

[1] Lu, Cong, Philip Ball, Yee Whye Teh, and Jack Parker-Holder. "Synthetic experience replay." Advances in Neural Information Processing Systems 36 (2024).
[2] Sun, Yihao, Jiaji Zhang, Chengxing Jia, Haoxin Lin, Junyin Ye, and Yang Yu. "Model-Bellman inconsistency for model-based offline reinforcement learning." In International Conference on Machine Learning, pp. 33177-33194. PMLR, 2023.
[3] Rigter, Marc, Bruno Lacerda, and Nick Hawes. "Rambo-rl: Robust adversarial model-based offline reinforcement learning." Advances in neural information processing systems 35 (2022): 16082-16097.

Meta-review: This paper addresses the issue of overestimation in value estimates that preclude high update-to-data ratio. The paper proposes to use a world model and update the value using on-policy data generated by the world model to correct this overestimation. The effectiveness of their approach is shown on hard continuous control Dog tasks from the DM control suite.

The approach is intuitive, and the results are quite promising. The paper contributes significantly to understanding and addressing the challenges of high update-to-data ratios in RL. Likewise, I recommend accepting the paper.

Decision: Accept (Spotlight)