TurboTrain: Towards Efficient and Balanced Multi-Task Learning for Multi-Agent Perception and Prediction

TL;DR

TurboTrain introduces an efficient training framework for multi-agent perception and prediction that leverages spatiotemporal pretraining and gradient conflict suppression, reducing manual tuning and improving performance.

Contribution

It presents a novel training framework combining spatiotemporal pretraining and balanced multi-task learning to streamline multi-agent system training.

Findings

Improves multi-agent perception and prediction performance on V2XPnP-Seq dataset.

Reduces training time and manual tuning requirements.

Enhances detection and prediction accuracy through balanced learning.

Abstract

End-to-end training of multi-agent systems offers significant advantages in improving multi-task performance. However, training such models remains challenging and requires extensive manual design and monitoring. In this work, we introduce TurboTrain, a novel and efficient training framework for multi-agent perception and prediction. TurboTrain comprises two key components: a multi-agent spatiotemporal pretraining scheme based on masked reconstruction learning and a balanced multi-task learning strategy based on gradient conflict suppression. By streamlining the training process, our framework eliminates the need for manually designing and tuning complex multi-stage training pipelines, substantially reducing training time and improving performance. We evaluate TurboTrain on a real-world cooperative driving dataset, V2XPnP-Seq, and demonstrate that it further improves the performance of state-of-the-art multi-agent perception and prediction models. Our results highlight that pretraining effectively captures spatiotemporal multi-agent features and significantly benefits downstream tasks. Moreover, the proposed balanced multi-task learning strategy enhances detection and prediction.