Summary: TD-MPC (Hansen et al., 2022) combines temporal difference learning with Model Predictive Control (MPC) in latent space. It learns an ensemble of latent dynamics models and plans with MPPI (Model Predictive Path Integral) — achieving strong performance on continuous control with per-domain hyperparameter tuning.


Core Architecture

┌─────────────────────────────────────────────────────────────┐
│                    TD-MPC AGENT                             │
├─────────────────────────────────────────────────────────────┤
│  Encoder:         x → z (CNN for vision, MLP for states)   │
│  Dynamics Ensemble: {f_θ₁, f_θ₂, ..., f_θ₅}                │
│                   z' = f_θ(z, a)  (deterministic per head)  │
│  Reward Model:    r̂ = R_φ(z, a)                              │
│  Value Model:     V_ψ(z)  (TD target)                       │
│  Policy:          π_φ(z)  (actor)                           │
└─────────────────────────────────────────────────────────────┘
                              ↓
                    MPPI PLANNING (online)
                    - Sample K action sequences
                    - Rollout with ensemble
                    - Score by Σ γ^t r̂ + γ^H V(z_H)
                    - Select best, update belief
                    - Return first action

Key Components

Component Design
Latent Space 50-100 dim (deterministic)
Dynamics 5-head ensemble MLP (2-layer, 512 hidden)
Planning MPPI with 512 samples, 5-10 iterations
Value TD target: r + gamma V(z') (detached)
Actor Squashed Gaussian, updated via DPG on imagined trajectories

TD-MPC vs DreamerV3

Aspect TD-MPC DreamerV3
Planning Online MPPI (test-time) Offline actor (amortized)
Dynamics Deterministic ensemble Stochastic RSSM
Hyperparameters Per-domain tuning Fixed (breakthrough)
Asymptotic Perf Higher (online search) Slightly lower
Sample Efficiency Good Excellent
Vision Works (CNN encoder) Native (CNN encoder)
Implementation Complex (MPPI + ensemble) Cleaner (actor-critic)

MPPI Planning Details

Model Predictive Path Integral (MPPI):

  1. Sample K action sequences from N(mu, Sigma)
  2. Rollout each through ensemble (mean dynamics)
  3. Compute cost: J(tau) = sum_t gamma^t c(z_t, a_t) + gamma^H V(z_H)
  4. Weight samples: w_i ∝ exp(-(1)/(lambda) J(tau_i))
  5. Update mean: mu ≥ ts sum w_i tau_i
  6. Iterate 5-10 times
  7. Execute a_0 = mu_0

Ensemble helps: Epistemic uncertainty → natural exploration.


When to Use TD-MPC

Scenario Recommendation
Max asymptotic performance TD-MPC (with tuning)
Limited tuning budget DreamerV3
Vision + proprioception Both work
Sparse rewards DreamerV3 (better credit assignment)
Real robot (sample budget) DreamerV3 / Diffusion Policy
Research: understand planning TD-MPC (explicit planning)

TD-MPC2 (2024) Improvements

Improvement Effect
Simpler architecture Removed policy network, pure MPPI
Better value target Uses V directly (no actor-critic)
Faster planning Fewer MPPI iterations needed
State-based Matches DreamerV3 on state tasks

Related Concepts


Sources

  • Hansen et al., "TD-MPC: Temporal Difference Learning for Model Predictive Control" (ICLR 2023)
  • Hansen et al., "TD-MPC2: Scalable, Robust World Models" (2024)
  • model-based-rl — Comparison table