TaskLoom: Weaving Knowledge Across Tasks in World Models
Abstract
World models have significantly improved the sample efficiency of model-based reinforcement learning (MBRL) by enabling policy learning in imagination, thereby reducing the need for direct interaction with the real environment. However, most existing world model methods are trained independently for each task or perform multi-task learning using offline datasets, failing to fully exploit the latent relationships among tasks in online interactive scenarios. To address this limitation, we propose TaskLoom, a knowledge-sharing world model architecture for online RL. TaskLoom adopts a grouped two-stage training paradigm: first, the tasks are divided into several groups based on the similarity of world model gradients, and fine-grained knowledge is shared among tasks within each group; second, coarse-grained knowledge is exchanged across groups, enabling hierarchical knowledge transfer and reuse. Experimental results show that TaskLoom outperforms baseline methods on widely used benchmarks such as Proprio Control, Visual Control and Meta-World, validating the effectiveness of the proposed knowledge-sharing mechanism for both low-dimensional state and high-dimensional visual inputs.
Lay Summary
(1) Problem: When training AI for tasks like robotics, agents learn efficiently by building an internal “world model” to practice actions in their imagination. However, when trying to master multiple skills in real time, current systems learn each task completely in isolation. They miss the chance to share valuable insights between related activities. (2) Solution: To solve this, we introduce TaskLoom, an AI architecture designed to share knowledge while actively learning. TaskLoom groups similar tasks together so the AI can share detailed, specific skills within each group. Then, it exchanges broader, general concepts across these different groups. (3) Impact: This layered sharing approach allows the AI to learn multiple skills much more effectively. Our tests show TaskLoom significantly outperforms standard methods on challenges ranging from basic physical movements to complex visual tasks. This brings us closer to adaptable AI that efficiently masters skills in the real world.