Pith. sign in

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 →

arxiv 2507.01823 v1 pith:E3T6Q6SH submitted 2025-07-02 cs.LG cs.RO

classification cs.LGcs.RO
keywords model-basedreinforcementlearningknowledgedistillationmulti-taskTD-MPC2MT30benchmarkmodelcompressionFP16quantizationcontinuouscontrol
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

This paper argues that a large model-based reinforcement learning agent can be compressed into a 1M-parameter student that outperforms the original compact model on the MT30 multi-task benchmark. The method adds a reward-distillation term to the standard TD-MPC2 loss, training the student to match the teacher's reward predictions on the same state-action inputs. With a batch size of 256 and 1M training steps, the distilled student reaches a normalized score of 28.45 after FP16 quantization, compared with 18.93 for the original 1M checkpoint. The paper also finds that smaller batch sizes help distillation, that reward-only distillation beats attempts to also match next-state latents, and that FP16 quantization roughly halves model size with no performance loss. If these results hold, they point to a practical recipe for turning large world models into deployable multi-task controllers.

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.

Watch

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 extensions of the paper, not claims the author makes directly.

  • 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.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 6 minor

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)
  1. [Abstract and Section 5, Discussion] This is the load-bearing comparison for the paper's central claim, so it must be presented fairly.
  2. [Tables 2 and 4; Section 4.2] This issue directly affects reproducibility of the central claim.
  3. [Section 4.2, Table 2] This is a substantive inconsistency in the results that needs to be addressed.
  4. [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)
  1. [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.
  2. [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.
  3. [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.
  4. [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.
  5. [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.
  6. [Throughout] The symbol 'd coef' is sometimes written with a space and sometimes as 'd_coef'; use a single notation consistently.

Circularity Check

0 steps flagged · score 1.0 of 10

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 1 free parameters · 3 assumptions · 0 invented entities

The central claim rests on one tuned hyperparameter (d_coef), the assumed validity of the TD-MPC2 framework and teacher checkpoint, and the empirical hypothesis that reward-prediction matching transfers knowledge. No new entities are introduced.

free parameters (1)
  • d_coef (distillation coefficient) = 0.4 (0.45 in Figure 3)
    Selected by sweeping 0.05-0.9 on MT30 normalized score (Table 1) and taking the best value; this is tuning on the evaluation metric.
assumptions (3)
  • domain assumption The TD-MPC2 loss functions and architecture are effective for the 1M-parameter student.
    Section 3.3 states the student retains the original TD-MPC2 architecture and planning hyperparameters without re-validation.
  • domain assumption The pretrained 317M teacher checkpoint is a competent multi-task world model.
    Section 3.3 uses a checkpoint from Hansen et al. (2024) as the distillation target; the teacher's own MT30 score is not reported.
  • ad hoc to paper MSE between teacher and student reward predictions transfers useful task knowledge.
    This is the core method in Section 3.1. The paper's own negative results with latent-state distillation (Section 4.6) show the reward-only choice is empirical, not derived.

how reviews work

0 comments
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

Figures reproduced from arXiv: 2507.01823 by the authors.

Figure 1
Figure 1. Task visualizations. Visualization of a random initial state for 12 of the 30 tasks that we consider (MT30). The tasks have different action spaces, degrees of freedom, and objectives. Parisotto et al.26 introduced the Actor-Mimic method, which uses policy distillation to transfer knowledge from multiple task-specific teacher networks to a single multi-task network. Czarnecki et al.4 proposed an improved policy dist… view at source ↗
Figure 2
Figure 2. Our distillation approach consists of two main loss function components: the original loss from TD-MPC2 (in red), calculated as a linear combination of consistency, reward, and value losses (α denotes a scaling coefficient for each respective loss); and the distillation loss (in teal) that calculates MSE between student’s (green) and teacher’s (blue) rewards produced from inferring the same state-action pair. The to… view at source ↗
Figure 3
Figure 3. Performance comparison between a knowledge-distilled model (d coef of 0.45) and a model trained from scratch. Both models have 1M parameters and were trained with a batch size of 256. The normalized scores are shown for five checkpoints, taken at 200K-step intervals up to 1M steps. 4.3 Batch Size Study We investigated how batch size influences the quality of knowledge transfer, convergence speed, and computational e… view at source ↗

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

36 extracted references · 22 canonical work pages

  1. [18]

    Knowledge transfer in model-based reinforcement learning agents for efficient multi-task learning, 2025

    Dmytro Kuzmenko and Nadiya Shvai. Knowledge transfer in model-based reinforcement learning agents for efficient multi-task learning, 2025. 10

  2. [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

  3. [2]

    Multitask learning

    Rich Caruana. Multitask learning. Machine learning, 28(1):41–75, 1997. 9

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

Show all 36 references
  1. [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

  2. [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

  3. [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

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

  9. [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...

  10. [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

  11. [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

  12. [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

  13. [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

  14. [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

  15. [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

  16. [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

  17. [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

  18. [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

  19. [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

  20. [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

  21. [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

  22. [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

  23. [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

  24. [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

  25. [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

  26. [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

  27. [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

  28. [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

Pith tools

Reviewed August 6, 2026 · model on record in the stance chip above.