Kolaps27/UI-layout-optimizer
0
1"""2ui_env.py3---------4Environment Engine for an Adaptive UI Layout Optimization system.5"""6 7from __future__ import annotations8import random9from typing import Literal, Optional10from pydantic import BaseModel, Field, model_validator11 12# ---------------------------------------------------------------------------13# Task Class (Required by OpenEnv Validator)14# ---------------------------------------------------------------------------15 16EPS = 1e-617 18def safe_grader(fn):19 def wrapper(x=None):20 try:21 if x is None:22 val = fn()23 else:24 try:25 val = fn(x)26 except TypeError:27 val = fn()28 29 if not isinstance(val, (int, float)):30 val = 0.531 32 val = float(val)33 34 val = max(min(val, 1.0 - EPS), EPS)35 36 return val37 38 except Exception:39 return 0.540 41 return wrapper42 43def clamp_score(raw: float) -> float:44 # Retained as a stub just in case other direct internal refs exist,45 # though safe_grader is now the strict master gate46 if not isinstance(raw, (int, float)):47 return 0.548 val = float(raw)49 return max(min(val, 1.0 - EPS), EPS)50 51def normalize(x: float, lo: float, hi: float) -> float:52 """Min-max normalize x from [lo, hi] to [0, 1], clamped."""53 if hi <= lo:54 return 0.555 return max(0.0, min(1.0, (x - lo) / (hi - lo)))56 57 58def grade_easy(x=None) -> float:59 try:60 return clamp_score(UIEnv.grade_easy(x))61 except Exception:62 return 0.563 64 65def grade_medium(x=None) -> float:66 try:67 return clamp_score(UIEnv.grade_medium(x))68 except Exception:69 return 0.570 71 72def grade_hard(x=None) -> float:73 try:74 return clamp_score(UIEnv.grade_hard(x))75 except Exception:76 return 0.577 78 79class Task:80 def __init__(self, name: str, grader):81 self.name = name82 self.grader = grader83 84# ---------------------------------------------------------------------------85# Constants86# ---------------------------------------------------------------------------87 88BUTTON_SIZE_MIN: float = 0.589BUTTON_SIZE_MAX: float = 2.090FORM_LENGTH_MIN: int = 191FORM_LENGTH_MAX: int = 1092STEPS_MIN: int = 193STEPS_MAX: int = 1094 95BUTTON_SIZE_DELTA: float = 0.196FORM_LENGTH_DELTA: int = 197STEPS_DELTA: int = 198 99INVALID_ACTION_REWARD: float = -0.1100MAX_STEPS_PER_EPISODE: int = 20101 102BUTTON_SWEET_LOW: float = 0.9103BUTTON_SWEET_HIGH: float = 1.3104 105# ---------------------------------------------------------------------------106# Data Models107# ---------------------------------------------------------------------------108 109class Layout(BaseModel):110 """Represents the current UI layout configuration."""111 button_size: float = Field(112 default=1.0,113 ge=BUTTON_SIZE_MIN,114 le=BUTTON_SIZE_MAX,115 description="Size multiplier for UI buttons (0.5 - 2.0).",116 )117 form_length: int = Field(118 default=5,119 ge=FORM_LENGTH_MIN,120 le=FORM_LENGTH_MAX,121 description="Number of fields in the form (1 - 10).",122 )123 steps: int = Field(124 default=3,125 ge=STEPS_MIN,126 le=STEPS_MAX,127 description="Number of wizard / checkout steps (1 - 10).",128 )129 130class Observation(BaseModel):131 """Full observable state returned to the agent after every transition."""132 device: Literal["mobile", "desktop"] = Field(133 description="Device type the user is on.",134 )135 layout: Layout = Field(136 description="Current layout configuration.",137 )138 progress: float = Field(139 ge=0.0,140 le=1.0,141 description="User's task-completion progress in [0, 1].",142 )143 last_action: Optional[str] = Field(144 default=None,145 description="String name of the most recently applied action, or None.",146 )147 reward: float = Field(default=0.0, description="Step reward")148 done: bool = Field(default=False, description="Is episode done")149 info: dict = Field(default_factory=dict, description="Extra info")150 151class Action(BaseModel):152 """An action the agent can submit to the environment."""153 type: Literal[154 "increase_button",155 "decrease_form",156 "increase_steps",157 "decrease_steps",158 "reorder_sections",159 "set_button_size",160 "noop",161 ] = Field(description="Discrete action type.")162 value: Optional[float] = Field(163 default=None,164 description="Optional scalar payload (used by set_button_size).",165 )166 167 @model_validator(mode="after")168 def _value_required_for_set_button_size(self) -> "Action":169 """Ensure `value` is provided when action type requires it."""170 if self.type == "set_button_size" and self.value is None:171 raise ValueError("'value' must be provided for action type 'set_button_size'.")172 return self173 174# ---------------------------------------------------------------------------175# Environment Engine176# ---------------------------------------------------------------------------177 178class UIEnv:179 """Adaptive UI Layout Optimization - Environment Engine."""180 181 def __init__(self, seed: int = 42, task: str = "easy") -> None:182 self._seed: int = seed183 self.task: str = task184 self._rng: random.Random = random.Random(seed)185 186 # OpenEnv task list with graders — NO dynamic generation, NO conditionals187 self.tasks = [188 Task(name="easy", grader=safe_grader(self.grade_easy)),189 Task(name="medium", grader=safe_grader(self.grade_medium)),190 Task(name="hard", grader=safe_grader(self.grade_hard)),191 ]192 self.task_dict = {t.name: t for t in self.tasks}193 194 self._layout: Layout = Layout()195 self._device: Literal["mobile", "desktop"] = "desktop"196 self._progress: float = 0.0197 self._last_action: Optional[str] = None198 self._step_count: int = 0199 self._prev_score: float = 0.0200 201 self._prefers_short_forms: bool = False202 self._prefers_large_buttons: bool = False203 self._user_type: str = "new"204 205 self._ready: bool = False206 207 def reset(self) -> Observation:208 if self.task == "easy":209 steps = self._rng.randint(2, 3)210 form_length = self._rng.randint(2, 4)211 button_size = self._rng.uniform(0.9, 1.2)212 elif self.task == "medium":213 steps = self._rng.randint(3, 5)214 form_length = self._rng.randint(4, 6)215 button_size = self._rng.uniform(0.7, 1.5)216 elif self.task == "hard":217 steps = self._rng.randint(5, 8)218 form_length = self._rng.randint(6, 10)219 button_size = self._rng.uniform(0.5, 2.0)220 else:221 steps = self._rng.randint(3, 5)222 form_length = self._rng.randint(4, 6)223 button_size = 1.0224 225 self._layout = Layout(226 button_size=button_size,227 form_length=form_length,228 steps=steps,229 )230 self._clamp_layout()231 232 self._device = self._rng.choice(("mobile", "desktop"))233 self._progress = 0.0234 self._last_action = None235 self._step_count = 0236 237 self._prefers_short_forms = self._rng.choice([True, False])238 self._prefers_large_buttons = self._rng.choice([True, False])239 self._user_type = self._rng.choice(["impatient", "careful", "new"])240 241 self._ready = True242 # Initialize prev_score for delta-based reward shaping243 task_obj = next((t for t in self.tasks if t.name == self.task), self.tasks[0])244 self._prev_score = task_obj.grader(None)245 return self._get_observation()246 247 def step(self, action: Action) -> tuple[Observation, float, bool, dict]:248 if not self._ready:249 self.reset()250 251 action_reward_offset: float = self._apply_action(action)252 self._step_count += 1253 254 outcome, user_reward = self._simulate_user()255 done = False256 257 if outcome == "drop":258 done = True259 elif outcome == "distrust":260 pass261 else:262 self._progress += 1.0 / max(1, self._layout.steps)263 if self._progress >= 0.999:264 self._progress = 1.0265 outcome = "complete"266 done = True267 268 reward = user_reward + action_reward_offset269 if outcome == "complete":270 reward += 2.0271 elif outcome == "continue":272 reward += 0.1273 274 reward -= 0.05275 276 # Delta-based grader-aligned shaping: reward reflects improvement277 task_obj = next((t for t in self.tasks if t.name == self.task), self.tasks[0])278 current_score = task_obj.grader(None)279 score_delta = current_score - self._prev_score280 281 alpha = 10.0282 reward += alpha * score_delta # reward improvement, penalize degradation283 284 self._prev_score = current_score285 286 if self._step_count >= MAX_STEPS_PER_EPISODE:287 done = True288 289 info: dict = {290 "completed": (outcome == "complete"),291 "outcome": outcome,292 "progress": self._progress,293 "step_count": self._step_count,294 "user_type": self._user_type,295 }296 297 if done:298 info["score"] = current_score299 # Terminal grader alignment boost300 reward += current_score301 302 return self._get_observation(), reward, done, info303 304 def state(self) -> Observation:305 if not self._ready:306 raise RuntimeError("Call reset() before state().")307 return self._get_observation()308 309 def close(self) -> None:310 pass311 312 async def reset_async(self) -> Observation:313 return self.reset()314 315 async def step_async(self, action: Action) -> Observation:316 obs, reward, done, info = self.step(action)317 obs.reward = reward318 obs.done = done319 obs.info = info320 return obs321 322 def _simulate_user(self) -> tuple[str, float]:323 if self._step_count <= 3:324 return "continue", 0.0325 326 layout = self._layout327 drop_chance = 0.0328 distrust_chance = 0.0329 330 if layout.steps > 3:331 drop_chance += 0.05 * (layout.steps - 3)332 if layout.form_length > 5:333 drop_chance += 0.04 * (layout.form_length - 5)334 if self._prefers_short_forms and layout.form_length > 4:335 drop_chance += 0.05336 if layout.steps < 2:337 distrust_chance += 0.20338 if layout.button_size < 0.9 or layout.button_size > 1.3:339 distrust_chance += 0.10340 drop_chance += 0.02341 342 if self._user_type == "impatient":343 drop_chance += 0.06344 elif self._user_type == "careful":345 distrust_chance += 0.08346 347 if self.task == "hard":348 drop_chance += 0.04349 elif self.task == "easy":350 drop_chance -= 0.05351 distrust_chance -= 0.05352 353 drop_chance = max(0.0, min(1.0, drop_chance))354 distrust_chance = max(0.0, min(1.0 - drop_chance, distrust_chance))355 356 roll = self._rng.random()357 358 if roll < drop_chance:359 return "drop", -1.0360 elif roll < drop_chance + distrust_chance:361 return "distrust", -0.2362 else:363 return "continue", 0.0364 365 def _apply_action(self, action: Action) -> float:366 reward: float = 0.0367 match action.type:368 case "increase_button":369 self._layout.button_size += BUTTON_SIZE_DELTA370 case "decrease_form":371 self._layout.form_length -= FORM_LENGTH_DELTA372 case "increase_steps":373 self._layout.steps += STEPS_DELTA374 case "decrease_steps":375 self._layout.steps -= STEPS_DELTA376 case "set_button_size":377 proposed: float = action.value378 if not (BUTTON_SIZE_MIN <= proposed <= BUTTON_SIZE_MAX):379 reward = INVALID_ACTION_REWARD380 self._layout.button_size = proposed381 case "reorder_sections" | "noop":382 pass383 384 self._clamp_layout()385 self._last_action = action.type386 return reward387 388 def _clamp_layout(self) -> None:389 self._layout.button_size = max(BUTTON_SIZE_MIN, min(BUTTON_SIZE_MAX, self._layout.button_size))390 self._layout.form_length = max(FORM_LENGTH_MIN, min(FORM_LENGTH_MAX, self._layout.form_length))391 self._layout.steps = max(STEPS_MIN, min(STEPS_MAX, self._layout.steps))392 393 def _get_observation(self) -> Observation:394 return Observation(395 device=self._device,396 layout=self._layout.model_copy(),397 progress=self._progress,398 last_action=self._last_action,399 )400 401 # ---------------------------------------------------------------------------402 # Graders (deterministic · partial-credit · strictly bounded in (0,1))403 # ---------------------------------------------------------------------------404 405 def grade_easy(self, *args, **kwargs) -> float:406 """Easy task — single objective: maximize completion progress.407 408 Sub-metrics (weighted sum):409 80 % completion progress410 20 % button-size proximity to sweet spot (1.1)411 A baseline agent can reach moderate scores easily.412 """413 progress = getattr(self, '_progress', 0.0)414 layout = getattr(self, '_layout', Layout())415 416 # --- sub-metric 1: completion progress ---417 m_progress = normalize(progress, 0.0, 1.0)418 419 # --- sub-metric 2: button in sweet spot (peak at 1.1) ---420 bs_err = abs(layout.button_size - 1.1)421 m_button = 1.0 - normalize(bs_err, 0.0, 1.6) # 1.6 = max possible error422 423 score = 0.80 * m_progress + 0.20 * m_button424 return clamp_score(score)425 426 def grade_medium(self, *args, **kwargs) -> float:427 """Medium task — multiple objectives, weighted sum, mild interactions.428 429 Sub-metrics:430 40 % completion progress431 25 % button-size proximity432 20 % form-length optimality (ideal ≈ 3)433 15 % step-count optimality (ideal ≈ 2)434 Requires coordinated improvements across dimensions.435 """436 progress = getattr(self, '_progress', 0.0)437 layout = getattr(self, '_layout', Layout())438 439 m_progress = normalize(progress, 0.0, 1.0)440 441 bs_err = abs(layout.button_size - 1.1)442 m_button = 1.0 - normalize(bs_err, 0.0, 1.6)443 444 fl_err = abs(layout.form_length - 3)445 m_form = 1.0 - normalize(fl_err, 0.0, 9.0) # range 1-10, ideal 3446 447 st_err = abs(layout.steps - 2)448 m_steps = 1.0 - normalize(st_err, 0.0, 9.0) # range 1-10, ideal 2449 450 score = 0.40 * m_progress + 0.25 * m_button + 0.20 * m_form + 0.15 * m_steps451 return clamp_score(score)452 453 def grade_hard(self, *args, **kwargs) -> float:454 """Hard task — conflicting objectives, geometric mean.455 456 Trade-offs (CANNOT maximise all simultaneously):457 Conversion — wants short forms + few steps458 Data quality — wants more fields + more steps459 Usability — device-dependent button sweet-spot460 Progress — completion rate461 462 Scoring: weighted geometric mean in log-space.463 Final score min-max stretched so worst ≈ 0.05, best ≈ 0.95.464 """465 import math466 467 progress = getattr(self, '_progress', 0.0)468 layout = getattr(self, '_layout', Layout())469 device = getattr(self, '_device', 'desktop')470 471 FLOOR = 0.05 # sub-metric floor (keeps log finite)472 473 # --- conversion: wants form_length ≤ 3 and steps ≤ 2 ---474 conv_raw = 1.0 - 0.08 * max(0, layout.form_length - 3) \475 - 0.10 * max(0, layout.steps - 2)476 m_conv = max(FLOOR, min(1.0, conv_raw))477 478 # --- data quality: wants form_length ≥ 6 and steps ≥ 5 ---479 qual_raw = 0.10 * min(layout.form_length, 10) \480 + 0.07 * min(layout.steps, 10)481 m_qual = max(FLOOR, normalize(qual_raw, 0.0, 1.7))482 483 # --- usability: device-dependent button sweet-spot ---484 optimal_bs = 1.3 if device == 'mobile' else 1.0485 usab_raw = 1.0 - abs(layout.button_size - optimal_bs) / 1.5486 m_usab = max(FLOOR, min(1.0, usab_raw))487 488 # --- progress ---489 m_prog = max(FLOOR, normalize(progress, 0.0, 1.0))490 491 # --- weighted geometric mean (log-space) ---492 log_score = (0.25 * math.log(m_prog)493 + 0.30 * math.log(m_conv)494 + 0.25 * math.log(m_qual)495 + 0.20 * math.log(m_usab))496 raw = math.exp(log_score)497 498 # stretch so empirical range [~0.05, ~0.85] maps to [~0.05, ~0.95]499 score = normalize(raw, 0.05, 0.90)500 return clamp_score(score)501 502if __name__ == "__main__":503 env = UIEnv()504 print("\n--- self-test ---")505 for t in env.tasks:506 # test with 0 args, None, and {}507 for label, call_args in [("None", (None,)), ("dict", ({},)), ("text", ("test",))]:508 val = t.grader(*call_args)509 ok = isinstance(val, float) and 0.0 < val < 1.0510 print(f" {t.name:8s} grader({label:6s}) = {val:.6f} {'OK' if ok else 'FAIL'}")511 # variation: change state and re-grade512 env._progress = 0.8513 env._layout = Layout(button_size=1.1, form_length=3, steps=2)514 print("\n--- after optimized state ---")515 for t in env.tasks:516 val = t.grader(None)517 print(f" {t.name:8s} = {val:.6f}")518