server/nfl_playcall_env_environment.py

import pandas as pd
from pathlib import Path
from uuid import uuid4

from openenv.core.env_server.interfaces import Environment
from openenv.core.env_server.types import State

from models import NflPlaycallAction, NflPlaycallObservation

DATA_PATH = Path(__file__).resolve().parent.parent / "data.csv"
FORMATION = "Nickel (4-2-5)"


class NflPlaycallEnvironment(Environment):
    SUPPORTS_CONCURRENT_SESSIONS: bool = True

    def __init__(self):
        df = pd.read_csv(DATA_PATH)

        # Derive play_type: run / pass / play_action
        df["play_type"] = "pass"
        df.loc[df["isDropback"] == False, "play_type"] = "run"
        df.loc[(df["isDropback"] == True) & (df["playAction"] == True), "play_type"] = (
            "play_action"
        )

        # Filter to constant formation, dedupe to play level
        df = df[df["defFormation"] == FORMATION].drop_duplicates(
            subset=["gameId", "playId"]
        )
        self._df = df[["play_type", "yardsGained"]].dropna().reset_index(drop=True)

        self._state = State(episode_id=str(uuid4()), step_count=0)

    def reset(self) -> NflPlaycallObservation:
        self._state = State(episode_id=str(uuid4()), step_count=0)
        return NflPlaycallObservation(
            defensive_formation=FORMATION,
            yards_gained=0.0,
            reward=0.0,
            done=False,
        )

    def step(self, action: NflPlaycallAction) -> NflPlaycallObservation:
        self._state.step_count += 1
        subset = self._df[self._df["play_type"] == action.play_type]
        yards = float(subset.sample(min(30, len(subset)))["yardsGained"].mean())
        return NflPlaycallObservation(
            defensive_formation=FORMATION,
            yards_gained=yards,
            reward=yards,
            done=True,
        )

    @property
    def state(self) -> State:
        return self._state