NeurIPS 2025
Xiaoyuan Zhang1,2*, Yizhe Huang1,2*, Chengdong Ma1, Zhixun Chen3, Long Ma1,
Yali Du4, Song-Chun Zhu1,2†, Yaodong Yang1†, Xue Feng2†
1PKU 2BIGAI 3HKUSTGZ 4KCL
*Equal contribution †Equal corresponding authors
Designing adaptive mechanisms to align individual and collective interests remains a central challenge in artificial social intelligence. Existing methods often struggle with modeling heterogeneous agents possessing persistent latent traits (e.g., skills, preferences) and dealing with complex multi-agent system dynamics. These challenges are compounded by the critical need for high sample efficiency due to costly real-world interactions. World Models, by learning to predict environmental dynamics, offer a promising pathway to enhance mechanism design in heterogeneous and complex systems. In this paper, we introduce a novel method named SWM-AP (Social World Model-Augmented Mechanism Design Policy Learning), which learns a social world model hierarchically modeling agents' behavior to enhance mechanism design. Specifically, the social world model infers agents' traits from their interaction trajectories and learns a trait-based model to predict agents' responses to the deployed mechanisms. The mechanism design policy collects extensive training trajectories by interacting with the social world model, while concurrently inferring agents' traits online during real-world interactions to further boost policy learning efficiency. Experiments in diverse settings (tax policy design, team coordination, and facility location) demonstrate that SWM-AP outperforms established model-based and model-free RL baselines in cumulative rewards and sample efficiency.
The AI mechanism designer (Principal) faces the challenge of making decisions within complex social systems. Direct interaction with real social systems is costly and risky, presenting two major challenges:
1. Complex System Dynamics: Individual interactions and environmental changes give rise to intricate and evolving system behaviors.
2. Heterogeneous Agent Traits: Individuals possess diverse latent preferences and needs that impact optimal decision-making.
Our proposed Social World Model aids the Principal by simulating interactions in a secure and efficient manner.
The Social World Model (SWM) utilizes a Posterior Trait Tracker to infer latent agent traits from full trajectories, which subsequently aid in predicting future state of background agents.
Concurrently, the Mechanism Design Policy employs a Prior Trait Tracker for real-time inference of agent traits based on partial history, informing its mechanism design. Interactions between the policy and the learned SWM, leveraging imagined trajectories, enhance sample efficiency for policy learning.
We implement a task where team structures are dynamically changed to optimize collective performance.
Example
Performance
Confusion matrices for trait inference in the AdaSociety task.
Left: Standard training on mixed-ambiguity states, showing moderate diagonal dominance due to modeling shortcuts in low-ambiguity late-episode states. Right: High-ambiguity training on initial states, demonstrating improved trait disentanglement with stronger diagonal dominance.
In this experimental setup(AI-Economist), low-level agents are trained using reinforcement learning. As a result, the planner interacts with a continually adapting and improving population of agents, making the environment increasingly complex over time.
We assess the effectiveness of our proposed method by comparing it with two baseline algorithms: the model-based MBPO and the model-free PPO. The evaluation is based on three social metrics recorded after 1,000 agent steps: Equality, Productivity, and their product, Equality*Productivity. The product metric serves as an indicator of overall social welfare within the simulated environment. As shown, our method exhibits effective control over taxation, enabling the planner to optimize social welfare.
Overall Performance
Equality
Productivity
This paper proposes a social world model-enhanced mechanism design policy learning approach, which employs unsupervised learning to infer the hidden trait of background agents, thereby enhancing the prediction of group dynamics and mechanism design policy learning. Experimental results in facility location, team structure optimization, and taxation setting tasks demonstrate that our method outperforms both model-based and model-free reinforcement learning approaches. Our approach inspires new directions for world model research towards more complex and realistic social world.
Current limitations include the scalability of the Social World Model to extremely large-scale systems and the challenge of ensuring the direct interpretability of all inferred latent traits. Future work will focus on developing more scalable SWM architectures, potentially leveraging structured priors, and on enhancing trait interpretability through techniques like disentangled representation learning.
@inproceedings{zhang2025swmap,
title={Social World Model-Augmented Mechanism Design Policy Learning},
author={Zhang, Xiaoyuan and Huang, Yizhe and Ma, Chengdong and Chen, Zhixun and Ma, Long and Du, Yali and Zhu, Song-Chun and Yang, Yaodong and Feng, Xue},
booktitle={Advances in Neural Information Processing Systems},
year={2025}
}