""" Logistics Shipment RL Environment =================================== Meta PyTorch OpenEnv Hackathon — Real-World Task Simulation This implements the strict OpenEnv `Environment` interface with pure Pydantic types for Observation, Action, and State. No "simulated" MCP abstractions are used here. """ import copy import random from typing import Any, Dict, List, Literal, Optional, Union from uuid import uuid4 from pydantic import BaseModel, Field try: from openenv.core.env import Environment from openenv.core.env_server.types import Action, Observation, State except ImportError: from openenv.core.env_server.interfaces import Environment from openenv.core.env_server.types import Action, Observation, State # --------------------------------------------------------------------------- # Shared route/carrier data # --------------------------------------------------------------------------- ROUTES = { "R1": {"name": "NH-48 Express (Mumbai–Pune)", "origin": "Mumbai", "destination": "Pune", "hours": 3.5, "cost": 120, "congestion": "heavy", "available": True}, "R2": {"name": "Western Highway Alt (Mumbai–Pune)", "origin": "Mumbai", "destination": "Pune", "hours": 4.0, "cost": 105, "congestion": "light", "available": True}, "R3": {"name": "NH-44 North Corridor (Delhi–Agra)", "origin": "Delhi", "destination": "Agra", "hours": 4.5, "cost": 160, "congestion": "moderate", "available": True}, "R4": {"name": "Yamuna Expressway (Delhi–Agra Alt)", "origin": "Delhi", "destination": "Agra", "hours": 3.2, "cost": 175, "congestion": "clear", "available": True}, "R5": {"name": "Chennai–Bangalore NH-48", "origin": "Chennai","destination": "Bangalore","hours": 5.0,"cost": 200,"congestion": "heavy", "available": True}, "R6": {"name": "Bangalore Alt Bypass", "origin": "Chennai","destination": "Bangalore","hours": 5.5,"cost": 185,"congestion": "light", "available": True}, } CARRIERS = ["FastFreight", "SpeedLane", "IndiaFreight", "CoastCargo", "NorthStar", "BlueLine"] # --------------------------------------------------------------------------- # Task Definitions (Easy / Medium / Hard) # --------------------------------------------------------------------------- TASKS = { "TASK-EASY": { "name": "Port Backlog Clearance", "description": "Single disruption at JNPT. Clear 2 delayed shipments within 3 turns.", "max_turns": 3, "baseline_delay": 5.5, "disruptions": ["Port congestion at JNPT (Mumbai): 4h backlog on R1"], "shipments": [ {"id": "SHIP-001", "cargo": "Fresh Vegetables (perishable)", "origin": "Mumbai", "destination": "Pune", "carrier": "FastFreight", "route": "R1", "sla_buffer_h": -1.0, "delay_h": 3.0, "value": 12000, "priority": False, "status": "DELAYED", "notes": "Stuck at gate"}, {"id": "SHIP-002", "cargo": "Auto Parts", "origin": "Mumbai", "destination": "Pune", "carrier": "SpeedLane", "route": "R1", "sla_buffer_h": 0.5, "delay_h": 2.5, "value": 31000, "priority": False, "status": "IN_TRANSIT", "notes": "Moving slowly"}, ], }, "TASK-MEDIUM": { "name": "Mumbai Crisis Coordination", "description": "Port congestion + accident + carrier strike. Manage 4 shipments over 5 turns.", "max_turns": 5, "baseline_delay": 11.0, "disruptions": ["Port congestion at JNPT: 6h backlog", "Khopoli accident: R1 +2.5h delay", "Carrier strike (FastFreight): 40% loss"], "shipments": [ {"id": "SHIP-001", "cargo": "Fresh Pharmaceuticals (perishable)", "origin": "Mumbai", "destination": "Pune", "carrier": "FastFreight", "route": "R1", "sla_buffer_h": -2.0, "delay_h": 3.5, "value": 45000, "priority": False, "status": "DELAYED", "notes": "Reefer stuck"}, {"id": "SHIP-002", "cargo": "Consumer Electronics", "origin": "Delhi", "destination": "Agra", "carrier": "NorthStar", "route": "R3", "sla_buffer_h": 1.5, "delay_h": 0.0, "value": 28000, "priority": False, "status": "IN_TRANSIT", "notes": "On time"}, {"id": "SHIP-003", "cargo": "Server Hardware (high-value)", "origin": "Mumbai", "destination": "Pune", "carrier": "SpeedLane", "route": "R1", "sla_buffer_h": -4.0, "delay_h": 5.0, "value": 180000, "priority": True, "status": "DELAYED", "notes": "Customs blocked"}, {"id": "SHIP-004", "cargo": "Industrial Chemicals (hazmat)", "origin": "Mumbai", "destination": "Pune", "carrier": "FastFreight", "route": "R1", "sla_buffer_h": -1.0, "delay_h": 2.5, "value": 22000, "priority": False, "status": "DELAYED", "notes": "Queued"}, ], }, "TASK-HARD": { "name": "Multi-Port Network Collapse", "description": "Simultaneous failures at 3 ports + weather event. 7 shipments, 7 turns.", "max_turns": 7, "baseline_delay": 28.0, "disruptions": ["JNPT CLOSED", "Chennai Port: 50% capacity", "BlueLine bankruptcy: stranded shipments", "Cold chain failure"], "shipments": [ {"id": "SHIP-001", "cargo": "COVID Vaccines", "origin": "Mumbai", "destination": "Pune", "carrier": "BlueLine", "route": "R1", "sla_buffer_h": -6.0, "delay_h": 8.0, "value": 950000, "priority": True, "status": "DELAYED", "notes": "Stranded"}, {"id": "SHIP-002", "cargo": "Election Ballots", "origin": "Delhi", "destination": "Agra", "carrier": "BlueLine", "route": "R3", "sla_buffer_h": -3.0, "delay_h": 4.0, "value": 0, "priority": True, "status": "DELAYED", "notes": "CRITICAL"}, {"id": "SHIP-003", "cargo": "Surgical Equipment", "origin": "Chennai", "destination": "Bangalore", "carrier": "CoastCargo", "route": "R5", "sla_buffer_h": -2.0, "delay_h": 5.0, "value": 340000, "priority": False, "status": "DELAYED", "notes": "Suspended soon"}, {"id": "SHIP-004", "cargo": "Petroleum (hazmat)", "origin": "Mumbai", "destination": "Pune", "carrier": "FastFreight", "route": "R1", "sla_buffer_h": -1.0, "delay_h": 3.0, "value": 88000, "priority": False, "status": "DELAYED", "notes": "Hazmat required"}, {"id": "SHIP-005", "cargo": "Consumer Electronics", "origin": "Delhi", "destination": "Agra", "carrier": "NorthStar", "route": "R3", "sla_buffer_h": 2.0, "delay_h": 2.0, "value": 120000, "priority": False, "status": "IN_TRANSIT", "notes": "Hub cyber incident"}, {"id": "SHIP-006", "cargo": "Blood Bank Supplies", "origin": "Chennai", "destination": "Bangalore", "carrier": "IndiaFreight", "route": "R5", "sla_buffer_h": -4.0, "delay_h": 6.0, "value": 75000, "priority": False, "status": "DELAYED", "notes": "Reefer failure"}, {"id": "SHIP-007", "cargo": "Agricultural Seeds", "origin": "Mumbai", "destination": "Pune", "carrier": "SpeedLane", "route": "R1", "sla_buffer_h": -8.0, "delay_h": 0.0, "value": 15000, "priority": False, "status": "CRITICAL", "notes": "Spoils in 24h"}, ], }, } # --------------------------------------------------------------------------- # Strict Pydantic Models for OpenEnv Compliance # --------------------------------------------------------------------------- class LogisticsAction(Action): """The strict typed action model representing what the AI can do.""" action_type: Literal["get_network_status", "reroute_shipment", "set_priority", "communicate_eta", "escalate", "end_turn"] = Field( ..., description="Which function to execute" ) shipment_id: Optional[str] = Field(None, description="Shipment ID if applicable") new_route: Optional[str] = Field(None, description="New route ID (for reroute)") new_carrier: Optional[str] = Field(None, description="New carrier (for reroute)") reason: Optional[str] = Field(None, description="Reason for action") message: Optional[str] = Field(None, description="Customer ETA message") priority_ids: Optional[List[str]] = Field(None, description="Shipment IDs to prioritize") class LogisticsObservation(Observation): """The strict typed observation model representing what the AI sees.""" task: str = Field(..., description="Current task ID") turn: int = Field(..., description="Current turn number") max_turns: int = Field(..., description="Maximum turns for this task") disruptions: List[str] = Field(default_factory=list, description="Active disruptions") shipments: List[Dict[str, Any]] = Field(default_factory=list, description="All shipments") feedback: Optional[str] = Field(None, description="Feedback from last action") incremental_reward: float = Field(0.0, description="Reward gained on the exact last step") turn_reward: Optional[float] = Field(None, description="Total reward for completed turn") cumulative_reward: float = Field(0.0, description="Total running reward") reward_breakdown: Optional[Dict[str, float]] = Field(None, description="Detailed score split") route_load: Dict[str, float] = Field(default_factory=dict, description="Current background traffic load (0.0 - 1.0)") class LogisticsState(State): """The strict typed state model representing the internal environment tracking.""" task_id: str = "TASK-MEDIUM" turn: int = 0 cumulative_reward: float = 0.0 incremental_reward: float = 0.0 actions_this_turn: int = 0 turn_committed: bool = False # Internal arrays shipments: List[Dict[str, Any]] = Field(default_factory=list) disruptions: List[str] = Field(default_factory=list) priority_set: List[str] = Field(default_factory=list) communications: Dict[str, str] = Field(default_factory=dict) escalations: List[str] = Field(default_factory=list) reroutings: Dict[str, Dict[str, str]] = Field(default_factory=dict) # Theme #1: Multi-Agent / Resource Scarcity route_load: Dict[str, float] = Field(default_factory=dict, description="Current background load per route (0.0 to 1.0)") # --------------------------------------------------------------------------- # Scoring Helpers # --------------------------------------------------------------------------- def _score_message(message: str) -> float: txt = message.lower() score = 0.0 if any(w in txt for w in ["sorry", "apologis", "apolog", "regret"]): score += 0.20 if any(w in txt for w in ["eta", "arrive", "delivery", "reschedule", "expect", "pm", "am", "hour"]): score += 0.40 if any(w in txt for w in ["due to", "because", "weather", "port", "delay", "congestion", "strike", "customs"]): score += 0.30 if len(message) > 80: score += 0.10 return min(1.0, score) def _message_feedback(score: float) -> str: if score >= 0.9: return "Excellent empathetic message with clear ETA." elif score >= 0.7: return "Good, but lacks either apology or specific cause/ETA." else: return "Poor. Include apology, cause of delay, and specific ETA next time." # --------------------------------------------------------------------------- # Environment Engine # --------------------------------------------------------------------------- class LogisticsShipmentEnvironment(Environment[LogisticsAction, LogisticsObservation, LogisticsState]): """ Pure Python explicit RL Environment honoring the Hackathon `openenv` spec strictly. No "MCP wrapper" translation - direct Action models to Observation models. """ SUPPORTS_CONCURRENT_SESSIONS = False def __init__(self): super().__init__() self._task_def = TASKS["TASK-MEDIUM"] self._state = LogisticsState( episode_id=str(uuid4()), step_count=0, task_id="TASK-MEDIUM", shipments=copy.deepcopy(self._task_def["shipments"]), disruptions=list(self._task_def["disruptions"]) ) self._simulate_background_traffic() @property def state(self) -> LogisticsState: return self._state def reset( self, seed: Optional[int] = None, episode_id: Optional[str] = None, task_id: str = "TASK-MEDIUM", **kwargs: Any, ) -> LogisticsObservation: self._task_def = TASKS.get(task_id, TASKS["TASK-MEDIUM"]) print(f"DEBUG: Resetting environment with task {task_id}. Shipments found: {len(self._task_def['shipments'])}") # Deepcopy the data so we can mutate it this episode initial_shipments = copy.deepcopy(self._task_def["shipments"]) if seed is not None: random.seed(seed) for s in initial_shipments: s["sla_buffer_h"] += round(random.uniform(-0.5, 0.5), 1) self._state = LogisticsState( episode_id=episode_id or str(uuid4()), step_count=0, task_id=task_id, shipments=initial_shipments, disruptions=list(self._task_def["disruptions"]) ) self._simulate_background_traffic() print(f"DEBUG: State created with {len(self._state.shipments)} shipments and {len(self._state.route_load)} route loads.") return self._build_observation("Environment reset complete.") def step(self, action: LogisticsAction, timeout_s: Optional[float] = None, **kwargs: Any) -> LogisticsObservation: self._state.incremental_reward = 0.0 # reset instantaneous counter self._state.step_count += 1 feedback = "" done = False # Theme #1: Simulate other agents' actions (background traffic) # Only update traffic at the end of the turn, so it doesn't jump constantly during intermediate actions if action.action_type == "end_turn": self._simulate_background_traffic() if action.action_type == "get_network_status": self._state.actions_this_turn += 1 self._state.incremental_reward = 0.01 feedback = "Network status fetched." elif action.action_type == "reroute_shipment": feedback = self._handle_reroute(action) self._state.actions_this_turn += 1 elif action.action_type == "set_priority": feedback = self._handle_priority(action) self._state.actions_this_turn += 1 elif action.action_type == "communicate_eta": feedback = self._handle_communication(action) self._state.actions_this_turn += 1 elif action.action_type == "escalate": feedback = self._handle_escalate(action) self._state.actions_this_turn += 1 elif action.action_type == "end_turn": feedback, done = self._handle_end_turn() else: feedback = f"Unknown action: {action.action_type}" self._state.cumulative_reward += self._state.incremental_reward obs = self._build_observation(feedback) obs.done = done obs.reward = self._state.incremental_reward return obs def _build_observation(self, feedback: str) -> LogisticsObservation: return LogisticsObservation( task=self._state.task_id, turn=self._state.turn, max_turns=self._task_def["max_turns"], disruptions=self._state.disruptions, shipments=self._state.shipments, feedback=feedback, incremental_reward=round(self._state.incremental_reward, 3), cumulative_reward=round(self._state.cumulative_reward, 3), route_load=self._state.route_load, done=False, reward=0.0 ) # ------------------------------------------------------- # Action Handlers # ------------------------------------------------------- def _handle_reroute(self, action: LogisticsAction) -> str: s_id = action.shipment_id new_r = action.new_route if not s_id or not new_r: return "Error: Missing shipment_id or new_route" shipment = next((s for s in self._state.shipments if s["id"] == s_id), None) if not shipment: return f"Error: Shipment {s_id} not found." if new_r not in ROUTES: return f"Error: Route {new_r} not valid." if shipment["route"] == new_r: return "Error: Already on that route." # Theme #1: Check Capacity current_load = self._state.route_load.get(new_r, 0.2) if current_load > 0.85: self._state.incremental_reward -= 0.05 return f"Error: Route {new_r} is at critical capacity ({(current_load*100):.1f}%). Reroute failed. Other agents are saturating this corridor." old_cong = ROUTES.get(shipment["route"], {}).get("congestion", "unknown") new_cong = ROUTES[new_r]["congestion"] savings_map = {("heavy", "light"): 2.5, ("heavy", "clear"): 3.0, ("heavy", "moderate"): 1.5, ("moderate", "light"): 1.0, ("moderate", "clear"): 1.5} savings = savings_map.get((old_cong, new_cong), 0.5) shipment["route"] = new_r if action.new_carrier: shipment["carrier"] = action.new_carrier shipment["delay_h"] = max(0.0, shipment["delay_h"] - savings) if shipment["delay_h"] == 0: shipment["status"] = "IN_TRANSIT" urgency_bonus = 0.05 if shipment["sla_buffer_h"] < 0 else 0.0 step_reward = min(0.15, savings / 20.0) + urgency_bonus self._state.incremental_reward += step_reward return f"Rerouted {s_id} to {new_r}. Delay saved: {savings}h. Immediate reward: {step_reward:.3f}." def _handle_priority(self, action: LogisticsAction) -> str: s_ids = action.priority_ids if not s_ids: return "Error: priority_ids missing." if len(s_ids) > 3: return "Error: Max 3 priority shipments allowed." self._state.priority_set = s_ids for s in self._state.shipments: s["priority"] = s["id"] in s_ids correct = [sid for sid in s_ids if any(s["id"] == sid and (s["value"] > 50000 or s["sla_buffer_h"] < 0) for s in self._state.shipments)] reward = len(correct) * 0.03 self._state.incremental_reward += reward return f"Priorities assigned to {s_ids}. Immediate reward: {reward:.3f}." def _handle_communication(self, action: LogisticsAction) -> str: if not action.shipment_id or not action.message: return "Error: missing shipment_id or message." # ANTI-HACK: Penalize duplicate messages to same shipment this turn if action.shipment_id in self._state.communications: self._state.incremental_reward -= 0.5 return (f"PENALTY: Duplicate message to {action.shipment_id} this turn. " f"Reward: -0.500. Do not spam communications — focus on uncontacted delayed shipments.") self._state.communications[action.shipment_id] = action.message score = _score_message(action.message) shipment = next((s for s in self._state.shipments if s["id"] == action.shipment_id), None) bonus = 0.02 if shipment and shipment["sla_buffer_h"] < 0 else 0.0 step_rew = (score * 0.10) + bonus self._state.incremental_reward += step_rew return f"Message logged for {action.shipment_id}. Reward: {step_rew:.3f}. Feedback: {_message_feedback(score)}" def _handle_escalate(self, action: LogisticsAction) -> str: if not action.shipment_id: return "Error: missing shipment_id." if action.shipment_id not in self._state.escalations: self._state.escalations.append(action.shipment_id) self._state.incremental_reward -= 0.1 return f"{action.shipment_id} escalated to human. Penalty -0.1 applied." return "Already escalated." def _handle_end_turn(self) -> tuple[str, bool]: if self._state.turn_committed: return "Turn already committed.", False self._state.turn_committed = True # Compute multi-dimensional turn reward total_delay = sum(s["delay_h"] for s in self._state.shipments) baseline = self._task_def["baseline_delay"] delay_saved = max(0.0, baseline - total_delay) delay_score = min(1.0, delay_saved / (baseline * 0.8)) on_time = sum(1 for s in self._state.shipments if s["sla_buffer_h"] >= 0) sla_score = on_time / len(self._state.shipments) delayed = [s for s in self._state.shipments if s["sla_buffer_h"] < 0] comm_delayed = {sid for sid in self._state.communications if any(s["id"] == sid and s["sla_buffer_h"] < 0 for s in self._state.shipments)} coverage = len(comm_delayed) / len(delayed) if delayed else 1.0 quality = sum(_score_message(m) for m in self._state.communications.values()) / len(self._state.communications) if self._state.communications else 0.0 comm_score = (0.5 * coverage) + (0.5 * quality) escalation_penalty = len(self._state.escalations) * 0.1 esc_score = max(0.0, 1.0 - escalation_penalty) act_bonus = 0.05 if self._state.actions_this_turn >= 3 else 0.0 turn_rew = min(1.0, (0.40 * delay_score + 0.30 * sla_score + 0.20 * comm_score + 0.10 * esc_score + act_bonus)) self._state.incremental_reward = turn_rew self._state.turn += 1 for s in self._state.shipments: s["sla_buffer_h"] -= 1.0 if s["sla_buffer_h"] < 0 and s["status"] == "IN_TRANSIT": s["status"] = "DELAYED" done = self._state.turn >= self._task_def["max_turns"] # Reset turn state self._state.reroutings.clear() self._state.priority_set.clear() self._state.communications.clear() self._state.escalations.clear() self._state.actions_this_turn = 0 self._state.turn_committed = False msg = f"Turn committed! Score: {turn_rew:.3f} | Delay: {delay_score:.2f}, SLA: {sla_score:.2f}, Comm: {comm_score:.2f}, Esc: {esc_score:.2f}" if done: msg += f" | 🏁 Episode Complete!" return msg, done def _simulate_background_traffic(self): """Simulate the actions of other agents in the network (Theme #1).""" # Routes have different base loads and volatility for r_id in ROUTES: base = 0.3 if ROUTES[r_id]["congestion"] == "clear" else 0.6 # Random volatility to simulate other agents' spikes self._state.route_load[r_id] = min(1.0, max(0.0, base + random.uniform(-0.2, 0.4)))