RonyForAI/Mirage_DB_RL
0
1"""2Mirage_RL_environment.py — Production query join-order optimisation environment.3 4Core design:5 - Episodes are drawn from a pool of enterprise query scenarios (e-commerce,6 analytics, financial) sampled randomly or by seed for reproducibility.7 - Cardinality estimation noise (log-normal) is applied to table_rows in the8 observation — the agent sees estimated stats, not ground truth, matching9 real planner conditions.10 - Cost is computed on TRUE row counts (hidden from the agent's observation),11 reflecting actual query execution cost rather than planner estimates.12 - Rewards are normalised to [0.0, 1.0] at every step.13"""14 15from __future__ import annotations16 17import math18import random19from typing import List, Optional20 21from openenv.core.env_server.interfaces import Environment22 23try:24 from Mirage_RL.models import QueryAction, QueryObservation, QueryState25 from Mirage_RL.server.tasks import (26 TaskConfig, TableSpec, Scenario, TASKS,27 apply_estimation_noise, compute_cost_bounds,28 compute_step_cost_bounds, compute_order_quality, grade,29 )30except ImportError:31 from models import QueryAction, QueryObservation, QueryState # type: ignore32 from server.tasks import ( # type: ignore33 TaskConfig, TableSpec, Scenario, TASKS,34 apply_estimation_noise, compute_cost_bounds,35 compute_step_cost_bounds, compute_order_quality, grade,36 )37 38# Join-type cost multipliers (mirrors the agent's cost model exactly)39_JOIN_MULTIPLIER = {0: 1.0, 1: 2.0, 2: 0.8}40 41 42class QueryEnv(Environment):43 """44 Multi-task, multi-scenario database query join-order optimisation environment.45 46 Episode lifecycle:47 1. reset(task_id, seed) — sample a scenario, apply cardinality noise48 2. step(action) — choose next table + join strategy; get [0,1] reward49 3. repeat until all tables joined (done=True)50 51 Reward design:52 Per-step : normalised quality of this join decision → [0.0, 1.0]53 Final : overall grader score from grade() → [0.0, 1.0]54 Invalid : action on already-joined table → 0.0 reward (no progress, no crash)55 """56 57 def __init__(self) -> None:58 self._task: TaskConfig = TASKS["medium"]59 self._scenario: Scenario = TASKS["medium"].scenarios[0]60 61 # True rows (used for cost), estimated rows (shown to agent)62 self._true_rows: List[int] = []63 self._est_rows: List[int] = []64 self._cost_bounds: tuple[float, float] = (1.0, 0.0) # (worst, best)65 66 # Intermediate result size tracking:67 # true — used internally for cost accuracy68 # est — shown to agent (product of est_rows × sel for joined tables)69 self._true_running_size: float = 1.070 self._est_running_size: float = 1.071 72 self._state = QueryState()73 self._step_count: int = 074 self._done: bool = False75 self._rng = random.Random()76 77 # ──────────────────────────────────────────────────────────────────────────78 # OpenEnv API79 # ──────────────────────────────────────────────────────────────────────────80 81 def reset(self, task_id: str = "medium", seed: Optional[int] = None, **kwargs) -> QueryObservation:82 """83 Reset to a new episode.84 85 Args:86 task_id: "easy" | "medium" | "hard"87 seed: Optional integer seed for reproducible episode sampling.88 The same seed + task_id always produces the same scenario89 and cardinality noise, enabling reproducible baselines.90 """91 self._task = TASKS.get(task_id, TASKS["medium"])92 93 # Seed the RNG: deterministic when seed given, random otherwise94 effective_seed = seed if seed is not None else random.randint(0, 2**31)95 self._rng.seed(effective_seed)96 97 # Sample a scenario from the pool98 self._scenario = self._rng.choice(self._task.scenarios)99 100 # Apply cardinality estimation noise to get the agent-visible estimates101 self._true_rows = [t.true_rows for t in self._scenario.tables]102 self._est_rows = [103 apply_estimation_noise(t.true_rows, t.noise_sigma, self._rng)104 for t in self._scenario.tables105 ]106 107 # Precompute cost bounds using TRUE rows (for normalisation)108 self._cost_bounds = compute_cost_bounds(self._scenario.tables)109 110 # Reset episode state111 n = len(self._scenario.tables)112 self._state.chosen_order = []113 self._state.remaining_tables = list(range(n))114 self._state.current_cost = 0.0115 self._state.final_cost = 0.0116 self._state.scenario_name = self._scenario.name117 self._step_count = 0118 self._done = False119 self._true_running_size = 1.0120 self._est_running_size = 1.0121 122 return self._make_obs(reward=0.0)123 124 def step(self, action: QueryAction) -> QueryObservation:125 """126 Execute one join-order decision.127 128 Returns:129 QueryObservation with reward in [0.0, 1.0].130 """131 # ── Guard: invalid table selection ────────────────────────────────────132 if action.next_table not in self._state.remaining_tables:133 # Zero reward; episode does NOT terminate (agent must recover)134 return self._make_obs(reward=0.0)135 136 # ── Apply action ──────────────────────────────────────────────────────137 idx = action.next_table138 self._state.chosen_order.append(idx)139 self._state.remaining_tables.remove(idx)140 self._step_count += 1141 142 # Cost computed on TRUE rows (not the noisy estimate the agent sees)143 added_cost = self._true_step_cost(idx, action)144 self._state.current_cost += added_cost145 146 # ── Normalised per-step reward ────────────────────────────────────────147 # MUST evaluate bounds using the running_size BEFORE it gets multiplied148 table_spec = self._scenario.tables[idx]149 worst_step, best_step = compute_step_cost_bounds(table_spec, self._true_running_size)150 if worst_step == best_step:151 step_reward = 1.0152 else:153 step_reward = (worst_step - added_cost) / (worst_step - best_step)154 step_reward = float(max(0.0, min(1.0, step_reward)))155 156 # Update intermediate size trackers for the NEXT step157 sel = self._scenario.tables[idx].selectivity158 self._true_running_size *= self._true_rows[idx] * sel159 self._est_running_size *= self._est_rows[idx] * sel160 161 # ── Episode complete? ─────────────────────────────────────────────────162 if not self._state.remaining_tables:163 self._done = True164 self._state.final_cost = self._state.current_cost165 final_score = grade(166 self._scenario.tables,167 self._state.final_cost,168 chosen_order=list(self._state.chosen_order),169 )170 return self._make_obs(reward=final_score)171 172 return self._make_obs(reward=step_reward)173 174 @property175 def state(self) -> QueryState:176 return self._state177 178 # ──────────────────────────────────────────────────────────────────────────179 # Internal helpers180 # ──────────────────────────────────────────────────────────────────────────181 182 def _true_step_cost(self, table_idx: int, action: QueryAction) -> float:183 """184 Compute join cost using TRUE row count (not the noisy estimate), 185 adding an intermediate size penalty for cascading costs.186 This strongly encourages joining highly selective, small tables early.187 """188 true_rows = self._true_rows[table_idx]189 sel = self._scenario.tables[table_idx].selectivity190 has_idx = self._scenario.tables[table_idx].has_index191 192 base_rows = true_rows * 0.5 if (action.use_index and has_idx) else true_rows193 mult = _JOIN_MULTIPLIER.get(action.join_type, 1.0)194 195 # Additive penalty: the size of the accumulated running result196 penalty = max(0.0, self._true_running_size - 1.0)197 return (base_rows * sel * mult) + penalty198 199 def _make_obs(self, reward: float) -> QueryObservation:200 """Build observation — agent sees ESTIMATED rows and intermediate size."""201 return QueryObservation(202 done = self._done,203 reward = round(reward, 6),204 tables = list(t.name for t in self._scenario.tables),205 table_rows = list(self._est_rows),206 selectivities = [t.selectivity for t in self._scenario.tables],207 has_index = [t.has_index for t in self._scenario.tables],208 chosen_order = list(self._state.chosen_order),209 remaining_tables = list(self._state.remaining_tables),210 step_number = self._step_count,211 current_cost = round(self._state.current_cost, 6),212 query_context = self._scenario.query_context,213 intermediate_size = round(self._est_running_size, 2),214 )