grid_env/client.py

"""
Client wrapper for the warehouse fulfillment environment.
"""

from __future__ import annotations

from .env import WarehouseFulfillmentEnv
from .models import WarehouseAction, WarehouseObservation, WarehouseState

try:
    from openenv.core.env_client import EnvClient  
except ImportError:  
    EnvClient = object


class WarehouseEnvClient(EnvClient):
    """
    Local client facade used when OpenEnv transport is not available.
    """

    def __init__(self, task_id: str = "easy_single_pick", seed: int = 7) -> None:
        self.env = WarehouseFulfillmentEnv(task_id=task_id, seed=seed)

    def reset(self, task_id: str | None = None, seed: int | None = None) -> WarehouseObservation:
        return self.env.reset(task_id=task_id, seed=seed)

    def step(self, action: WarehouseAction | str):
        return self.env.step(action)

    def state(self) -> WarehouseState:
        return self.env.state()