Spaces:
Running
Running
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |
| # All rights reserved. | |
| # | |
| # This source code is licensed under the BSD-style license found in the | |
| # LICENSE file in the root directory of this source tree. | |
| from __future__ import annotations | |
| try: | |
| from openenv.core.client_types import StepResult | |
| from openenv.core.env_client import EnvClient | |
| from openenv.core.env_server.types import State | |
| except ImportError: | |
| from core.client_types import StepResult | |
| from core.env_client import EnvClient | |
| from core.env_server.types import State | |
| from .models import GridWorldAction, GridWorldObservation, MoveAction | |
| class GridWorldEnv(EnvClient[GridWorldAction, GridWorldObservation, State]): | |
| """ | |
| A WebSocket-based client for interacting with the GridWorld environment. | |
| This client inherits from EnvClient and is configured with the | |
| GridWorld Pydantic models for automatic (de)serialization. | |
| """ | |
| def step_move(self, move: MoveAction) -> StepResult[GridWorldObservation]: | |
| """ | |
| Helper method to send a simple move action. | |
| Args: | |
| move: The MoveAction enum (e.g., MoveAction.UP) | |
| """ | |
| action_payload = GridWorldAction(action=move) | |
| # 'super().step' comes from the base HTTPEnvClient | |
| return super().step(action_payload) | |
| # --- REQUIRED ABSTRACT METHODS (The Missing Pieces) --- | |
| def _step_payload(self, action: GridWorldAction) -> dict: | |
| """Convert the Pydantic action model to a dictionary.""" | |
| # Uses Pydantic v2 'model_dump'. If this fails, try 'action.dict()' | |
| return action.model_dump() | |
| def _parse_result(self, data: dict) -> StepResult[GridWorldObservation]: | |
| """Convert the raw dictionary response into a typed StepResult.""" | |
| return StepResult( | |
| observation=GridWorldObservation( | |
| **data["observation"], reward=data["reward"], done=data["done"] | |
| ), | |
| reward=data["reward"], | |
| done=data["done"], | |
| metadata=data.get("metadata"), | |
| ) | |
| def _parse_state(self, data: dict) -> State: | |
| """Convert the raw state dictionary into a State object.""" | |
| return State(**data) | |