Agent-Capability-Aware Advantage Estimation

• Agent-capability-aware advantage estimation is a framework that calibrates policy gradients using capability-sensitive baselines to achieve unbiased value estimation.
• It employs reweighting and importance correction strategies to address agent heterogeneity and enhance cross-agent credit assignment.
• The approach enables robust, stable learning in collaborative environments by adapting learning signals to each agent's unique capabilities.

Agent-capability-aware advantage estimation is a framework within multi-agent and heterogeneous-agent reinforcement learning (MARL/HARL) for computing policy gradients that correct for variations in individual agent capabilities while maximizing effective sample sharing. This methodology explicitly addresses heterogeneity in agent skill, architecture, or policy distribution, yielding unbiased and stable optimization in collaborative training regimes that feature shared but independently generated rollouts. Key advances include the capability-sensitive construction of advantage baselines, reweighting mechanisms that modulate cross-agent credit assignment, rigorous importance correction for policy mismatches, and theoretical guarantees on estimator unbiasedness and policy improvement.

1. Problem Setting and Motivation

Heterogeneous-agent learning arises when a collection of $n$ agents, potentially with distinct policy classes, parameterizations, or modeling assumptions, seek to maximize returns from a common reward source. For agent $k \in \{1, \dots, n\}$, let policy $\pi_k$ generate candidate responses $Y_k(q)=\{y_{k,1}, \dots, y_{k,G}\}$ per query $q$. All agent-generated responses are pooled, with the shared reward function $R(\cdot)$ assigning a scalar reward to each $y$. Heterogeneity is manifest in factors such as model size, initialization, or architecture.

Traditional on-policy methods compute advantages independently and ignore inter-agent calibration, leading to misaligned baselines and sample inefficiency in the presence of differing agent abilities. Agent-capability-aware advantage estimation remedies this by (a) jointly utilizing all agents' rollouts, and (b) adapting baselines and learning signals in a capability-sensitive manner, which is central to Heterogeneous Agent Collaborative Policy Optimization (HACPO) (Zhang et al., 3 Mar 2026).

2. Construction of the Agent-Capability-Aware Advantage

In the absence of cross-agent correction, the typical single-agent estimator for policy gradient methods such as PPO/GSPO takes the form

$$A_{t,i}^{(k)}(\text{single}) = \frac{R(y_{t,i}^{(k)}) - \frac{1}{G}\sum_{i} R(y_{t,i}^{(k)})}{\sigma_t^{(k)}}$$

where $\sigma_t^{(k)}$ is the empirical standard deviation over the batch for agent $k$. This approach neglects the comparative strength of peer agents and fails under capability disparity.

HACPO introduces a centralized, capability-weighted baseline for agent $k$: $$\hat{\mu}_t^{(k)} = \frac{1}{nG} \sum_{j=1}^n \sum_{i=1}^G \omega_t^{(k,j)} R(y_{t,i}^{(j)})$$ with weighting

$$\omega_t^{(k,j)} = \frac{\hat{P}_t^{(k)}}{\hat{P}_t^{(j)}}$$

where $\hat{P}_t^{(k)}$ is a sliding-window mean of recent batch rewards for agent $k$. The global empirical variance $\sigma_{t,\text{joint}}$ is computed over the joint rollout pool. The agent-capability-aware advantage is then

$$A_{t,i}^{(k)} = \frac{R(y_{t,i}^{(k)}) - \hat{\mu}_t^{(k)}}{\sigma_{t,\text{joint}}}$$

This baseline adapts to both the scale and relative strength of each agent’s output, ensuring that stronger and weaker models calibrate their value functions appropriately (Zhang et al., 3 Mar 2026).

3. Guarantees and Statistical Properties

Formal analysis demonstrates that under weak independence assumptions for the capability weights (specifically, independence from the stochasticity of the current batch), the modified baseline $\hat{\mu}_t^{(k)}$ is an unbiased estimator of the on-policy value $\mathbb{E}_{y \sim \pi_k}[R(y)]$. This implies that the centered advantage estimator $A_{t,i}^{(k)}$ is zero-mean. Theorems in HACPO establish:

• Unbiasedness: $$\mathbb{E}[\hat{\mu}_t^{(k)}] = \mathbb{E}_{y \sim \pi_k(\cdot|q_t)}[R(y)]$$.
• Gradient Alignment: The cross-agent/heterogeneous gradient update direction remains positively aligned with the agent’s own return-maximizing direction, i.e., $$\left\langle \nabla_\theta J_{\text{hete}}^{(k)}, \nabla_\theta J_{\text{homo}}^{(k)} \right\rangle > 0$$.

These properties guarantee that sharing and reweighting rollouts across heterogeneous agents does not bias the learning update and that learning signals remain correctly oriented toward improvement of the individual agent's policy (Zhang et al., 3 Mar 2026).

4. Discrepancy Mitigation Mechanisms

To robustly enable cross-agent rollout reuse, HACPO deploys four technical mechanisms:

1. Agent-Capability-Aware Advantage Estimation: Joint reward statistics and capability-weighted baselines (as above).
2. Model Capabilities Discrepancy Coefficient: Cross-agent samples are scaled by capability discrepancy,

$$\tilde{A}_{t,i}^{(k)} = \begin{cases} A_{t,i}^{(k)}, & j=k \ \omega_t^{(j,k)} A_{t,i}^{(j)}, & j \neq k \end{cases}$$

with $\omega_t^{(j,k)} = \hat{P}_t^{(j)}/\hat{P}_t^{(k)}$, amplifying stronger agents' signals and attenuating weaker ones.

1. Exponential Importance Sampling: For off-policy/cross-distribution samples, use sequence-level importance ratios

$$s_{t,i}^{(k,j)} = \left(\frac{\pi_k(y_{t,i}^{(j)})}{\pi_j(y_{t,i}^{(j)})}\right)^{1/|y|}$$

and reweight further by a function of $\mathrm{sg}(s)$ and a stability parameter $\alpha$.

2. Stepwise Clipping: Apply an asymmetric, tightening lower-bound clip to $s_{t,i}^{(k,j)}$ to ensure that only downweighting (never amplification) of cross-agent signals is permitted, with the lower bound increased at each mini-batch.

The aggregate update is given by

$$J^{(k)} = \underbrace{\frac{1}{G} \sum_{i=1}^G \min(s_{t,i}^{(k,k)}, \mathrm{clip}(s_{t,i}^{(k,k)})) A_{t,i}^{(k)}}_{\text{on-policy}} + \underbrace{\frac{1}{G} \sum_{j \neq k} \sum_{i=1}^G \mathrm{clip}(\tilde{s}_{t,i}^{(k,j)}) \omega_t^{(j,k)} A_{t,i}^{(j)}}_{\text{cross-agent}}$$

correcting for both capability and policy discrepancies (Zhang et al., 3 Mar 2026).

5. Illustrative Example: Two-Agent Rollout Sharing

Consider agents A (strong) and B (weak), each generating $G=4$ responses with rewards:

• A: $\{1,0,1,1\}$ ($\hat{P}_t^A = 0.75$), B: $\{0,1,0,0\}$ ($\hat{P}_t^B = 0.25$).

Capability coefficients: $\omega^{(A,B)}=3.0$, $\omega^{(B,A)}=1/3$. The mixed baseline for A: $$\hat{\mu}_t^{(A)} = \frac{1}{8}[3 \cdot 1.0 + 3 \cdot 3.0] = 0.75$$ with a joint empirical $\sigma=0.43$. Each advantage for A (for its own responses) is thus: $$A_{\text{A}}(a_1) = (1-0.75)/0.43 \approx +0.58,\quad A_{\text{A}}(a_2) = (0-0.75)/0.43 \approx -1.74, \text{ etc.}$$ For B’s samples, before policy update, the raw advantage is further scaled by $\omega^{(B,A)}=1/3$. Stepwise importance weighting and clipping are applied to all cross-agent samples.

This allows performant, unbiased, and stable learning where stronger models benefit from occasional correct outputs by weaker agents, while weak agents are guided and safely regularized by the strong agent’s solutions (Zhang et al., 3 Mar 2026).

6. Extensions: Generalized Per-Agent Advantage Estimation

In other MARL settings, explicit agent capability factors $\kappa_i(s) \in [0,1]$ (e.g., skill, actuator strength, or confidence) may be incorporated multiplicatively into per-agent value and advantage estimation. In the Generalized Per-Agent Advantage Estimator (GPAE), the per-agent value is modified as

$$V^i(s, \boldsymbol{a}^{-i}) \leftarrow \kappa_i(s) \mathbb{E}_{a^i \sim \pi^i}[ Q(s, a^i, \boldsymbol{a}^{-i}) ]$$

and per-agent trace weights for off-policy corrections are contracted for agents with low capability: $$c_t^{i,\text{DT\_cap}} = \lambda \min(1, \rho_t^i \min(\eta, \boldsymbol{\rho}_t^{-i})) \times (1-\alpha(1-\kappa_i(s_t)))$$ where $\alpha \in [0,1]$ regulates the downweighting effect. This capability-sensitized mechanism preserves credit assignment fidelity and contraction properties in heterogeneous teams (Kim et al., 3 Mar 2026).

7. Impact and Significance

Agent-capability-aware advantage estimation enables robust, sample-efficient, and theoretically grounded learning in heterogeneous agent populations. By reconciling policy improvement steps across agents with differing abilities, it maximizes the utility of available rollouts and ensures unbiased value learning. These mechanisms are central to recent advances in collaborative RL, including HACRL and the GPAE framework, and they represent a marked improvement over isolated optimization and naive cross-agent rollout reuse. They also provide a foundation for principled capability adaptation in future multi-agent learning systems (Zhang et al., 3 Mar 2026, Kim et al., 3 Mar 2026).