Knowledge Transfer in Model-Based Reinforcement Learning Agents for Efficient Multi-Task Learning

2025-01-09

AI Summary

To address the poor generalization and deployment challenges of multi-task model-based reinforcement learning (MBRL) in resource-constrained settings, this paper proposes a knowledge transfer framework tailored for lightweight agents. We first compress a high-capacity multi-task world model (317M parameters) into an ultra-compact agent (1M parameters) via knowledge distillation, then apply FP16 post-training quantization to further reduce model size by 50%. Evaluated on the MT30 multi-task benchmark, our method achieves a normalized score of 28.45—improving upon the original 1M baseline by 50.5% and establishing a new state-of-the-art (SOTA). The model size is reduced by 317× with zero performance degradation. Our core contribution lies in establishing an efficient, fidelity-preserving knowledge distillation and compression paradigm that transfers capabilities from large-scale multi-task world models to extremely lightweight MBRL agents—significantly enhancing multi-task generalization and deployment feasibility on edge devices.

Technology Category

• Machine Learning: Learning on the Edge & Model Compression
• Multiagent Systems: Multiagent Learning
• Data Mining & Knowledge Management: Data Compression

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

We propose an efficient knowledge transfer approach for model-based reinforcement learning, addressing the challenge of deploying large world models in resource-constrained environments. Our method distills a high-capacity multi-task agent (317M parameters) into a compact 1M parameter model, achieving state-of-the-art performance on the MT30 benchmark with a normalized score of 28.45, a substantial improvement over the original 1M parameter model's score of 18.93. This demonstrates the ability of our distillation technique to consolidate complex multi-task knowledge effectively. Additionally, we apply FP16 post-training quantization, reducing the model size by 50% while maintaining performance. Our work bridges the gap between the power of large models and practical deployment constraints, offering a scalable solution for efficient and accessible multi-task reinforcement learning in robotics and other resource-limited domains.

Problem

• Model-based Reinforcement Learning
• Knowledge Transfer
• Resource-constrained Environment

Innovation

• Model Compression
• Performance Enhancement
• Efficient Multitasking