sasikumarM/detraff-env

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the BSD-style license found in the5# LICENSE file in the root directory of this source tree.6 7"""Detraff Env Environment Client."""8 9from typing import Dict10 11from openenv.core import EnvClient12from openenv.core.client_types import StepResult13from openenv.core.env_server.types import State14 15from .models import DetraffAction, DetraffObservation16 17 18class DetraffEnv(19    EnvClient[DetraffAction, DetraffObservation, State]20):21    """22    Client for the Detraff Env Environment.23 24    This client maintains a persistent WebSocket connection to the environment server,25    enabling efficient multi-step interactions with lower latency.26    Each client instance has its own dedicated environment session on the server.27 28    Example:29        >>> # Connect to a running server30        >>> with DetraffEnv(base_url="http://localhost:8000") as client:31        ...     result = client.reset()32        ...     print(result.observation.echoed_message)33        ...34        ...     result = client.step(DetraffAction(message="Hello!"))35        ...     print(result.observation.echoed_message)36 37    Example with Docker:38        >>> # Automatically start container and connect39        >>> client = DetraffEnv.from_docker_image("detraff_env-env:latest")40        >>> try:41        ...     result = client.reset()42        ...     result = client.step(DetraffAction(message="Test"))43        ... finally:44        ...     client.close()45    """46 47    def _step_payload(self, action: DetraffAction) -> Dict:48        """49        Convert DetraffAction to JSON payload for step message.50 51        Args:52            action: DetraffAction instance53 54        Returns:55            Dictionary representation suitable for JSON encoding56        """57        return {58            "message": action.message,59        }60 61    def _parse_result(self, payload: Dict) -> StepResult[DetraffObservation]:62        """63        Parse server response into StepResult[DetraffObservation].64 65        Args:66            payload: JSON response data from server67 68        Returns:69            StepResult with DetraffObservation70        """71        obs_data = payload.get("observation", {})72        observation = DetraffObservation(73            echoed_message=obs_data.get("echoed_message", ""),74            message_length=obs_data.get("message_length", 0),75            done=payload.get("done", False),76            reward=payload.get("reward"),77            metadata=obs_data.get("metadata", {}),78        )79 80        return StepResult(81            observation=observation,82            reward=payload.get("reward"),83            done=payload.get("done", False),84        )85 86    def _parse_state(self, payload: Dict) -> State:87        """88        Parse server response into State object.89 90        Args:91            payload: JSON response from state request92 93        Returns:94            State object with episode_id and step_count95        """96        return State(97            episode_id=payload.get("episode_id"),98            step_count=payload.get("step_count", 0),99        )100