REVIEW 4 major objections 7 minor 1 cited by
TD-M(PC)$^2$: Improving Temporal Difference MPC Through Policy Constraint
T0 review · 4 major / 7 minor · reviewed 2026-08-09 · deepseek-v4-flash
Pith's one-line read Adding a policy-regularization term that keeps the learned policy close to the MPC planner's behavior reduces persistent value overestimation and substantially improves TD-MPC2 on high-dimensional continuous control.
desk verdict A simple policy-constraint patch to TD-MPC2 that gets large empirical gains on high-dimensional control, but the theory is sloppier than the experiments and the HumanoidBench baseline needs a controlled re-run. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
What carries the argument
The central object is the policy-regularization term $\beta\log\mu(a|s)$ added to the maximum-entropy policy objective in Eq. 10, where $\mu$ is the behavior policy, a weighted mixture of past planner policies $\pi_{H,k}$. Because exact $\log\mu$ is unavailable, the implementation maximizes $\mathbb{E}_{\mu'\sim\{\mu\}}[\log\mu']$ over the planner actions stored in the replay buffer, treating that as a lower bound. This term is the Lagrangian relaxation of the KL-constrained policy improvement problem in Eq. 8, and it works by penalizing the learned policy for selecting actions the planner would not take, keeping value-function queries inside the support of the data distribution.
What would settle it
Compare TD-M(PC)2 against a version that uses the exact planner log-density $\log\mu(a|s)$ computed from the MPC Gaussian parameters instead of the stored-action surrogate. If performance and value error change materially, the lower-bound surrogate is doing the work; if they do not, the claimed mechanism is not the source of the gains.
Extended reading notes
Core claim
The paper claims that in the TD-MPC2 pipeline the data-generating policy is the H-step lookahead planner $\pi_H$, while value targets bootstrap the nominal policy $\pi$, and the resulting distributional mismatch makes value errors invisible to the training distribution. Theorems 3.1, 4.1, and 4.2 formalize how approximation error accumulates through policy iteration, widens the performance gap between $\pi_H$ and $\pi$, and forces the two policies to diverge, so the overestimated regions are never revisited and corrected. The remedy is distribution-constrained policy improvement: replace the unconstrained policy update with the KL-constrained problem of Eq. 8 and implement its Lagrangian form as the policy loss $L_\pi = -\mathbb{E}_{a\sim\pi}[Q(s,a) - \alpha\log\pi(a|s) + \beta\log\mu(a|s)]$, using the stored planner actions as a surrogate for the behavior policy $\mu$. Empirically this reduces value approximation error sharply on 4-DoF, 36-DoF, and 61-DoF tasks and yields large performance gains, especially on HumanoidBench locomotion.
Load-bearing premise
The load-bearing premise is that maximizing the log-likelihood of stored planner actions is a faithful lower bound on the true planner policy's density, so the extra term really constrains the learned policy toward in-distribution actions rather than merely imitating the planner.
Editorial extensions
If this is right
- Value estimates in TD-M(PC)2 track Monte Carlo returns far more closely than those of TD-MPC2, so the planner's terminal bootstrapping becomes more trustworthy in high-dimensional tasks.
- On the 14-task HumanoidBench locomotion suite, the method is reported to improve over TD-MPC2 by over 100% on average, with the largest gains on Run, Slide, and Pole.
- On DMControl's seven high-dimensional tasks, the method slightly outperforms the baseline overall, with clear gains on the three dog tasks and comparable performance on humanoid tasks.
- The modification costs no additional computation and no per-task hyperparameter tuning, so it can be applied to existing TD-MPC-style implementations with a few lines of code.
- Because the constraint is applied to the policy prior rather than the planner, online exploration is not restricted; conservatism enters only where values are bootstrapped.
Reading between the lines
- If the mismatch diagnosis generalizes, the same one-term constraint should transfer to other planning-plus-value algorithms, including diffusion-based planners, since the failure mode is structural rather than specific to TD-MPC2.
- The ablation showing near-parity of a pure behavior-cloning variant suggests that conservatism itself, not improved value learning, may be the active ingredient; a $\beta$-annealing schedule could separate imitation from value-based improvement.
- A direct confirmation would be to track the total-variation divergence between $\pi_H$ and $\pi$ over training; the paper predicts this divergence should shrink under TD-M(PC)2 while value error falls.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper claims that TD-MPC2-style algorithms suffer from persistent value overestimation caused by a structural mismatch between the planner-driven data-collection policy π_H and the learned nominal policy π. It proposes TD-M(PC)2, a minimalist modification that adds a policy regularization term to the actor loss so that the learned policy stays close to the behavior policy, thereby reducing out-of-distribution value queries. The paper supports this claim with theoretical bounds (Theorems 3.1, 4.1, 4.2), empirical demonstrations of overestimation and its mitigation (Figures 1, 5, 7), benchmark results on HumanoidBench and DMControl, and ablation studies. The headline empirical result is that TD-M(PC)2 improves over TD-MPC2 by over 100% on the HumanoidBench locomotion average.
Significance. If the stated mechanism is correct, the paper offers a simple, general, and computationally cheap improvement to a prominent MBRL pipeline, with particularly large gains in very high-dimensional control tasks. The theoretical analysis connecting policy mismatch to error accumulation is a valuable addition to the literature, and the paper is transparent about the looseness of its bounds. However, the link between the theoretical constrained policy-improvement formulation (Eq. 8) and the implemented loss (Eq. 10) is not presently coherent, and the ablation results suggest that the observed gains may be driven more by conservatism than by the claimed value-learning improvement. The empirical comparison against HumanoidBench baselines that were not re-run under the paper's own protocol also weakens the headline quantitative claim. These issues are fixable, but they are central to the paper's contribution as currently stated.
major comments (4)
- [Section 5.2, Equations (8), (10), (23)] The claimed relationship between the constrained policy improvement and the implemented loss is not correct as written. The Lagrangian of (8) with an entropy bonus αH(π) is E_{a∼π}[Q(s,a)] − (α+β) E_{a∼π}[log π(a|s)] + β E_{a∼π}[log μ(a|s)], whereas Eq. (10) has coefficient α on the log π term. With β=1.0 and α=1e−4 (Table 1), the missing β log π term is the dominant entropy contribution. Moreover, the sentence 'we can maximize E_{μ′∼{μ}}[log μ′] as the lower bound of log(μ)' is not mathematically meaningful: for a fixed μ, E_{μ′∼μ}[log μ′(a)] is a constant and is not a lower bound for log μ(a|s) at actions sampled from π. Equation (23) is coherent only if μ_t is interpreted as the planner density evaluated at π-sampled actions, but that interpretation conflicts with the stored-action wording in Algorithm 1 and with the behavior-cloning variant in Eq. (25). Because the paper's central claim is that the added term implements a policy constraint that reduces out-of-distribution value queries, the authors must state precisely what loss is implemented (ideally with the actual code), and either derive it from (8) or explicitly present it as a heuristic cross-entropy regularizer.
- [Section 6.2, Figure 3] The HumanoidBench comparison relies on author-provided TD-MPC2 scores from an external repository (humanoid-bench) rather than on runs under the paper's own protocol. The headline claim that TD-M(PC)2 'improves TD-MPC2 by over 100%' therefore depends on an uncontrolled baseline; differences in seeds, environment versions, or evaluation procedures could change the comparison. Please either re-run TD-MPC2 under identical conditions for all tasks or clearly restrict the claim to the externally reported numbers and report the baseline's seed count and protocol.
- [Figures 1, 5, 7 and Appendix C] The ground-truth value estimates are Monte Carlo averages over 100 episodes and are plotted without confidence intervals or per-seed traces. The claim of 'persistent value overestimation' is the paper's motivation, and the proposed reduction in overestimation is a central result; without error bars it is hard to assess whether the Hopper-Stand case (15% error) or the convergence behavior in Figure 7 is significant. Please add confidence intervals or per-seed curves, and state the number of episodes and seeds used for the true-value estimator.
- [Section 6.3, Equation (25)] The 'pure behavior cloning' variant is not behavior cloning in the standard sense. Equation (25) maximizes E_{s∼B} E_{a∼π(·|s)} log μ(a|s), i.e., it minimizes KL(π∥μ), while BC normally maximizes the log-likelihood of stored actions under π, E_{(s,a)∼B} log π(a|s). As written, the ablation shows that a conservative policy update (whether or not it uses the Q-function) is sufficient for most of the gain, but it does not isolate the mechanism of improved value learning through reduced OOD queries. Please correct the description of the variant and, if the mechanism claim is to be retained, add an ablation that holds the policy update fixed and measures value error (e.g., Figure 5) for the full method versus the BC variant.
minor comments (7)
- [References, [19]] Reference [19] is not Hafner et al.'s Dreamer paper; it is a database paper for emotion recognition. Please correct the citation for Dreamer.
- [Theorem 3.1, Eq. (4)] The phrase 'for any policy µ' after the displayed equation is dangling and does not connect to the theorem statement.
- [Figures 1 and 5] The x-axis label 'Training Iterations (x10)' is ambiguous; it should read 10^5 or 1e5 steps.
- [Algorithm 1] The line 'Initialize ... by pertaining on uniformly sampled data' contains a typo; it should be 'pretraining'.
- [Figure 8 caption] The caption 'Constrained policy update through A W ACMean' is missing a space, and 'ramdom' is a typo for 'random'.
- [Appendix B, Eq. (23) and (25)] The text calls this a 'TD3-BC style policy constraint,' but TD3-BC uses a squared-distance penalty on actions, not a log-density term; please clarify the analogy.
- [Appendix A.7, Theorem 4.1 proof] In the displayed chain of equalities, the term '+ γ^H E_{ˆτ_k}[V^{π_{k−1}}(s_H) + \hat V_{k−1}(s_H)]' should have a minus sign before the value function to be consistent with the preceding line.
Circularity Check
No significant circularity: the derivation is not forced by its inputs; the Section 5.2 Lagrangian mismatch is a correctness/fidelity issue, not a circular reduction.
full rationale
The paper's central claim is that SAC-style policy iteration with planner-generated data causes persistent value overestimation and that a policy-regularization term mitigates it. This is not circular: the overestimation is measured against Monte Carlo returns (Figure 1), the policy mismatch is defined independently as a KL/performance gap (Theorems 4.1 and 4.2), and the proposed loss (Eq. 10, implemented as Eq. 23) is a modification of a known TD3-BC-style constraint (cited to [6]), not a fitted parameter that is then reported as a prediction. The empirical validation is against external benchmarks (TD-MPC2, DreamerV3, SAC) with no task-specific tuning, so the claimed improvements are not constructed from the method's own outputs. There are no self-citations by the present authors, and Theorem 3.1 is explicitly adopted from LOOP [34], which is independent prior work. The most serious flaw is in Section 5.2: Eq. 10 is not the Lagrangian of Eq. 8 (the beta log pi term is absent), and the sentence 'we can maximize E_{mu'~{mu}}[log mu'] as the lower bound of log(mu)' is mathematically incoherent as written. This is a real implementation-fidelity and correctness concern, and the BC ablation in Figure 6 suggests the gains may come from imitation rather than the stated KL mechanism. However, an unproved or even incorrect surrogate does not make the argument circular: the method is fully specified, the value-learning improvement is tested empirically, and no conclusion is assumed by construction. The approximation weakens the theoretical link between Eq. 8 and the implemented loss, but it does not reduce the paper's claims to their inputs. Under the hard rule that circularity requires exhibiting a specific reduction of a claimed result to a fit or to a self-citation chain, no such step exists here.
Assumptions & free parameters
free parameters (2)
- Prior constraint coefficient β =
1.0
- Scale threshold S =
2.0
assumptions (4)
- domain assumption The behavior policy μ can be represented as a weighted sum of past H-step lookahead policies π_H (Section 5.1).
- ad hoc to paper Stored planner actions a_t are a valid empirical surrogate for the behavior policy μ in the constraint (Section 5.2).
- ad hoc to paper Maximizing E_{μ'∼{μ}}[log μ'] is a valid lower bound for log μ(a|s) (Section 5.2).
- standard math Prior bounds from LOOP (Theorem 1) and Bertsekas (Lemma 6.1) are correct and applicable (Appendix A).
Cite this review
Pith. "Pith review of TD-M(PC)$^2$: Improving Temporal Difference MPC Through Policy Constraint." pith.science (2026). https://pith.science/paper/LGMCAM66
@misc{pith2026250203550,
author = {Pith},
title = {Pith review of: TD-M(PC)$^2$: Improving Temporal Difference MPC Through Policy Constraint},
year = {2026},
howpublished = {\url{https://pith.science/paper/LGMCAM66}},
note = {Machine review of arXiv:2502.03550}
}
read the original abstract
Model-based reinforcement learning algorithms that combine model-based planning and learned value/policy prior have gained significant recognition for their high data efficiency and superior performance in continuous control. However, we discover that existing methods that rely on standard SAC-style policy iteration for value learning, directly using data generated by the planner, often result in \emph{persistent value overestimation}. Through theoretical analysis and experiments, we argue that this issue is deeply rooted in the structural policy mismatch between the data generation policy that is always bootstrapped by the planner and the learned policy prior. To mitigate such a mismatch in a minimalist way, we propose a policy regularization term reducing out-of-distribution (OOD) queries, thereby improving value learning. Our method involves minimum changes on top of existing frameworks and requires no additional computation. Extensive experiments demonstrate that the proposed approach improves performance over baselines such as TD-MPC2 by large margins, particularly in 61-DoF humanoid tasks. View qualitative results at https://darthutopian.github.io/tdmpc_square/.
Figures
Figures from the paper (5 more)
Forward citations
Cited by 1 Pith paper
-
DoublyAware: Dual Planning and Policy Awareness for Temporal Difference Learning in Humanoid Locomotion
DoublyAware combines conformal trajectory filtering with a group-relative policy constraint to improve sample efficiency of TD-MPC for simulated humanoid locomotion.
Reference graph
Works this paper leans on
-
[1]
A. Argenson and G. Dulac-Arnold. Model-based of- fline planning. arXiv preprint arXiv:2008.05556, 2020
arXiv 2008
- [2]
-
[3]
M. Bhardwaj, A. Handa, D. Fox, and B. Boots. In- formation theoretic model predictive q-learning. In Learning for Dynamics and Control, pages 840–850. PMLR, 2020
work page 2020
-
[4]
A. Chan, H. Silva, S. Lim, T. Kozuno, A. R. Mah- mood, and M. White. Greedification operators for policy optimization: Investigating forward and reverse kl divergences. Journal of Machine Learning Research, 23(253):1–79, 2022
work page 2022
-
[5]
K. Chua, R. Calandra, R. McAllister, and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31, 2018
2018
-
[6]
S. Fujimoto and S. S. Gu. A minimalist approach to offline reinforcement learning. Advances in neural in- formation processing systems, 34:20132–20145, 2021
work page 2021
-
[7]
S. Fujimoto, H. Hoof, and D. Meger. Addressing func- tion approximation error in actor-critic methods. In International conference on machine learning, pages 1587–1596. PMLR, 2018
work page 2018
-
[8]
S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052–2062. PMLR, 2019
work page 2019
Show all 41 references
-
[9]
D. Garg, J. Hejna, M. Geist, and S. Ermon. Extreme q-learning: Maxent rl without entropy. arXiv preprint arXiv:2301.02328, 2023
2023 arXiv
-
[10]
Haarnoja, A
T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep rein- forcement learning with a stochastic actor. In Inter- national conference on machine learning, pages 1861–
-
[11]
Hafner, T
D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555–2565. PMLR, 2019
2019
-
[12]
Hafner, T
D. Hafner, T. Lillicrap, M. Norouzi, and J. Ba. Mas- tering atari with discrete world models. arXiv preprint arXiv:2010.02193, 2020
2010 arXiv
-
[13]
Hafner, J
D. Hafner, J. Pasukonis, J. Ba, and T. Lillicrap. Mas- tering diverse domains through world models. arXiv preprint arXiv: 2301.04104, 2023
2023 arXiv
-
[14]
Hansen, H
N. Hansen, H. Su, and X. Wang. Td-mpc2: Scalable, robust world models for continuous control. arXiv preprint arXiv:2310.16828, 2023
2023 arXiv
-
[15]
Hansen, X
N. Hansen, X. Wang, and H. Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.04955, 2022
2022 arXiv
-
[16]
Hansen-Estruch, I
P. Hansen-Estruch, I. Kostrikov, M. Janner, J. G. Kuba, and S. Levine. Idql: Implicit q-learning as an actor- critic method with diffusion policies. arXiv preprint arXiv:2304.10573, 2023
2023 arXiv
-
[17]
Janner, J
M. Janner, J. Fu, M. Zhang, and S. Levine. When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019
2019
-
[18]
Kabzan, L
J. Kabzan, L. Hewing, A. Liniger, and M. N. Zeilinger. Learning-based model predictive control for autonomous racing. IEEE Robotics and Automa- tion Letters, 4(4):3363–3370, 2019
2019
-
[19]
Katsigiannis and N
S. Katsigiannis and N. Ramzan. Dreamer: A database for emotion recognition through eeg and ecg signals from wireless low-cost off-the-shelf devices. IEEE journal of biomedical and health informatics, 22(1):98– 107, 2017
2017
-
[20]
Kostrikov, A
I. Kostrikov, A. Nair, and S. Levine. Offline reinforce- ment learning with implicit q-learning. arXiv preprint arXiv:2110.06169, 2021
2021 arXiv
-
[21]
Kumar, J
A. Kumar, J. Fu, M. Soh, G. Tucker, and S. Levine. Sta- bilizing off-policy q-learning via bootstrapping error reduction. Advances in neural information processing systems, 32, 2019
2019
-
[22]
Kumar, A
A. Kumar, A. Zhou, G. Tucker, and S. Levine. Con- servative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1179–1191, 2020
2020
-
[23]
Levine, A
S. Levine, A. Kumar, G. Tucker, and J. Fu. Offline rein- forcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643, 2020
2005 arXiv
-
[24]
C. Li, A. Krause, and M. Hutter. Robotic world model: A neural network simulator for robust policy optimiza- tion in robotics. arXiv preprint arXiv:2501.10100 , 2025
2025
-
[25]
Littman and A
M. Littman and A. Moore. Reinforcement learning: A survey, journal of artificial intelligence research 4, 1996
1996
-
[26]
Lowrey, A
K. Lowrey, A. Rajeswaran, S. Kakade, E. Todorov, and I. Mordatch. Plan online, learn offline: Efficient learn- ing and exploration via model-based control. arXiv preprint arXiv:1811.01848, 2018. 9
2018 arXiv
-
[27]
Y . Lu, J. Fu, G. Tucker, X. Pan, E. Bronstein, R. Roelofs, B. Sapp, B. White, A. Faust, S. Whiteson, et al. Imitation is not enough: Robustifying imitation with reinforcement learning for challenging driving scenarios. In 2023 IEEE/RSJ International Confer- ence on Intelligen...
2023
-
[28]
R. Munos. Performance bounds in l p-norm for ap- proximate value iteration. SIAM journal on control and optimization, 46(2):541–561, 2007
2007
-
[29]
A. Nair, A. Gupta, M. Dalal, and S. Levine. Awac: Accelerating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.09359, 2020
2006 arXiv
-
[30]
Nakamoto, S
M. Nakamoto, S. Zhai, A. Singh, M. Sobol Mark, Y . Ma, C. Finn, A. Kumar, and S. Levine. Cal-ql: Calibrated offline rl pre-training for efficient online fine-tuning. Advances in Neural Information Process- ing Systems, 36, 2024
2024
-
[31]
X. B. Peng, A. Kumar, G. Zhang, and S. Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. arXiv preprint arXiv:1910.00177, 2019
1910 arXiv
-
[32]
Schulman
J. Schulman. Trust region policy optimization. arXiv preprint arXiv:1502.05477, 2015
2015 arXiv
-
[33]
Sferrazza, D.-M
C. Sferrazza, D.-M. Huang, X. Lin, Y . Lee, and P. Abbeel. Humanoidbench: Simulated humanoid benchmark for whole-body locomotion and manipula- tion. arXiv preprint arXiv:2403.10506, 2024
2024 arXiv
-
[34]
Sikchi, W
H. Sikchi, W. Zhou, and D. Held. Learning off-policy with online planning. InConference on Robot Learning, pages 1622–1633. PMLR, 2022
2022
-
[35]
S. P. Singh and R. C. Yee. An upper bound on the loss from approximate optimal-value functions. Machine Learning, 16:227–233, 1994
1994
-
[36]
R. S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4):160–163, 1991
1991
-
[37]
R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018
2018
-
[38]
Tassa, Y
Y . Tassa, Y . Doron, A. Muldal, T. Erez, Y . Li, D. d. L. Casas, D. Budden, A. Abdolmaleki, J. Merel, A. Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv:1801.00690, 2018
2018 arXiv
-
[39]
Thrun and A
S. Thrun and A. Schwartz. Issues in using function approximation for reinforcement learning. In Proceed- ings of the 1993 connectionist models summer school, pages 255–263. Psychology Press, 2014
1993
-
[40]
Williams, N
G. Williams, N. Wagener, B. Goldfain, P. Drews, J. M. Rehg, B. Boots, and E. A. Theodorou. Information theoretic mpc for model-based reinforcement learning. In 2017 IEEE international conference on robotics and automation (ICRA), pages 1714–1721. IEEE, 2017
2017
-
[41]
G. Zhou, S. Swaminathan, R. V . Raju, J. S. Guntupalli, W. Lehrach, J. Ortiz, A. Dedieu, M. L´azaro-Gredilla, and K. Murphy. Diffusion model predictive control. arXiv preprint arXiv:2410.05364, 2024. 10 A Theory and Discussion A.1 Useful Lemma Lemma A.1. [35] Suppose πk+1 is 1...
2024 arXiv
Reviewed August 9, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.