rishithayanidhi/datacenter-cooling-optimization

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""8Baseline Rule-Based Agent for Data Center Cooling Optimization.9 10This agent implements simple heuristic-based cooling management:11- Monitor temperature relative to safe limits12- Increase cooling if temperature is too high13- Decrease cooling if temperature is too low14- Balance cooling across zones15"""16 17import random18from typing import List, Tuple19import numpy as np20 21try:22    from ..models import CoolingAction, CoolingObservation23except ImportError:24    from models import CoolingAction, CoolingObservation25 26 27class BaselineAgent:28    """29    Simple rule-based agent for data center cooling management.30    31    Strategy:32    1. Monitor each zone's temperature33    2. If temperature > SAFE_MAX: increase cooling34    3. If temperature < SAFE_MIN: decrease cooling35    4. If temperature in range: maintain current cooling36    5. Prefer balanced cooling across all zones37    """38    39    SAFE_TEMPERATURE_MIN = 18.0  # °C40    SAFE_TEMPERATURE_MAX = 42.0  # °C41    WARNING_TEMPERATURE = 40.0   # °C42    CRITICAL_TEMPERATURE = 48.0  # °C43    44    def __init__(self, zone_count: int = 4, strategy: str = "conservative"):45        """46        Initialize the baseline agent.47        48        Args:49            zone_count: Number of cooling zones50            strategy: Type of control strategy51                - "reactive": Simple threshold-based52                - "conservative": Preemptive cooling increase53                - "aggressive": Fast response to changes54        """55        self.zone_count = zone_count56        self.strategy = strategy57        self.previous_temps = [None] * zone_count58        self.step_count = 059    60    def select_action(self, observation: CoolingObservation) -> CoolingAction:61        """62        Select cooling action based on observation.63        64        Args:65            observation: Current environment state66            67        Returns:68            CoolingAction to apply69        """70        self.step_count += 171        72        # Select zone with largest temperature deviation73        zone_id, adjustment = self._select_zone_and_adjustment(observation)74        75        return CoolingAction(76            zone_id=zone_id,77            cooling_adjustment=adjustment,78            duration=179        )80    81    def _select_zone_and_adjustment(82        self, observation: CoolingObservation83    ) -> Tuple[int, float]:84        """85        Determine which zone to adjust and by how much.86        87        Args:88            observation: Current environment state89            90        Returns:91            Tuple of (zone_id, adjustment_value)92        """93        temps = observation.zone_temperatures94        current_cooling = observation.zone_cooling_levels95        96        # Find zone with most problematic temperature97        zone_issues = []98        for i, temp in enumerate(temps):99            # Measure deviation from safe range100            if temp > self.SAFE_TEMPERATURE_MAX:101                deviation = temp - self.SAFE_TEMPERATURE_MAX102                priority = 2.0 + deviation  # Highest priority103            elif temp < self.SAFE_TEMPERATURE_MIN:104                deviation = self.SAFE_TEMPERATURE_MIN - temp105                priority = -1.0 - deviation  # Decrease cooling106            else:107                # In safe range - maintain108                priority = 0.5109            110            zone_issues.append((i, priority, temp))111        112        # Apply strategy113        if self.strategy == "reactive":114            zone_id, _, temp = max(zone_issues, key=lambda x: abs(x[1]))115            adjustment = self._calculate_adjustment_reactive(temp)116        117        elif self.strategy == "conservative":118            zone_id, _, temp = max(zone_issues, key=lambda x: abs(x[1]))119            adjustment = self._calculate_adjustment_conservative(temp, current_cooling[zone_id])120        121        else:  # aggressive122            zone_id, _, temp = max(zone_issues, key=lambda x: abs(x[1]))123            adjustment = self._calculate_adjustment_aggressive(temp, current_cooling[zone_id])124        125        return zone_id, adjustment126    127    def _calculate_adjustment_reactive(self, temperature: float) -> float:128        """Simple threshold-based adjustment."""129        if temperature > self.CRITICAL_TEMPERATURE:130            return 0.8  # Maximum increase131        elif temperature > self.WARNING_TEMPERATURE:132            return 0.5  # Significant increase133        elif temperature > self.SAFE_TEMPERATURE_MAX:134            return 0.3  # Moderate increase135        elif temperature < self.SAFE_TEMPERATURE_MIN:136            return -0.3  # Decrease cooling137        else:138            return 0.0  # Maintain139    140    def _calculate_adjustment_conservative(self, temperature: float, current_cooling: float) -> float:141        """Preemptive cooling increase to prevent overheating."""142        if temperature > self.CRITICAL_TEMPERATURE:143            return 1.0  # Maximum increase144        elif temperature > self.WARNING_TEMPERATURE:145            return 0.6  # Preemptive increase146        elif temperature > self.SAFE_TEMPERATURE_MAX:147            return 0.4  # Early response148        elif temperature < self.SAFE_TEMPERATURE_MIN - 2:149            return -0.5  # Aggressive decrease150        elif temperature < self.SAFE_TEMPERATURE_MIN:151            return -0.2  # Gradual decrease152        else:153            return 0.1  # Slight preemptive increase154    155    def _calculate_adjustment_aggressive(self, temperature: float, current_cooling: float) -> float:156        """Fast response to temperature changes."""157        temp_diff = temperature - (self.SAFE_TEMPERATURE_MAX + self.SAFE_TEMPERATURE_MIN) / 2158        159        # Response magnitude proportional to deviation160        if abs(temp_diff) > 10:161            adjustment = 0.9 if temp_diff > 0 else -0.7162        elif abs(temp_diff) > 5:163            adjustment = 0.6 if temp_diff > 0 else -0.4164        elif abs(temp_diff) > 2:165            adjustment = 0.3 if temp_diff > 0 else -0.2166        else:167            adjustment = 0.0168        169        return adjustment170    171    def get_name(self) -> str:172        """Get agent name for logging."""173        return f"BaselineAgent({self.strategy})"174 175 176class SmartBaselineAgent(BaselineAgent):177    """178    Enhanced baseline agent with zone balancing.179    180    Improvements over basic agent:181    - Considers temperature variance across zones182    - Tries to balance cooling distribution183    - Tracks long-term trends184    """185    186    def __init__(self, zone_count: int = 4):187        """Initialize smart baseline agent."""188        super().__init__(zone_count, strategy="conservative")189        self.temp_history = [[] for _ in range(zone_count)]190        self.history_size = 10191    192    def _select_zone_and_adjustment(193        self, observation: CoolingObservation194    ) -> Tuple[int, float]:195        """196        Select zone and adjustment considering balancing.197        198        Args:199            observation: Current environment state200            201        Returns:202            Tuple of (zone_id, adjustment_value)203        """204        temps = observation.zone_temperatures205        current_cooling = observation.zone_cooling_levels206        207        # Update history208        for i, temp in enumerate(temps):209            self.temp_history[i].append(temp)210            if len(self.temp_history[i]) > self.history_size:211                self.temp_history[i].pop(0)212        213        # Find zone with highest temperature214        max_zone = np.argmax(temps)215        max_temp = temps[max_zone]216        217        # Consider temperature trend218        if len(self.temp_history[max_zone]) > 1:219            trend = self.temp_history[max_zone][-1] - self.temp_history[max_zone][0]220        else:221            trend = 0222        223        # Adjust based on current state and trend224        if max_temp > self.CRITICAL_TEMPERATURE or (max_temp > self.WARNING_TEMPERATURE and trend > 0):225            adjustment = 0.7226        elif max_temp > self.SAFE_TEMPERATURE_MAX:227            adjustment = 0.4228        elif max_temp < self.SAFE_TEMPERATURE_MIN and trend < 0:229            adjustment = -0.4230        elif max_temp < self.SAFE_TEMPERATURE_MIN:231            adjustment = -0.2232        else:233            adjustment = 0.1234        235        return max_zone, adjustment236    237    def get_name(self) -> str:238        """Get agent name for logging."""239        return "SmartBaselineAgent"240 241 242def run_baseline_evaluation(243    env_client,244    num_episodes: int = 5,245    agent_type: str = "smart",246) -> dict:247    """248    Run baseline agent for evaluation.249    250    Args:251        env_client: Environment client instance252        num_episodes: Number of episodes to run253        agent_type: Type of baseline agent ("smart" or "reactive")254        255    Returns:256        Dictionary with evaluation results257    """258    if agent_type == "smart":259        agent = SmartBaselineAgent(zone_count=4)260    else:261        agent = BaselineAgent(zone_count=4, strategy="conservative")262    263    results = {264        "agent": agent.get_name(),265        "episodes": [],266        "avg_reward": 0.0,267        "avg_violations": 0,268        "avg_energy": 0.0,269    }270    271    episode_rewards = []272    episode_violations = []273    episode_energies = []274    275    for episode_idx in range(num_episodes):276        obs_result = env_client.reset()277        observation = obs_result.observation278        state = obs_result279        280        episode_reward = 0.0281        episode_violations = 0282        episode_energy = 0.0283        284        for step in range(500):  # Max 500 steps per episode285            # Agent selects action286            action = agent.select_action(observation)287            288            # Step environment289            result = env_client.step(action)290            observation = result.observation291            step_reward = result.reward or 0.0292            293            episode_reward += step_reward294            episode_violations += 1 if max(observation.zone_temperatures) > 50.0 else 0295            episode_energy += observation.total_energy_consumption296            297            if result.done:298                break299        300        results["episodes"].append({301            "episode": episode_idx,302            "reward": episode_reward,303            "violations": episode_violations,304            "energy": episode_energy,305        })306        307        episode_rewards.append(episode_reward)308        episode_violations.append(episode_violations)309        episode_energies.append(episode_energy)310    311    # Calculate averages312    if episode_rewards:313        results["avg_reward"] = np.mean(episode_rewards)314        results["avg_violations"] = np.mean(episode_violations)315        results["avg_energy"] = np.mean(episode_energies)316    317    return results318