aparekh02/overflow-openenv

1"""2Data models for the Overflow Environment.3 4An autonomous vehicle fleet oversight environment where an LLM agent5controls one car on a 2D road grid while other cars follow scripted rules.6 7Structured observation fields (cars, proximities, lane_occupancies) are8compatible with the Overflow frontend's CarState / AnomalyObservation types.9"""10 11from typing import Any, Dict, List, Optional12 13from pydantic import BaseModel, Field14 15try:16    from openenv.core.env_server.types import Action, Observation, State17except ImportError:18    class Action(BaseModel): pass19    class Observation(BaseModel):20        done: bool = False21        reward: float = 0.022    class State(BaseModel):23        episode_id: str = ""24        step_count: int = 025 26# ── Structured sub-models (frontend-compatible) ─────────────────────────27 28 29class Position(BaseModel):30    """2D position on the road. x = longitudinal, y = lateral."""31 32    x: float = 0.033    y: float = 0.034 35 36class CarStateData(BaseModel):37    """38    Structured per-car snapshot — matches the frontend CarState interface.39 40    Frontend type:41        interface CarState {42            carId: number; lane: number;43            position: { x: number; y: number };44            speed: number; acceleration: number;45        }46    """47 48    carId: int49    lane: int50    position: Position51    speed: float52    acceleration: float = 0.053 54 55class ProximityData(BaseModel):56    """Pairwise distance between two cars."""57 58    carA: int59    carB: int60    distance: float61 62 63class LaneOccupancyData(BaseModel):64    """Which cars are in a given lane."""65 66    lane: int67    carIds: List[int]68 69 70# ── OpenEnv core models ─────────────────────────────────────────────────71 72 73class OverflowAction(Action):74    """75    Action for the Overflow environment.76 77    The LLM agent outputs a driving decision and optional reasoning.78    """79 80    decision: str = Field(81        default="maintain",82        description="Driving decision: accelerate, brake, lane_change_left, lane_change_right, maintain",83    )84    reasoning: str = Field(85        default="",86        description="The LLM's chain-of-thought reasoning for this decision",87    )88 89 90class OverflowObservation(Observation):91    """92    Observation from the Overflow environment.93 94    Contains both:95    - Text fields (scene_description, incident_report) for the LLM to read.96    - Structured fields (cars, proximities, lane_occupancies) for the frontend97      to render, matching the Overflow frontend AnomalyObservation shape.98    """99 100    # ── Text (for the LLM) ──101    scene_description: str = Field(102        default="", description="Text description of the traffic scene"103    )104    incident_report: str = Field(105        default="", description="Observer's incident report, empty if no incident"106    )107 108    # ── Structured (for the frontend / viz) ──109    cars: List[CarStateData] = Field(110        default_factory=list, description="Structured state of every car"111    )112    proximities: List[ProximityData] = Field(113        default_factory=list, description="Pairwise proximity measurements"114    )115    lane_occupancies: List[LaneOccupancyData] = Field(116        default_factory=list, description="Per-lane vehicle occupancy"117    )118 119 120class OverflowState(State):121    """122    Internal state for the Overflow environment.123    """124 125    crash_count: int = Field(default=0, description="Number of crashes this episode")126    near_miss_count: int = Field(127        default=0, description="Number of near misses this episode"128    )129    cars_reached_goal: int = Field(130        default=0, description="Number of cars that reached their goal"131    )132    total_cars: int = Field(133        default=5, description="Total number of cars in the simulation"134    )135