---
title: 'MAD-TD: Stabilizing Off-Policy TD Learning'
url: https://www.emergentmind.com/topics/model-augmented-data-for-temporal-difference-learning-mad-td
type: topic
---

# MAD-TD: Stabilizing Off-Policy TD Learning

Model-Augmented Data for Temporal Difference Learning (MAD-TD) is a reinforcement learning (RL) methodology designed to stabilize off-policy temporal-difference (TD) learning under high update-to-data (UTD) regimes by blending real and model-generated transitions. It addresses instability and Q-function overestimation issues that arise when function approximators are updated many times per environment sample—scenarios common in deep RL with replay buffers and modern control tasks. MAD-TD integrates a learned one-step world model to generate on-policy synthetic rollouts, mixing these with environment transitions during value updates to close the generalization gap to unseen actions and eliminate the need for parameter resets or ensembles [2410.08896].

## 1. Motivation and Problem Formulation

The core problem addressed by MAD-TD is the instability observed in high-UTD off-policy TD learning: updating neural value function approximators multiple times per collected transition can lead to misgeneralization, Q-function overestimation, and catastrophic divergence. In off-policy deep RL algorithms such as DDPG, TD3, and SAC, the value function is trained using transitions $(s, a, r, s')$ from a replay buffer $D_{\mathrm{env}}$. The UTD ratio $K$ denotes the number of critic gradient updates per new environment step. While higher $K$ can enhance sample efficiency, it creates a scenario where the target policy $\pi$ evolves much faster than the data in the buffer, resulting in TD updates computed on $(s', a')$ pairs (with $a' \sim \pi(\cdot | s')$) that may not be present in collected data. This breaks the “coverage” necessary for stable value learning, leading to overestimation and critical instability [2410.08896].

## 2. World Model Construction and Supervision

To counter the generalization gap, MAD-TD employs a one-step world model $p_\phi(s', r | s, a)$. Raw states $s$ are encoded via a neural embedding $\phi(s) \in \mathbb{R}^d$ with a SimNorm nonlinearity. The reward $\hat{r}$ is predicted by a shallow MLP $Rw(\phi(s), a)$, while the next latent state $\hat{s}'$ is predicted by a categorical/softmax output head $p_\phi(\cdot | \phi(s), a)$. Model training minimizes the negative log-likelihood of observed transitions from $D_{\mathrm{env}}$:
\[
L_{\mathrm{model}}(\phi) = \mathbb{E}_{(s, a, r, s') \sim D_{\mathrm{env}}} [ -\log p_\phi(s', r | s, a) ]
\]
which decomposes into reward MSE and state prediction cross-entropy. An additional VAML-style penalty aligns the value predicted at $\phi(s')$ with model rollouts, further regularizing the learned environment dynamics [2410.08896].

## 3. Model-Augmented Data Mixing in TD Learning

For each update, MAD-TD generates a batch by sampling $(1 - \alpha)N$ real transitions from $D_{\mathrm{env}}$ and $\alpha N$ synthetic transitions from the model ($\alpha \approx 5\%$ in practice). Model rollouts are constructed by sampling states $s$ (typically from $B_{\mathrm{env}}$ or uniformly from $D_{\mathrm{env}}$), computing $a_{\mathrm{model}} = \pi(s)$, and sampling $(\hat{r}, \hat{s}') \sim p_\phi(\cdot | s, a_{\mathrm{model}})$. The critic (value function) is trained to minimize TD error over the combined batch $D = D_{\mathrm{env}} \cup D_{\mathrm{model}}$:
\[
L_{\mathrm{TD}}(\theta) = \mathbb{E}_{(s, a, r, s') \sim D}\left[ (r + \gamma V_\theta(s') - V_\theta(s))^2 \right]
\]
or, for Q-functions, using the standard clipped double-Q target. This Dyna-style augmentation allows the critic to see transitions for current on-policy actions that have not yet been observed from the real environment, closing the coverage gap and improving stability [2410.08896].

## 4. Algorithmic Workflow and Pseudocode

The MAD-TD workflow is as follows:

- Initialize replay buffer $D_{\mathrm{env}}$, world model $p_\phi$, value network $V_\theta$, policy network $\pi_\psi$, target networks.
- For each environment time step:
  - Collect an (s, a, r, s') transition and append to $D_{\mathrm{env}}$.
  - For $K$ critic updates (where $K$ is the UTD ratio):
    - Sample $N(1 - \alpha)$ transitions from $D_{\mathrm{env}}$.
    - For $\alpha N$ model transitions: for $s$ in $S_{\mathrm{batch}}$, compute $a_{\mathrm{model}} = \pi_\psi(s)$; sample $(\hat{r}, \hat{s}') \sim p_\phi(\cdot | s, a_{\mathrm{model}})$.
    - Form batch $B = B_{\mathrm{env}} \cup B_{\mathrm{model}}$.
    - For each $(s, a, r, s') \in B$, compute $\delta(s, a) = r + \gamma V_\theta^-(s') - V_\theta(s)$ (or Q-value variant), minimize $L_{\mathrm{TD}}(\theta)$.
    - Update the actor via policy gradient on $V_\theta$.
    - Perform model update by minimizing $L_{\mathrm{model}}(\phi)$.
    - Execute soft target network updates as required [2410.08896].

This regime maintains a high UTD ratio while stabilizing value learning.

## 5. Theoretical Underpinnings of Stability

The necessity for coverage of on-policy state-action pairs is rooted in classical TD learning theory. Given a parametric value function $V_\theta$, TD learning converges stably only if the data distribution $D_{\mathrm{data}}$ matches the occupancy measure of the target policy. The “key matrix” 
\[
M = \sum_i D^{\pi_i}(I - \gamma P \Pi)
\]
is positive definite when data policies $\pi_i$ coincide with the target, but may acquire negative eigenvalues as policy drift increases, precipitating divergence. The “remainder” term
\[
\gamma \sum_i D^{\pi_i} P (\Pi_i - \Pi) 
\]
is zero only under perfect coverage; otherwise, unseen actions break positivity. A single on-policy occupancy measure $D^\pi$ ensures positive definiteness and stable TD. By supplementing each batch with even a small amount of on-policy, model-generated data, MAD-TD restores the requisite positive definiteness, directly repairing the main cause of instability at high UTD [2410.08896].

## 6. Empirical Results and Analysis

Experiments on challenging DeepMind Control Suite tasks—specifically, dog (walk, trot, run) and humanoid (stand, walk, run) environments—demonstrate several phenomena:

- As UTD increases, standard TD3 suffers from markedly increased Q-overestimation, especially for unseen on-policy actions; by UTD=8–16, this causes critical instability.
- Adding as little as $5\%$ model-generated on-policy data reduces on-policy TD error to the held-out validation error level, eliminates initial Q-overshoot, and preserves stability up to UTD=16 and beyond.
- Baselines such as BRONet (which requires periodic network resets to avoid collapse at high UTD) fail catastrophically without resets, whereas MAD-TD achieves equivalent or superior performance continuously, removing the need for resets.
- With action repeat 2 and 2M steps, MAD-TD slightly outperforms both BRONet and TD-MPC2 on mean and interquartile-mean metrics; with 1M steps it outperforms both.
- Ablation studies show that using random, rather than on-policy, actions in the model eliminates the stability gains; shrinking the model below 64 hidden units degrades performance below the model-free baseline, confirming the importance of model accuracy.
- The Q-surface produced by MAD-TD is smoother and more robust to small perturbations of $\pi(s)$, indicating mitigated value exploitation by the actor [2410.08896].

## 7. Connections to Related Approaches

MAD-TD is a modern deep RL analog of Dyna-style methods, combining model-based and model-free updates by using synthetic rollouts to assist TD learning. The method is related to classical Dyna with imagined transitions, experience replay, and successor representation–based backups. Unlike prior works, MAD-TD focuses on the stabilization provided by even a small proportion of on-policy model data in high-UTD regimes. Notably, it avoids the need for network resets, critic ensembles, or complex regularizers. The empirical findings validate the hypothesis that the requisite coverage of on-policy actions—a problem also noted in the analysis of classic off-policy TD divergence—is the central factor behind instability under function approximation, and that this coverage can be restored efficiently through simple model-augmented data mixtures [2410.08896][1902.02907].

---

MAD-TD thus establishes a simple, theoretically sound, and empirically validated paradigm for leveraging model-generated on-policy data to stabilize deep TD learning in the presence of high update-to-data ratios, closing the value generalization gap and enabling robust, efficient learning in high-dimensional continuous-control environments.

Source: https://www.emergentmind.com/topics/model-augmented-data-for-temporal-difference-learning-mad-td