StarTripper/ticket_ordering

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 8"""Ticket Ordering Environment Client."""9 10 11from typing import Dict12 13from openenv.core import EnvClient14from openenv.core.client_types import StepResult15 16from models import TicketOrderingAction, TicketOrderingObservation, TicketOrderingState17 18 19class TicketOrderingEnv(20    EnvClient[TicketOrderingAction, TicketOrderingObservation, TicketOrderingState]21):22    """23    Client for the Ticket Ordering Environment.24 25    This client maintains a persistent WebSocket connection to the environment server,26    enabling efficient multi-step interactions with lower latency.27    Each client instance has its own dedicated environment session on the server.28    """29 30 31    def _step_payload(self, action: TicketOrderingAction) -> Dict:32        """33        Convert TicketOrderingAction to JSON payload for step message.34 35        Args:36            action: TicketOrderingAction instance37 38        Returns:39            Dictionary representation suitable for JSON encoding40        """41        return action.model_dump()42 43 44    def _parse_result(self, payload: Dict) -> StepResult[TicketOrderingObservation]:45        """46        Parse server response into StepResult[TicketOrderingObservation].47 48        Args:49            payload: JSON response data from server50 51        Returns:52            StepResult with TicketOrderingObservation53        """54 55        observation = TicketOrderingObservation.model_validate(payload.get("observation"))56 57        return StepResult(58            observation=observation,59            reward=payload.get("reward"),60            done=payload.get("done", False),61        )62 63 64    def _parse_state(self, payload: Dict) -> TicketOrderingState:65        """66        Parse server response into State object.67 68        Args:69            payload: JSON response from state request70 71        Returns:72            State object with:73                episode_id,74                step_count,75                ordering_criteria: Criteria by which the ordering is to take place.76                optimally_ordered_ticket_ids: IDs of the tickets assuming them in their optimal positions.77                tickets: The tickets that the agent operates on during an episode. Initially shuffled randomly.78                optimality: 1.0 - Normalized Spearman footrule distance between optimal ticket IDs and actual tickets. Should always be between 0.0 and 1.0.79        """80 81        state = TicketOrderingState.model_validate(payload)82 83        return state84