File size: 4,230 Bytes
02f4a63 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | """OpenRA-RL environment client.
Provides the EnvClient subclass for connecting to the OpenRA-RL
environment server over WebSocket.
"""
import os
from typing import Any, Dict
from openenv.core.client_types import StepResult
from openenv.core.env_client import EnvClient
from websockets.asyncio.client import connect as ws_connect
from openra_env.models import (
BuildingInfoModel,
EconomyInfo,
MapInfoModel,
MilitaryInfo,
OpenRAAction,
OpenRAObservation,
OpenRAState,
ProductionInfoModel,
UnitInfoModel,
)
class OpenRAEnv(EnvClient[OpenRAAction, OpenRAObservation, OpenRAState]):
"""WebSocket client for the OpenRA-RL environment.
Usage:
async with OpenRAEnv(base_url="http://localhost:8000") as env:
result = await env.reset()
while not result.done:
action = OpenRAAction(commands=[...])
result = await env.step(action)
"""
async def connect(self) -> "OpenRAEnv":
"""Connect with ping keepalive disabled.
OpenRA operations (especially reset) can take 60-120+ seconds
with software rendering. The default websockets ping_interval=20s
would kill the connection before the server responds.
"""
if self._ws is not None:
return self
ws_url_lower = self._ws_url.lower()
is_localhost = "localhost" in ws_url_lower or "127.0.0.1" in ws_url_lower
old_no_proxy = os.environ.get("NO_PROXY")
if is_localhost:
current_no_proxy = old_no_proxy or ""
if "localhost" not in current_no_proxy.lower():
os.environ["NO_PROXY"] = (
f"{current_no_proxy},localhost,127.0.0.1"
if current_no_proxy
else "localhost,127.0.0.1"
)
try:
self._ws = await ws_connect(
self._ws_url,
open_timeout=self._connect_timeout,
max_size=self._max_message_size,
ping_interval=None,
)
except Exception as e:
raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e
finally:
if is_localhost:
if old_no_proxy is None:
os.environ.pop("NO_PROXY", None)
else:
os.environ["NO_PROXY"] = old_no_proxy
return self
def _step_payload(self, action: OpenRAAction) -> Dict[str, Any]:
"""Convert action to JSON for WebSocket transport."""
return action.model_dump()
def _parse_result(self, data: Dict[str, Any]) -> StepResult[OpenRAObservation]:
"""Parse server response into StepResult."""
obs_data = data.get("observation", data)
observation = OpenRAObservation(
tick=obs_data.get("tick", 0),
economy=EconomyInfo(**obs_data.get("economy", {})),
military=MilitaryInfo(**obs_data.get("military", {})),
units=[UnitInfoModel(**u) for u in obs_data.get("units", [])],
buildings=[BuildingInfoModel(**b) for b in obs_data.get("buildings", [])],
production=[ProductionInfoModel(**p) for p in obs_data.get("production", [])],
visible_enemies=[UnitInfoModel(**u) for u in obs_data.get("visible_enemies", [])],
visible_enemy_buildings=[BuildingInfoModel(**b) for b in obs_data.get("visible_enemy_buildings", [])],
map_info=MapInfoModel(**obs_data.get("map_info", {})),
available_production=obs_data.get("available_production", []),
done=obs_data.get("done", False),
reward=obs_data.get("reward"),
result=obs_data.get("result", ""),
spatial_map=obs_data.get("spatial_map", ""),
spatial_channels=obs_data.get("spatial_channels", 0),
)
return StepResult(
observation=observation,
reward=data.get("reward", obs_data.get("reward")),
done=data.get("done", obs_data.get("done", False)),
)
def _parse_state(self, data: Dict[str, Any]) -> OpenRAState:
"""Parse state response into OpenRAState."""
return OpenRAState(**data)
|