Making Universal Policies Universal

TL;DR

This paper introduces a universal policy framework for multi-agent sequential decision tasks, leveraging a diffusion-based planner and inverse dynamics model to enable positive transfer and generalization across diverse agents.

Contribution

It proposes a novel training method for a universal policy that pools data from multiple agents, improving transferability and generalization in complex environments.

Findings

Achieved up to 42.20% improvement in task accuracy over single-agent training.

Demonstrated positive transfer across different agents in BabyAI environment.

Showed the planner's ability to generalize to unseen agents.

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.