models.py

from typing import Tuple
from core.env_server import Action, Observation, State

class GridAction(Action):
    direction: str

class GridObservation(Observation):
    agent_pos: Tuple[int, int]
    goal_pos: Tuple[int, int]
    reward: float
    done: bool

class GridState(State):
    episode_id: str
    steps_taken: int
    grid_size: int = 5