""" environment.py — Core MetaQuestOSEnv engine. Bridges the database, task handlers, and API layer. """ from __future__ import annotations import time import random from typing import Any, Dict, List, Optional, Tuple import database as db from models import ( TaskID, SensorStatus, KernelStatus, TrackingMode, ThermalState, SensorState, ResourceState, SpatialAnchorState, KernelState, SystemObservation, ActionResult, ResetResult, StateResult, TaskInfo, ) from tasks import ( TASK_INITIALIZERS, TASK_HANDLERS, TASK_ACTIONS, TASK_MAX_STEPS, TASK_DESCRIPTIONS, ) class MetaQuestOSEnv: """ Simulates the system reliability layer of a Meta Quest-class MR device. All state is persisted in SQLite so it survives across HTTP requests. Real psutil metrics are injected into resource readings. """ def __init__(self) -> None: db.init_db() # ------------------------------------------------------------------ # Public API methods (called by FastAPI endpoints) # ------------------------------------------------------------------ def reset(self, task_id: TaskID, seed: Optional[int] = None) -> ResetResult: if seed is not None: random.seed(seed) initializer = TASK_INITIALIZERS[task_id] initializer() obs = self._build_observation(task_id, step=0) return ResetResult( observation=obs, task_id=task_id, task_description=TASK_DESCRIPTIONS[task_id], available_actions=TASK_ACTIONS[task_id], max_steps=TASK_MAX_STEPS[task_id], ) def step(self, action: str, task_id: Optional[TaskID] = None) -> ActionResult: # Resolve task_id tid_str = db.get_state("task_id") if tid_str is None: raise ValueError("Environment not initialized. Call reset() first.") tid = TaskID(tid_str) # Check if already done if db.get_state("done", False): obs = self._build_observation(tid) return ActionResult( observation=obs, reward=obs.cumulative_reward, done=True, info={"message": "Episode already complete. Call reset() to start a new episode."}, ) # Increment step current_step = db.get_state("step", 0) + 1 db.set_state("step", current_step) # Validate action valid_actions = TASK_ACTIONS[tid] if action not in valid_actions: # Partial match — try to help agent close = [a for a in valid_actions if action.split(" ")[0] in a] reward, success, message = 0.0, False, ( f"Invalid action '{action}'. " f"Did you mean: {close[:3]}?" if close else f"Invalid action. Valid actions: {valid_actions[:5]}..." ) else: handler = TASK_HANDLERS[tid] reward, success, message = handler(action) # Update cumulative reward prev_cum = db.get_state("cumulative_reward", 0.0) new_cum = prev_cum + reward db.set_state("cumulative_reward", new_cum) # Log action db.log_action(current_step, tid.value, action, reward, success, message) # Check max steps max_steps = TASK_MAX_STEPS[tid] done = db.get_state("done", False) if current_step >= max_steps and not done: done = True db.set_state("done", True) # Build observation obs = self._build_observation(tid, step=current_step, last_action=action, last_action_success=success, last_action_message=message, reward=reward, cumulative_reward=new_cum, done=done) return ActionResult( observation=obs, reward=reward, done=done, info={ "step": current_step, "max_steps": max_steps, "message": message, "action_valid": success, }, ) def state(self) -> StateResult: tid_str = db.get_state("task_id") if tid_str is None: raise ValueError("Environment not initialized. Call reset() first.") tid = TaskID(tid_str) step = db.get_state("step", 0) done = db.get_state("done", False) obs = self._build_observation(tid, step=step) return StateResult(observation=obs, task_id=tid, step=step, done=done) def list_tasks(self) -> List[TaskInfo]: difficulties = { TaskID.THERMAL_MITIGATION: "easy", TaskID.SENSOR_RECOVERY: "medium", TaskID.KERNEL_PANIC_RECOVERY: "hard", } return [ TaskInfo( task_id=tid, name=tid.value.replace("_", " ").title(), description=TASK_DESCRIPTIONS[tid], difficulty=difficulties[tid], max_steps=TASK_MAX_STEPS[tid], success_threshold=1.0, available_actions=TASK_ACTIONS[tid], ) for tid in TaskID ] # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _build_observation( self, task_id: TaskID, step: Optional[int] = None, last_action: Optional[str] = None, last_action_success: bool = True, last_action_message: str = "", reward: float = 0.0, cumulative_reward: Optional[float] = None, done: Optional[bool] = None, ) -> SystemObservation: s = db.get_all_state() if step is None: step = s.get("step", 0) if done is None: done = s.get("done", False) if cumulative_reward is None: cumulative_reward = s.get("cumulative_reward", 0.0) thermals = ThermalState( cpu_temp_c=s.get("cpu_temp_c", 65.0), gpu_temp_c=s.get("gpu_temp_c", 68.0), optical_pod_temp_c=s.get("optical_pod_temp_c", 42.0), throttle_active=s.get("throttle_active", False), refresh_rate_hz=s.get("refresh_rate_hz", 90), fan_speed_pct=s.get("fan_speed_pct", 40.0), ) sensors = SensorState( imu_status=SensorStatus(s.get("imu_status", SensorStatus.ACTIVE.value)), imu_latency_ms=s.get("imu_latency_ms", 2.0), lidar_status=SensorStatus(s.get("lidar_status", SensorStatus.ACTIVE.value)), lidar_latency_ms=s.get("lidar_latency_ms", 5.0), hand_tracking_status=SensorStatus(s.get("hand_tracking_status", SensorStatus.ACTIVE.value)), hand_tracking_latency_ms=s.get("hand_tracking_latency_ms", 8.0), tracking_mode=TrackingMode(s.get("tracking_mode", TrackingMode.SLAM.value)), tracking_stability=s.get("tracking_stability", 1.0), ) resources = ResourceState( battery_pct=s.get("battery_pct", 80.0), ram_used_mb=s.get("ram_used_mb", 2048.0), ram_total_mb=s.get("ram_total_mb", 8192.0), cpu_usage_pct=s.get("cpu_usage_pct", 30.0), frame_buffer_healthy=s.get("frame_buffer_healthy", True), background_processes=s.get("background_processes", 5), ) spatial = SpatialAnchorState( anchor_count=s.get("anchor_count", 12), anchor_stability=s.get("anchor_stability", 1.0), world_lock_active=s.get("world_lock_active", True), session_data_intact=s.get("session_data_intact", True), cached_anchors_cleared=s.get("cached_anchors_cleared", False), ) kernel = KernelState( status=KernelStatus(s.get("kernel_status", KernelStatus.ACTIVE.value)), crashed_app_id=s.get("crashed_app_id"), safe_mode_initiated=s.get("safe_mode_initiated", False), memory_lock_cleared=s.get("memory_lock_cleared", False), uptime_seconds=time.time() - s.get("uptime_seconds", time.time()), ) return SystemObservation( step=step, task_id=task_id, thermals=thermals, sensors=sensors, resources=resources, spatial_anchors=spatial, kernel=kernel, last_action=last_action, last_action_success=last_action_success, last_action_message=last_action_message, reward=reward, cumulative_reward=cumulative_reward, done=done, info={ "task_description": TASK_DESCRIPTIONS[task_id], "available_actions": TASK_ACTIONS[task_id], "max_steps": TASK_MAX_STEPS[task_id], }, ) # Singleton instance _env_instance: Optional[MetaQuestOSEnv] = None def get_env() -> MetaQuestOSEnv: global _env_instance if _env_instance is None: _env_instance = MetaQuestOSEnv() return _env_instance