models.py

"""Pydantic models for the BrainRL region selection environment."""

from pydantic import Field

try:
    from openenv.core.env_server.types import Action, Observation
except ImportError:  # pragma: no cover
    from pydantic import BaseModel

    class Action(BaseModel):  # type: ignore[no-redef]
        pass

    class Observation(BaseModel):  # type: ignore[no-redef]
        pass


class BrainRegionAction(Action):
    """Action selecting the next brain region to acquire."""

    region_id: str = Field(
        ...,
        description="Candidate region identifier to select, for example 'pSTS'.",
    )


class BrainRegionObservation(Observation):
    """Observation returned by the BrainRL environment."""

    selection_state: str = Field(
        default="",
        description="JSON string containing selected regions, budget, score, and candidates.",
    )
    task_name: str = Field(default="roi_selection", description="Current task name")
    timestep: int = Field(default=0, ge=0, description="Current selection step")
    max_timesteps: int = Field(default=0, ge=0, description="Selection budget")
    feedback: str = Field(default="", description="Feedback from the latest region action")
    score: float = Field(default=0.0, description="Current cumulative R2 estimate")