Yalpha/AGRITECH-META
0
1import random2from models import Action, Observation, Reward, StepInfo3 4CROP_NITROGEN_COST = {"rice": 0.20, "wheat": 0.10, "none": 0.0}5CROP_BASE_YIELD = {"rice": 0.85, "wheat": 0.70, "none": 0.0}6CROP_WATER_NEED = {"rice": 0.25, "wheat": 0.12, "none": 0.0}7 8WEATHER_SEQUENCE = ["normal", "normal", "rainy", "drought", "normal",9 "rainy", "normal", "drought", "normal", "normal"]10 11SCENARIOS = {12 "fertile": dict(nitrogen=0.75, moisture=0.65, groundwater=0.90, budget=150.0),13 "drought": dict(nitrogen=0.45, moisture=0.20, groundwater=0.35, budget=100.0),14 "degraded": dict(nitrogen=0.25, moisture=0.45, groundwater=0.60, budget=80.0),15 "default": dict(nitrogen=0.50, moisture=0.50, groundwater=0.70, budget=120.0),16}17 18EPISODE_LENGTH = 519FERTILIZER_THRESHOLD = 0.620IRRIGATION_THRESHOLD = 0.6521BUDGET_PER_STEP = 20.022EARLY_STOP_SOIL = 0.1023EARLY_STOP_BUDGET = -30.024 25 26def _clamp(v, lo=0.0, hi=1.0):27 """Hard clamp — applied to every computed float value."""28 return max(lo, min(hi, v))29 30 31class AgriEnv:32 """33 AgriDecisionEnv v3 – Sustainable Farming RL Environment.34 35 API (OpenEnv compliant):36 reset() -> Observation37 step(action) -> (Observation, float, bool, dict)38 state() -> Observation39 40 - reward is a plain float, always in [0.0, 1.0]41 - info is a plain dict (not Pydantic)42 - No randomness in step() — fully deterministic via WEATHER_SEQUENCE43 """44 45 def __init__(self, scenario="default", seed=42):46 self._scenario = scenario47 self._seed = seed48 self._nitrogen = 0.549 self._moisture = 0.550 self._groundwater = 0.751 self._budget = 120.052 self._season = 053 self._last_crop = "none"54 self._done = False55 self._history = []56 self._fertilizer_window = []57 58 def reset(self):59 preset = SCENARIOS.get(self._scenario, SCENARIOS["default"])60 self._nitrogen = preset["nitrogen"]61 self._moisture = preset["moisture"]62 self._groundwater = preset["groundwater"]63 self._budget = preset["budget"]64 self._season = 065 self._last_crop = "none"66 self._done = False67 self._history = []68 self._fertilizer_window = []69 return self.state()70 71 def step(self, action):72 """73 Returns: (Observation, float reward [0,1], bool done, dict info)74 """75 if self._done:76 raise RuntimeError("Episode finished. Call reset().")77 78 crop = action.crop if action.crop in CROP_NITROGEN_COST else "none"79 fertilizer = _clamp(round(float(action.fertilizer), 2))80 irrigation = _clamp(round(float(action.irrigation), 2))81 82 # Weather — pure lookup, zero randomness83 weather = WEATHER_SEQUENCE[self._season % len(WEATHER_SEQUENCE)]84 85 # Nitrogen86 new_nitrogen = _clamp(87 self._nitrogen - CROP_NITROGEN_COST[crop] + fertilizer * 0.3088 )89 90 # Delayed fertilizer penalty (3-step rolling avg)91 self._fertilizer_window.append(fertilizer)92 if len(self._fertilizer_window) > 3:93 self._fertilizer_window.pop(0)94 avg_fert = sum(self._fertilizer_window) / len(self._fertilizer_window)95 delayed_fert_penalty = _clamp(max(0.0, avg_fert - FERTILIZER_THRESHOLD) * 0.4)96 97 # Moisture + groundwater98 weather_m = {"rainy": +0.12, "normal": 0.0, "drought": -0.10}[weather]99 new_moisture = _clamp(self._moisture + irrigation * 0.28 + weather_m - 0.08)100 new_groundwater = _clamp(self._groundwater - irrigation * 0.18 - CROP_WATER_NEED[crop])101 102 # Budget103 new_budget = self._budget - (BUDGET_PER_STEP + fertilizer * 15.0 + irrigation * 10.0)104 105 # Derived metrics106 soil_quality = _clamp(new_nitrogen * 0.6 + new_moisture * 0.4)107 yield_score = _clamp(CROP_BASE_YIELD[crop] * (new_nitrogen + new_moisture) / 2.0)108 109 # Penalties — each individually clamped110 fert_penalty = _clamp(max(0.0, fertilizer - FERTILIZER_THRESHOLD) * 0.35)111 irrig_penalty = _clamp(max(0.0, irrigation - IRRIGATION_THRESHOLD) * 0.30)112 groundwater_penalty = _clamp(max(0.0, 0.20 - new_groundwater) * 0.50)113 budget_penalty = _clamp(max(0.0, -new_budget / 100.0) * 0.40)114 monocrop_penalty = 0.12 if (crop == self._last_crop and crop != "none") else 0.0115 soil_bonus = 0.15 if new_nitrogen >= 0.45 and new_moisture >= 0.35 else 0.0116 117 # Reward — HARD CLAMP guarantees [0.0, 1.0]118 reward: float = round(_clamp(119 yield_score + soil_bonus120 - fert_penalty - delayed_fert_penalty121 - irrig_penalty - monocrop_penalty122 - groundwater_penalty - budget_penalty123 ), 4)124 125 # History126 self._history.append({127 "season": self._season, "crop": crop, "weather": weather,128 "nitrogen": new_nitrogen, "moisture": new_moisture,129 "groundwater": new_groundwater, "soil_quality": soil_quality,130 "budget": new_budget, "reward": reward,131 })132 133 # Update state134 self._nitrogen = round(new_nitrogen, 4)135 self._moisture = round(new_moisture, 4)136 self._groundwater = round(new_groundwater, 4)137 self._budget = round(new_budget, 4)138 self._last_crop = crop139 self._season += 1140 141 terminated = soil_quality < EARLY_STOP_SOIL or new_budget < EARLY_STOP_BUDGET142 self._done = self._season >= EPISODE_LENGTH or terminated143 144 info = {145 "yield_score": round(yield_score, 4),146 "soil_health": round(soil_quality, 4),147 "water_used": round(irrigation + CROP_WATER_NEED[crop], 4),148 "budget_remaining": round(new_budget, 2),149 "weather": weather,150 "penalties": {151 "fertilizer": round(fert_penalty, 4),152 "delayed_fert": round(delayed_fert_penalty, 4),153 "irrigation": round(irrig_penalty, 4),154 "monocrop": round(monocrop_penalty, 4),155 "groundwater": round(groundwater_penalty, 4),156 "budget": round(budget_penalty, 4),157 },158 }159 return self.state(), reward, self._done, info160 161 def state(self):162 weather = WEATHER_SEQUENCE[self._season % len(WEATHER_SEQUENCE)]163 soil_quality = _clamp(self._nitrogen * 0.6 + self._moisture * 0.4)164 return Observation(165 nitrogen = self._nitrogen,166 moisture = self._moisture,167 soil_quality= round(soil_quality, 4),168 last_crop = self._last_crop,169 season = self._season,170 weather = weather,171 groundwater = self._groundwater,172 budget = self._budget,173 )174 175 @property176 def history(self): return self._history177 178 @property179 def done(self): return self._done180