TanujInsane/document-classification-env
3
1import gymnasium as gym2from gymnasium import spaces3import numpy as np4import time5import random6from tasks import TaskDataGenerator, get_task_config7from constants import CATEGORIES8 9 10class DocumentClassificationEnv(gym.Env):11 metadata = {12 "render_modes": ["human"],13 "name": "DocumentClassification-v1"14 }15 16 # Different difficulty levels = different number of categories17 # Easy = 5, Medium = 10, Hard = 22 (gets gnarly fast)18 # Category Maps are now centralized in constants.py19 CATEGORY_MAPS = CATEGORIES20 21 def __init__(self, task_difficulty="easy", seed=None):22 super().__init__()23 self.task_difficulty = task_difficulty24 if task_difficulty not in ["easy", "medium", "hard"]:25 raise ValueError(f"Pick easy, medium, or hard. Got: {task_difficulty}")26 if seed is not None:27 np.random.seed(seed)28 random.seed(seed)29 self.task_config = get_task_config(task_difficulty)30 self.num_categories = self.task_config["num_categories"]31 self.feature_dim = self.task_config["feature_dim"]32 self.action_space = spaces.Discrete(self.num_categories + 2) # +2 for Tools33 self.observation_space = spaces.Dict({34 "document_id": spaces.Text(max_length=20),35 "content": spaces.Text(max_length=10000),36 "word_count": spaces.Box(low=0, high=10000, shape=(1,), dtype=np.int32),37 "has_urgency_markers": spaces.MultiBinary(1),38 "features": spaces.Box(low=-1.0, high=1.0, shape=(self.feature_dim,), dtype=np.float32),39 "document_index": spaces.Box(low=0, high=10000, shape=(1,), dtype=np.int32),40 "total_documents": spaces.Box(low=0, high=10000, shape=(1,), dtype=np.int32),41 "metadata_response": spaces.Text(max_length=1000),42 "sla_remaining": spaces.Box(low=0, high=300, shape=(1,), dtype=np.float32),43 })44 self.data_generator = TaskDataGenerator(task_difficulty, seed)45 self.sla_seconds = self.task_config.get("sla_seconds", None)46 self.documents = None47 self.labels = None48 self.current_document_index = 049 self.episode_rewards = []50 self.correct = 051 self.total = 052 self.episode_times = []53 self.current_metadata_response = "Not requested. Use action to request customer history."54 self.doc_start_time = None # Tracks time spent on current document55 self.sla_breaches = 0 # Count of SLA breaches56 self.start_time = None57 self.action_history = []58 self.reward_history = []59 self.processing_times = []60 61 def reset(self, seed=None, options=None):62 if seed is not None:63 np.random.seed(seed)64 random.seed(seed)65 self.documents, self.labels = self.data_generator.generate_task_data()66 self.current_document_index = 067 self.episode_rewards = []68 self.correct = 069 self.total = 070 self.episode_times = []71 self.action_history = []72 self.reward_history = []73 self.processing_times = []74 self.current_metadata_response = "Not requested. Use action to request customer history."75 self.doc_start_time = time.time()76 self.sla_breaches = 077 self.start_time = time.time()78 return self._obs(), {}79 80 def _compute_sla_penalty(self):81 """Compute SLA penalty based on time spent on current document.82 The longer the agent deliberates, the higher the penalty."""83 if self.sla_seconds is None or self.doc_start_time is None:84 return 0.085 elapsed = time.time() - self.doc_start_time86 if elapsed > self.sla_seconds:87 self.sla_breaches += 188 # Penalty scales with how much over SLA (capped at -0.3)89 overage_ratio = min((elapsed - self.sla_seconds) / self.sla_seconds, 1.0)90 return -0.3 * overage_ratio91 return 0.092 93 def step(self, action):94 t0 = time.time()95 96 # Enforce per-episode time limits97 limit_per_ep = self.task_config.get("time_limit_per_episode")98 if limit_per_ep is not None and self.start_time is not None:99 if time.time() - self.start_time > limit_per_ep:100 info = {"episode_complete": True, "time_limit_exceeded": True, "episode_accuracy": self.correct / self.total if self.total > 0 else 0.0}101 info["episode_summary"] = self._summary()102 return self._obs(), 0.0, True, False, info103 104 if self.current_document_index >= len(self.documents):105 return self._obs(), 0.0, True, False, {"episode_complete": True}106 107 true_label = self.labels[self.current_document_index]108 true_cat_str = self.CATEGORY_MAPS[self.task_difficulty][true_label]109 sla_penalty = self._compute_sla_penalty()110 111 # Tool: Request Metadata112 if action == self.num_categories:113 if random.random() < 0.15:114 self.current_metadata_response = "[SYSTEM: Database timeout. No history available.]"115 else:116 self.current_metadata_response = f"[SYSTEM: User history suggests high correlation with {true_cat_str.split('-')[0]}]"117 118 tool_reward = -0.05 + sla_penalty # Tool cost + SLA pressure119 self.episode_rewards.append(tool_reward)120 self.episode_times.append(time.time() - t0)121 122 info = {123 "is_correct": False,124 "true_category": true_cat_str,125 "predicted_category": "TOOL_METADATA",126 "processing_time_ms": (time.time() - t0) * 1000,127 "episode_accuracy": self.correct / self.total if self.total > 0 else 0.0,128 "sla_penalty": sla_penalty,129 }130 return self._obs(), tool_reward, False, False, info131 132 # Tool: Escalate to Human133 elif action == self.num_categories + 1:134 self.current_metadata_response = "Not requested. Use action to request customer history."135 esc_reward = 0.0 + sla_penalty # Escalation + SLA pressure136 self.episode_rewards.append(esc_reward)137 self.total += 1138 self.episode_times.append(time.time() - t0)139 140 self.current_document_index += 1141 self.doc_start_time = time.time() # Reset SLA timer for next doc142 done = self.current_document_index >= len(self.documents)143 144 info = {145 "is_correct": False,146 "true_category": true_cat_str,147 "predicted_category": "ESCALATED",148 "processing_time_ms": (time.time() - t0) * 1000,149 "episode_accuracy": self.correct / self.total if self.total > 0 else 0.0,150 "sla_penalty": sla_penalty,151 }152 if done: info["episode_summary"] = self._summary()153 return self._obs(), esc_reward, done, False, info154 155 # Standard Classification156 self.current_metadata_response = "Not requested. Use action to request customer history."157 is_correct = action == true_label158 159 processing_time = time.time() - t0160 is_vip = self.documents[self.current_document_index].get("is_vip", False)161 error_penalty = -0.5 if is_vip else -0.4162 base_reward = 1.0 if is_correct else error_penalty163 164 # Speed Bonus: Scale reward if classification is correct and fast165 speed_bonus = 0.0166 if is_correct and self.task_difficulty in ["medium", "hard"]:167 # Bonus if under 1.0s (medium) or 0.5s (hard)168 threshold = 1.0 if self.task_difficulty == "medium" else 0.5169 if processing_time < threshold:170 speed_bonus = 0.2 * (1.0 - (processing_time / threshold))171 172 # Processing Cost: Fixed deduction for resource usage (Hard only)173 proc_cost = -0.01 if self.task_difficulty == "hard" else 0.0174 175 reward = base_reward + sla_penalty + speed_bonus + proc_cost176 177 self.episode_rewards.append(reward)178 self.action_history.append(action)179 self.reward_history.append(reward)180 self.processing_times.append(processing_time)181 182 if is_correct:183 self.correct += 1184 self.total += 1185 self.episode_times.append(processing_time)186 187 self.current_document_index += 1188 self.doc_start_time = time.time() # Reset SLA timer for next doc189 done = self.current_document_index >= len(self.documents)190 obs = self._obs()191 192 try:193 pred_cat = self.CATEGORY_MAPS[self.task_difficulty][action]194 except KeyError:195 pred_cat = f"INVALID_ACTION_{action}"196 197 info = {198 "is_correct": is_correct,199 "true_category": true_cat_str,200 "predicted_category": pred_cat,201 "processing_time_ms": processing_time * 1000,202 "episode_accuracy": self.correct / self.total if self.total > 0 else 0.0,203 "sla_penalty": sla_penalty,204 "speed_bonus": speed_bonus,205 "proc_cost": proc_cost,206 }207 if done:208 info["episode_summary"] = self._summary()209 return obs, reward, done, False, info210 211 def state(self):212 """Return the full internal environment state with rich diagnostic history."""213 # Ground-truth label for the current document (hidden from agent)214 current_true_label = None215 current_true_category = None216 if self.documents and self.current_document_index < len(self.documents):217 true_label_idx = self.labels[self.current_document_index]218 current_true_label = int(true_label_idx)219 current_true_category = self.CATEGORY_MAPS[self.task_difficulty].get(current_true_label)220 221 return {222 # ── Progress ──────────────────────────────────────────────────223 "document_index": self.current_document_index,224 "total_documents": len(self.documents) if self.documents else 0,225 # ── Hidden ground truth (not in observation) ──────────────────226 "current_true_label": current_true_label,227 "current_true_category": current_true_category,228 # ── Accumulated metrics ───────────────────────────────────────229 "total_reward": sum(self.episode_rewards),230 "correct": self.correct,231 "total_classified": self.total,232 "accuracy": self.correct / self.total if self.total > 0 else 0.0,233 "avg_time_ms": float(np.mean(self.episode_times) * 1000) if self.episode_times else 0.0,234 "sla_breaches": self.sla_breaches,235 # ── History Traces (NEW) ──────────────────────────────────────236 "action_history": list(self.action_history),237 "reward_history": [float(r) for r in self.reward_history],238 "processing_times_ms": [float(t * 1000) for t in self.processing_times],239 # ── Environment config ────────────────────────────────────────240 "difficulty": self.task_difficulty,241 "num_categories": self.num_categories,242 "sla_seconds": self.sla_seconds,243 }244 245 def _get_sla_remaining(self):246 """Calculate SLA time remaining for current document."""247 if self.sla_seconds is None or self.doc_start_time is None:248 return 999.0 # No SLA pressure249 elapsed = time.time() - self.doc_start_time250 return max(0.0, self.sla_seconds - elapsed)251 252 def _obs(self):253 sla_remaining = self._get_sla_remaining()254 if self.current_document_index >= len(self.documents):255 return {256 "document_id": "",257 "content": "",258 "word_count": np.array([0], dtype=np.int32),259 "has_urgency_markers": np.array([0], dtype=np.int8),260 "features": np.zeros(self.feature_dim, dtype=np.float32),261 "document_index": np.array([self.current_document_index], dtype=np.int32),262 "total_documents": np.array([len(self.documents)], dtype=np.int32),263 "metadata_response": self.current_metadata_response,264 "sla_remaining": np.array([0.0], dtype=np.float32),265 }266 doc = self.documents[self.current_document_index]267 return {268 "document_id": doc["id"],269 "content": doc["content"],270 "word_count": np.array([doc["word_count"]], dtype=np.int32),271 "has_urgency_markers": np.array([int(doc["has_urgency_markers"])], dtype=np.int8),272 "features": np.array(doc["features"], dtype=np.float32),273 "document_index": np.array([self.current_document_index], dtype=np.int32),274 "total_documents": np.array([len(self.documents)], dtype=np.int32),275 "metadata_response": self.current_metadata_response,276 "sla_remaining": np.array([sla_remaining], dtype=np.float32),277 }278 279 def _summary(self):280 total_reward = sum(self.episode_rewards)281 avg_reward = total_reward / len(self.episode_rewards) if self.episode_rewards else 0.0282 acc = self.correct / self.total if self.total > 0 else 0.0283 avg_time = np.mean(self.episode_times) * 1000 if self.episode_times else 0.0284 elapsed = time.time() - self.start_time285 return {286 "total_reward": total_reward,287 "avg_reward": avg_reward,288 "accuracy": acc,289 "total_classified": self.total,290 "avg_time_ms": avg_time,291 "total_time_s": elapsed,292 }293 294 def render(self):295 pass296 297 def close(self):298 pass299 300 301