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"""8Data Center Cooling Optimization Environment Implementation.9 10This environment simulates a realistic data center cooling management task where11agents learn to balance thermal stability and energy efficiency.12"""13 14import logging15import math16import os17import sys18from pathlib import Path19from uuid import uuid420from typing import List, Dict, Tuple21 22import numpy as np23from openenv.core.env_server.interfaces import Environment24from openenv.core.env_server.types import State25 26# Add parent directory to path for imports27sys.path.insert(0, str(Path(__file__).parent.parent))28 29try:30    from models import CoolingAction, CoolingObservation, CoolingState31except ImportError:32    try:33        from ..models import CoolingAction, CoolingObservation, CoolingState34    except ImportError:35        from models import CoolingAction, CoolingObservation, CoolingState36 37log = logging.getLogger("environment")38 39 40class DataCenterCoolingEnvironment(Environment):41    """42    Simulates a data center cooling management task.43 44    The environment models a configurable-zone data center where an agent must manage45    cooling levels to maintain safe temperatures while minimizing energy consumption.46 47    Physics constants are read from environment variables at class-definition time so48    they can be overridden without modifying source code.49    """50 51    SUPPORTS_CONCURRENT_SESSIONS: bool = True52 53    # Physics constants — all overridable via environment variables54    NUM_ZONES: int = int(os.getenv("NUM_ZONES", "4"))55    AMBIENT_TEMPERATURE: float = float(os.getenv("AMBIENT_TEMPERATURE", "20.0"))56    SAFE_TEMPERATURE_MIN: float = float(os.getenv("SAFE_TEMPERATURE_MIN", "15.0"))57    SAFE_TEMPERATURE_MAX: float = float(os.getenv("SAFE_TEMPERATURE_MAX", "45.0"))58    CRITICAL_TEMPERATURE: float = float(os.getenv("CRITICAL_TEMPERATURE", "50.0"))59 60    # Thermal dynamics61    THERMAL_CAPACITANCE: float = float(os.getenv("THERMAL_CAPACITANCE", "10.0"))62    COOLING_EFFICIENCY: float = float(os.getenv("COOLING_EFFICIENCY", "5.0"))63    WORKLOAD_TO_HEAT: float = float(os.getenv("WORKLOAD_TO_HEAT", "20.0"))64 65    # Energy66    BASE_POWER: float = float(os.getenv("BASE_POWER", "50.0"))67    COOLING_POWER_PER_UNIT: float = float(os.getenv("COOLING_POWER_PER_UNIT", "8.0"))68 69    # Time scales70    TIME_STEP_DURATION: float = float(os.getenv("TIME_STEP_DURATION", "1.0"))71 72    def __init__(self, task_type: str = "easy") -> None:73        """74        Initialize the data center cooling environment.75 76        Args:77            task_type: Type of task — "easy", "medium", or "hard".78                       Can also be set via TASK_TYPE env var.79        """80        self.task_type = os.getenv("TASK_TYPE", task_type).lower()81        if self.task_type not in ("easy", "medium", "hard"):82            log.warning("Unknown task_type '%s'; defaulting to 'easy'", self.task_type)83            self.task_type = "easy"84 85        self._episode_id = str(uuid4())86        self._step_count = 087 88        # Episode length — judges run 50 steps per task; overridable for local testing89        self._max_steps = int(os.getenv("MAX_EPISODE_STEPS", "50"))90 91        log.info(92            "DataCenterCoolingEnvironment init — task=%s max_steps=%d zones=%d",93            self.task_type, self._max_steps, self.NUM_ZONES,94        )95 96        try:97            # Initialize state98            self._zone_temperatures: List[float] = [25.0] * self.NUM_ZONES99            self._zone_cooling_levels: List[float] = [0.5] * self.NUM_ZONES100            self._zone_workload_intensity: List[float] = [101                self._get_initial_workload() for _ in range(self.NUM_ZONES)102            ]103        except Exception as exc:104            log.error("Failed to initialise environment state: %s", exc)105            raise106 107        # Tracking108        self._thermal_violations = 0109        self._cumulative_energy = 0.0110        self._cumulative_reward = 0.0111        112        # For physics113        self._step_index = 0114        115    def _get_initial_workload(self) -> float:116        """Get initial workload based on task type."""117        if self.task_type == "easy":118            return 0.5  # Constant moderate load119        elif self.task_type == "medium":120            return 0.6  # Slightly higher121        else:  # hard122            return 0.4  # Variable123    124    def _generate_workload(self, step: int) -> List[float]:125        """126        Generate workload for each zone based on task type and time.127        Creates dynamic heat generation to test agent's cooling management.128        129        Args:130            step: Current simulation step131            132        Returns:133            List of workload intensities [0, 1] for each zone134        """135        workloads = []136        137        for zone in range(self.NUM_ZONES):138            if self.task_type == "easy":139                # Moderate constant workload - requires agent to maintain cooling140                base = 0.7  # Increased from 0.5141                # Add gentle variation by zone for realistic behavior142                base += 0.05 * math.sin(step * 0.02 + zone * 0.5)143                144            elif self.task_type == "medium":145                # Fluctuating workload with clear peaks and valleys146                base = 0.65 + 0.2 * math.sin(step * 0.05 + zone * 0.5)147                # Occasional activity spikes148                if step % 50 == 0:149                    base += 0.2150                    151            else:  # hard152                # Challenging: unpredictable spikes and difficult patterns153                base = 0.5154                # Random spikes that agents must respond to quickly155                if step % 40 == 0 and step > 0:156                    base += 0.4  # Significant thermal spike157                # Multiple harmonic patterns to prevent simple solutions158                base += 0.2 * math.sin(step * 0.1 + zone * 0.3)159                base += 0.1 * math.sin(step * 0.03 + zone * 1.5)160            161            # Clamp to [0, 1]162            workloads.append(max(0.0, min(1.0, base)))163        164        return workloads165    166    def _update_temperatures(self, actions_applied: List[Tuple[int, float]]) -> None:167        """168        Update zone temperatures based on workload, cooling, and physics.169        170        Args:171            actions_applied: List of (zone_id, cooling_adjustment) tuples172        """173        # Get current workload174        workloads = self._generate_workload(self._step_index)175        self._zone_workload_intensity = workloads176        177        # Apply cooling adjustments178        for zone_id, cooling_adj in actions_applied:179            zone_id = int(zone_id)180            change = cooling_adj * 0.1  # Max 10% change per step181            self._zone_cooling_levels[zone_id] = max(0.0, min(1.0, 182                self._zone_cooling_levels[zone_id] + change))183        184        # Update temperatures for each zone185        for zone in range(self.NUM_ZONES):186            # Heat generation from workload187            heat_generated = workloads[zone] * self.WORKLOAD_TO_HEAT188            189            # Cooling effect190            cooling_effect = self._zone_cooling_levels[zone] * self.COOLING_EFFICIENCY191            192            # Net heat flow193            net_heat = heat_generated - cooling_effect194            195            # Temperature change (simplified thermal dynamics)196            temp_change = net_heat / self.THERMAL_CAPACITANCE197            198            # Update temperature with thermal dissipation toward ambient199            T = self._zone_temperatures[zone]200            ambient_diff = self.AMBIENT_TEMPERATURE - T201            dissipation = 0.1 * ambient_diff  # Natural cooling202            203            self._zone_temperatures[zone] = T + temp_change + dissipation204            205            # Clamp to reasonable range206            self._zone_temperatures[zone] = max(10.0, 207                min(70.0, self._zone_temperatures[zone]))208    209    def _calculate_reward(self) -> Tuple[float, Dict]:210        """211        Calculate reward based on environment state.212        213        CRITICAL: Reward must vary meaningfully to distinguish good from bad actions214        Returns:215            Tuple of (reward_value, reward_breakdown)216        """217        reward_breakdown = {}218        total_reward = 0.0219        220        # Get current state metrics221        mean_temp = np.mean(self._zone_temperatures)222        max_temp = max(self._zone_temperatures)223        min_temp = min(self._zone_temperatures)224        temp_variance = np.var(self._zone_temperatures)225        226        # ===== THERMAL PERFORMANCE (PRIMARY REWARD) =====227        # This is the main signal: how well you're keeping temps in safe range228        229        if max_temp > self.CRITICAL_TEMPERATURE:230            # MASSIVE PENALTY for critical overheat (>50°C)231            overheat_severity = min(20.0, (max_temp - self.CRITICAL_TEMPERATURE) / 2.0)232            thermal_reward = -1.0 * min(1.0, overheat_severity)  # -1.0 to 0.0233            self._thermal_violations += 1234            235        elif max_temp > self.SAFE_TEMPERATURE_MAX:236            # LARGE PENALTY for overheating (45-50°C)237            excess_heat = (max_temp - self.SAFE_TEMPERATURE_MAX) / 5.0  # 0-1238            thermal_reward = -0.8 * excess_heat  # -0.8 to 0.0239            240        elif max_temp < self.SAFE_TEMPERATURE_MIN:241            # MEDIUM PENALTY for overcooling (<15°C)242            undercool = (self.SAFE_TEMPERATURE_MIN - max_temp) / 10.0243            thermal_reward = -0.3 * undercool  # -0.3 to 0.0244            245        else:246            # GOOD ZONE: Reward based on distance from ideal247            # Ideal = 35°C (middle of safe range)248            ideal_temp = 35.0249            temp_error = abs(mean_temp - ideal_temp)250            251            # Normalize error (0-10°C error -> 0-1)252            normalized_error = min(1.0, temp_error / 10.0)253            254            # Thermal reward: 0.7 max (when perfect), 0.0 at boundaries255            thermal_reward = 0.7 * (1.0 - normalized_error)256            257            # Bonus for low variance (zones stable with each other)258            # Variance should be < 5 for good thermal distribution259            normalized_variance = min(1.0, temp_variance / 10.0)260            variance_bonus = 0.2 * (1.0 - normalized_variance)261            thermal_reward += variance_bonus262        263        reward_breakdown["thermal"] = thermal_reward264        total_reward += thermal_reward265        266        # ===== ENERGY EFFICIENCY (SECONDARY REWARD) =====267        # Reward for efficient cooling: don't over-cool268        avg_cooling = np.mean(self._zone_cooling_levels)269        total_cooling = sum(self._zone_cooling_levels)270        271        # Only apply energy penalty if temps are well-managed272        if thermal_reward > -0.5:273            # Penalty for excessive cooling (>0.7 average is wasteful)274            if avg_cooling > 0.7:275                excess_cooling = (avg_cooling - 0.7) / 0.3276                energy_penalty = -0.15 * min(1.0, excess_cooling)277            else:278                # Reward for efficient cooling (between 0.3-0.7)279                efficiency = 1.0 - abs(avg_cooling - 0.5) / 0.5280                energy_penalty = 0.1 * max(0.0, efficiency - 0.5)281            282            reward_breakdown["energy"] = energy_penalty283            total_reward += energy_penalty284        285        # ===== BALANCED COOLING BONUS =====286        # Reward for keeping all zones similar temp (good coordination)287        if temp_variance < 3.0:288            coordination_bonus = 0.1289            reward_breakdown["coordination"] = coordination_bonus290            total_reward += coordination_bonus291        292        # ===== FINAL CLIPPING =====293        # Ensure reward is in [-1, 1] and normalize to [0, 1] for judges294        total_reward = np.clip(total_reward, -1.0, 1.0)295        296        # Normalize to [0.0, 1.0] range for judge compatibility297        # Map: -1.0 -> 0.0, 0.0 -> 0.5, 1.0 -> 1.0298        normalized_reward = (total_reward + 1.0) / 2.0299        300        reward_breakdown["total"] = normalized_reward301        return normalized_reward, reward_breakdown302    303    def reset(self) -> CoolingObservation:304        """305        Reset the environment for a new episode.306 307        Returns:308            CoolingObservation with initial state309        """310        try:311            self._episode_id = str(uuid4())312            self._step_count = 0313            self._step_index = 0314            self._thermal_violations = 0315            self._cumulative_energy = 0.0316            self._cumulative_reward = 0.0317 318            # Initial temperatures vary by task difficulty319            initial_temp_map = {320                "easy": float(os.getenv("INITIAL_TEMP_EASY", "30.0")),321                "medium": float(os.getenv("INITIAL_TEMP_MEDIUM", "35.0")),322                "hard": float(os.getenv("INITIAL_TEMP_HARD", "38.0")),323            }324            initial_temp = initial_temp_map.get(self.task_type, 30.0)325 326            self._zone_temperatures = [327                initial_temp + np.random.uniform(-2, 2) for _ in range(self.NUM_ZONES)328            ]329            self._zone_cooling_levels = [0.4] * self.NUM_ZONES330            self._zone_workload_intensity = [331                self._get_initial_workload() for _ in range(self.NUM_ZONES)332            ]333 334            log.info(335                "reset — episode=%s task=%s initial_temps=%s",336                self._episode_id, self.task_type,337                [round(t, 1) for t in self._zone_temperatures],338            )339            return self._get_observation()340        except Exception as exc:341            log.error("reset failed: %s", exc, exc_info=True)342            raise343    344    def step(self, action: CoolingAction) -> CoolingObservation:  # type: ignore[override]345        """346        Execute one step of the environment.347 348        Args:349            action: CoolingAction with zone and cooling adjustment350 351        Returns:352            CoolingObservation with updated state353        """354        try:355            self._step_count += 1356            self._step_index += 1357 358            actions_applied = [(action.zone_id, action.cooling_adjustment)]359            self._update_temperatures(actions_applied)360 361            reward, breakdown = self._calculate_reward()362            self._cumulative_reward += reward363 364            done = (365                self._step_count >= self._max_steps366                or max(self._zone_temperatures) > 70.0367            )368 369            log.debug(370                "step=%d zone=%d adj=%.2f reward=%.4f done=%s max_temp=%.1f",371                self._step_count, action.zone_id, action.cooling_adjustment,372                reward, done, max(self._zone_temperatures),373            )374 375            obs = self._get_observation()376            obs.done = done377            obs.reward = reward378            return obs379 380        except Exception as exc:381            log.error("step %d failed: %s", self._step_count, exc, exc_info=True)382            raise383    384    def _get_observation(self) -> CoolingObservation:385        """386        Generate observation from current state.387        388        Returns:389            CoolingObservation with current metrics390        """391        max_temp = max(self._zone_temperatures)392        min_temp = min(self._zone_temperatures)393        temp_variance = np.var(self._zone_temperatures)394        395        return CoolingObservation(396            zone_temperatures=self._zone_temperatures.copy(),397            zone_workload_intensity=self._zone_workload_intensity.copy(),398            zone_cooling_levels=self._zone_cooling_levels.copy(),399            total_energy_consumption=self.BASE_POWER + sum(self._zone_cooling_levels) * self.COOLING_POWER_PER_UNIT / 10.0,400            ambient_temperature=self.AMBIENT_TEMPERATURE,401            timestamp=self._step_count,402            task_name=self.task_type,403            max_temperature=max_temp,404            min_temperature=min_temp,405            temperature_variance=temp_variance,406        )407    408    @property409    def state(self) -> CoolingState:410        """411        Get the current episode state.412        413        Returns:414            CoolingState with episode metadata415        """416        return CoolingState(417            episode_id=self._episode_id,418            step_count=self._step_count,419            task_type=self.task_type,420            max_steps=self._max_steps,421            total_reward=self._cumulative_reward,422            thermal_violations=self._thermal_violations,423            energy_consumed=self._cumulative_energy,424            workload_profile={425                "easy": "constant",426                "medium": "fluctuating", 427                "hard": "spike"428            }.get(self.task_type, "constant"),429            initial_temperatures=[25.0] * self.NUM_ZONES,430        )431