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"""8Ticket Ordering Environment Implementation.9"""10 11 12import heapq13import logging14import numpy as np15from uuid import uuid416from copy import deepcopy17 18from openenv.core.env_server.interfaces import Any, Environment, Optional19 20 21logger = logging.getLogger("uvicorn.error")22logger.setLevel(logging.DEBUG)23 24 25try:26    from models import Ticket, TicketHeuristic, TicketOrderingConfig, TicketOrderingState, TicketOrderingObservation, TicketOrderingAction27    from problem_generator import generate_problem_statement, GenerationDifficulty28except ImportError:29    from ..models import Ticket, TicketHeuristic, TicketOrderingConfig, TicketOrderingState, TicketOrderingObservation, TicketOrderingAction30    from ..problem_generator import generate_problem_statement, GenerationDifficulty31 32 33class TicketOrderingEnvironment(Environment):34    """35    An environment that orders tickets.36    """37 38    # Enable concurrent WebSocket sessions.39    # Set to True if your environment isolates state between instances.40    # When True, multiple WebSocket clients can connect simultaneously, each41    # getting their own environment instance (when using factory mode in app.py).42    SUPPORTS_CONCURRENT_SESSIONS: bool = True43 44    def __init__(self):45        """Initialize the ticket_ordering environment."""46 47        self._config = TicketOrderingConfig()48 49        self._reset_count = 050 51        self.rng = np.random.default_rng(42)52 53        self._state = TicketOrderingState(54            episode_id=str(uuid4()),55            step_count=0,56            ordering_criteria="",57            optimally_ordered_ticket_ids=[],58            tickets=[59                Ticket(id=0, thread=[], heuristic=TicketHeuristic()),60                Ticket(id=1, thread=[], heuristic=TicketHeuristic())61            ],62            optimality=0.5,63        )64        self._state_id_index_map: dict[int, int] = {}65 66        self._current_candidate: Ticket = Ticket(67            id=0,68            thread=[],69            heuristic=TicketHeuristic(priority=0.0, summary="")70        )71        self._current_references: list[Ticket] = []72        self._current_heuristics: dict[int, TicketHeuristic] = {}73 74 75    def reset(76        self,77        seed: Optional[int] = 42,78        episode_id: Optional[str] = str(uuid4()),79        difficulty: Optional[int] = GenerationDifficulty.Medium.value,80        **kwargs: Any,81    ) -> TicketOrderingObservation:82        """83        Reset the environment.84 85        Returns:86            TicketOrderingObservation87        """88 89        self._reset_count += 190 91        self.rng = np.random.default_rng(seed if seed is not None else 42)92 93        generated_criteria, generated_tickets = generate_problem_statement(94            GenerationDifficulty(difficulty) if difficulty is not None else GenerationDifficulty.Medium95        )96        shuffled_tickets = deepcopy(generated_tickets)97        self.rng.shuffle(shuffled_tickets)98        optimal_ticket_ids = [ticket.id for ticket in generated_tickets]99 100        self._state = TicketOrderingState(101            episode_id=episode_id or str(uuid4()),102            step_count=0,103            ordering_criteria=generated_criteria,104            optimally_ordered_ticket_ids=optimal_ticket_ids,105            tickets=shuffled_tickets,106            optimality=self.compute_optimality(shuffled_tickets, optimal_ticket_ids),107        )108        self._state_id_index_map = self._make_id_index_map(self._state.tickets)109 110        self._current_candidate = self._state.tickets[0]111        self._current_references = self._state.tickets[0:1]112        self._current_heuristics = self.select_heuristics(self._state)113 114 115        return TicketOrderingObservation(116            done=False,117            reward=0.0,118            metadata={},119 120            ordering_criteria=self._state.ordering_criteria,121            reference_tickets=self._current_references,122            candidate_ticket=self._current_candidate,123            ticket_heuristics=self._current_heuristics,124            total_tickets=len(self._state.tickets),125            completed_iterations=self._state.step_count126        )127 128 129    def step(self, action: TicketOrderingAction) -> TicketOrderingObservation:  # type: ignore[override]130        """131        Execute a step in the environment.132 133        Args:134            action: TicketOrderingAction containing the assigned priority and summary for the candidate ticket,135            next reference ticket ids, next candidate ticket id and whether or not ordering should stop (in case136            the agent decides ordering has "reached optimality" at any point in time)137 138        Returns:139            TicketOrderingObservation140        """141 142        previous_state = deepcopy(self._state)143 144        updated_state = self.get_updated_state(previous_state, self._current_candidate, action)145 146        self._current_candidate = self.select_candidate(updated_state, action)147        self._current_references = self.select_references(updated_state, action)148        self._current_heuristics = self.select_heuristics(updated_state)149 150        self._state = deepcopy(updated_state)151 152        observation = self.construct_observation(previous_state, updated_state, action)153 154        return observation155 156 157    @property158    def state(self) -> TicketOrderingState:159        """160        Get the current environment state.161 162        Returns:163            Current State with episode_id and step_count164        """165        return self._state    166 167 168    def select_candidate(self, state: TicketOrderingState, action: TicketOrderingAction) -> Ticket:169        id_index_map = self._make_id_index_map(state.tickets)170 171        if action.next_candidate_id in id_index_map:172            candidate_ticket_index = id_index_map[action.next_candidate_id]173        else:174            candidate_ticket_index = id_index_map[state.tickets[0].id]175 176        return deepcopy(state.tickets[candidate_ticket_index])177 178 179    def select_references(self, state: TicketOrderingState, action: TicketOrderingAction) -> list[Ticket]:180        id_index_map = self._make_id_index_map(state.tickets)181        182        reference_ticket_indices = []183        for id in action.next_reference_ids[:self._config.max_reference_tickets]:184            if id in id_index_map:185                reference_ticket_indices.append(id_index_map[id])186            else:187                reference_ticket_indices.append(id_index_map[state.tickets[0].id])188 189        return deepcopy(190            [state.tickets[index] for index in reference_ticket_indices]191        )192 193 194    def select_heuristics(self, state: TicketOrderingState) -> dict[int, TicketHeuristic]:195        heuristics = {}196 197        assert self._config.max_heurestics >= 2198        assert self._config.max_heurestics % 2 == 0199 200        n = self._config.max_heurestics // 2201        largest = heapq.nlargest(n, state.tickets, key=lambda x: (x.heuristic.times_assigned, self.rng.random()))202        smallest = heapq.nsmallest(n, state.tickets, key=lambda x: (x.heuristic.times_assigned, self.rng.random()))203 204        heuristics.update(205            {ticket.id: ticket.heuristic for ticket in largest}206        )207        heuristics.update(208            {ticket.id: ticket.heuristic for ticket in smallest}209        )210 211        return heuristics212 213 214    def construct_observation(215        self,216        previous_state: TicketOrderingState,217        new_state: TicketOrderingState,218        action: TicketOrderingAction,219    ) -> TicketOrderingObservation:220        observation = TicketOrderingObservation(221            done=action.end_ordering or (new_state.step_count >= self._config.max_steps),222            reward=self.construct_reward(previous_state, new_state, action),223 224            ordering_criteria=new_state.ordering_criteria,225            reference_tickets=self._current_references,226            candidate_ticket=self._current_candidate,227            ticket_heuristics=self._current_heuristics,228 229            total_tickets=len(new_state.tickets),230            completed_iterations=new_state.step_count231        )232 233        return observation234 235 236    def construct_reward(237        self,238        previous_state: TicketOrderingState,239        new_state: TicketOrderingState,240        action: TicketOrderingAction,241    ) -> float:242        reward = 0.0243 244        improvement = new_state.optimality - previous_state.optimality245 246        # Actually really good reward metric because huge reorders are rewarded / penalized hugely247        # but small ones are rewarded / penalized slightly (proportional, smooth reward metric).248        # RL algorithms are usually built with a value decay anyways (gamma) this means early-large reorders are prioritized more than late-large ones249        reward += improvement250 251        # Ending ordering comes on purpose comes with consequences! Ending early with an optimality of more than 0.5 is rewarded252        # while ending it with an optimality of less than 0.5 is punished.253        #254        # More steps taken results in lessened importance to this reward signal overall, a more careful agent that takes more steps is to be255        # rewarded and punished less than an overly aggressive and unreliable one. Note that an aggressive but reliably high performing256        # agent is still rewarded greatly.257        if action.end_ordering:258            reward += (new_state.optimality * 2.0 - 1.0) * (1.5 - new_state.step_count / self._config.max_steps)259 260        return reward261 262 263    def get_updated_state(self, previous_state: TicketOrderingState, candidate: Ticket, action: TicketOrderingAction) -> TicketOrderingState:264        id_index_map = self._make_id_index_map(previous_state.tickets)265        updated_state = deepcopy(previous_state)266 267        updated_state.step_count += 1268 269        candidate_ticket_index = id_index_map[candidate.id]270        updated_state.tickets[candidate_ticket_index].heuristic.priority = action.candidate_priority271        updated_state.tickets[candidate_ticket_index].heuristic.summary = action.candidate_summary272        updated_state.tickets[candidate_ticket_index].heuristic.times_assigned += 1273 274        updated_state.tickets = self.reorder_tickets(updated_state.tickets)275 276        updated_state.optimality = self.compute_optimality(updated_state.tickets, updated_state.optimally_ordered_ticket_ids)277 278        return updated_state279 280 281    def compute_optimality(self, tickets: list[Ticket], optimal_ids: list[int]):282        normalized_distance = self.normalized_spearman_footrule_distance(283            [ticket.id for ticket in tickets],284            [id for id in optimal_ids],285        )286        return 1.0 - normalized_distance287 288 289    def _make_id_index_map(self, tickets: list[Ticket]) -> dict[int, int]:290        id_index_map = {ticket.id: index for index, ticket in enumerate(tickets)}291        return id_index_map292 293 294    def reorder_tickets(self, tickets: list[Ticket]) -> list[Ticket]:295        copied_tickets = deepcopy(tickets)296        copied_tickets.sort(key=lambda ticket: ticket.heuristic.priority)297        return copied_tickets298 299 300    def normalized_spearman_footrule_distance(self, a: list[int], b: list[int]) -> float:301        assert len(a) == len(b)302        assert len(a) > 0303 304        a_value_index_map = {305            value: index for index, value in enumerate(a)306        }307 308        accumulated_distance = 0.0309        for index, value in enumerate(b):310            a_index = a_value_index_map[value]311            distance = float(abs(index - a_index))312            accumulated_distance += distance313 314        n = len(a)315        if n > 1:316            max_distance = (n * n) / 2.0 if n % 2 == 0 else (n * n - 1) / 2.0317        else:318            max_distance = 1.0319 320        normalized_distance = accumulated_distance / max_distance321 322        return normalized_distance323