Skip to content

Encouraging metric-aware diversity in contrastive representation space

🕒 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.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

WCD (Wasserstein Contrastive Diversity) addresses the collapse of Wasserstein-based diversity signals in parameter-sharing MARL by learning a contrastive trajectory representation space where agent policies are distinguishable, then maximizing pairwise Wasserstein distances in that latent space as intrinsic exploration rewards. It outperforms mutual-information and prior Wasserstein-based baselines across Pac-Men, SMAC, and SMACv2.

Problem

In cooperative MARL with parameter sharing, agents develop nearly identical policies, causing prior Wasserstein distance-based diversity methods (MAPD, DiCo) to fail: when initial policy distributions are nearly identical, their Wasserstein distance collapses toward zero, producing no meaningful diversity gradient. Mutual information / KL-divergence-based methods are metric-agnostic and saturate with only slight trajectory differences, insufficient for genuine exploration coverage.

Method

WCD operates in three stages integrated with QMIX:

  1. Contrastive trajectory encoding: Each agent's observation-action sequence (o_t^a, u_t^a) is encoded via MLP encoder g_θe → latent z_t^a, then summarized by a GRU autoregressive model g_θg into trajectory representation c_t^a. Training minimizes an InfoNCE loss that contrasts the next-step embedding z_{t+1}^a against embeddings from other agents' trajectories (used as noise), forcing c_t^a to predict its own future while repelling other agents' futures.

  2. Efficient Wasserstein computation: Instead of parameterizing dual functions with neural networks (expensive for all agent pairs), dual functions μ and ν are represented as linear combinations of Gaussian random feature maps φ(x) = (1/√m) cos(Gx + b). Optimization reduces to learning dual vectors λ_μ, λ_ν ∈ ℝ^m via SGD over the smoothed Fenchel-Rockafellar dual form.

  3. Nearest-neighbor intrinsic reward: For agent a, the intrinsic reward is r_w = min_{a'≠a} W(p^{π_a}, p^{π_{a'}}), maximizing distance to the closest neighbor in representation space. A separate intrinsic utility network Q^w_a(c_t^a, Q_a) is trained via TD loss to propagate this signal; total loss is L_TD + αL^w_TD.

Key Contributions

  • Identifies the root cause of Wasserstein distance collapse in parameter-sharing MARL and proposes a CPC-based contrastive representation space to restore its effectiveness.
  • Novel InfoNCE-trained trajectory encoder using cross-agent noise samples to produce metrically separated representations.
  • Gaussian kernel / random feature map approximation for efficient Wasserstein dual function optimization, avoiding per-pair neural network parameterization.
  • Nearest-neighbor Wasserstein intrinsic reward that scales to arbitrary numbers of agents.
  • Integration with both value-based (QMIX) and policy-gradient (MAPPO) MARL frameworks.

Results

  • Pac-Men: WCD agents spread across all four edge rooms; MAPD agents cluster, failing to discover the long-path room.
  • SMAC: Outperforms all baselines (QMIX, QTRAN, RODE, MAVEN, EOI, SCDS, PMIC, LIPO, FoX, MAPD, DiCo) in all 6 scenarios; largest gains on super-hard maps (6h_vs_8z, corridor, 3s5z_vs_3s6z) where QMIX and MAPD near-zero.
  • SMACv2 (stochastic): Superior performance in all three race scenarios (terran, protoss, zerg 5v5); visitation heatmaps show broader state coverage vs. baselines confined to partial regions.
  • Ablations: Removing the autoregressive model, replacing InfoNCE with identity prediction or generative reconstruction, or using average (vs. nearest-neighbor) Wasserstein distance each causes significant degradation; all Wasserstein variants still beat QMIX alone, confirming the representation learning is load-bearing.

Limitations

  • The Euclidean cost function for the Wasserstein distance is used for simplicity; optimal cost function selection for specific multi-agent task geometries remains an open problem.
  • The approach is validated only within the CTDE / parameter-sharing paradigm; applicability to heterogeneous-parameter or fully decentralized settings is not studied.
  • Computational cost of maintaining and updating Gaussian kernel dual functions and the additional intrinsic utility network Q^w_a per agent adds overhead not quantified against baselines.

Relevance to Agentic AI / LLM Agents

This work is directly relevant to multi-agent systems research, particularly the challenge of inducing behavioral diversity in agents that share underlying model parameters—a problem that arises in any multi-agent LLM or tool-use setting where a single model is instantiated multiple times. The insight that metric-aware (Wasserstein) diversity signals require a well-structured latent space to be effective connects to broader questions about how to encourage specialization among homogeneous LLM-based agents in collaborative tasks. The CPC-based trajectory representation approach offers a transferable technique for making distributional divergence objectives meaningful when agent policies are nearly identical at initialization.