TD-MPC Agent: Data-Efficient Model-Based RL

• TD-MPC Agent is a model-based reinforcement learning method that leverages an implicit latent-space world model to integrate MPC planning with temporal difference value learning.
• It employs a decoder-free architecture with stochastic dynamics, reward prediction, and policy priors to mitigate policy mismatch and out-of-distribution estimation errors.
• TD-MPC algorithms have demonstrated scalability and improved sample efficiency in complex continuous control tasks across simulated and robotic high-dimensional environments.

TD-MPC Agent refers to a class of model-based reinforcement learning (MBRL) algorithms that integrate online trajectory optimization—model-predictive control (MPC)—in the latent space of a learned world model, with temporal difference (TD) value learning to produce highly data-efficient agents for continuous control. The TD-MPC framework and its modern descendants, such as TD-MPC2 and TD-M(PC)$^2$, operate by iteratively updating both a stochastic world model and a value-predictive policy in a learned embedding space, enabling fast planning and improved robustness to high-dimensional dynamics. Recent work emphasizes the importance of controlling policy mismatch and out-of-distribution errors in the value estimation process, introducing explicit policy constraints to regularize actor updates. TD-MPC algorithms have been empirically validated to scale to tens of millions of parameters and excel at both single-task and multi-task continuous-control domains, including high-DoF simulated or robotic environments (Hansen et al., 2023, Lin et al., 5 Feb 2025).

1. Architectural Foundations

At the core of TD-MPC is an implicit, decoder-free world model trained entirely in a learned latent space. This model consists of:

• State encoder $h_e$: maps observed state $s_t$ to latent vector $z_t = h_e(s_t) \in \mathbb{R}^n$.
• Latent dynamics $d_\psi$: predicts $z_{t+1} = d_\psi(z_t, a_t)$ for action $a_t$.
• Reward model $R_\theta$: outputs predicted reward $R_\theta(z_t, a_t)$.
• Action-value ensemble $Q_\phi$: approximates $Q^\pi(z_t, a_t)$, typically as an ensemble of $K$ categorical critics.
• Policy-prior network $\pi_\omega$: diagonal Gaussian $\pi_\omega(a|z_t)$ used in value bootstrapping and as a proposal in MPC planning.

The architecture relies on Simplicial Normalization (SimNorm) in latent layers to control norm explosion, LayerNorm+Mish activations for stability, and a shared latent space for state, transition, and reward/value functions (Hansen et al., 2023, Lin et al., 5 Feb 2025).

2. Model-Predictive Planning in Latent Space

At each environment interaction step, the agent solves a short-horizon trajectory optimization (typically $H=3$) entirely in latent space to select the next action: $$J(\tau) = \sum_{i=0}^{H-1} \gamma^i R_\theta(z_i, a_i) + \gamma^H V_\phi(z_H)$$ subject to $z_0 = z_t$, $z_{i+1} = d_\psi(z_i, a_i)$. This is achieved using a sample-based optimizer such as Model Predictive Path Integral (MPPI) or Cross-Entropy Method (CEM), where candidate action sequences are proposed (with a subset drawn from $\pi_\omega$), unrolled in the latent world model, and weighted by their return. Elite sequences are used to update the trajectory distribution, and the first action of the best plan is executed (Hansen et al., 2023, Lin et al., 5 Feb 2025).

Bootstrapping the terminal value with $$V_\phi(z_H) = \mathbb{E}_{a \sim \pi_\omega(\cdot|z_H)}[Q_\phi(z_H, a)]$$ ensures planning incorporates expected future value beyond the planning horizon.

3. Temporal Difference Value Learning

TD-MPC algorithms train their critic ensembles via TD bootstrapping using both real and model rollouts: $$\mathcal{T}^{\pi_\omega}Q(z, a) = R_\theta(z, a) + \gamma\,\mathbb{E}_{z' \sim d_\psi(z, a),\, a' \sim \pi_\omega(\cdot|z')}[Q_\phi(z', a')]$$ The supervised loss for $Q_\phi$ adopts categorical or quantile target regression—formulated as cross-entropy between predicted and bootstrapped target distributions. Critic targets are computed via ensemble averaging, often with the "min of two" trick to mitigate overestimation.

The world model optimizes a summed reconstruction (prediction) loss, reward loss, and value loss weighted by coefficients $c_d, c_r, c_q$, respectively (Hansen et al., 2023, Lin et al., 5 Feb 2025). Model and value losses are aggregated over $H$-step rollouts drawn from the replay buffer.

4. Policy Mismatch and KL-Regularized Actor Updates

A principal challenge in TD-MPC is policy mismatch: the planner ($\pi_H$) generating buffer experience may systematically differ from the policy prior ($\pi_\omega$) used in value learning. As a consequence, value targets are often evaluated on out-of-distribution (OOD) state-action pairs, leading to persistent value overestimation and error accumulation. Theoretical analysis shows the performance gap $$\delta_k = \|V^{\pi_{H,k}} - V^{\pi_\omega,k}\|_\infty$$ grows with both model errors and planner–prior divergence, unless sufficiently corrected (Lin et al., 5 Feb 2025).

TD-M(PC)$^2$ introduces KL-divergence regularization between the learned policy $\pi$ and the buffer’s mixture of planner policies $\mu$, enforcing

$$D_{KL}[\pi(\cdot|z)\;\|\;\mu(\cdot|z)] \le \epsilon$$

within policy improvement. The actor loss becomes

$$\mathcal{L}_\pi = -\mathbb{E}_{z\sim\mathcal{B},\,a\sim\pi} \Big[Q_\phi(z,a) - \alpha\log\pi(a|z) + \beta\log\mu(a|z)\Big]$$

where $\alpha$ controls entropy regularization and $\beta$ the KL penalty. This suppresses OOD action queries, stabilizing value learning, and can be introduced smoothly via curriculum on $\beta$ (Lin et al., 5 Feb 2025).

5. Training Protocol and Hyperparameterization

TD-M(PC)$^2$ is trained by alternating model/policy updates and online environment interaction:

• Data collection: At each step, encode observation, compute action via MPC planning (using $\pi_\omega$ as initialization), step in environment, store transition and planner action distribution in buffer.
• Model/value updates: Every $N$ steps, sample $H$-step batch, compute bootstrapped targets, minimize $\mathcal{L}_\text{model}$ over dynamics, reward, critic heads, and update target critics via Polyak averaging.
• Policy update: Minimize $\mathcal{L}_\pi$ including both entropy and KL penalty. Adam optimizer with learning rate $3\,\times\,10^{-4}$, batch size $256$, gradient norm clip $20$.

Canonical hyperparameters include: $H=3$, MPPI iterations $6$, $512$ samples ($64$ elites), replay buffer $10^6$, ensemble $K=5$ with $101$-bin categorical critics on $[-10,10]$, $\alpha \approx 1\text{e}{-4}$, $\beta$ ramped in as value quality improves (Lin et al., 5 Feb 2025). Architecturally, the encoder is a $2\times256$-unit MLP, and policy/value heads use $512$-unit Mish + LayerNorm networks.

6. Empirical Performance and Comparative Analysis

TD-M(PC)$^2$ demonstrates significant improvement over baseline TD-MPC2 and related methods across high-degree-of-freedom tasks, such as 61-DoF simulated humanoid control. The addition of the KL constraint term in actor updates eliminates persistent value overestimation, particularly in areas of state space only visited by the planner. Empirically, performance gains appear most pronounced in scenarios where extrapolation errors otherwise dominate, highlighting the benefit of explicit OOD penalty in actor optimization (Lin et al., 5 Feb 2025).

When extended further (e.g., TD-GRPC (Nguyen et al., 19 May 2025)), group-based policy constraints and trust-region KL penalties can enhance stability and performance in highly unstable or distribution-shift-prone settings, such as humanoid locomotion.

7. Related Methods and Evolution

The TD-MPC family has evolved through:

• Original TD-MPC: trajectory optimization in latent space using an implicit world model, with value learning via TD backups (Hansen et al., 2023, Kuzmenko et al., 2 Jul 2025).
• TD-MPC2: architectural improvements (SimNorm, larger ensembles, robust categorical value heads), large-scale multitask capabilities (Hansen et al., 2023).
• Policy-constrained variants: TD-M(PC)$^2$ (minimal KL actor regularization to address policy mismatch) (Lin et al., 5 Feb 2025), and TD-GRPC (group-relative weighted actor and strict trust regions) (Nguyen et al., 19 May 2025).
• Extensions to hierarchical RL (IQL-TD-MPC) and knowledge distillation (TD-MPC-Opt) (Chitnis et al., 2023, Kuzmenko et al., 2 Jul 2025).

A plausible implication is that the KL-regularized actor update introduced in TD-M(PC)$^2$ is now regarded as a required component to achieve consistent value estimation and sample efficiency in off-policy, model-based MPC agents operating at scale. This regularization aligns the policy learning distribution with the data-collection policy, bounding extrapolation errors, and providing improved learning stability and asymptotic performance.