Source code for examples.golf.schemas

from typing import Final

import numpy as np
import pandas as pd
import pandera.pandas as pa
from pandera.typing.pandas import Index, Series

from sara.oar import OARSchema




class GolfSchema(OARSchema):
    """Pandera OAR Schema for golf env."""

    position: Series[pd.Float64Dtype] = pa.Field(
        alias=("obs0", "position"),  # use an alias for the name of the multicolumn
        metadata={"description": "the ball position"},  # you can use metadata to
        # hold description of the variables
    )
    move: Series[pd.Float64Dtype] = pa.Field(
        alias=("act0", "move"),
        metadata={"description": "the ball variation of position"},
    )
    distance: Series[pd.Float64Dtype] = pa.Field(
        alias=("rew1", "distance"),
        le=0.0,  # you can use validators to impose conditions. Here, the
        #         distance is always negative thanks to `le`
        metadata={"description": "the negative distance from the hole"},
    )
    episodes: Index[pd.Int64Dtype] = pa.Field(
        metadata={"description": "the episode number"},
    )
    date: Index[np.datetime64] = pa.Field(
        metadata={"description": "date of the move"},
    )

    class Config:
        metadata: Final[dict] = {
            "description": (
                "A simple golf environment. "
                "The objective is to shot a ball in a hole"
            ),
            "process": (
                "1. Agent observes ball position.\n"
                "2. Agent shots the ball.\n"
                "3. Ball moves.\n"
            ),
        }