overflow_env/models.py

from typing import List, Optional
from pydantic import BaseModel, Field


class Position(BaseModel):
    x: float = 0.0
    y: float = 0.0


class CarStateData(BaseModel):
    carId: int
    lane: int
    position: Position
    speed: float
    acceleration: float = 0.0


class ProximityData(BaseModel):
    carA: int
    carB: int
    distance: float


class LaneOccupancyData(BaseModel):
    lane: int
    carIds: List[int]


class OverflowAction(BaseModel):
    decision: str = Field(default="maintain")
    reasoning: str = Field(default="")


class OverflowObservation(BaseModel):
    done: bool = False
    reward: float = 0.0
    scene_description: str = ""
    incident_report: str = ""
    cars: List[CarStateData] = Field(default_factory=list)
    proximities: List[ProximityData] = Field(default_factory=list)
    lane_occupancies: List[LaneOccupancyData] = Field(default_factory=list)


class OverflowState(BaseModel):
    episode_id: str = ""
    step_count: int = 0
    crash_count: int = 0
    near_miss_count: int = 0
    cars_reached_goal: int = 0
    total_cars: int = 5