Social World Model-Augmented Mechanism Design Policy Learning¶

🕒 Published (v1): 2025-01-01 · Source: NeurIPS · Venue: NeurIPS 2025 · link

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

TL;DR¶

SWM-AP introduces a model-based RL framework for AI-driven mechanism design that hierarchically infers heterogeneous agents' latent traits (skills, preferences) from interaction trajectories and uses a learned social world model to simulate rollouts, dramatically improving sample efficiency. The mechanism design policy trains primarily on imagined trajectories and performs real-time trait inference online. Tested on tax policy, team coordination, and facility location, it outperforms Dreamer, MBPO, and PPO baselines on both final reward and sample efficiency.

Problem¶

Existing RL-based mechanism design methods treat background agents as homogeneous, ignoring persistent latent traits (skill, risk tolerance, preferences) that drive heterogeneous responses to incentives. Standard MBRL world models cannot capture these trait-conditioned social dynamics, and direct real-world interaction with social systems (tax policy, resource allocation) is costly and risky. Information asymmetry — the principal never observes agents' latent traits — compounds the challenge.

Method¶

SWM-AP couples two components:

Social World Model (SWM): A Posterior Trait Tracker infers latent trait vectors \(\hat{m}^{\text{post}}\) from complete offline trajectories via unsupervised learning (ELBO objective with KL regularization toward a uniform prior). A trait-conditioned dynamics model then predicts next observations and social welfare: \(\hat{s}^{\text{obs}}_{t+1}(s^{\text{obs}}_t, \pi_t, \hat{m}^{\text{post}})\).
Mechanism Design Policy: A Prior Trait Tracker performs online, causal inference of \(\hat{m}^{\text{prior}}\) from partial history up to timestep \(t\), supervised by the Posterior Tracker's outputs (cross-entropy or MSE loss). The PPO-trained policy \(\Pi(\pi_t | s^{\text{obs}}_t, \hat{m}^{\text{prior}})\) trains on a mix of real environment rollouts and imagined trajectories generated by SWM, using the inferred traits to condition both planning and action selection.

Training alternates: collect real trajectories → jointly train Posterior Tracker + dynamics model → distill into Prior Tracker → generate imagined rollouts → update policy on combined data.

Key Contributions¶

A hierarchical Social World Model that jointly learns unsupervised latent trait inference and trait-conditioned multi-agent dynamics via an ELBO-grounded objective.
A Prior/Posterior Trait Tracker architecture that bridges offline (full-trajectory) trait inference with online (causal) real-time adaptation during policy deployment.
Demonstrated that trait-aware world models yield lower state and reward prediction error than trait-agnostic baselines (Dreamer, MBPO), translating to higher final reward and fewer real environment steps.
Validated across three structurally distinct mechanism design domains: facility location (rule-based agents), team structure optimization (AdaSociety), and tax policy (AI-Economist with RL-trained agents).

Results¶

Facility location (5 facilities, 8 agents, 8×8 grid): SWM-AP achieves highest final converged reward and requires the fewest training steps to reach PPO's final performance level among PPO, MBPO, and Dreamer.
Team structure optimization (AdaSociety): SWM-AP final group reward ≈ 40.225 vs. MBPO ≈ 40.181 vs. PPO ≈ 39.580; SWM-AP requires fewest steps to reach target (bars in Figure 4a; exact step counts not extracted but SWM-AP bar is shortest).
Tax policy (AI-Economist, 4 agents, 25×25 map): SWM-AP matches PPO's productivity while significantly improving equality, yielding higher equality × productivity (social welfare product) than both PPO and MBPO; MBPO raises equality at the cost of large productivity drop.
SWM achieves lower state prediction loss and reward prediction loss than MBPO across all three domains.

Limitations¶

Scalability of the Posterior Trait Tracker to very large agent populations is not addressed; experiments use at most 8 agents.
Latent trait representations are not interpretable by design; the paper acknowledges that the inferred \(\hat{m}\) vectors lack semantic transparency.
All background agents must have fixed (non-evolving) traits per episode; the framework does not handle agents whose traits drift over time.
Experiments use relatively small, grid-world environments; generalization to continuous high-dimensional social systems is unverified.

Relevance to Agentic AI / LLM Agents¶

SWM-AP is a concrete instantiation of a principal agent that must model, predict, and steer a population of heterogeneous background agents — a core sub-problem in multi-agent agentic systems where a coordinating agent must reason about others' latent intentions and capabilities. The Posterior/Prior Trait Tracker duality directly operationalizes Theory-of-Mind-style belief inference within a world model, which is increasingly central to agentic architectures that must cooperate or compete with black-box agents. The work also demonstrates how MBRL's sample-efficiency gains transfer to social system control, relevant to researchers building agentic systems that must learn with minimal costly real-world interaction. Its hierarchical principal–agent framing connects directly to orchestrator–subagent patterns in LLM agent frameworks.