""" server/environment.py — Core RL Environment Logic ================================================== PURPOSE: The "Kitchen" of our restaurant analogy. All game rules live here. No HTTP, no FastAPI — just pure Python logic. OpenEnv standard interface: reset() → initialises a new episode, returns first Observation step() → takes an action, returns StepResponse(Observation, reward, done, info) state() → returns a DEEP COPY of full internal State (for grader/debug) close() → optional cleanup (added for spec completeness) REWARD FUNCTION: RESPOND (1): email_value = base_priority × urgency_multiplier (range: 1–10+) normalized_cost = estimated_response_time / 120.0 (range: 0–1.5 for 180-min emails) reward = (email_value - 5.0 × normalized_cost) × (relationship_health / 100) IGNORE (0): if Spam → reward = 0 (correct decision, no time penalty, no relationship penalty) else → health_penalty = 15 × importance_weight reward = -1 × (email_value × health_penalty / 100) reading_cost = max(MIN_IGNORE_TIME, int(estimated_response_time × IGNORE_READING_TIME_FACTOR)) time_budget_remaining -= reading_cost Episode end — Sunset Penalty (only if task config features.sunset_penalty = true): penalty = -Σ(base_priority × relationship_health / 100) for remaining emails Episode end — Time Bonus (if inbox fully cleared before time runs out): bonus = (time_remaining / original_time_budget) × 10 TASK FEATURE FLAGS (read from task config): sunset_penalty → Task 1: false, Task 2: true, Task 3: true dynamic_followups → Task 1: false, Task 2: true, Task 3: true repeat_ignore_penalty→Task 1: false, Task 2: false, Task 3: true time_traps → Task 1: false, Task 2: false, Task 3: true OWNER: Algorithm Engineer """ import copy import json import os import random from typing import Dict, List, Optional, Tuple from models import ( RELATIONSHIP_CONFIG, Action, Email, EmailObservation, Relationship, State, StepResponse, ) # --------------------------------------------------------------------------- # CONSTANTS # --------------------------------------------------------------------------- DEFAULT_TIME_BUDGET = 480 # 8-hour workday in minutes INITIAL_RELATIONSHIP_HEALTH = 75.0 # All senders start at 75/100 # Agents pay a reading cost even when they ignore an email. # Without this, ignoring is zero-cost and the env is trivially exploitable. IGNORE_READING_TIME_FACTOR = 0.2 # Reading before ignoring costs 20% of response time MIN_IGNORE_TIME = 2 # Minimum 2 minutes to read any email before ignoring class EmailTriageEnv: """ The Email Triage RL Environment. One episode = one simulated workday. Each step = agent decides IGNORE (0) or RESPOND (1) to one email. Episode ends when: 1. Inbox is empty (current_email_index >= len(inbox)) — clean finish 2. Time budget runs out (time_budget_remaining <= 0) — forced end Task feature flags are read from task_config["features"] and gate mechanics: - dynamic_followups: whether ignoring VIPs injects follow-up emails - sunset_penalty: whether running out of time incurs a penalty - repeat_ignore_penalty:whether ignoring same sender twice adds extra health damage - time_traps: whether long-body VIP emails get 120-180 min response times """ def __init__(self, task_config: Optional[dict] = None, seed: Optional[int] = None): """ Args: task_config: Dict from tasks/task_X.json. Controls num_emails, distribution, time_budget, and features. If None, uses Task 1 defaults (no followups, no sunset penalty). seed: Random seed for reproducibility. When provided, the FIRST episode is seeded deterministically. Subsequent reset() calls advance the RNG normally so that repeated runs produce different episodes (variance check passes). To replay the exact same episode, pass the same seed to a fresh EmailTriageEnv() instance — do NOT re-seed on reset(). """ self.task_config = task_config or { "num_emails": 20, "vip_count": 5, "normal_count": 10, "spam_count": 5, "time_budget": DEFAULT_TIME_BUDGET, "features": { "action_cost_asymmetry": True, "dynamic_followups": False, "sunset_penalty": False, "repeat_ignore_penalty": False, "time_traps": False, }, } # Seed once at construction time only. if seed is not None: random.seed(seed) # Internal state — private, only exposed via state() deep copy self._inbox: List[Email] = [] self._relationships: Dict[str, Relationship] = {} self._current_email_index: int = 0 self._current_timestep: int = 0 self._time_budget_remaining: int = DEFAULT_TIME_BUDGET self._total_time_spent: int = 0 self._emails_handled: int = 0 self._ignored_senders: Dict[str, int] = {} # Load email templates self._email_bank = self._load_email_bank() # ----------------------------------------------------------------------- # FEATURE FLAG HELPERS # ----------------------------------------------------------------------- def _feature(self, name: str) -> bool: """Read a boolean feature flag from task_config['features'], defaulting to False.""" features = self.task_config.get("features", {}) return bool(features.get(name, False)) # ----------------------------------------------------------------------- # PUBLIC INTERFACE # ----------------------------------------------------------------------- def reset(self) -> EmailObservation: """ Start a new episode. 1. Sample emails from bank (VIP/Normal/Spam mix per task config) 2. Shuffle inbox order 3. Initialise all relationship healths at 75 4. Reset all counters 5. Return first email as Observation NOTE: Does NOT re-seed the RNG. Each call to reset() advances the RNG from where it left off, producing a different episode each time. This is intentional — Phase 2 variance checks require different episodes to produce different scores. """ cfg = self.task_config vip_pool = self._email_bank.get("vip_emails", []) normal_pool = self._email_bank.get("normal_emails", []) spam_pool = self._email_bank.get("spam_emails", []) vip_emails = random.sample(vip_pool, min(cfg["vip_count"], len(vip_pool))) normal_emails = random.sample(normal_pool, min(cfg["normal_count"], len(normal_pool))) spam_emails = random.sample(spam_pool, min(cfg["spam_count"], len(spam_pool))) if len(vip_emails) < cfg["vip_count"]: print(f"[WARNING] VIP pool too small: requested {cfg['vip_count']}, got {len(vip_emails)}", flush=True) if len(normal_emails) < cfg["normal_count"]: print(f"[WARNING] Normal pool too small: requested {cfg['normal_count']}, got {len(normal_emails)}", flush=True) if len(spam_emails) < cfg["spam_count"]: print(f"[WARNING] Spam pool too small: requested {cfg['spam_count']}, got {len(spam_emails)}", flush=True) all_emails = vip_emails + normal_emails + spam_emails random.shuffle(all_emails) # Build Email objects — assign sequential IDs for this episode self._inbox = [] for i, template in enumerate(all_emails): email = Email( email_id = i, sender = template["sender"], sender_domain = template.get("sender_domain", "external"), subject = template["subject"], body = template["body"], sender_importance = template["sender_importance"], base_priority = template["base_priority"], estimated_response_time = self._sample_time_cost(template), is_followup = template.get("is_followup", False), parent_email_id = template.get("parent_email_id", None), received_at_timestep = 0, ) email.bank_email_id = template.get("email_id") self._inbox.append(email) # Initialise relationships for all unique senders self._relationships = {} for email in self._inbox: if email.sender not in self._relationships: cfg_rel = RELATIONSHIP_CONFIG[email.sender_importance] self._relationships[email.sender] = Relationship( sender_email = email.sender, health = INITIAL_RELATIONSHIP_HEALTH, importance = email.sender_importance, importance_weight = cfg_rel["importance_weight"], degradation_rate = abs(cfg_rel["ignore_penalty"]), ) # Reset counters self._current_email_index = 0 self._current_timestep = 0 self._time_budget_remaining = self.task_config.get("time_budget", DEFAULT_TIME_BUDGET) self._total_time_spent = 0 self._emails_handled = 0 self._ignored_senders = {} return self._build_observation() def step(self, action: int) -> StepResponse: """ Execute one action and advance the environment. Args: action: 0 = IGNORE, 1 = RESPOND Returns: StepResponse(observation, reward, done, info) """ # Guard against stepping after episode ended if self._current_email_index >= len(self._inbox): return StepResponse( observation=self._build_final_observation(), reward=0.0, done=True, info={"error": "step() called after episode ended"}, ) current_email = self._inbox[self._current_email_index] relationship = self._relationships[current_email.sender] # ------ Calculate reward and time cost ------ if action == Action.RESPOND: reward, time_cost = self._reward_respond(current_email, relationship) # Update time budget self._time_budget_remaining -= time_cost self._total_time_spent += time_cost # Boost relationship health boost = RELATIONSHIP_CONFIG[current_email.sender_importance]["respond_boost"] relationship.health = min(100.0, relationship.health + boost) relationship.interaction_count += 1 info = { "action": "respond", "time_cost": time_cost, "relationship_delta": boost, "reward": reward, } else: # IGNORE reward = self._reward_ignore(current_email, relationship) # Agents still pay a reading cost on ignore (except Spam). if current_email.sender_importance != "Spam": reading_cost = max( MIN_IGNORE_TIME, int(current_email.estimated_response_time * IGNORE_READING_TIME_FACTOR), ) self._time_budget_remaining -= reading_cost self._total_time_spent += reading_cost # Apply relationship penalty penalty = RELATIONSHIP_CONFIG[current_email.sender_importance]["ignore_penalty"] relationship.health = max(0.0, relationship.health + penalty) # penalty is negative # Track ignore escalation sender = current_email.sender self._ignored_senders[sender] = self._ignored_senders.get(sender, 0) + 1 current_email.times_ignored += 1 current_email.sender_urgency_multiplier += 0.5 # Extra repeat-ignore penalty — ONLY if task enables it (Task 3) if self._feature("repeat_ignore_penalty") and self._ignored_senders[sender] > 1: relationship.health = max(0.0, relationship.health - 10) # Dynamic follow-up injection — ONLY if task enables it (Tasks 2 & 3) if self._feature("dynamic_followups") and \ current_email.sender_importance in ("VIP", "Normal"): relationship.is_angry = True # Guard prevents injecting a follow-up from a follow-up if not current_email.is_followup: already_has_followup = any( e.parent_email_id == current_email.email_id for e in self._inbox ) if not already_has_followup: followup = Email( email_id = len(self._inbox), sender = current_email.sender, sender_domain = current_email.sender_domain, subject = f"FOLLOW UP: {current_email.subject}", body = ( f"I still need a response on my previous email:\n\n" f"---\n{current_email.body[:500]}\n---\n\n" f"Please reply as soon as possible." ), sender_importance = current_email.sender_importance, base_priority = min(10, current_email.base_priority + 2), estimated_response_time = max(5, current_email.estimated_response_time // 2), is_followup = True, parent_email_id = current_email.email_id, received_at_timestep = self._current_timestep, ) insert_pos = min(self._current_email_index + 2, len(self._inbox)) self._inbox.insert(insert_pos, followup) elif not self._feature("dynamic_followups") and \ current_email.sender_importance in ("VIP", "Normal"): # Still mark angry (affects relationship score display) but no followup injected relationship.is_angry = True info = { "action": "ignore", "time_cost": reading_cost, "ignore_count": self._ignored_senders.get(current_email.sender, 0), "relationship_delta": penalty, "reward": reward, } else: # Spam: correct ignore — zero time cost, zero relationship damage info = { "action": "ignore", "time_cost": 0, "ignore_count": 0, "relationship_delta": 0, "reward": reward, } self._emails_handled += 1 self._current_email_index += 1 self._current_timestep += 1 inbox_empty = self._current_email_index >= len(self._inbox) time_up = self._time_budget_remaining <= 0 done = inbox_empty or time_up original_budget = self.task_config.get("time_budget", DEFAULT_TIME_BUDGET) if done: if time_up and not inbox_empty: # Sunset penalty — ONLY if task config enables it if self._feature("sunset_penalty"): sunset_penalty = self._calculate_sunset_penalty() reward += sunset_penalty info["sunset_penalty"] = sunset_penalty else: info["sunset_penalty"] = 0.0 elif inbox_empty: # Time bonus — always applied when inbox is cleared time_ratio = self._time_budget_remaining / original_budget time_bonus = round(time_ratio * 10, 4) reward += time_bonus info["time_bonus"] = time_bonus next_obs = self._build_final_observation() if done else self._build_observation() return StepResponse( observation=next_obs, reward=round(reward, 4), done=done, info=info, ) def state(self) -> State: """ Returns the full internal state as a DEEP COPY. Deep copy ensures the grader records frozen snapshots, not live references. """ return State( inbox = copy.deepcopy(self._inbox), current_email_index = self._current_email_index, relationships = copy.deepcopy(self._relationships), current_timestep = self._current_timestep, time_budget_remaining = self._time_budget_remaining, total_time_spent = self._total_time_spent, emails_handled = self._emails_handled, ) def close(self) -> None: """Clean up any resources. No-op — extend if file handles or sockets are added.""" pass # ----------------------------------------------------------------------- # PRIVATE HELPERS # ----------------------------------------------------------------------- def _reward_respond(self, email: Email, relationship: Relationship) -> Tuple[float, int]: """ Reward formula for RESPOND action. reward = (email_value - 5.0 × normalized_cost) × (relationship_health / 100) email_value = base_priority × urgency_multiplier normalized_cost = estimated_response_time / 120.0 (120-min email → cost 1.0, 60-min → 0.5, 5-min → ~0.04) Returns: (reward, time_cost_minutes) """ email_value = email.base_priority * email.sender_urgency_multiplier action_cost = email.estimated_response_time normalized_cost = action_cost / 120.0 reward = (email_value - 5.0 * normalized_cost) * (relationship.health / 100) return round(reward, 4), action_cost def _reward_ignore(self, email: Email, relationship: Relationship) -> float: """ Reward formula for IGNORE action. Spam: 0.0 — correct decision, no penalty. Others: -1 × (email_value × health_penalty / 100) """ if email.sender_importance == "Spam": return 0.0 health_penalty = 15 * relationship.importance_weight email_value = email.base_priority * email.sender_urgency_multiplier reward = -1 * (email_value * health_penalty / 100) return round(reward, 4) def _calculate_sunset_penalty(self) -> float: """ Penalty when time runs out with emails remaining. FinalPenalty = -Σ(base_priority × relationship_health / 100) for remaining emails """ penalty = 0.0 for i in range(self._current_email_index, len(self._inbox)): email = self._inbox[i] rel = self._relationships.get(email.sender) if rel: penalty -= email.base_priority * (rel.health / 100) return round(penalty, 4) def _build_observation(self) -> EmailObservation: """Build the agent's limited, partial view of the current email.""" if self._current_email_index >= len(self._inbox): return self._build_final_observation() email = self._inbox[self._current_email_index] rel = self._relationships.get(email.sender) rel_score = rel.health if rel else 50.0 emails_remaining = len(self._inbox) - self._current_email_index - 1 return EmailObservation( email_id = email.email_id, sender = email.sender, subject = email.subject, body = email.body, sender_importance = email.sender_importance, email_length = len(email.body), # Proxy — NOT actual time cost relationship_score = rel_score, time_budget_remaining = self._time_budget_remaining, emails_remaining = emails_remaining, ) def _build_final_observation(self) -> EmailObservation: """Dummy observation returned when episode is done.""" return EmailObservation( email_id = -1, sender = "done@done.com", subject = "Episode Complete", body = "", sender_importance = "Spam", email_length = 0, relationship_score = 0.0, time_budget_remaining = self._time_budget_remaining, emails_remaining = 0, ) def _sample_time_cost(self, template: dict) -> int: """ Derives estimated_response_time from body length so the agent can actually learn the proxy signal it's given in observations. Task 3 time_traps feature enables the 120-180 min bucket for long emails, making even VIP emails potentially not worth responding to. Buckets: < 150 chars → 5–15 min (quick reply) 150–300 chars → 15–35 min (standard) 300–500 chars → 35–70 min (deep work) > 500 chars → 70–120 min (black hole) > 500 chars + time_traps → 70–180 min (time trap — Task 3 only) """ body_length = len(template.get("body", "")) if body_length < 150: return random.randint(5, 15) elif body_length < 300: return random.randint(15, 35) elif body_length < 500: return random.randint(35, 70) else: # time_traps extends the upper bound to 180 min (Task 3 only) upper = 180 if self._feature("time_traps") else 120 return random.randint(70, upper) def _load_email_bank(self) -> dict: """Load pre-written email templates from data/email_bank.json.""" possible_paths = [ os.path.join(os.path.dirname(__file__), "..", "data", "email_bank.json"), os.path.join(os.path.dirname(__file__), "data", "email_bank.json"), "data/email_bank.json", ] for path in possible_paths: if os.path.exists(path): with open(path, "r") as f: return json.load(f) # Minimal fallback — fires loud warnings in reset() when pool is too small return { "vip_emails": [ { "sender": "boss@company.com", "sender_domain": "internal", "subject": "Q4 Report Review Needed", "body": "Please review the attached Q4 report and share feedback by EOD.", "sender_importance": "VIP", "base_priority": 9, "estimated_response_time": 45, "is_followup": False, "parent_email_id": None, }, { "sender": "cto@company.com", "sender_domain": "internal", "subject": "Architecture Decision Needed", "body": "We need a decision on the new microservices architecture. Please review and respond.", "sender_importance": "VIP", "base_priority": 8, "estimated_response_time": 60, "is_followup": False, "parent_email_id": None, }, ], "normal_emails": [ { "sender": "teammate@company.com", "sender_domain": "internal", "subject": "PR Review Request", "body": "Hey, can you review my pull request when you get a chance? It's a small fix.", "sender_importance": "Normal", "base_priority": 5, "estimated_response_time": 15, "is_followup": False, "parent_email_id": None, }, ], "spam_emails": [ { "sender": "newsletter@random.com", "sender_domain": "external", "subject": "Top 10 Productivity Hacks!", "body": "Unsubscribe at any time. Limited offer today only!", "sender_importance": "Spam", "base_priority": 1, "estimated_response_time": 5, "is_followup": False, "parent_email_id": None, }, ], }