🤖 AI

ARMS: Automatic Reward Shaping for Sparse-Reward Multi-Agent Reinforcement Learning

The paper proposes ARMS, a self-supervised framework for multi-agent reinforcement learning that automatically generates dense reward signals from sparse rewards via trajectory ranking while theoretically guaranteeing the preservation of Nash equilibria through conditional best-response reasoning, thereby improving sampling efficiency and stability in sparse-reward environments.

Original authors: Elie Abboud, Oren Gal

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Elie Abboud, Oren Gal

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a group of dogs to run a relay race. In a perfect world, every time a dog takes a step in the right direction, you give it a treat. But in the real world (and in the complex world of Artificial Intelligence), you often can't give a treat for every step. Instead, you only give a treat when the dog finally crosses the finish line.

This is the problem of Sparse Rewards. If the dog has to run a long distance and only gets a treat at the very end, it has no idea which steps were good and which were bad. It's like trying to learn a new language by only being told "Correct!" or "Wrong!" once a year after you've finished a sentence.

In the world of Multi-Agent Reinforcement Learning (MARL), this is even harder. Imagine not just one dog, but a whole pack. They all have to learn at the same time, and they have to coordinate with each other. If one dog changes its strategy, it changes the environment for all the other dogs. This creates a chaotic, shifting landscape where learning is incredibly difficult.

The Problem: The "Silent" Coach

The paper argues that when you have many agents (like our dogs) learning together with sparse rewards, they often get stuck. They might start running in circles, bumping into each other, or just standing still because they can't figure out what to do without constant feedback.

Traditional methods try to fix this by having a human expert design a "cheat sheet" (called Reward Shaping) that gives the agents little hints along the way. But this is risky. If the human gives the wrong hints, the agents might learn to "game the system"—they might find a shortcut that gives them lots of points but doesn't actually solve the problem (like a dog running in circles to get treats instead of running the race).

The Solution: ARMS (The Self-Learning Coach)

The authors propose a new system called ARMS (Automatic Reward-shaping in Multi-agent Systems). Instead of a human designing the hints, ARMS acts like a smart coach that learns how to give hints automatically.

Here is how it works, using a simple analogy:

The Observation: The coach watches the dogs run the race many times. Even though the dogs only get a treat at the finish line, the coach can see the whole race.
The Ranking: The coach looks at two different race attempts. "Look," the coach says, "In Race A, the dogs finished faster and didn't crash. In Race B, they crashed and took longer. Therefore, Race A is better than Race B."
The Lesson: The coach doesn't just say "Good job" or "Bad job." It uses these rankings to invent a new reward system. It creates a dense signal (a constant stream of small hints) that tells the dogs, "This step was good because it looks like the steps in the winning race," or "This step was bad because it looks like the steps in the losing race."
The Safety Net: The most important part of ARMS is that it is mathematically proven to be safe. The authors show that even though the coach is inventing new rules to give hints, it never changes the ultimate goal of the game. It guarantees that the "winning strategies" (called Nash Equilibria in the paper) remain the same. The agents might learn faster, but they won't learn the wrong thing.

The "Oscillation" Trap

The paper discovered a funny and dangerous glitch that happens when there are too many agents and not enough exploration.

Imagine the dogs are so confused that they all start doing the exact same silly dance. They bump into each other, stop, dance again, and bump again. Because they are all doing the same thing, the "coach" (the reward system) sees this pattern over and over and thinks, "Oh, this is the best way to move!" It reinforces the bad behavior, and the dogs get stuck in an endless loop of crashing and dancing.

The paper found that if you tell the agents to be a little more curious (increase "exploration"), they will try random, weird moves. This breaks the loop, allowing the coach to see that the "dance" is actually a bad idea and to start teaching them the right way to run the race.

The Results

The authors tested this in a virtual world where agents (like little robots) had to navigate a grid, avoid obstacles, and reach targets without crashing into each other.

Faster Learning: When the rewards were very sparse (hard to learn), ARMS helped the agents learn much faster than if they had no help or if they used old, hand-designed hints.
Better Teamwork: The agents learned to coordinate better, crashing into each other less often.
Generalization: The agents trained with ARMS were better at navigating new maps they had never seen before, proving they actually learned the concept of "running a race" rather than just memorizing one specific track.

Summary

In short, ARMS is a system that lets a group of AI agents teach themselves how to work together when they don't get much feedback. It does this by watching their own past attempts, ranking the good ones against the bad ones, and automatically generating a helpful "cheat sheet" to guide them. Crucially, it does this without breaking the rules of the game, ensuring they learn the right solution, not just a clever trick.

Technical Summary: ARMS – Automatic Reward Shaping for Sparse-Reward Multi-Agent Reinforcement Learning

1. Problem Statement

Multi-Agent Reinforcement Learning (MARL) faces a critical bottleneck when dealing with sparse rewards. In sparse environments, agents receive feedback infrequently, which substantially slows learning. This issue is exacerbated in MARL compared to Single-Agent RL (SARL) due to non-stationarity and co-adaptation: as agents simultaneously update their policies, the environment appears non-stationary from any single agent's perspective.

While reward shaping (augmenting or replacing sparse signals with dense ones) is a common solution, it is particularly risky in MARL. In SARL, improper shaping can alter the set of optimal policies. In MARL, the risk is more severe: shaping can alter the set of Nash equilibria, fundamentally changing the game's strategic structure and leading to suboptimal or unintended collective behaviors. Existing automatic reward shaping methods, such as Self-Supervised Online Reward Shaping (SORS), rely on global trajectory ranking guarantees that do not directly transfer to MARL due to the lack of a global optimality notion and the interdependence of agent policies.

2. Methodology: ARMS Framework

The authors propose ARMS (Automatic Reward-shaping in Multi-agent Systems), a self-supervised framework designed to learn dense shaping signals from sparse environmental rewards while preserving the game-theoretic equilibrium structure.

Theoretical Foundation: Conditional Best-Response Invariance

The core theoretical contribution is a reformulation of policy invariance for multi-agent settings.

Single-Agent Limitation: Standard trajectory-ranking guarantees (e.g., from SORS) assume that if two reward functions induce the same total ordering over all trajectory pairs, they yield the same optimal policies. This relies on a global notion of optimality absent in MARL.
ARMS Formulation: The authors prove that if the policies of all other agents are fixed, an individual agent faces an induced single-agent environment. Under this conditional best-response view, if a shaping reward function is total-order equivalent to the original reward over the trajectories of this induced environment, it preserves the agent's set of best responses.
Equilibrium Preservation: By extension, if this condition holds for all agents, the set of Nash equilibria remains invariant. This provides a game-theoretic justification for using trajectory ranking in MARL without altering the underlying strategic solution.

Algorithmic Implementation

ARMS operates in an alternating loop between two phases (Algorithm 1):

Reinforcement Learning Phase: Agents interact with the environment using their current shaped rewards ( $R_{\theta}$ ) and a base MARL algorithm (e.g., IPPO or MAPPO). They collect joint trajectories, which are stored in a buffer.
Reward-Shaping Phase: The system samples pairs of trajectory segments from the buffer. For each agent, a parameterized reward network ( $F_\theta$ $F_{θ}$ ) is updated to minimize a binary classification loss. This loss enforces that the shaped reward's preference ordering over trajectory pairs matches the preference ordering induced by the original sparse environmental reward.
- Shared Parameters: To improve efficiency and scalability, all agents share the same reward-shaping parameters ( $\theta$ ). The network is trained on agent-wise trajectory orderings, aggregating signals across all agents.
- Input: The reward network takes local observations and actions as input, modeling a decentralized information flow.

3. Key Contributions

The paper outlines three primary contributions:

Multi-Agent Policy Invariance Formulation: The authors establish that total-order-equivalent shaping rewards preserve each agent's best-response set under fixed opponent policies, thereby preserving the set of Nash equilibria. This addresses the theoretical gap preventing the direct application of single-agent trajectory-ranking methods to MARL.
The ARMS Framework: A practical, self-supervised online reward learning method that leverages trajectory ranking while accounting for multi-agent non-stationarity. It learns a shared shaping signal that is agnostic to the underlying RL backbone.
Empirical Analysis of MARL Dynamics: The study identifies a specific MARL failure mode where limited exploration coupled with policy-reward dynamics induces oscillatory behavior and convergence to suboptimal cyclic patterns. The authors demonstrate that increasing exploration mitigates this effect.

4. Experimental Results

Experiments were conducted in a partially observable multi-agent pathfinding domain (Lifelong Pathfinding) using the POGEMA environment. The authors evaluated ARMS against a baseline with no shaping and Potential-Based Reward Shaping (PBRS) using IPPO and MAPPO backbones.

Sampling Efficiency: ARMS significantly outperformed baselines in terms of cumulative original reward, particularly as the number of agents increased (8, 16, 32) and as rewards became sparser (delays of 10, 20, 30 timesteps).
Generalization: Policies trained with ARMS generalized effectively to 50 unseen maps (including random and maze-like environments), accumulating higher rewards and exhibiting better coordination than baselines.
Coordination and Collisions: ARMS resulted in substantially fewer normalized collisions compared to PBRS and no shaping, indicating that the learned signal promotes coordinated behavior rather than just reward maximization.
Exploration Sensitivity: A critical finding was that ARMS could fail to converge to optimal policies if exploration was too low (e.g., exploration coefficient $\alpha = 0.023$ ), leading to agents "hacking" the reward via oscillatory behaviors. Increasing the exploration coefficient (e.g., $\alpha = 0.35$ ) stabilized learning and restored high throughput.

5. Significance and Claims

The paper claims that ARMS is the first automatic reward shaping framework for MARL whose design is explicitly motivated by a game-theoretic equilibrium-preservation result.

Unlike previous methods that risk altering the strategic structure of multi-agent problems, ARMS provides a theoretical guarantee that the set of Nash equilibria is preserved under specific conditions. The authors emphasize that while trajectory ranking is effective in SARL, it requires a fundamentally different formulation in MARL to handle non-stationarity and the lack of global optimality. The framework demonstrates that automatic, dense reward signals can be learned to accelerate MARL in sparse environments without compromising the integrity of the game's solution space, provided that sufficient exploration is maintained to prevent suboptimal feedback loops.