HATRPO-W: Adaptive Trust Region Optimization

• The paper introduces HATRPO-W, a multi-agent trust-region policy optimization method that reallocates KL budgets based on per-agent natural-gradient geometry to boost learning efficiency.
• It employs a KKT-based approach to solve a global budget allocation problem, offering adaptive trust region sizes for heterogeneous agents in cooperative MARL settings.
• Empirical results show that HATRPO-W delivers lower variance, faster convergence, and higher returns compared to uniform KL threshold methods.

HATRPO-W is a multi-agent trust-region policy optimization method for cooperative multi-agent reinforcement learning (MARL) that extends Heterogeneous-Agent Trust Region Policy Optimization (HATRPO) by allocating the Kullback-Leibler (KL) divergence threshold across agents under a single global KL budget. Introduced in “Multi-Agent Trust Region Policy Optimisation: A Joint Constraint Approach” (Shek et al., 14 Aug 2025), it is a Karush-Kuhn-Tucker-based (KKT-based) method that optimizes threshold assignment under global KL constraints. Its motivating premise is that assigning each agent the same KL threshold can be suboptimal in heterogeneous settings, producing slow or locally optimal updates; HATRPO-W instead redistributes trust-region budgets according to per-agent natural-gradient geometry, while preserving the sequential update structure of HATRPO (Shek et al., 14 Aug 2025).

1. Research lineage and nomenclature

HATRPO-W is not the original HATRPO. The 2021 work “Trust Region Policy Optimisation in Multi-Agent Reinforcement Learning” introduced HATRPO and HAPPO, established the multi-agent advantage decomposition lemma and the sequential policy update scheme, and justified the monotonic improvement property of HATRPO/HAPPO; that work does not define a variant named “HATRPO-W” (Kuba et al., 2021). The 2023 work “Heterogeneous-Agent Reinforcement Learning” likewise develops HATRL, derives HATRPO and HAPPO by tractable approximations, and introduces HAML as a general template for cooperative MARL algorithmic designs, but it also does not define a variant named “HATRPO-W” (Zhong et al., 2023).

Within that lineage, HATRPO-W should be understood as a later refinement of trust-region budgeting rather than a replacement for the sequential trust-region framework itself. The underlying algorithmic substrate remains the HATRPO regime of agent-by-agent policy improvement in heterogeneous-agent settings; the innovation is the replacement of uniform per-agent KL thresholds by an adaptive allocation under a joint constraint (Shek et al., 14 Aug 2025).

2. Cooperative-game formulation and the HATRPO baseline

HATRPO-W is formulated in a cooperative Markov (stochastic) game with agents $i \in \{1,\dots,N\}$, global state $s \in S$, and joint action $a=(a_1,\dots,a_N) \in A=\prod_i A_i$. The joint policy factorizes as

$\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$

while agents may be heterogeneous, with different observation modalities, action spaces, and policy parameterizations $\theta_i$. Centralized training with decentralized execution is assumed (Shek et al., 14 Aug 2025).

Let $d^{\pi}(s)$ denote the discounted state visitation distribution under the joint policy $\pi$, and define the joint advantage

$A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s).$

For agent $i$, holding other agents at their old policies, the per-agent advantage used in HATRPO is

$$A_i^{\pi^{\text{old}}}(s,a_i) \;\triangleq\; \mathbb{E}_{a_{-i} \sim \pi_{-i}^{\text{old}}(\cdot|s)} \big[\,A^{\pi^{\text{old}}}(s,(a_i,a_{-i}))\,\big].$$

The standard HATRPO surrogate for agent $s \in S$0 updating from $s \in S$1 to $s \in S$2, with other agents fixed at $s \in S$3, is

$s \in S$4

subject to the per-agent KL trust region

$s \in S$5

Standard HATRPO assigns each agent the same KL threshold. HATRPO-W replaces that uniform assignment by an adaptive allocation under a single global budget (Shek et al., 14 Aug 2025).

The natural-gradient quantities are

$s \in S$6

and

$s \in S$7

Under the TRPO quadratic approximation, the predicted improvement for agent $s \in S$8 with KL budget $s \in S$9 is

$a=(a_1,\dots,a_N) \in A=\prod_i A_i$0

Writing $a=(a_1,\dots,a_N) \in A=\prod_i A_i$1 and $a=(a_1,\dots,a_N) \in A=\prod_i A_i$2 makes the threshold-allocation problem explicit (Shek et al., 14 Aug 2025).

3. Global KL allocation and the KKT solution

The central idea of HATRPO-W is to replace uniform per-agent trust regions with a single weighted global constraint. Given weights $a=(a_1,\dots,a_N) \in A=\prod_i A_i$3, the joint budget is

$a=(a_1,\dots,a_N) \in A=\prod_i A_i$4

Under the standard TRPO per-agent approximation,

$a=(a_1,\dots,a_N) \in A=\prod_i A_i$5

so the global constraint reduces to $a=(a_1,\dots,a_N) \in A=\prod_i A_i$6 (Shek et al., 14 Aug 2025).

HATRPO-W chooses $a=(a_1,\dots,a_N) \in A=\prod_i A_i$7 by solving

$a=(a_1,\dots,a_N) \in A=\prod_i A_i$8

The corresponding Lagrangian is

$a=(a_1,\dots,a_N) \in A=\prod_i A_i$9

with dual variables $\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$0 and $\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$1. For each $\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$2 with $\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$3,

$\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$4

Substituting into the active budget constraint yields

$\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$5

The resulting closed-form allocation is

$\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$6

In the special case $\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$7,

$\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$8

The paper characterizes this as a separable “water-filling–like” problem and interprets the optimum through the leveling rule

$\pi(\cdot|s)=\prod_{i=1}^N \pi_i(\cdot|s),$9

so that agents with larger $\theta_i$0 receive a larger share of the global KL budget, moderated by $\theta_i$1 (Shek et al., 14 Aug 2025).

4. Per-agent TRPO step and sequential training procedure

Once the KKT-optimal thresholds are computed, each agent performs the standard TRPO step with its personalized budget: $\theta_i$2 HATRPO-W therefore changes the step size allocation, not the local trust-region mechanism itself. Agents with higher $\theta_i$3 make larger steps, but all remain within personalized trust regions (Shek et al., 14 Aug 2025).

A training iteration proceeds in four stages. First, trajectories are collected under the current joint policy $\theta_i$4, and per-agent gradients $\theta_i$5 together with Fisher-vector products for $\theta_i$6 are estimated, typically by conjugate gradient. Second, the thresholds $\theta_i$7 are computed in closed form from the KKT solution. Third, agents are updated sequentially, each using its natural-gradient step and an optional line search to ensure the realized average KL satisfies

$\theta_i$8

Fourth, already-updated agents are kept fixed while proceeding to the next agent, and new trajectories are collected for the next outer iteration (Shek et al., 14 Aug 2025).

The overhead of the allocation itself is $\theta_i$9, since it only requires a single pass to compute $d^{\pi}(s)$0. The dominant cost remains per-agent conjugate gradient for $d^{\pi}(s)$1, reported as $d^{\pi}(s)$2 per agent and iteration, using only Fisher-vector products (Shek et al., 14 Aug 2025).

5. Guarantees and theoretical interpretation

The theoretical status of HATRPO-W is inherited from the trust-region structure of HATRPO. Under standard TRPO assumptions, when the surrogate is optimized under a KL trust region and the line search or damping ensures that the realized KL respects $d^{\pi}(s)$3, HATRPO’s sequential per-agent updates retain the joint monotonic improvement guarantee; HATRPO-W does not alter the per-agent TRPO guarantee, but only reallocates the KL budgets $d^{\pi}(s)$4 under a global cap (Shek et al., 14 Aug 2025). The allocation subproblem is strictly concave in $d^{\pi}(s)$5 over $d^{\pi}(s)$6, its constraints are linear, and the KKT system has a unique solution when $d^{\pi}(s)$7 and $d^{\pi}(s)$8; feasibility holds for any $d^{\pi}(s)$9 (Shek et al., 14 Aug 2025).

This guarantee structure rests on earlier HATRPO theory. The 2021 formulation establishes the multi-agent advantage decomposition lemma, the sequential policy update scheme, and a monotonic improvement theorem stating that a sequence of joint policies updated by the theoretical algorithm satisfies $\pi$0 for all $\pi$1, with convergence to Nash equilibria under randomized update order in cooperative Markov games (Kuba et al., 2021). The 2023 HARL framework generalizes this perspective through HATRL and HAML, and proves that algorithms derived from HAML inherently enjoy monotonic improvement of joint return and convergence to Nash Equilibrium (Zhong et al., 2023).

A common misconception is to treat HATRPO-W as a weighted objective over agents. The defining construction is narrower: it is a weighted or globally budgeted allocation of KL thresholds, not a redefinition of the cooperative reward. The objective being optimized in the allocation stage is the sum of TRPO-predicted improvements $\pi$2 under $\pi$3 (Shek et al., 14 Aug 2025).

6. Empirical behavior, comparison, and limitations

HATRPO-W is evaluated on a matrix game, a differential Gaussian game, and Multi-Agent MuJoCo benchmarks including Ant, HalfCheetah, Hopper, and Walker with heterogeneous agent decompositions such as $\pi$4 Ant, $\pi$5 Ant, and $\pi$6 Ant. Baselines are MADDPG, MAPPO, HAPPO, and HATRPO, with centralized training, sequential agent updates, consistent hardware, 5 seeds on MuJoCo, and at least 10 seeds on the smaller games (Shek et al., 14 Aug 2025). The reported findings are that both HATRPO-W and HATRPO-G improve final performance over HATRPO, with gains exceeding $\pi$7, while HATRPO-G can reach approximately $\pi$8. HATRPO-W is reported to demonstrate lower variance and more stable learning dynamics across tasks, to reduce steps to $\pi$9 of maximum reward in the 4-agent matrix game under tighter KL budgets, to escape local optima in the differential game by adaptive allocation to the high-impact agent, and to achieve notably higher returns on Ant and HalfCheetah variants with KL usage heatmaps showing meaningful, non-uniform allocation over time (Shek et al., 14 Aug 2025).

The comparison with HATRPO-G clarifies the role of the weighted KKT solution. HATRPO-W solves a principled KKT system under a global budget and produces a closed-form $A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s).$0 proportional to $A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s).$1, whereas HATRPO-G greedily prioritizes agents with the highest local gain per KL and allocates budget iteratively. The paper characterizes HATRPO-W as typically yielding smoother, lower-variance training, while HATRPO-G can be more aggressive and may produce faster initial improvement at the cost of higher variance (Shek et al., 14 Aug 2025).

The principal limitations are also explicit. Performance is sensitive to $A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s).$2: budgets that are too small can overly constrain learning, while budgets that are too large can break the surrogate’s guarantees. Misallocation can occur if $A^{\pi}(s,a)=Q^{\pi}(s,a)-V^{\pi}(s).$3 is noisy because of poor advantage estimates or inaccurate Fisher approximations. Computational overhead remains that of TRPO/HATRPO, since conjugate gradient and Fisher-vector products are still required per agent. The method remains on-policy, so off-policy corrections must be handled carefully to maintain KL control. The formulation is compatible with discrete and continuous policies, but accurate KL estimation is required; Gaussian policies should use exact KL, and categorical policies use a sum over logits (Shek et al., 14 Aug 2025).