| from openenv.core.env_client import EnvClient |
| from openenv.core.client_types import StepResult |
| from openenv.core.env_server.types import State |
| from models import ContainmentAction, CityObservation |
|
|
|
|
| class CascadeContainmentEnv(EnvClient[ContainmentAction, CityObservation, State]): |
| """ |
| Client for the Cascade Containment environment. |
| |
| Async: |
| async with CascadeContainmentEnv(base_url="http://localhost:7860") as env: |
| obs = await env.reset("easy") |
| result = await env.step(ContainmentAction(action_type="allocate", district_id=0)) |
| |
| Sync: |
| with CascadeContainmentEnv(base_url="http://localhost:7860").sync() as env: |
| obs = env.reset("easy") |
| result = env.step(ContainmentAction(action_type="allocate", district_id=0)) |
| """ |
|
|
| def _step_payload(self, action: ContainmentAction) -> dict: |
| return { |
| "action_type": action.action_type, |
| "district_id": action.district_id, |
| } |
|
|
| def _parse_result(self, result: dict) -> StepResult: |
| observation = CityObservation(**result["observation"]) |
| return StepResult( |
| observation = observation, |
| reward = result.get("reward", 0.0), |
| done = result.get("done", False), |
| ) |
|
|
| def _parse_state(self, result: dict) -> State: |
| return State( |
| episode_id = result.get("episode_id", ""), |
| step_count = result.get("step_count", 0), |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| |
| with CascadeContainmentEnv(base_url="http://localhost:7860").sync() as env: |
| obs = env.reset() |
| print(f"✓ Connected") |
| print(f" Districts: {len(obs.observation.districts)}") |
| print(f" Resources: {obs.observation.available_resources}") |
| print(f" Max steps: {obs.observation.max_steps}") |
|
|
| result = env.step(ContainmentAction(action_type="allocate", district_id=0)) |
| print(f" Step reward: {result.reward}") |
| print(f"✓ End-to-end OK") |
|
|
