TD-MPC-Opt: Distilling Model-Based Multi-Task Reinforcement Learning Agents
This addresses practical deployment limitations for multi-task reinforcement learning in robotics and other resource-constrained applications, though it is incremental as it builds on existing distillation and model-based methods.
They tackled the challenge of deploying large model-based reinforcement learning agents in resource-constrained environments by distilling a high-capacity multi-task agent into a compact model, achieving a state-of-the-art normalized score of 28.45 on the MT30 benchmark, up from 18.93, and reducing model size by ~50% with quantization.
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.