| """Client for the SUPERBEAM OpenEnv environment (WebSocket session).""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Dict |
|
|
| from openenv.core.client_types import StepResult |
| from openenv.core.env_client import EnvClient |
|
|
| from .models import SuperBeamAction, SuperBeamObservation, SuperBeamState |
|
|
|
|
| class SuperBeamEnv(EnvClient[SuperBeamAction, SuperBeamObservation, SuperBeamState]): |
| """Typed client. |
| |
| Example: |
| >>> with SuperBeamEnv(base_url="http://localhost:8000").sync() as env: |
| ... r = env.reset(level="classic") |
| ... r = env.step(SuperBeamAction(action="right")) |
| ... print(r.observation.facts, r.reward, r.done) |
| """ |
|
|
| def _step_payload(self, action: SuperBeamAction) -> Dict[str, Any]: |
| return {"action": action.action} |
|
|
| def _parse_result(self, payload: Dict[str, Any]) -> StepResult[SuperBeamObservation]: |
| obs_data = dict(payload.get("observation", {})) |
| obs_data["done"] = payload.get("done", False) |
| obs_data["reward"] = payload.get("reward") |
| obs_data["metadata"] = payload.get("metadata") or obs_data.get("metadata", {}) |
| obs = SuperBeamObservation(**obs_data) |
| return StepResult(observation=obs, reward=obs.reward, done=obs.done) |
|
|
| def _parse_state(self, payload: Dict[str, Any]) -> SuperBeamState: |
| return SuperBeamState(**payload) |
|
|
