TD-M(PC)²: Improving Temporal-Difference MPC Through Policy Constraint

Plan-based model-based RL (TD-MPC family) integrates sampling-based MPC with a learned value prior. Existing practice relies on online planning to collect training data, so value learning is entirely off-policy — and we find that this incurs persistent value over-estimation in high-dimensional tasks. We trace the cause to a structural policy mismatch between the planner (exploration policy) and the nominal actor (exploitation policy). The fix is conservative: a soft-constrained policy update that mitigates out-of-distribution queries. The result is a minimalist drop-in modification of TD-MPC2 that improves performance by large margins, over 100 % on 61-DoF humanoid locomotion, with no additional computational overhead.

§ 01The Problem

Plan-based MBRL — PETS, LOOP, TD-MPC, TD-MPC2 — couples an MPC planner with a learned actor-critic value prior. The planner $\pi_H$ performs $H$-step lookahead in a learned latent dynamics model and uses the value $\hat V$ as a terminal cost. The actor $\pi$, in turn, is what the critic $\hat Q$ is evaluated against. This combination is highly sample-efficient on low-dimensional continuous control — but its scaling to high-dimensional control (humanoid, dexterous-hand) stalls, and the bottleneck is value learning, not the planner.

Because online planning is leveraged for exploration, the data in the replay buffer is generated entirely by the planner $\pi_H$. The critic $\hat Q$, however, is bootstrapped from the nominal actor $\pi$. Whenever $\pi$ disagrees with $\pi_H$ — and in high-dimensional control they disagree a lot — Bellman targets are evaluated at out-of-distribution actions: actions the planner does not visit. With function approximation, that extrapolation error compounds across policy iterations rather than self-correcting, producing the persistent value over-estimation we observe in practice (the off-policy "deadly triad").

§ 02An Intuition From First Principles

The mechanism and our fix fit in one composite picture. Along the action axis $a$, the upper plot shows $Q^{\pi}$ (truth, blue solid) and $\hat Q$ (learned, red dashed); the lower plot shows the action-density of the buffer policy $\mu$ and of the actor $\pi$. The red dot marks the Bellman target $\hat Q(s', \pi(s'))$ — the value the critic is asked to back up.

§ 03The Minimalist Approach

Our remedy is conservative: a soft, state-conditional KL between the actor $\pi$ and the behaviour policy $\mu$ stored in the replay buffer. In practice this reduces to a single log-likelihood term in the actor loss, evaluated under the stored planner Gaussian:

The boxed term is the only delta from TD-MPC2. No new networks, no separate behaviour-density model — we re-use the stored Gaussian planner $\mu \approx \pi_H$:

# vanilla TD-MPC2 actor loss
loss_pi = -( Q(s, a)/S_q - α * log_pi )

# TD-M(PC)² — one extra term
log_mu  = stored_planner.log_prob(a)
loss_pi = -( Q(s,a)/S_q + β * log_mu / S_q - α * log_pi )

§ 04Theory

The off-policy mismatch is not a quirk of training — it is structural. Three results together close the loop:

Policy mismatch is not noise to be averaged away — it is a structural feedback loop. Cut the loop by keeping $\pi$ close to the buffer, not by closing the planner / actor gap. — § 3.3

§ 05Results

Standard benchmarks · aggregate

Three suites: 14 HumanoidBench locomotion tasks (61-DoF, dexterous hands), all 23 DMControl tasks, and the 7 hardest high-dimensional DMControl tasks (Dog 36-DoF, Humanoid 21-DoF). TD-M(PC)² tops SAC, PPO, DreamerV3 and TD-MPC2 on every suite.

HumanoidBench Results

On the 61-DoF HumanoidBench suite — dexterous whole-body control with hand-conditioned manipulation — TD-M(PC)² beats every baseline on the large majority of tasks. The biggest gains are on the hardest: run, slide, stair, pole, where the baseline produces exaggerated motion driven by over-estimated values.

The mechanism · value-bias diagnostics

The simplest way to see why ours wins is the value-bias diagnostic. We compute $\mathbb{E}_{s_0}[\hat V(s_0) - V^\pi(s_0)]$ along training. The top two rows are baselines that diverge wildly; the bottom row is ours — bias collapses to near-zero on all four tasks.

The bias on the baseline grows steeply with action-dim, exactly as Thms 2 + 3 predict — and ours collapses each to near-zero:

Qualitative roll-outs · TD-MPC2 vs. ours

The diagnostics manifest as behaviour. The baseline acts on over-estimated values: motions are aggressive, exaggerated, and unstable. Our agent acts cautiously, recovering a stable, smooth gait. Three side-by-side roll-outs on HumanoidBench (61-DoF, dexterous hands):

§ 06Takeaways

1. Policy mismatch is the bottleneck in plan-based MBRL — it scales with action-dim and is hidden by low-dim benchmarks.
2. Constrain the actor, not the planner. A state-conditional soft KL kills OOD bootstrapping without collapsing exploration.
3. Simple wins. ~10 LoC on TD-MPC2, no new networks, no extra compute, no new hyper-parameters.
4. The fix is orthogonal to double-Q / TD3 clipping — those address variance; we address the missing data.

§ 07BibTeX

@article{lin2025td,
  title={TD-M (PC) $\^{} 2$: Improving Temporal Difference MPC Through Policy Constraint},
  author={Lin, Haotian and Wang, Pengcheng and Schneider, Jeff and Shi, Guanya},
  journal={arXiv preprint arXiv:2502.03550},
  year={2025}
}