COALA-PG: Learning-Aware Policy Gradients

• COALA-PG is a learning-aware policy gradient framework for multi-agent reinforcement learning that decomposes global gradients over asynchronous coagent networks.
• It leverages execution path semantics and adaptive credit assignment to boost sample efficiency and coordination in hierarchical, nonstationary environments.
• The framework offers theoretical guarantees for convergence and robust performance in diverse applications, from StarCraft II micromanagement to continuous control.

Co-agent Learning-Aware Policy Gradients (COALA-PG) represent a principled framework for multi-agent reinforcement learning (MARL) in which the policy gradient estimators and credit assignment mechanisms are explicitly constructed to be “learning-aware.” That is, COALA-PG and its variants ensure that policy updates in complex agent networks take into account the learning dynamics and parameterization structure of other participating agents (coagents) without assuming static environments or policy stationarity. The framework is most closely associated with coagent networks—arbitrary, potentially asynchronous, modular arrangements of stochastic agents whose joint behavior is coordinated through flexible execution rules. COALA-PG algorithms decompose the overall policy gradient and exploit execution-path semantics, parameter sharing, and asynchronous learning to achieve efficient, scalable, and robust multi-agent learning.

1. Coagent Network Foundations and Policy Gradient Decomposition

COALA-PG is grounded in the theory of coagent networks, which model a reinforcement learning agent as a composite of multiple stochastic coagents. Each coagent $\pi_o$ operates on its local state space $S_o$ and selects sub-actions that may be internal (contributing to downstream coagents) or external (primitive actions affecting the environment). The overall agent policy $\Pi$ thus comprises execution paths—ordered sequences of coagent activations and outputs:

$$P = ((x_{o_1}, u_{o_1}), (x_{o_2}, u_{o_2}), \ldots, (x_{o_k}, u_{o_k}))$$

A central theoretical result is that the global policy gradient can be written as a sum over coagent gradients:

$$\nabla_{\theta} J_{\Pi} = \sum_{x_0 \in S_{\text{init}}} d(x_0) \sum_{o} \sum_{x, x_o \in S_o} d(x_o, x|x_0) \sum_{u \in A_{x_o}} \frac{d\pi_o}{d\theta}(u|x_o) Q_{\pi_o}(x_o, u)$$

This decomposition is preserved regardless of parameter sharing schemes, asynchronous execution schedules, or hierarchical policy structures (Zini et al., 2020, Ullah et al., 2020, Kostas et al., 2023).

2. Asynchronous and Hierarchical Policy Gradient Extensions

COALA-PG generalizes naturally to asynchronous coagent networks, where each node executes and learns at its own rate. Through the augmentation of the state and action space, the asynchronous network can be mapped to an equivalent synchronous acyclic network, preserving the dynamics and expected return. The asynchronous policy gradient is:

$$\Delta_i (\theta_i) = \mathbb{E} \left[ \sum_t E_{t}^{i} \gamma^t G_t \frac{\partial}{\partial \theta_i} \log \pi_i (X_t, U_t, \theta_i) \right]$$

where $E_{t}^{i}$ is the binary indicator for coagent $i$'s execution at timestep $t$. The summation over coagent-local gradients recovers the full policy gradient for the network.

COALA-PG thus enables direct derivation of gradient rules for hierarchical architectures, including option-critic models and multi-level temporal abstraction. Policy gradient updates for termination functions, intra-option policies, and hierarchical advisors are unified under this framework, eliminating the need for custom derivations typically required in semi-Markov or SMDP settings (Kostas et al., 2019, Zini et al., 2020, Kostas et al., 2023, Jin et al., 2022).

3. Learning-Aware Gradient Conditioning and Multi-Agent Credit Assignment

A primary challenge in MARL is credit assignment—mapping joint outcomes to informative learning signals for individual agents. Naïve approaches either share the global advantage signal, risking suboptimal convergence, or use local returns, risking loss of coordination. COALA-PG improves credit assignment by making gradients “learning-aware,” i.e., conditioning policy updates on the actions, states, and adaptive policy dynamics of other agents.

This includes:

• Conditioning critic gradients via hypernetwork mixing with latent state representations, improving signal richness and implicit credit assignment (see (Zhou et al., 2020)).
• Adaptive entropy regularization, dynamically rescaling entropy gradients to ensure consistent exploration and avoid premature convergence (Zhou et al., 2020).
• Difference rewards techniques, in which each agent’s policy gradient is scaled by the difference between realized team reward and the expected reward from that agent’s average policy (Castellini et al., 2020).
• Polarization gradients, which reshape the advantage landscape to suppress suboptimal joint actions and eliminate centralized-decentralized mismatch (Chen et al., 2022).
• Coalition-based credit assignment frameworks such as CORA, which decompose advantages according to the contributions of agent coalitions, ensuring rational and fair distribution of learning incentives (Ji et al., 3 Jun 2025).
• Utilizing optimal baseline techniques to reduce variance in multi-agent gradient estimation, further stabilizing and harmonizing policy updates (Kuba et al., 2021).

4. Execution Path Semantics and Modular Algorithm Design

A distinguishing feature of COALA-PG is its reliance on execution path semantics to structure learning, exploration, and credit assignment. An execution path identifies the specific sequence of coagent activations that produce a primitive action in the environment. By updating only along executed paths (i.e., updating parameters only for those coagents called in a timestep), COALA-PG increases runtime and sample efficiency, particularly in hierarchical or asynchronous settings. Modular design is achieved: coagents can operate with unique or shared parameters and can be updated in parallel without violating convergence guarantees.

Recent refinements address termination bias in hierarchical policies (option-critic), improve learning efficiency in nonstationary environments (see experiments in (Zini et al., 2020)), and permit integration of non-differentiable components (e.g., pure sampling-based modules) due to the independence of coagent learning rules (Kostas et al., 2023).

5. Applications: Efficient Coordination and Emergent Cooperation

Empirical validation demonstrates that COALA-PG frameworks outperform conventional multi-agent policy gradient methods in coordinated, hierarchical, and temporally-extended tasks. Applications include:

• Cooperative navigation and socio-economic matrix games, illustrating suppression of relative overgeneralization and improved convergence to optimal joint actions (Wei et al., 2018, Chen et al., 2022).
• StarCraft II micromanagement scenarios, where explicit credit assignment and polarization yield faster convergence and higher cumulative reward relative to baselines (Chen et al., 2022, Jin et al., 2022).
• Social dilemma environments, including iterated prisoner’s dilemma and CleanUp-like tasks, where learning-aware agents shape the dynamics of naive learners and achieve stable cooperation (Meulemans et al., 24 Oct 2024, Li et al., 19 Feb 2024).
• Large-scale agent populations and continuous control domains (MuJoCo), where generalized coagent networks scale efficiently, learning effective policies even with high-dimensional state/action spaces (Kostas et al., 2023, Ji et al., 3 Jun 2025).

Sample efficiency and coordination are improved via asynchronous option-based joint policies and modular generative adversarial learning between agents and hierarchical advisors (Lyu et al., 2022, Jin et al., 2022).

6. Theoretical Guarantees and Convergence Properties

COALA-PG frameworks are supported by rigorous theoretical proofs of unbiasedness, convergence, and, under standard step-size conditions, attainment of locally optimal policies. The general asynchronous coagent policy gradient theorem ensures that local updates aggregate to a global policy gradient, even under shared parameters or hierarchical abstraction. Coalition-based allocation ensures rationality and stability of credit signals (the core assignment property) (Ji et al., 3 Jun 2025). Learning-aware meta-MARL, as formulated in (Kim et al., 2020, Meulemans et al., 24 Oct 2024), integrates own-learning and peer-learning gradient terms, capturing nonstationary adaptation and influencing opponent learning trajectories. Altruistic gradient adjustment (AgA) aligns individual gradients toward stable collective optima with explicit second-order theoretical control (Li et al., 19 Feb 2024).

7. Implications, Limitations, and Future Directions

COALA-PG establishes a unified formalism for cooperative MARL that directly addresses the deficiencies of vanilla policy gradient estimators—overgeneralization, centralized-decentralized mismatch, poorly conditioned learning signals, and lack of scaling. While modularity enables heterogeneity, high variance in decentralized updates can emerge for large networks, sometimes necessitating hybridization with backpropagation or critic-based variance reduction (Kostas et al., 2023, Kuba et al., 2021). Random coalition sampling provides practical tractability but can introduce slack in the rationality guarantees (Ji et al., 3 Jun 2025). The framework paves the way for principled integration with game-theoretic credit assignment, differentiated meta-learning, and deep sequence models that exploit learning traces. Open questions remain on optimal coalition sampling, interplay with differentiable game dynamic modeling, and robustification in highly adversarial or partially observable environments.

---

COALA-PG constitutes a rigorous, general, and flexible foundation for multi-agent policy gradient algorithms that are both learning-aware and credit-aware, enabling robust coordination, hierarchical decision making, and scalable learning across diverse MARL domains.