File size: 1,883 Bytes
64b4cae 1f7a59c 64b4cae | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | # models.py
# Pydantic models for the study optimizer environment
# Spent way too long deciding whether to use dataclasses or pydantic - went with pydantic
# because the FastAPI integration is just too good to pass up
from pydantic import BaseModel, Field
from typing import Optional
class Observation(BaseModel):
"""What the agent can see about the current state of the student."""
energy: float = Field(..., ge=0.0, le=100.0, description="Energy level 0-100")
focus: float = Field(..., ge=0.0, le=100.0, description="Focus score 0-100")
time_studied: float = Field(..., ge=0.0, description="Total minutes studied this session")
session_minutes: float = Field(..., ge=0.0, description="Total minutes elapsed in session")
consecutive_study_minutes: float = Field(
..., ge=0.0, description="Minutes studied without a break"
)
break_count: int = Field(..., ge=0, description="How many breaks taken so far")
productive_minutes: float = Field(
..., ge=0.0, description="Minutes where focus was above 60"
)
class Action(BaseModel):
"""The action the agent wants to take."""
# 0 = study, 1 = short break, 2 = long break
# took me a while to figure out the right granularity here
# originally had 5 actions but it made training way slower
action: int = Field(..., ge=0, le=2, description="0=study, 1=short_break, 2=long_break")
class StepResult(BaseModel):
"""What the environment returns after an action."""
observation: Observation
reward: float
done: bool
info: dict # extra stuff for debugging and grading
# TODO: maybe add a 'truncated' field later for gymnasium compatibility
class TaskInfo(BaseModel):
"""Info about the current task being evaluated."""
task_id: str
score: float = Field(..., gt=0.0, lt=1.0)
done: bool
extra: Optional[dict] = None
|