Hierarchical Multi-Agent Reinforcement Learning for Carbon-Aware AI Data Centers in Power Distribution Systems

Abstract: Eco-friendly energy management for artificial intelligence data centers (AIDCs) is crucial because of the significant increase in energy consumption-induced carbon emissions from AIDCs resulting from the rapid expansion of AI applications. This paper proposes a hierarchical carbon-aware multi-agent reinforcement learning (CA-MARL) framework for robust and efficient operations of AIDCs under uncertainties while ensuring low-carbon operation of power distribution systems. The framework comprises a workload manager (WM) agent and multiple local AIDC agents trained using a multi-agent transformer method, corresponding to a global AIDC aggregator and a local AIDC operator, respectively. Leveraging AIDC operation data along with nodal carbon intensity (NCI) calculated from the carbon emission flow-integrated distribution system operator problem, the WM agent spatially allocates AI training and inference jobs among all AIDCs. Based on the jobs allocated from the WM agent and NCI information, each AIDC agent schedules economical and eco-friendly operations of the AIDC by performing the following tasks: i) temporal shifting of training jobs, ii) spatial allocation of training graphics processing unit (GPU) blocks and inference GPUs within the AIDC, and iii) control of the supply air temperature of the cooling system. The effectiveness of the proposed framework was assessed using an IEEE 33-node power distribution system.

• The paper introduces a CA-MARL framework integrating carbon emission flow models to optimize spatio-temporal workload and cooling scheduling in distributed AI data centers.
• It demonstrates that hierarchical MAT-based coordination can cut carbon emissions by up to 15–20% and improve cooling efficiency compared to non-hierarchical methods.
• Scalability is achieved by decomposing the action space into hierarchical decision layers, enabling real-time, grid-aware scheduling under variable workloads and energy signals.

Hierarchical Multi-Agent Reinforcement Learning for Carbon-Aware AI Data Center Coordination

Introduction

The deployment of large generative models and AI applications has precipitated a dramatic increase in the energy demands and resultant carbon emissions of Artificial Intelligence Data Centers (AIDCs). Integrating AIDCs into distribution grids introduces complex operational and environmental challenges due to their high power density, workload heterogeneity, and scheduling flexibility. Traditional data center management frameworks are increasingly insufficient for managing the spatio-temporal flexibility demanded by AI workloads while simultaneously initiating grid-aware, low-carbon operations.

This paper addresses these limitations by formulating a hierarchical, carbon-aware multi-agent reinforcement learning (CA-MARL) framework, targeting distributed power distribution systems containing AIDCs. The approach incorporates dynamic nodal carbon intensity (NCI) traced via carbon emission flow (CEF) models, hierarchical agents, and a multi-agent transformer (MAT)-based learning architecture for the coordinated scheduling of AIDC workloads.

Figure 1: A reference hierarchical system integrated with a DSO, a WM agent, and AIDC agents.

System Model and Carbon Trace Integration

Unlike conventional approaches, the proposed framework fully integrates the CEF model into the distribution system operator (DSO) environment. This enables the computation of power-flow-aware, real-time NCI signals, reflecting the physical distribution of both active and virtual carbon flows.

The system comprises three core entities:

1. DSO (Environment): Solves the CEF-augmented optimal power flow for node-level NCI calculation, providing closed-loop feedback to agents.
2. Workload Manager (WM) Agent: Implements global AIDC aggregator responsibilities, performing spatial allocation of training and inference jobs based on grid, workload, and carbon signals.
3. AIDC Agents: Local AIDC operators schedule jobs temporally and spatially within each data center, allocate GPUs, and optimize CRAC-based cooling parameters.

This design enables joint optimization across spatial and temporal axes, leveraging both job deferral and site-specific allocation to minimize system-wide emissions.

Figure 2: CRAC-based cooling system for AIDC.

Hierarchical CA-MARL Framework

To address the exponential complexity of joint workload-cooling-resource scheduling across multi-site AIDCs, a hierarchical MAT-based MARL architecture is adopted.

• Architecture: The WM agent (top layer) produces job allocation decisions, followed by conditioned local policies executed by AIDC agents (bottom layer).
• Coordination: No explicit communication among AIDC agents; hierarchical information flow is mediated by the WM, which significantly curtails combinatorial action-space explosion.
• Learning: MAT's encoder-decoder with random order permutation ensures centralized training and decentralized inference, maintaining agent heterogeneity while enabling scalable credit assignment and advantage estimation.

  Figure 3: Structure of the proposed hierarchical CA-MARL framework comprising environment, agents, and MAT.

Spatio-Temporal Workload and Resource Scheduling

The integrated framework allows the following multilayer decision process:

• The WM agent allocates workloads among geographically distributed AIDCs, based on current NCI, workload arrivals, and grid price/carbon signals.
• AIDC agents perform:
  • Temporal job shifting of delay-tolerant training workloads.
  • Spatial allocation of available GPU resources for both training and inference.
  • Optimal control of CRAC supply air temperature, balancing IT and cooling energy under thermal constraints.

This enables adaptive shifting of workloads from high-NCI to low-NCI sites and from periods of grid carbon intensity peaks to troughs, concurrently optimizing economic and environmental objectives.

Figure 4: Spatio-temporal workload scheduling of WM and AIDC agents.

Action Space Scalability and Complexity Analysis

By explicit hierarchical decomposition, the policy search complexity transitions from exponential to additive, enabling stable training as the number of AIDCs and control dimensions scale. Quantitative results demonstrate a reduction of action-space cardinality by many orders of magnitude in the hierarchical framework compared to non-hierarchical baselines.

Numerical Results and Empirical Performance

Simulations leverage the IEEE 33-node test feeder, augmented with grid-connected AIDCs exhibiting realistic, stochastic AI workload arrivals. Four methods are compared: MAPPO, Transformer, decentralized MAT (MAT-Dec), and the proposed MAT.

Key results include:

• The hierarchical MAT method achieves the lowest total carbon emissions and job drop rates, outperforming all baselines in both economic and environmental metrics.
• Carbon-aware operation (joint mode) reduces emissions by up to 15–20% relative to power-only (economic) operation.
• The joint MAT policy dynamically shifts job allocations—both spatially and temporally—in response to time-varying NCI and electricity price signals.
• Significant reduction in PUE, indicating improved cooling efficiency due to coordinated IT-cooling optimization.

  Figure 5: Comparison of rewards among the MAPPO, Transformer, MAT-Dec, and proposed MAT methods in a joint mode.

  

  Figure 6: Spatio-temporal training/inference job shifting for three AIDCs under varying NCIs during a day.

Ablation Studies and Sensitivity Analyses

• The superiority of dynamic MAT-based workload shifting over static allocation is empirically validated, with static policies leading to higher carbon emissions and increased job loss.
• Sensitivity analyses on reward weight ratios and emission penalty parameters surface explicit trade-offs between cost and emissions, and identify practical operating points for robust grid and AIDC coordination.

  Figure 7: Performance comparison of the MAT method between static and dynamic job shifting of the WM agent.

  

  Figure 8: Comparison of rewards for the proposed MAT method with different numbers of AIDCs (3, 4, and 5) in the joint mode.

Implications, Limitations, and Future Directions

This work demonstrates that integrating carbon flow models (NCI/CEF) into hierarchical MARL frameworks facilitates truly grid-aware, low-carbon operation of distributed AIDCs, addressing physical, economic, and environmental constraints in a unified learning framework. The empirical results validate both the scalability and efficacy of the MAT-based architecture in large-scale, multi-site coupled power-cyber systems.

However, application in real-world deployments requires further extension:

• Inclusion of heterogeneous renewable generation and storage co-optimization.
• Expansion to transactive energy models with competitive, non-collaborative AIDC operators.
• Integration with uncertainty-aware and safe RL mechanisms for robustness under adversarial conditions and partial observability.

Conclusion

The hierarchical CA-MARL methodology introduced in this work represents a significant advancement in coordinated power-AIDC scheduling. By integrating CEF-based NCI signals and leveraging MAT-based hierarchical MARL, the framework achieves scalable, coordinated, and substantially lower-carbon AI data center operations, with robust real-time adaptivity to both grid and workload uncertainties. These findings have strong implications for sustainable, future-scaled AI infrastructure planning, and for reinforcement learning applications in coupled power-cyber-physical networks.