REVIEW 4 major objections 6 minor 36 references
TD-MPC-Opt: Distilling Model-Based Multi-Task Reinforcement Learning Agents
T0 review · 4 major / 6 minor · reviewed 2026-08-06 · deepseek-v4-flash
Pith's one-line read Distilled 1M-parameter RL agent scores 28.45 on MT30, up from 18.93
desk verdict The +50.2% headline is an apples-to-oranges comparison; the same-budget distillation gain is small and unquantified, but the empirical sweep is worth a careful revision. 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 load-bearing mechanism is the reward distillation loss $L_{\mathrm{distill}} = \mathrm{MSE}(R_{\mathrm{teacher}}(s,a), R_{\mathrm{student}}(s,a))$, added to the original TD-MPC2 loss as $L_{\mathrm{total}} = L_{\mathrm{orig}} + d_{\mathrm{coef}} \cdot L_{\mathrm{distill}}$. The teacher is a frozen 317M-parameter TD-MPC2 checkpoint from [12]; the student is the 1M-parameter TD-MPC2 backbone. The distillation coefficient $d_{\mathrm{coef}}$ (optimal around 0.4) controls how much the teacher's reward knowledge shapes the student, and the paper combines this with a batch size of 256 and extended 1M-step training. FP16 post-training quantization halves the model size to 3.9 MiB while slightly improving the score, making the mechanism a full pipeline from large teacher to deployable compact agent.
What would settle it
Train the distilled 1M/256/1M student and the from-scratch 1M/256 baseline across at least five seeds each and compare the distributions of MT30 normalized scores; if the confidence intervals overlap substantially or the distilled mean does not exceed the from-scratch mean, the claimed distillation benefit and the 28.45 state-of-the-art are not established.
Extended reading notes
Core claim
The central discovery is that matching only the reward predictions of a frozen 317M-parameter TD-MPC2 teacher is enough to transfer a large part of its multi-task competence to a 1M-parameter TD-MPC2 student. The student is trained with the original TD-MPC2 objective (consistency, reward, and value losses) plus an MSE penalty between teacher and student reward outputs, scaled by a coefficient d_coef, with values near 0.4 performing best. After 1M steps with batch size 256, the distilled student scores 28.12; applying FP16 post-training quantization pushes the normalized MT30 score to 28.45, which the paper reports as the new state of the art for this benchmark and a +50.2% improvement over the previously reported 18.93 for the 1M model trained from scratch. The paper also reports that a 200K-step distilled run with batch size 256 reaches 17.85 versus 14.04 from scratch, that a 1024-batch run does not benefit from distillation, and that adding next-state latent matching via linear projection or PCA degrades scores to 7.69 or 8.78, respectively.
Load-bearing premise
The single-run scores for the distilled and from-scratch 1M-parameter models are treated as reliable point estimates; if run-to-run variance on MT30 is as large as a few points, the claimed edge of distillation and the 28.45 state-of-the-art could disappear.
Editorial extensions
If this is right
- A 1M-parameter world model can match or exceed the MT30 performance of a 317M-parameter teacher, so multi-task continuous control competence is highly compressible when distillation uses reward predictions.
- Reward-only distillation is sufficient for strong transfer on MT30; attempts to also match next-state latent representations currently hurt performance because of the teacher-student dimension gap.
- FP16 post-training quantization can cut model size by about half without degrading (and in this study slightly improving) normalized score, making FP16 a safe default compression step for these agents.
- Training with batch size 256 outperforms batch size 1024 for distillation, suggesting that frequent weight updates matter more than large batches in this low-resource regime.
- Longer distillation (1M steps) lets the student overtake a from-scratch baseline trained under identical conditions, whereas short distillation (200K steps) only helps at small batch sizes.
Reading between the lines
- Editorial inference: The reported 0.76-point gap between the distilled 1M/256 model (28.12) and the from-scratch 1M/256 baseline (27.36) is based on single runs; with typical MT30 run-to-run variance, this difference may not be robust, so the most convincing evidence for distillation is the 200K-step regime, where the distilled model beats from scratch by 3.81 points.
- Editorial inference: The surprising FP16 gain (28.12 to 28.45) could be a seed effect rather than a real benefit of quantization; a multi-seed comparison is needed before treating quantized models as superior to their float32 counterparts.
- Editorial inference: A testable extension would be to distill from the same teacher into a student with a larger latent dimension (closer to the teacher's 1376) and add a latent-matching term; if the dimension-mismatch explanation is right, such a student should beat reward-only distillation on tasks that require next-state prediction, such as acrobot-swingup and cheetah-jump.
- Editorial inference: The finding that batch size 256 beats 1024 may depend on the teacher's own training batch size; matching the student's batch to the teacher's could be an alternative explanation for the effect, and it can be tested directly.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper proposes TD-MPC-Opt, a method for distilling a 317M-parameter TD-MPC2 teacher into a 1M-parameter student on the MT30 benchmark, using an additional reward-prediction MSE loss (L_distill) and post-training FP16 quantization. The central claim is that the resulting 1M-parameter student achieves a state-of-the-art normalized score of 28.45, surpassing the previously reported 1M-parameter model score of 18.93 by 50.2%, and that distillation improves over from-scratch training in some regimes. Experiments compare distillation against from-scratch baselines at several batch sizes and training lengths, and explore quantization (FP16, mixed precision, INT8). The paper also reports limitations including no physical-system validation and evaluation only on MT30.
Significance. If the central claim were robustly supported, the paper would make a useful practical contribution: reward-level distillation from a large world model to a 1M-parameter student, combined with FP16 quantization, could enable deployment of multi-task model-based RL agents in resource-constrained settings. The paper's strengths include a simple, clearly described distillation objective, evaluation on a standard benchmark (MT30), and a code release. However, the current experimental evidence does not establish the headline improvements: the main comparison is confounded by different training budgets, all key results are single runs without error bars, and one regime (1024/337K) shows distillation underperforming from-scratch training. The significance of the claimed state-of-the-art result therefore remains unverified.
major comments (4)
- [Abstract and Section 5, Discussion] This is the load-bearing comparison for the paper's central claim, so it must be presented fairly.
- [Tables 2 and 4; Section 4.2] This issue directly affects reproducibility of the central claim.
- [Section 4.2, Table 2] This is a substantive inconsistency in the results that needs to be addressed.
- [Section 3.1, Table 1, and Section 5] This is a methodological caveat that affects the strength of the state-of-the-art claim.
minor comments (6)
- [Abstract and Section 1 contribution list] The abstract reports a '+50.2%' improvement, while the contribution bullet in Section 1 reports '+48.5%' for the same comparison; these numbers should be reconciled.
- [Section 3.4 and Section 4.2] The dataset size is given as 345,690,000 transitions in Section 3.4 but as '349 million transitions' in Section 4.2; these should be made consistent.
- [Table 1] Table 1 is titled 'Impact of d_coef' but its first two rows are not d_coef values; they are ablation variants (reward + next-state with linear projection or PCA). The table should be restructured or retitled to avoid confusion.
- [Figure 3] Figure 3 shows training trajectories at five checkpoints but provides no error bars or multiple-run information; since the underlying scores are single runs, the visual 'consistently superior performance' claim is unsupported. This should be clarified in the caption or text.
- [Section 4.1] The text says 'reward + next-state (linear projection) 7.69' and 'reward + next-state (PCA) 8.78' in Table 1, but Section 4.6 describes these as attempts at latent next-state distillation; the terminology ('next-state latent' vs 'next-state') should be harmonized.
- [Throughout] The symbol 'd coef' is sometimes written with a space and sometimes as 'd_coef'; use a single notation consistently.
Circularity Check
No construction-level circularity: the distillation target is a frozen teacher's reward output, not the benchmark score, and no fitted parameter is renamed as a prediction.
full rationale
The paper's central claim is that a 1M-parameter student trained with an additional MSE reward-distillation loss against a frozen 317M teacher reaches 28.45 on MT30. That score is not defined in terms of the distillation loss, the student weights, or the teacher output: the normalized score is defined independently as the average of per-task 1-1000 scores, and the distillation loss only shapes the student's training objective. There is no equation in which the predicted quantity equals a fitted input by construction. The d_coef is tuned on MT30, but tuning a hyperparameter on the evaluation benchmark is an overfitting/selection concern, not a circular derivation, and the paper reports scores for several d_coef values. The FP16 result (28.45) exceeding its unquantized source (28.12) and the absence of error bars are reproducibility and statistical-evidence concerns, not circularity. The only self-citation, [18], appears in the introduction as 'This work extends our preliminary results described in 18' and does not carry the load-bearing argument, so it does not make the reasoning circular. No self-definitional, fitted-input-renamed-as-prediction, or citation-imported uniqueness step is identifiable from the paper's equations and text.
Assumptions & free parameters
free parameters (1)
- d_coef (distillation coefficient) =
0.4 (0.45 in Figure 3)
assumptions (3)
- domain assumption The TD-MPC2 loss functions and architecture are effective for the 1M-parameter student.
- domain assumption The pretrained 317M teacher checkpoint is a competent multi-task world model.
- ad hoc to paper MSE between teacher and student reward predictions transfers useful task knowledge.
Cite this review
Pith. "Pith review of TD-MPC-Opt: Distilling Model-Based Multi-Task Reinforcement Learning Agents." pith.science (2026). https://pith.science/paper/E3T6Q6SH
@misc{pith2026250701823,
author = {Pith},
title = {Pith review of: TD-MPC-Opt: Distilling Model-Based Multi-Task Reinforcement Learning Agents},
year = {2026},
howpublished = {\url{https://pith.science/paper/E3T6Q6SH}},
note = {Machine review of arXiv:2507.01823}
}
abstract
We present a novel approach to knowledge transfer in model-based reinforcement learning, addressing the critical challenge of deploying large world models in resource-constrained environments. Our method efficiently distills a high-capacity multi-task agent (317M parameters) into a compact model (1M parameters) on the MT30 benchmark, significantly improving performance across diverse tasks. Our distilled model achieves a state-of-the-art normalized score of 28.45, surpassing the original 1M parameter model score of 18.93. This improvement demonstrates the ability of our distillation technique to capture and consolidate complex multi-task knowledge. We further optimize the distilled model through FP16 post-training quantization, reducing its size by $\sim$50\%. Our approach addresses practical deployment limitations and offers insights into knowledge representation in large world models, paving the way for more efficient and accessible multi-task reinforcement learning systems in robotics and other resource-constrained applications. Code available at https://github.com/dmytro-kuzmenko/td-mpc-opt.
Figures
Reference graph
Works this paper leans on
-
[18]
Dmytro Kuzmenko and Nadiya Shvai. Knowledge transfer in model-based reinforcement learning agents for efficient multi-task learning, 2025. 10
work page 2025
-
[1]
Learning dexterous in-hand manipulation
OpenAI Andrychowicz, Bowen Baker, Maciek Chociej, Rafal Jozefowicz, Bob McGrew, Jakub Pachocki, Arthur Petron, Matthias Plappert, Glenn Powell, Alex Ray, et al. Learning dexterous in-hand manipulation. The International Journal of Robotics Research , 39(1):3– 20, 2020
work page 2020
-
[2]
Rich Caruana. Multitask learning. Machine learning, 28(1):41–75, 1997. 9
work page 1997
-
[3]
Deep reinforce- ment learning in a handful of trials using probabilistic dynamics models
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforce- ment learning in a handful of trials using probabilistic dynamics models. In Advances in Neural Information Processing Systems , pages 4754–4765, 2018
work page 2018
-
[4]
Distilling policy distillation
Wojciech M Czarnecki, Razvan Pascanu, Simon Osindero, Siddhant Jayakumar, Grzegorz Swirszcz, and Max Jaderberg. Distilling policy distillation. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 1331–1340. PMLR, 2019
work page 2019
-
[5]
Model-agnostic meta-learning for fast adaptation of deep networks
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International Conference on Machine Learning , pages 1126–1135. PMLR, 2017
work page 2017
-
[6]
Model predictive control: Theory and practice—a survey
Carlos E Garcia, David M Prett, and Manfred Morari. Model predictive control: Theory and practice—a survey. Automatica, 25(3):335–348, 1989
1989
-
[7]
Pwm: Policy learning with large world models
Ignat Georgiev, Varun Giridhar, Nicklas Hansen, and Animesh Garg. Pwm: Policy learning with large world models. arXiv preprint arXiv:2407.02466 , 2024
arXiv 2024
Show all 36 references
-
[8]
A survey of quantization methods for efficient neural network inference
Amir Gholami, Sehoon Kim, Zhen Dong, Zhewei Yao, Michael W Mahoney, and Kurt Keutzer. A survey of quantization methods for efficient neural network inference. arXiv preprint arXiv:2103.13630, 2021
2021 arXiv
-
[9]
Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. International conference on machine learning , pages 1861–1870, 2018
2018
-
[10]
Dream to control: Learning behaviors by latent imagination
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:1912.01603 , 2020
1912 arXiv
-
[11]
Mastering diverse domains through world models
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301.04104 , 2023
2023 arXiv
-
[12]
Td-mpc2: Scalable, robust world models for continuous control
Nicklas Hansen, Hao Su, and Xiaolong Wang. Td-mpc2: Scalable, robust world models for continuous control. arXiv preprint arXiv:2310.16828 , 2024
2024 arXiv
-
[13]
Temporal difference learning for model predictive control
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. In International Conference on Machine Learning , pages 8385–8406. PMLR, 2022
2022
-
[14]
Distilling the knowledge in a neural network
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. In NIPS Deep Learning and Representation Learning Workshop , 2015
2015
-
[15]
When to trust your model: Model-based policy optimization
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. In Advances in Neural Information Processing Systems , pages 12519–12530, 2019
2019
-
[16]
Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation
Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, et al. Qt-opt: Scalable deep reinforcement learning for vision-based robotic manipulation. arXiv preprint arXiv:1806.1...
2018 arXiv
-
[17]
Model- ensemble trust-region policy optimization
Thanard Kurutach, Ignasi Clavera, Yan Duan, Aviv Tamar, and Pieter Abbeel. Model- ensemble trust-region policy optimization. In International Conference on Learning Rep- resentations, 2018
2018
-
[19]
Multimodal reinforcement learning: A survey and taxonomy
Jaekyeom Lee, Yeong-Joon Jang, and Kee-Eung Cho. Multimodal reinforcement learning: A survey and taxonomy. IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022
2022
-
[20]
End-to-end training of deep visuomotor policies
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. The Journal of Machine Learning Research , 17(1):1334–1373, 2016
2016
-
[21]
Offline reinforcement learning: Tutorial, review, and perspectives on open problems
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.01643 , 2020
2005 arXiv
-
[22]
Mixed precision training
Paulius Micikevicius, Sharan Narang, Jonah Alben, Gregory Diamos, Erich Elsen, David Garcia, Boris Ginsburg, Michael Houston, Oleksii Kuchaiev, Ganesh Venkatesh, et al. Mixed precision training. arXiv preprint arXiv:1710.03740 , 2018
-
[23]
Playing atari with deep reinforcement learning, 2013
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing atari with deep reinforcement learning, 2013
2013
-
[24]
Human-level control through deep reinforcement learning
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. Nature, 518(7540):529–533, 2015
2015
-
[25]
Curriculum learning for reinforcement learning domains: A framework and survey
Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey. Journal of Machine Learning Research , 21(181):1–50, 2020
2020
-
[26]
Actor-mimic: Deep multitask and transfer reinforcement learning
Emilio Parisotto, Jimmy Lei Ba, and Ruslan Salakhutdinov. Actor-mimic: Deep multitask and transfer reinforcement learning. arXiv preprint arXiv:1511.06342 , 2015
2015 arXiv
-
[27]
Policy distillation
Andrei A Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell. Policy distillation. arXiv preprint arXiv:1511.06295 , 2015
2015 arXiv
-
[28]
Deep q-learning with quantized neural networks
Seunghyun Shin, Donnie H Ko, and Taehoon Kwon. Deep q-learning with quantized neural networks. In Proceedings of the 18th International Conference on Autonomous Agents and MultiAgent Systems, pages 1785–1787, 2019
2019
-
[29]
Decoupling representation learning from reinforcement learning
Adam Stooke, Kimin Lee, Pieter Abbeel, and Michael Laskin. Decoupling representation learning from reinforcement learning. arXiv preprint arXiv:2009.08319 , 2020
2009 arXiv
-
[30]
Learning to predict by the methods of temporal differences
Richard S Sutton. Learning to predict by the methods of temporal differences. Machine learning, 3(1):9–44, 1988
1988
-
[31]
Policy gradient methods for reinforcement learning with function approximation
Richard S Sutton, David A McAllester, Satinder P Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12, 2000
2000
-
[32]
Dm control: Software and tasks for continuous control
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Dm control: Software and tasks for continuous control. arXiv preprint arXiv:2006.12983 , 2020
2006 arXiv
-
[33]
Distral: Robust multitask reinforcement learning
Yee Whye Teh, Victor Bapst, Wojciech M Czarnecki, John Quan, James Kirkpatrick, Raia Hadsell, Nicolas Heess, and Razvan Pascanu. Distral: Robust multitask reinforcement learning. In Advances in Neural Information Processing Systems , pages 4496–4506, 2017. 11
2017
-
[34]
Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning
Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning, pages 1094–1100. PMLR, 2020
2020
-
[35]
Conservative q-learning for offline reinforcement learn- ing
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. Conservative q-learning for offline reinforcement learn- ing. Advances in Neural Information Processing Systems , 34:1179–1191, 2021
2021
-
[36]
A survey on multi-task learning
Yu Zhang and Qiang Yang. A survey on multi-task learning. IEEE Transactions on Knowledge and Data Engineering , 34(8):3555–3569, 2021. 12
2021
Reviewed August 6, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.