"""
InboxIQ RL Environment v3.0
- ML-powered email classification (TF-IDF + Naive Bayes)
- Confidence-weighted reward shaping
- Three difficulty tiers: easy medium hard
- Deterministic grading normalized to [0.0, 1.0]
"""
import os
import random
import csv
from typing import List, Dict, Optional, Any
try:
from server.classifier import EmailClassifier, analyze_sentiment, detect_urgency
from server.models import Action, Observation, State
from server.registry import TASK_REGISTRY
except ImportError:
from .classifier import EmailClassifier, analyze_sentiment, detect_urgency
from .models import Action, Observation, State
from .registry import TASK_REGISTRY
try:
from openenv.core.env_server.interfaces import Environment
except ImportError:
# Fallback if OpenEnv not installed
Environment = object
SENTIMENTS = ["Aggressive", "Professional", "Casual"]
SCORE_EPSILON = 1e-4
# Resolve dataset path try multiple locations
_possible_paths = [
os.path.join(
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
"dataset",
"spam_assassin.csv",
),
os.path.join(
os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
"dataset",
"spam_assassin.csv",
),
"/app/dataset/spam_assassin.csv", # Docker container path
"dataset/spam_assassin.csv",
"spam_assassin.csv",
]
_DEFAULT_DATASET = None
for path in _possible_paths:
if os.path.exists(path):
_DEFAULT_DATASET = path
break
class EmailEnv:
"""
InboxIQ OpenEnv-compliant RL environment for email triage.
Interface:
reset(task) observation
step(action) (observation, reward, done, info)
state() current observation
grader() normalized score [0.0, 1.0]
"""
# Enable concurrent sessions since state is isolated per instance
SUPPORTS_CONCURRENT_SESSIONS = True
def __init__(self, dataset_path: Optional[str] = None):
self.inbox: List[Dict] = []
self.total_reward: float = 0.0
self.steps: int = 0
self.max_steps: int = 10
self.current_task: str = "task1"
self.dataset_path = dataset_path or _DEFAULT_DATASET
self.raw_data: List[Dict] = []
self._episode_rewards: List[float] = []
self.initial_email_count: int = 0 # Track initial inbox size for progress
self.processed_emails: int = 0 # Track emails processed correctly
self.trajectory: List[Dict[str, Any]] = []
# ML classifier for intelligent grading
self.classifier = EmailClassifier(dataset_path=self.dataset_path)
self._load_dataset()
def _load_dataset(self):
"""Load dataset rows into memory."""
if not self.dataset_path:
print("[InboxIQ] No dataset path configured using simulation fallback.")
return
try:
print(f"[InboxIQ] Loading dataset from {self.dataset_path}...")
rows = []
with open(self.dataset_path, "r", encoding="utf-8", errors="ignore") as f:
reader = csv.DictReader(f)
for row in reader:
if row.get("text") and row.get("target") is not None:
rows.append(row)
if len(rows) > 500:
self.raw_data = random.sample(rows, 500)
else:
self.raw_data = rows
print(f"[InboxIQ] Dataset loaded. {len(self.raw_data)} samples available.")
except Exception as e:
print(f"[InboxIQ ERROR] Failed to load dataset: {e}")
def _get_random_email(self, is_train: bool = True) -> Dict:
"""Generate a random email, either from the dataset or via simulation."""
if is_train and self.raw_data:
row = random.choice(self.raw_data)
text = str(row["text"])
target = int(row["target"])
subject = "Dataset Message"
for line in text.split("\n"):
if line.lower().startswith("subject:"):
subject = line[8:].strip()
break
classification = self.classifier.classify(text, subject)
return {
"id": f"DB-{random.randint(1000, 9999)}",
"sender": "dataset@archive.org",
"subject": subject[:60] + ("..." if len(subject) > 60 else ""),
"type": "SPAM" if target == 1 else "WORK",
"urgency": classification["urgency"],
"sentiment": classification["sentiment"],
"sentiment_confidence": classification["sentiment_confidence"],
"spam_score": classification["spam_score"],
"confidence": classification["confidence"],
"wait": 0,
"ground_truth": target,
}
# Simulation fallback
sim_type = random.choice(["WORK", "SUPPORT", "SPAM"])
return {
"id": f"MSG-{random.randint(100, 999)}",
"sender": "bot@simulation.io",
"subject": "Simulated Email",
"type": sim_type,
"urgency": random.choice(["LOW", "MEDIUM", "HIGH"]),
"sentiment": random.choice(SENTIMENTS),
"sentiment_confidence": 0.5,
"spam_score": 0.8 if sim_type == "SPAM" else 0.1,
"confidence": 0.5,
"wait": 0,
}
def reset(
self,
seed: Optional[int] = None,
episode_id: Optional[str] = None,
task: Optional[str] = None,
**kwargs: Any,
) -> Observation:
"""Reset the environment for a new episode. Returns the initial observation."""
self.total_reward = 0.0
self.steps = 0
self.current_task = task or kwargs.get("task", "task1")
self._episode_rewards = []
if seed is not None:
random.seed(seed)
self.seed = seed
import uuid
self.episode_id = episode_id or str(uuid.uuid4())
task_key = self.current_task if self.current_task in TASK_REGISTRY else "task1"
self.current_task = task_key
task_info = TASK_REGISTRY[task_key]
self.max_steps = task_info.max_steps
count = task_info.config.get("count", 1)
self.inbox = [
self._get_random_email(is_train=(task_key in ["task1", "task2", "task3"]))
for _ in range(count)
]
self.initial_email_count = len(self.inbox)
self.processed_emails = 0
self.trajectory = []
# Reset rubric if present
if hasattr(self, "_reset_rubric"):
self._reset_rubric()
return self._create_observation()
def _is_correct_action(self, email: Dict, action_type: str) -> bool:
"""Determine if the action is correct for the given email."""
is_spam = email.get("type") == "SPAM" or (email.get("ground_truth", 0) == 1)
if is_spam:
return action_type == "delete"
else:
# For legitimate emails, opening or escalating is generally correct
return action_type in ["open", "escalate"]
def complex_grader(self, email: Dict, action_type: str) -> float:
"""
Native 0.0 1.0 reward function.
Every action maps to a reward strictly within [0.0, 1.0].
No negative values exist.
"""
is_spam = email.get("type") == "SPAM"
if "ground_truth" in email:
is_spam = email["ground_truth"] == 1
urgency = email.get("urgency", "MEDIUM")
sentiment = email.get("sentiment", "Professional")
# Action-specific reward mapping with reachable per-email optimum of 1.0
if action_type == "delete":
if is_spam:
reward = 1.0
else:
reward = 0.0 # CRITICAL: Deleting real mail = total failure
elif action_type == "open":
if not is_spam:
reward = 1.0
else:
reward = 0.05 # Opening spam = near-failure
elif action_type == "escalate":
if (urgency == "HIGH" or sentiment == "Aggressive") and not is_spam:
reward = 1.0
elif not is_spam:
reward = 0.55 # Unnecessary escalation on legitimate mail
else:
reward = 0.05 # Escalating spam
elif action_type == "defer":
if urgency == "HIGH":
reward = 0.1 # Deferring urgent = bad
else:
reward = 0.15 # Low-urgency defer = weak but non-zero
else:
reward = 0.0 # Unknown action
# Wait Decay (multiplier, never goes negative)
wait_steps = email.get("wait", 0)
reward *= 0.9**wait_steps
# Final clamp
reward = float(max(0.0, min(1.0, reward)))
self.classifier.record_reward(reward)
return round(reward, 4)
def step(self, action: Action, timeout_s: Optional[float] = None) -> Observation:
"""
Execute one step in the environment.
Returns: Observation with reward, done status, and metadata
"""
action_type = action.action_type.lower()
email_id = action.email_id
target_email = next((e for e in self.inbox if e["id"] == email_id), None)
step_reward = 0.0
# Track observation telemetry for this step
step_error_trace = None
current_step_feedback = None
if target_email:
# Get base reward from complex grading
step_reward = self.complex_grader(target_email, action_type)
if action_type in ["open", "delete", "escalate"]:
self.inbox = [e for e in self.inbox if e["id"] != email_id]
# Track correct processing
is_correct_action = self._is_correct_action(target_email, action_type)
if is_correct_action:
self.processed_emails += 1
current_step_feedback = (
f"Action '{action_type}' successfully executed on {email_id}."
)
else:
# Graceful System Degradation: intercept hallucinated IDs mathematically
if email_id and email_id.lower() != "none":
step_error_trace = f"InvalidTargetError: Action '{action_type}' failed. Identifier '{email_id}' does not exist in inbox."
step_reward = 0.0 # Invalid target should give 0 reward
else:
step_reward = 0.1 # Small reward for choosing to defer
current_step_feedback = "Deferred action; no target specified."
# CRITICAL: Final clamp to [0.0, 1.0] - NO INTERMEDIATE VALUES CAN EXCEED 1.0
step_reward = float(max(0.0, min(1.0, step_reward)))
# Tick wait counters for remaining emails
for email in self.inbox:
email["wait"] += 1
# Dynamic arrivals are restricted to task3 to keep task1/task2 deterministic.
if self.current_task == "task3" and random.random() < 0.3 and len(self.inbox) < 8:
self.inbox.append(
self._get_random_email(is_train=(self.current_task != "eval"))
)
self.total_reward += step_reward
self._episode_rewards.append(step_reward)
self.steps += 1
done = self.steps >= self.max_steps or len(self.inbox) == 0
# NOTE: Completion bonus is now handled in grader() via trajectory analysis
# to maintain proper normalization of individual step rewards [0.0, 1.0]
obs = self._create_observation(
reward=step_reward,
done=done,
error_trace=step_error_trace,
step_feedback=current_step_feedback,
)
# Record to trajectory for standardized grading
self.trajectory.append(
{
"action": action,
"observation": obs,
"reward": step_reward,
"done": done,
}
)
return obs
def grader(self) -> float:
"""
Normalized grader per OpenEnv/Scaler requirements.
Maps total_reward to a [0.0, 1.0] score using per-task
theoretical maxima. Negative totals are treated as failure (0.0).
"""
task_cfg = TASK_REGISTRY.get(self.current_task, TASK_REGISTRY["task1"])
if "reward_ceiling" in task_cfg.config:
r_max = float(task_cfg.config["reward_ceiling"])
else:
initial_count = task_cfg.config.get("count", 1)
# Theoretical max assumes perfect first-pass actions with unavoidable wait decay.
# Email i (0-indexed) is processed after i wait ticks, so best reward is 0.9**i.
r_max = float(sum(0.9**i for i in range(max(1, initial_count))))
if r_max <= 0.0:
return SCORE_EPSILON
normalized = float(max(0.0, min(1.0, self.total_reward / r_max)))
# Hackathon validator expects strict interval (0, 1), not inclusive bounds.
if normalized <= 0.0:
return SCORE_EPSILON
if normalized >= 1.0:
return 1.0 - SCORE_EPSILON
return normalized
def state(self) -> State:
"""Return the current environment state."""
if not hasattr(self, "episode_id") or self.episode_id is None:
import uuid
self.episode_id = str(uuid.uuid4())
return State(
episode_id=self.episode_id,
step_count=self.steps,
seed=getattr(self, "seed", None),
)
def _create_observation(
self,
reward: float = 0.0,
done: bool = False,
error_trace: Optional[str] = None,
step_feedback: Optional[str] = None,
) -> Observation:
"""Create an Observation object from current state."""
info = {
"total": round(self.total_reward, 2),
"classifier_stats": self.classifier.stats,
"episode_id": getattr(self, "episode_id", "N/A"),
}
if done:
gs = self.grader()
info["grader_score"] = round(gs, 4)
return Observation(
inbox=self.inbox,
steps=self.steps,
task=self.current_task,
max_steps=self.max_steps,
total_reward=round(self.total_reward, 4),
classifier_stats=self.classifier.stats,
reward=reward,
done=done,
info=info,
error_trace=error_trace,
step_feedback=step_feedback,
html_observation=self._generate_html_table(),
)
def _generate_html_table(self) -> str:
"""Generate a visual HTML representation of the current inbox."""
if not self.inbox:
return "Your inbox is currently empty.
"
html = ""
html += "| ID | Sender | Subject | Urgency | Sentiment | Spam Prob |
|---|
"
for email in self.inbox:
urgency = email.get("urgency", "LOW")
row_bg = "#ffffff"
if urgency == "HIGH":
row_bg = "#fff5f5"
elif email.get("type") == "SPAM":
row_bg = "#fcfcfc"
html += f""
html += f"| {email['id']} | "
html += f"{email['sender']} | "
html += f"{email['subject']} | "
html += f"{urgency} | "
html += f"{email.get('sentiment', 'Professional')} | "
html += f"{email.get('spam_score', 0.0):.2f} | "
html += "
"
html += "
"
return html