Prototypical Environment-aware Proxy with Coordinated Optimization for Multi-agent Dynamics Modeling

Yifan Wang, Fanchun Meng, Changhu Wang, Yuhang Pei, Tao Ren, Wei Ju, Dongjie Wang, Chong Chen, Xian-Sheng Hua, Xiao Luo

Keywords: Multi-agent dynamical systems, Graph ODEs, Prompt learning, Out-of-distribution generalization

Abstract: Modeling multi-agent dynamic systems is crucial for understanding collective behaviors in various scientific domains. While graph ordinary differential equation (ODE) approaches effectively capture continuous dynamics from irregular data, their performance struggles to generalize across temporal and parameter-induced shifts and degrades severely under potential out-of-distribution fluctuation. In this paper, we propose a novel approach named Prototypical Enviroment-aware Proxy with Coordinated Optimization (PEACE) for multi-agent dynamics modeling. The core of our PEACE is to learn a set of proxy models to simulate environment information while keeping the primary model fixed. In particular, our primary model utilizes temporal graph neural networks to extract invariant observation embeddings across different nodes. More importantly, a range of prototypical prompts are introduced to model temporal distribution shifts with graph ODEs, which are further incorporated with our observation embeddings to serve as proxy models. These proxy models would further generate diverse predictions of unseen trajectories, which are selected by a vision language model for data augmentation. To jointly learn the primary and proxy models, a bi-level strategy is adopted for alternative optimization. In the lower level, we update prompt parameters in the proxy models with our primary model frozen. In the upper level, we integrate all these proxy models and measure the gradient coordination to update our primary models. Extensive experiments on multiple real-world system dynamics datasets demonstrate the superiority of PEACE over state-of-the-art baselines, confirming its effectiveness and robustness.

Primary Area: learning on graphs and other geometries & topologies

Submission Number: 11503