🤖 AI Summary
This work addresses the generalizable policy learning problem in multi-agent settings where agents possess heterogeneous action spaces but share a common observation space. We propose the first diffusion-based planning framework trained on joint multi-agent trajectories, integrating inverse dynamics modeling with cross-agent trajectory alignment to enable forward transfer and zero-shot adaptation to unseen agents. Evaluated on the BabyAI benchmark, our method achieves up to a 42.20% improvement in task completion accuracy, significantly outperforming single-agent training and conventional imitation learning baselines. Key contributions include: (1) the first application of diffusion models to cross-agent policy generalization; (2) the construction of an action-agnostic implicit planning representation that decouples observation understanding from action execution; and (3) effective policy knowledge transfer across heterogeneous action spaces via joint trajectory training.
📝 Abstract
The development of a generalist agent capable of solving a wide range of sequential decision-making tasks remains a significant challenge. We address this problem in a cross-agent setup where agents share the same observation space but differ in their action spaces. Our approach builds on the universal policy framework, which decouples policy learning into two stages: a diffusion-based planner that generates observation sequences and an inverse dynamics model that assigns actions to these plans. We propose a method for training the planner on a joint dataset composed of trajectories from all agents. This method offers the benefit of positive transfer by pooling data from different agents, while the primary challenge lies in adapting shared plans to each agent's unique constraints. We evaluate our approach on the BabyAI environment, covering tasks of varying complexity, and demonstrate positive transfer across agents. Additionally, we examine the planner's generalisation ability to unseen agents and compare our method to traditional imitation learning approaches. By training on a pooled dataset from multiple agents, our universal policy achieves an improvement of up to $42.20%$ in task completion accuracy compared to a policy trained on a dataset from a single agent.