Task-agnostic Decision Transformer for Multi-type Agent Control with Federated Split Training

TL;DR

This paper tackles the challenge of coordinating multiple heterogeneous agents in offline reinforcement learning under privacy constraints, where state and action spaces vary across agents. It introduces the Federated Split Decision Transformer (FSDT), a two-stage training framework that keeps local embedding and prediction models on clients while employing a server-side Transformer decoder to synthesize cross-agent information for action prediction, effectively handling multi-type agents. The approach demonstrates superior performance on the D4RL benchmarks with reduced communication and computation overhead compared to centralized training, highlighting its potential for privacy-preserving, efficient training in autonomous driving decision systems. By leveraging split learning and a server-side transformer, FSDT enables scalable, type-agnostic decision-making across distributed offline data, with strong practical implications for real-world multi-agent control tasks.

Abstract

With the rapid advancements in artificial intelligence, the development of knowledgeable and personalized agents has become increasingly prevalent. However, the inherent variability in state variables and action spaces among personalized agents poses significant aggregation challenges for traditional federated learning algorithms. To tackle these challenges, we introduce the Federated Split Decision Transformer (FSDT), an innovative framework designed explicitly for AI agent decision tasks. The FSDT framework excels at navigating the intricacies of personalized agents by harnessing distributed data for training while preserving data privacy. It employs a two-stage training process, with local embedding and prediction models on client agents and a global transformer decoder model on the server. Our comprehensive evaluation using the benchmark D4RL dataset highlights the superior performance of our algorithm in federated split learning for personalized agents, coupled with significant reductions in communication and computational overhead compared to traditional centralized training approaches. The FSDT framework demonstrates strong potential for enabling efficient and privacy-preserving collaborative learning in applications such as autonomous driving decision systems. Our findings underscore the efficacy of the FSDT framework in effectively leveraging distributed offline reinforcement learning data to enable powerful multi-type agent decision systems.

Figures (5)

• Figure 1: Single-type agent learning task vs. Multi-type agent learning task.
• Figure 2: The core components of the Federated Split Decision Transformer (FSDT). Local agents train with their data, generating embeddings that reflect key features pertinent to decision-making tasks. These embeddings are sent to a server-side transformer decoder, synthesizing the information across all agents to predict actions.
• Figure 3: The framework of our algorithm. At the personalized client agent, we train the embedding and prediction models in stage 1. On the server side, we introduce the Transformer decoder without the embedding layer in stage 2.
• Figure 4: Performance comparison of different agents.
• Figure 5: The ablation experiment results of our proposed algorithm FSDT on Medium-Expert dataset.