Zero Shot Coordination for Sparse Reward Tasks with Diverse Reward Shapings

1. Problem Definition

The paper addresses a critical limitation in Multi-Agent Reinforcement Learning (MARL) known as Zero-Shot Coordination (ZSC).

• The Core Issue: Traditional MARL agents, trained via Self-Play, often develop specific conventions or strategies that fail when paired with agents trained under the same objectives but with different random seeds, algorithms, or training histories.
• The Gap: Existing ZSC solutions (e.g., Other-Play, TrajeDi, MAZE) assume that all agents share an identical reward function. However, in real-world scenarios (e.g., autonomous vehicles from different manufacturers), agents may share a sparse objective (e.g., "reach the destination") but utilize different reward shapings (different weights assigned to environmental features to densify the sparse reward).
• The Challenge: How to train an agent to cooperate effectively with unknown partners who have the same sparse goals but different internal reward structures, without additional training on the specific partner.

2. Methodology

The authors propose a framework that trains an ensemble of agents using diverse reward shapings to improve generalization.

A. Reward Shaping Formulation

The authors define reward shaping as a weighted combination of environmental features added to a base sparse reward:
$$R(s_t, a_t) = \text{sparse}(s_t, a_t) + \sum_{k=1}^{F} w_k f_k$$
Where $w_k$ are weights selected for specific features $f_k$ (e.g., distance to a pot, picking up a dish).

B. Four Selection Algorithms

To generate a population of diverse reward shapings ($P$), the paper introduces four selection methods:

1. LLM-Based: Uses an LLM (Claude Sonnet 4.5) prompted with examples of reward shapings and their performance to generate new, diverse sets of weights that balance performance and policy diversity.
2. Surrogate Network: Trains a Multi-Layer Perceptron (MLP) to predict the total sparse reward based on a given reward shaping. It then samples 1,000 random shapings and selects the top-performing ones based on the surrogate's predictions.
3. Stratified Grid: Uses Latin Hypercube Sampling to ensure that the selected weights cover the entire possible range (0–10) evenly. This guarantees statistical representativeness of the reward space.
4. Random: Uniformly samples weights from the possible range without constraints.

C. Ensemble Training Framework

Building upon the TrajeDi (Trajectory Diversity) algorithm:

1. Population Generation: For each selected reward shaping, a population of agents is trained using TrajeDi, which includes a Jensen-Shannon Divergence term to encourage policy diversity within the population.
2. Ensemble Construction: The final agent is an ensemble model composed of the "Best Response" (BR) agents from each of the $P$ trained populations.
3. Action Selection: The ensemble selects the action that is the mode (most frequent choice) among all BR agents in the ensemble.

3. Key Contributions

1. Novel Problem Formulation: Extends ZSC to scenarios where partners have identical sparse objectives but different reward shapings, a more realistic constraint than previous literature assumed.
2. Four Selection Methods: Introduces and evaluates four distinct algorithms (LLM, Surrogate, Stratified Grid, Random) for generating diverse reward shapings.
3. Ensemble Framework: Proposes a method to combine models trained on different reward shapings into a single robust policy.
4. Empirical Validation: Demonstrates that diversity in reward shapings, not just policy diversity, is crucial for generalization in ZSC.

4. Experimental Results

Experiments were conducted in the Overcooked environment (specifically Random0_Medium, Random3, and Unident_S) against partners trained with unknown reward shapings using standard MAPPO.

• Performance Gains:

  • The proposed methods significantly outperformed baseline ZSC algorithms (including standard TrajeDi and Hidden-Utility Self-Play).
  • Improvement Range: Sparse reward improvements ranged from 62.2% to 119.2% over baseline algorithms.
  • Top Performers: Stratified Grid and Surrogate Network consistently achieved the best results.
    • Stratified Grid: Achieved the highest generalization, likely due to its even coverage of the reward space.
    • Surrogate Network: Performed exceptionally well despite lower diversity, suggesting that selecting high-performing shapings is more critical than pure randomness.
  • Ensemble Effect: The ensemble approach itself provided a 40–70% improvement over non-ensembled baselines.
  • LLM-Based: Showed strong performance but was slightly less consistent than Stratified Grid.
  • Random: Outperformed baselines (+27% to +51%) but was the weakest of the proposed methods, highlighting that unguided diversity is insufficient.

• Diversity vs. Performance:

  • Stratified Grid had the highest range coverage (96.7%) and high standard deviation.
  • Surrogate Network had the lowest standard deviation and range coverage but achieved the highest performance in some environments. This indicates that quality of the reward shaping (finding high-performing weights) can sometimes outweigh pure statistical diversity.

5. Significance and Conclusion

• Realism: The paper bridges the gap between theoretical ZSC and real-world deployment, acknowledging that agents from different sources will likely optimize the same goal differently.
• Robustness: The ensemble approach creates agents that are resilient to the "convention mismatch" problem, capable of cooperating immediately with unknown partners.
• Future Directions: The authors suggest that future work should focus on hybridizing methods (e.g., using the Surrogate Network's performance focus with Stratified Grid's diversity constraints) and expanding to more complex environments.

In summary, the paper proves that training agents on a diverse set of reward shapings and combining them via an ensemble is a highly effective strategy for achieving robust Zero-Shot Coordination in sparse reward environments.