AIIR-MIX: Adaptive Intrinsic Reward Mixing

• AIIR-MIX is a multi-agent reinforcement learning architecture that integrates attention-driven intrinsic rewards with nonlinear, extrinsic-conditioned mixing for adaptive individual credit assignment.
• It employs a softmax-normalized attention mechanism to compute personalized intrinsic rewards and utilizes a hypernetwork-based mixing module to dynamically blend these with extrinsic signals.
• Empirical evaluations on StarCraft II micro-management benchmarks demonstrate significant performance gains, achieving win rates up to 96% and faster convergence than baseline methods.

AIIR-MIX is a multi-agent reinforcement learning (MARL) architecture designed to address the challenge of individual credit assignment and dynamic reward structuring in cooperative environments. Specifically, it integrates an attention-based intrinsic reward network with a non-linear, extrinsic-conditioned reward mixing module, enabling personalized per-agent feedback that adapts to team dynamics and environmental conditions. AIIR-MIX demonstrates empirical superiority over value-decomposition, policy-gradient, and other intrinsic reward MARL baselines on StarCraft II micro-management benchmarks (Li et al., 2023).

1. Architectural Composition

AIIR-MIX is composed of two primary, end-to-end differentiable modules:

• Attention Individual Intrinsic Reward (AIIR) Network: Computes for each agent $i$ a personalized intrinsic reward $r_i^{\rm in}$ that estimates the unique, context-dependent contribution of the agent to the team objective.
• Nonlinear Extrinsic-Conditioned Mixing Network (MIX): Combines all agents’ intrinsic rewards with the shared extrinsic reward $r^{\rm ex}$ into per-agent total rewards $r_i^{\rm total}$ using a hypernetwork whose parameters are adaptive functions of $r^{\rm ex}$.

The per-agent input at time $t$ consists of local state $s_i^t$ and last action $u_i^{t-1}$. The two-layer fully connected (FC) extractor produces a local embedding $\mathbf{v}_i^t$. Pairwise similarity $A^t_{ij}$ is calculated via cosine similarity, then normalized by softmax to obtain attention weights $r_i^{\rm in}$0. Agent $r_i^{\rm in}$1’s intrinsic reward embedding is $r_i^{\rm in}$2. A single FC projects $r_i^{\rm in}$3 to $r_i^{\rm in}$4.

The mixing network, parameterized by a hypernetwork $r_i^{\rm in}$5, produces two sets of weights $r_i^{\rm in}$6 conditioned on $r_i^{\rm in}$7, enabling the dynamic, nonlinear mixing:

$r_i^{\rm in}$8

$r_i^{\rm in}$9

This mechanism allows the network to up- or down-weight intrinsic versus extrinsic feedback on a per-timestep basis.

2. Learning Objective and Optimization

AIIR-MIX is trained with centralized training and decentralized execution (CTDE). The learning objective comprises several value and policy components:

• Extrinsic Critic ($r^{\rm ex}$0): Trained to minimize the mean squared Bellman residual based on the team’s extrinsic reward trace.
• Per-Agent Total Critic ($r^{\rm ex}$1): Trained on per-agent total rewards $r^{\rm ex}$2 with similar temporal-difference targets.
• Policy Gradient Update: Each agent’s policy $r^{\rm ex}$3 is updated via advantage estimation using extrinsic returns:

$r^{\rm ex}$4

where $r^{\rm ex}$5 is constructed from extrinsic returns.

No auxiliary regularization or penalty is imposed on the mixing network or attention weights; all network parameters are trained end-to-end with temporal-difference (TD) and policy gradient objectives.

3. Attention-based Credit Assignment and Dynamic Mixing

The core innovation of AIIR-MIX is its per-timestep, contribution-sensitive attention mechanism for intrinsic reward, coupled with extrinsic-conditioned reward mixing.

• Attention for Individual Credit: The softmax-normalized, similarity-based attention $r^{\rm ex}$6 lets each agent weigh the influence of other agents’ behaviors, actively targeting teammates whose actions most benefit team performance. The resulting intrinsic reward fluctuates responsively with in-episode events (e.g., peaking when agents coordinate attack or retreat on low HP).
• Dynamic Nonlinear Mixing: The mixing network’s weights and biases are adaptive functions of current extrinsic reward, enabling context-sensitive prioritization of intrinsic curiosity versus task progress. This is a departure from previous methods, which use a fixed linear combination.

This design achieves fine-grained, temporally-localized reward assignment, addressing deleterious credit diffusion and misattribution prevalent in scalar or static-mixed schemes.

4. Experimental Evaluation

AIIR-MIX was empirically validated on the SMAC (StarCraft Multi-Agent Challenge) micro-battle suite:

Map	COMA	QTRAN	QMIX	LIIR	AIIR-MIX
8m	45%	80%	75%	85%	96%
MMM	30%	90%	88%	65%	92%
2s3z	25%	70%	68%	55%	75%
3s5z	20%	65%	80%	50%	78%

All results are average test win rates over five random seeds. AIIR-MIX achieves the highest performance in every scenario and typically converges more rapidly than all baselines, including value-decomposition (QMIX, QTRAN), counterfactual policy gradient (COMA), and the linear-intrinsic reward model LIIR.

Ablation studies further separate the effect of the attention mechanism and nonlinear mixing. Substituting linear reward generators (RMIX) or reverting to linear mixing (LinearMix) significantly reduces win rates by up to 20 percentage points, demonstrating the unique advantage conferred by the full architecture (Li et al., 2023).

5. Interpretability, Behavioral Insights, and Limitations

Analysis of AIIR-MIX-trained agents indicates interpretability at both the attention and reward stages. Attention weights $r^{\rm ex}$7 noticeably spike when agents require close coordination (e.g., during joint attacks), and the intrinsic reward signal adapts to individual situations (e.g., falling health causing withdrawal). This confirms that the architecture dynamically discovers and reinforces team synergies and individual initiative.

Model complexity is increased due to the additional attention computation and reward mixing hypernetwork, leading to slower per-timestep training compared to linear baselines. The method relies on sufficiently informative extrinsic (team) rewards for effective training. The original work suggests future directions such as multi-head attention, more expressive hypernetworks, and exploration of sparse or delayed extrinsic signal regimes.

6. Comparison with Related MARL Approaches

In contrast to preceding approaches that sum environment- and hand-tuned intrinsic rewards or use static decomposition, AIIR-MIX uniquely combines:

• Learnable, per-agent, attention-driven intrinsic rewards for fine-grained credit
• Nonlinear, hypernetwork-conditioned reward mixing adaptive to environmental feedback

This fusion enables both efficient learning and robust credit assignment in cooperative MARL without explicit access to ground-truth contribution signals or reliance on manually weighted intrinsic terms. No explicit regularization on attention or mixing weights is required, simplifying implementation.

7. Implementation Summary and Practical Considerations

The architecture is realized with compact two-layer FC modules for embedding/feature extraction, cosine similarity for attention computation, and a parallel hypernetwork for reward mixing. Training leverages CTDE, five random seeds for statistical stability, and frequent evaluation. No additional regularization is employed.

Ablation and behavioral visualization methodologies enable diagnostic insight into the functioning of both intrinsic reward generation and its downstream effect on group behavior.

---

AIIR-MIX constitutes an integrated, end-to-end approach for adaptive, agent-centric credit assignment in cooperative multi-agent reinforcement learning, underpinned by attention-guided intrinsic reward estimation and extrinsic-conditioned, nonlinear reward mixing (Li et al., 2023).