openenv/models.py

from pydantic import BaseModel, ConfigDict, field_validator
import numpy as np

class Observation(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    soc: float; lmp: float; lmp_delta: float; regd: float; demand: float
    time: float; time_sin: float; time_cos: float; lmp_avg: float
    voltage: float; line_load: float

    def to_array(self):
        return np.array([
            self.soc, self.lmp, self.lmp_delta, self.regd, self.demand,
            self.time, self.time_sin, self.time_cos, self.lmp_avg,
            self.voltage, self.line_load
        ], dtype=np.float32)

class Action(BaseModel):
    power: float
    @field_validator("power", mode="before")
    @classmethod
    def coerce_to_float(cls, v):
        return float(np.asarray(v).reshape(-1)[0])