🤖 machine learning

CoFi-PGMA: Counterfactual Policy Gradients under Filtered Feedback for Multi-Agent LLMs

CoFi-PGMA is a unified framework designed to optimize multi-agent LLM systems by using counterfactual policy gradients to correct for "filtered feedback," such as selection-gated routing or obscured shared rewards in collaborative settings.

Original authors: Stela Tong, Elai Ben-Gal

Published 2026-04-28

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Stela Tong, Elai Ben-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 part of a high-stakes professional kitchen. To serve a perfect meal, you have two different ways of working, but both have a major problem: the feedback you get from the head chef is confusing or incomplete.

This paper, CoFi-PGMA, is a new mathematical "instruction manual" designed to help AI agents learn how to work better in these two messy kitchen scenarios.

Scenario 1: The "Winner-Takes-All" Audition (Competitive Routing)

Imagine three chefs (AI agents) each cook a different dish: one makes pasta, one makes steak, and one makes salad. The Head Chef (the Router) tastes them all, but only picks one to serve to the customer.

The Problem: If the Head Chef picks the steak, the pasta chef and the salad chef get zero feedback. They don't know if their dish was delicious or terrible; they just know they weren't picked. If they only learn from the "wins," they might all start making steak just to stay in the game, even if the customer actually wanted salad. This is called "Selection-Gated Feedback."
The CoFi-PGMA Solution: Instead of just telling the winner "Good job," this method uses a "Counterfactual" approach. It asks: "If we hadn't picked the steak, how much would the customer have liked the pasta?" It uses math to estimate the quality of the unchosen dishes. This way, the pasta chef learns how to improve even when they aren't the star of the show.

Scenario 2: The "Group Project" Mess (Collaborative Generation)

Now imagine a different setup: three chefs work together on a single, complex Beef Wellington. One prepares the meat, one makes the pastry, and one makes the sauce. At the end, the Head Chef tastes the final dish and gives the whole team a single score: "8 out of 10."

The Problem: This is the "Group Project" nightmare. If the dish was an 8, was it because the meat was perfect, or was the pastry so good it saved a mediocre meat? The chefs don't know their individual value. They might "free-ride" (do nothing and take the credit) or get discouraged because a teammate's mistake dragged their score down. This is called "Credit Assignment."
The CoFi-PGMA Solution: This method uses a "Leave-One-Out" strategy. It mathematically calculates: "What would the score have been if the pastry chef had just used store-bought dough instead of making it from scratch?" By measuring that specific difference, the system can pinpoint exactly how much each person contributed to the final success.

Why does this matter for AI?

Most AI today is trained like a single student studying alone. But the future of AI is Multi-Agent—groups of specialized AIs talking to each other, competing, or collaborating to solve complex problems (like coding an app or solving math).

If we train these "teams" using old methods, they become biased, lazy, or confused. CoFi-PGMA provides the "fair grading system" they need.

The Result: In the researchers' tests (using math problems), the AI agents trained with this new method were significantly smarter and more specialized than the ones using the old, "unfair" grading systems. They didn't just get the right answer; they actually learned how to play their specific roles better.

Technical Summary: CoFi-PGMA

CoFi-PGMA: Counterfactual Policy Gradients under Filtered Feedback for Multi-Agent LLMs

1. Problem Statement

As Large Language Model (LLM) deployments shift from monolithic policies to multi-agent architectures, standard Reinforcement Learning from Human Feedback (RLHF) objectives become misspecified. The authors identify that in multi-agent systems, the learning signal is "filtered" by the system mechanism in two primary ways:

Competitive Routing (Selection-Gated Feedback): A controller selects one expert from a pool of candidates. The system only observes the reward for the selected response, leaving the rewards of unselected candidates unobserved. Naive RL leads to statistical bias because the agent cannot learn from the "near misses" or the quality of unselected actions.
Collaborative Generation (Credit Assignment Problem): Multiple agents (e.g., planner, solver, critic) work together to produce a single output. The system provides a single shared reward. This obscures individual contributions, leading to high variance and "free-riding," where agents receive positive updates regardless of their specific marginal impact.

2. Methodology

The paper introduces CoFi-PGMA (Counterfactual Policy Gradients under Filtered Feedback for Multi-Agent LLMs), a unified framework that replaces naive reward signals with a Counterfactual Marginal Contribution ( $\Delta_{i,t}$ ). This signal measures the difference between the observed system reward and the reward that would have been obtained if agent $i$ had not participated (or had provided a baseline response).

The framework provides two distinct mathematical solutions based on the mechanism:

For Routing (Doubly Robust Estimator): To handle the missing data of unselected candidates, the authors use a Doubly Robust (DR) estimator. It combines a learned reward predictor (to estimate the value of unselected actions) with importance sampling (to correct the estimate of the selected action). This ensures that if either the router's probabilities or the reward model is accurate, the gradient remains unbiased.
For Collaboration (Leave-One-Out Estimator): To solve the credit assignment problem, the authors use a Leave-One-Out (LOO) approach. They construct a counterfactual trajectory by replacing an agent's contribution with a baseline (e.g., a null message) and resampling downstream agent responses. This isolates the causal effect of the specific agent on the final outcome.

3. Key Contributions

Theoretical Framework: A unified mathematical formulation for RLHF in multi-agent systems that accounts for both routing and collaboration.
Mathematical Proofs:
- Theorem 1 proves that counterfactual marginal contribution provides a valid policy-gradient signal for autoregressive multi-agent policies.
- Proposition 3 proves that the counterfactual approach reduces gradient variance by filtering out "action-independent noise" (fluctuations caused by other agents).
Practical Estimators: Derivation of the DR estimator for selection-gated feedback and the LOO estimator for shared rewards.
Algorithmic Integration: A training recipe that integrates these counterfactual estimators with modern optimization methods like GRPO (Group Relative Policy Optimization).

4. Results

The authors evaluated the framework on the GSM8K reasoning dataset using a competitive routing setup with three specialized LoRA adapters (Direct, Equation-first, and Final-answer-only) on a Llama-3.2-3B backbone.

Accuracy Improvement: The proposed DR-GRPO achieved 41% exact-match accuracy, significantly outperforming the "Winner-take-all" baseline (37%) and the frozen single-agent reference (35.2%).
Improved Routing Quality: DR-GRPO showed lower "regret" (the gap between the router's choice and the best possible candidate) and better reward-predictor calibration (lower Brier score).
Functional Specialization: The method encouraged higher "specialization scores," meaning agents actually learned to perform their assigned roles (e.g., the equation agent focused on math steps) rather than all agents converging to the same generic behavior.

5. Significance

This work is highly significant for the industrial evolution of LLMs. As AI systems move toward Mixture-of-Experts (MoE), agentic workflows, and multi-step reasoning pipelines, the "single-policy" assumption of standard RLHF fails. CoFi-PGMA provides a principled way to train these complex systems, ensuring that each component is optimized for its actual contribution to the final system performance, thereby improving both the efficiency and the specialized intelligence of multi-agent AI.