GradMAP: Gradient-Based Multi-Agent Proximal Learning for Grid-Edge Flexibility

Abstract: Coordinating large populations of grid-edge devices requires learning methods that remain fully decentralised in deployment while still respecting three-phase AC distribution-network physics. This paper proposes gradient-based multi-agent proximal learning (GradMAP) to address this challenge. GradMAP trains independent neural-network policies for each agent without any parameter sharing, and each agent uses only its own local observation for online decision-making without communication. During offline training, GradMAP embeds a differentiable three-phase AC power-flow model in a primal-dual learning loop and uses implicit differentiation to propagate exact network-constraint violations to update the policy parameters. To speed up training, GradMAP reuses expensive environment gradients through a proximal surrogate within a trust region defined in the more direct policy-output (action) space, instead of the probability distribution space used in other works, such as PPO. In case studies with 1,000 agents managing batteries, heat pumps, and controllable generators on the IEEE 123-bus feeder, GradMAP learns decentralised policies that minimise three-phase AC load-flow constraint violations within 15 minutes of training on a single workstation-class NVIDIA RTX PRO 5000 Blackwell 48GB GPU. This is a 3--5x training speed-up over gradient-based self-supervised learning benchmarks and substantially better training efficiency than multi-agent reinforcement-learning benchmarks. In out-of-sample tests, GradMAP also delivers among the lowest operating cost and constraint violations.

• The paper introduces an exact multi-agent gradient-based learning method that integrates differentiable AC power flow via implicit differentiation.
• It employs a proximal update strategy in the policy-output space, reducing computational costs and improving convergence efficiency in complex grid environments.
• GradMAP demonstrates superior voltage regulation and cost reduction on large-scale IEEE 123-bus tests, enabling decentralized control for thousands of agents.

GradMAP: Gradient-Based Multi-Agent Proximal Learning for Grid-Edge Flexibility

Introduction and Motivation

Efficient coordination of large-scale, heterogeneous grid-edge resources—such as batteries, heat pumps, and distributed generators—poses a fundamental challenge as distribution networks undergo decarbonization and electrification. Conventional centralized optimization is increasingly infeasible due to scalability and communication bottlenecks, while traditional distributed approaches face convergence and stability limitations. Recent learning-based strategies (including reinforcement learning and self-supervised policy optimization) have demonstrated some success but typically neglect exploitable grid physics, are restricted to the single-agent regime, and incur high computational costs when dealing with the nonlinearity and differentiability of AC power flow constraints.

The GradMAP (Gradient-based Multi-Agent Proximal learning) framework directly addresses these deficiencies by training fully decentralized, independent neural policies utilizing exact network-model gradients (via implicit differentiation) and efficiently reusing these gradients within a proximal update scheme defined in the policy-output space. This composite approach delivers highly scalable and computationally efficient learning for distributed flexibility, even in networks with exact three-phase unbalanced AC constraints and thousands of agents.

Methodological Advances

GradMAP introduces two principal innovations over prior multi-agent and gradient-based literature:

1. Exact Multi-Agent Gradient-Based Learning (GradMA): Each agent is trained with a separate neural policy, observing only local information and without any communication. Centralized training leverages full network knowledge for efficient policy optimization (centralized training, decentralized execution—CTDE paradigm), embedding a differentiable three-phase unbalanced AC power-flow solver into the learning loop. Network constraint violations are backpropagated to agents via implicit differentiation, enabling tight coordination across strongly coupled network constraints.

   Figure 1: Architecture diagram of the proposed GradMA and GradMAP framework, showing centralized training with differentiable network constraints and decentralized execution.

2. Gradient Reuse with Proximal Updates (GradMAP): To circumvent the high cost of environment-gradient evaluation—especially for large-scale AC power flows—GradMAP introduces a proximal update loop in the policy-output coordinates, allowing several low-cost parameter updates for each expensive rollout gradient computation. Unlike prior trust region methods (e.g., PPO) that use probabilistic KL divergence in parameter space, GradMAP's trust region is constructed directly in the (mean, standard deviation) output space, aligning with the coordinate system of the gradients and empirically improving performance in constrained grid environments.

Problem Formulation

The setting considered is a high-fidelity model of a residential 3-phase unbalanced distribution feeder (IEEE 123-bus topology), with up to 1,000 agents each controlling a heterogeneous device (battery, heat pump, generator). Agents are modeled with realistic device physics, including nonlinear battery efficiency and ramp constraints. The overall objective is social cost minimization subject to device dynamics and exact network power flow constraints:

• Each neural policy $\pi_i$ maps local agent observations $o_{i,t}$ to device power actions via normalized, clipped outputs.
• Constraints and their violations are aggregated in physically meaningful channels (voltages, energy/temperature end-states, ramps) and incorporated into a Lagrangian for primal-dual optimization.
• The constraints are enforced in learning via dual variables and explicit calculation of gradients with respect to agent outputs, including implicit differentiation through the AC network.

  Figure 2: IEEE123 feeder topology used in the main case study together with the connection locations of the 1,000 agents.

Implicit Differentiation for AC Power Flow

To model the exact physical coupling, GradMAP leverages implicit differentiation for steady-state AC power flow. Explicit backpropagation through iterative solvers is computationally infeasible for such large systems; instead, implicit differentiation at the fixed point enables fast calculation of the sensitivity of voltage violation terms with respect to nodal power injections:

• Gradients are evaluated via Krylov-subspace solvers (e.g., BiCGSTAB), requiring only a small memory overhead for large real/imaginary-valued Jacobians.
• The scalable computation of these environment gradients is a cornerstone of GradMAP's efficiency, permitting rapid repeated updates in policy space for each environment interaction.

Proximal Surrogate Loss and Trust Region

Gradient signals from environment rollouts are cached and re-applied in a trust region in the policy-output space. The trust metric is defined as the $\ell_2$ distance between old and updated policy means and standard deviations across all agents and timesteps. The proximal penalty is adaptively scaled to ensure updates remain within the region where linearization is a locally accurate surrogate for the nonlinear objective.

This direct policy-output trust region is demonstrated to be more effective than PPO-style KL or clipping-based trust regions, providing both theoretical alignment and empirical stability in high-dimensional, multi-agent, tightly-coupled network settings.

Empirical Evaluation

Benchmark Setup

The main evaluation is conducted on the IEEE 123-bus feeder (1,000 agents; heterogeneous batteries, heat pumps, and generators), with supporting results for 10- and 100-agent systems. Realistic time-series demand, PV, and temperature data drive the simulations, which include full AC network constraints at each 15-minute timestep (96 per episode, 31 days of test episodes). GradMAP is benchmarked against:

• GradMA (exact gradients without reuse);
• GradMA-RPO (probability-distribution space trust region);
• Black-box multi-agent RL (IPPO, MAPPO);
• Naive baseline (always-on or always-idle policies).

  Figure 3: Five-day visualization of the datasets used in the IEEE123 1,000-agent case study.

Constraint Satisfaction and Policy Coordination

GradMAP achieves significantly improved voltage regulation and power dispatch coordination relative to unconstrained learning, empirically reducing the maximum voltage deviation from 0.0362 p.u. to 0.0095 p.u. while agents maintain feasible device states and respect end-of-day constraints.

Figure 4: All 1,000 agents' normalized power on a held-out test day. Grey dashed lines indicate regulatory 0.95–1.05 p.u. voltage limits; GradMAP sharply reduces voltage violations compared to an unconstrained baseline.

Numerical Results and Convergence Analysis

Sample Efficiency and Wall-clock Time

GradMAP converges rapidly (within 15 minutes wall time on a single workstation-class NVIDIA GPU), outperforming all baselines on both cumulative cost and total constraint violations. It consistently achieves a 3–5$\times$ training speed-up compared to exact-gradient multi-agent self-supervised learning (GradMA), and at least a 10$\times$ speed-up against black-box MARL methods (IPPO/MAPPO).

Figure 5: Training convergence plots for the 1,000-agent case, showing both primal-steps (rollouts) and wall-clock time. GradMAP achieves the fastest and most reliable convergence across objectives and constraints.

Out-of-Sample and Scaling Properties

On 31 held-out test days, GradMAP policies maintain constraint satisfaction and achieve the lowest cumulative cost (12% below naive baseline, with all constraint violations sharply reduced versus all benchmarks).

Figure 6: Out-of-sample benchmark comparison for 1,000-agent case: GradMAP achieves lowest mean cost and maintains constraint violations at near-zero levels across all test days.

Scalability experiments up to 1,000 agents reveal that the computational cost scales favorably, with GradMAP maintaining the best performance and lowest runtime at all problem sizes.

Figure 7: Scaling comparison across 10-, 100-, and 1,000-agent systems: GradMAP's gradient reuse strategy results in the lowest training time and best cost/constraint trade-off.

Theoretical and Practical Implications

GradMAP demonstrates the feasibility of deploying fully decentralized, communication-free, independent neural policies for large agent populations in power systems by exploiting differentiability and proximal optimization strategies. The results contradict earlier assumptions on the necessity of probability-space trust regions for scalable gradient-based learning in multi-agent, nonlinear constrained domains and indicate that direct action-space trust regions yield both superior convergence and computational advantages.

Practically, the ability to train 1,000 independent agent policies (no parameter sharing) with real AC-grid physics and realistic device models—within minutes, on a single standard GPU—makes GradMAP directly applicable for DSOs, aggregators, and flexibility market operators. This framework is adaptable to broader distributed control problems with differentiable simulators, such as robotics and autonomous driving.

Future Directions

GradMAP's architecture admits several extensions:

• Incorporating higher-order derivatives (e.g., Hessian reuse) to further accelerate convergence.
• Adapting to discrete or hybrid action spaces (e.g., via bias-corrected straight-through estimators).
• Extending to broader differentiable environments with expensive or implicit simulators beyond grid applications.

Conclusion

GradMAP presents a robust, scalable, and computationally efficient solution for large-scale, fully decentralized grid-edge flexibility coordination, advancing the state-of-the-art in gradient-based multi-agent learning. By aligning the proximal update with policy output space and embedding exact differentiable AC network models, GradMAP achieves stronger performance—both in cost and constraint feasibility—than previous benchmarks and enables practical deployment at unprecedented scale and speed.