openenv/atari_env

1"""2core/runner_env.py3Minimal HTTP-based environment client.4- Talks to a single env worker exposing: POST /reset, POST /step5 6Future hooks (commented below) for:7- episode_id, seed on reset8- request_id on step9- custom headers (auth/trace)10"""11 12from __future__ import annotations13 14from abc import ABC, abstractmethod15from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar16from .containers.runtime import LocalDockerProvider17import requests18 19from .types import StepResult20 21if TYPE_CHECKING:22    from .containers.runtime import ContainerProvider23 24ActT = TypeVar("ActT")25ObsT = TypeVar("ObsT")26EnvClientT = TypeVar("EnvClientT", bound="HTTPEnvClient")27 28 29class HTTPEnvClient(ABC, Generic[ActT, ObsT]):30    def __init__(31        self,32        base_url: str,33        request_timeout_s: float = 15.0,34        default_headers: Optional[Dict[str, str]] = None,35        provider: Optional["ContainerProvider"] = None,36    ):37        self._base = base_url.rstrip("/")38        self._timeout = float(request_timeout_s)39        self._http = requests.Session()40        self._headers = default_headers or {}41        self._provider = provider42 43    @classmethod44    def from_docker_image(45        cls: Type[EnvClientT],46        image: str,47        provider: Optional["ContainerProvider"] = None,48    ) -> EnvClientT:49        """50        Create an environment client by spinning up a Docker container locally.51 52        This is a development utility that:53        1. Starts a Docker container from the specified image54        2. Waits for the server to be ready55        3. Creates and returns a client instance connected to the container56 57        Note: The container lifecycle management is left to the user or higher-level58        orchestration. The container will keep running until manually stopped.59 60        Args:61            image: Docker image name to run (e.g., "echo-env:latest")62            provider: Container provider to use (defaults to LocalDockerProvider)63 64        Returns:65            An instance of the client class connected to the running container66 67        Example:68            >>> from envs.coding_env.client import CodingEnv69            >>> from envs.coding_env.models import CodeAction70            >>>71            >>> # Create environment from image72            >>> env = CodingEnv.from_docker_image("coding-env:latest")73            >>>74            >>> # Use the environment75            >>> result = env.reset()76            >>> print(result.observation)77            >>>78            >>> step_result = env.step(CodeAction(code="print('hello')"))79            >>> print(step_result.observation.stdout)80            >>>81            >>> # Cleanup (optional)82            >>> env.close()83        """84 85        # Use default provider if none provided86        if provider is None:87            provider = LocalDockerProvider()88 89        # 1. Start container90        base_url = provider.start_container(image)91 92        # 2. Wait for server to be ready93        provider.wait_for_ready(base_url)94 95        # 3. Create and return client instance with provider reference96        return cls(base_url=base_url, provider=provider)97 98    @abstractmethod99    def _step_payload(self, action: ActT) -> dict:100        """Convert an Action object to the JSON body expected by the env server."""101        raise NotImplementedError102 103    @abstractmethod104    def _parse_result(self, payload: dict) -> StepResult[ObsT]:105        """Convert a JSON response from the env server to StepResult[ObsT]."""106        raise NotImplementedError107 108    @abstractmethod109    def _parse_state(self, payload: dict) -> Any:110        """Convert a JSON response from the state endpoint to a State object."""111        raise NotImplementedError112 113    # ---------- Environment Server Interface Methods ----------114    def reset(self) -> StepResult[ObsT]:115        body: Dict[str, Any] = {}116        # TODO: later:117        # body["seed"] = seed118        # body["episode_id"] = episode_id119        r = self._http.post(120            f"{self._base}/reset",121            json=body,122            headers=self._headers,123            timeout=self._timeout,124        )125        r.raise_for_status()126        return self._parse_result(r.json())127 128    def step(self, action: ActT) -> StepResult[ObsT]:129        body: Dict[str, Any] = {130            "action": self._step_payload(action),131            "timeout_s": int(self._timeout),132        }133        # TODO: later:134        # body["request_id"] = str(uuid.uuid4())135        # body["episode_id"] = current_episode_id136        r = self._http.post(137            f"{self._base}/step",138            json=body,139            headers=self._headers,140            timeout=self._timeout,141        )142        r.raise_for_status()143        return self._parse_result(r.json())144 145    def state(self) -> Any:146        """147        Get the current environment state from the server.148 149        Returns:150            State object with environment state information (e.g., episode_id, step_count)151 152        Example:153            >>> client = EchoEnv.from_docker_image("echo-env:latest")154            >>> result = client.reset()155            >>> state = client.state()156            >>> print(state.episode_id)157            >>> print(state.step_count)158        """159        r = self._http.get(160            f"{self._base}/state",161            headers=self._headers,162            timeout=self._timeout,163        )164        r.raise_for_status()165        return self._parse_state(r.json())166 167    def close(self) -> None:168        """169        Close the environment and clean up resources.170 171        If this client was created via from_docker_image(), this will stop172        and remove the associated container.173        """174        if self._provider is not None:175            self._provider.stop_container()176