Spaces:
Sleeping
Sleeping
| """ | |
| ContentModerationEnv — OpenEnv-compatible environment. | |
| Simulates a real-world content moderation queue where an AI agent must | |
| review user-generated content and decide on appropriate policy actions. | |
| API: | |
| reset(task_id) → ModerationObservation | |
| step(action) → (ModerationObservation, float, bool, dict) | |
| state() → dict | |
| """ | |
| from __future__ import annotations | |
| import copy | |
| import time | |
| from typing import Any, Dict, List, Optional, Tuple | |
| from env.grader import GRADERS, grade_single_action | |
| from env.models import ( | |
| EpisodeResult, | |
| ModerationAction, | |
| ModerationDecision, | |
| ModerationObservation, | |
| ModerationReward, | |
| Severity, | |
| ViolationCategory, | |
| ) | |
| from env.tasks import ALL_TASKS | |
| # --------------------------------------------------------------------------- | |
| # Reward constants | |
| # --------------------------------------------------------------------------- | |
| R_CORRECT_DECISION = 0.60 # full credit for correct approve/remove/etc. | |
| R_CORRECT_CATEGORY = 0.25 # full credit for correct violation category | |
| R_CORRECT_SEVERITY = 0.15 # full credit for correct severity | |
| R_EFFICIENCY_BONUS = 0.10 # bonus for not over-escalating benign content | |
| R_CONSISTENCY_BONUS = 0.05 # bonus for consistent policy application | |
| R_MISS_CRITICAL = -0.50 # penalty: approved content with CRITICAL severity | |
| R_MISS_SELF_HARM = -0.80 # extra penalty: missed self-harm | |
| R_OVER_REMOVE = -0.20 # penalty: removed clearly benign content | |
| R_CONFIDENCE_PENALTY = -0.05 # penalty per 0.1 below 0.5 confidence (low quality) | |
| class ContentModerationEnv: | |
| """ | |
| OpenEnv environment for content moderation. | |
| """ | |
| ENV_ID = "content-moderation-v1" | |
| VERSION = "1.0.0" | |
| MAX_STEPS = 50 # hard cap per episode | |
| def __init__(self) -> None: | |
| self._task_id: Optional[str] = None | |
| self._difficulty: Optional[str] = None | |
| self._items: List[Dict] = [] | |
| self._current_idx: int = 0 | |
| self._step_number: int = 0 | |
| self._done: bool = True | |
| self._decisions: List[Dict] = [] | |
| self._cumulative_reward: float = 0.0 | |
| self._recent_decisions: List[Dict[str, str]] = [] | |
| # Track consistency: same violation type should receive same decision | |
| self._violation_decision_history: Dict[str, str] = {} | |
| # ----------------------------------------------------------------------- | |
| # reset | |
| # ----------------------------------------------------------------------- | |
| def reset(self, task: str = "easy") -> ModerationObservation: | |
| """ | |
| Reset the environment and return the first observation. | |
| task: "easy" | "medium" | "hard" | |
| """ | |
| if task not in ALL_TASKS: | |
| raise ValueError(f"Unknown task '{task}'. Choose from: {list(ALL_TASKS)}") | |
| task_def = ALL_TASKS[task] | |
| self._task_id = task_def["task_id"] | |
| self._difficulty = task_def["difficulty"] | |
| self._items = copy.deepcopy(task_def["items"]) | |
| self._current_idx = 0 | |
| self._step_number = 0 | |
| self._done = False | |
| self._decisions = [] | |
| self._cumulative_reward = 0.0 | |
| self._recent_decisions = [] | |
| self._violation_decision_history = {} | |
| return self._make_observation() | |
| # ----------------------------------------------------------------------- | |
| # step | |
| # ----------------------------------------------------------------------- | |
| def step( | |
| self, action: ModerationAction | |
| ) -> Tuple[ModerationObservation, float, bool, Dict[str, Any]]: | |
| """ | |
| Process one moderation decision. | |
| Returns: (observation, reward, done, info) | |
| """ | |
| if self._done: | |
| raise RuntimeError("Episode is done. Call reset() first.") | |
| if self._current_idx >= len(self._items): | |
| raise RuntimeError("No more items in queue.") | |
| current_item_def = self._items[self._current_idx] | |
| self._step_number += 1 | |
| # Validate that the action references the correct content item | |
| expected_id = current_item_def["content"].content_id | |
| if action.content_id != expected_id: | |
| # Penalise wrong content ID (agent is confused) | |
| reward_obj = ModerationReward( | |
| decision_score=0.0, category_score=0.0, severity_score=0.0, | |
| efficiency_bonus=0.0, consistency_bonus=0.0, | |
| penalty=-0.30, | |
| total=-0.30, | |
| explanation=f"Wrong content_id. Expected {expected_id}, got {action.content_id}.", | |
| ) | |
| obs = self._make_observation() | |
| info = {"reward_breakdown": reward_obj.dict(), "error": "wrong_content_id"} | |
| return obs, reward_obj.total, self._done, info | |
| gt_decision = current_item_def["gt_decision"] | |
| gt_category = current_item_def["gt_category"] | |
| gt_severity = current_item_def["gt_severity"] | |
| # --- Compute component scores --- | |
| scores = grade_single_action(action, gt_decision, gt_category, gt_severity) | |
| # --- Efficiency bonus: don't over-escalate clearly benign content --- | |
| efficiency_bonus = 0.0 | |
| if (gt_decision == ModerationDecision.APPROVE | |
| and action.decision == ModerationDecision.APPROVE | |
| and gt_category == ViolationCategory.NONE): | |
| efficiency_bonus = R_EFFICIENCY_BONUS | |
| # --- Consistency bonus --- | |
| consistency_bonus = 0.0 | |
| vcat = action.violation_category.value | |
| if vcat != ViolationCategory.NONE.value: | |
| prev = self._violation_decision_history.get(vcat) | |
| if prev is None: | |
| self._violation_decision_history[vcat] = action.decision.value | |
| elif prev == action.decision.value: | |
| consistency_bonus = R_CONSISTENCY_BONUS | |
| # --- Penalties --- | |
| penalty = 0.0 | |
| # Missed critical self-harm | |
| if (gt_category == ViolationCategory.SELF_HARM | |
| and gt_severity in (Severity.HIGH, Severity.CRITICAL) | |
| and action.decision == ModerationDecision.APPROVE): | |
| penalty += R_MISS_SELF_HARM | |
| # Missed any critical violation | |
| elif (gt_severity == Severity.CRITICAL | |
| and action.decision == ModerationDecision.APPROVE): | |
| penalty += R_MISS_CRITICAL | |
| # Over-removed clean content | |
| if (gt_decision == ModerationDecision.APPROVE | |
| and action.decision == ModerationDecision.REMOVE): | |
| penalty += R_OVER_REMOVE | |
| # Low-confidence penalty | |
| if action.confidence < 0.5: | |
| shortfall = (0.5 - action.confidence) / 0.1 | |
| penalty += R_CONFIDENCE_PENALTY * shortfall | |
| # --- Assemble total reward --- | |
| raw_total = ( | |
| scores["decision_score"] * R_CORRECT_DECISION | |
| + scores["category_score"] * R_CORRECT_CATEGORY | |
| + scores["severity_score"] * R_CORRECT_SEVERITY | |
| + efficiency_bonus | |
| + consistency_bonus | |
| + penalty | |
| ) | |
| total_reward = round(max(-1.0, min(1.0, raw_total)), 4) | |
| self._cumulative_reward += total_reward | |
| # --- Record decision --- | |
| self._decisions.append({ | |
| "action": action, | |
| "gt_decision": gt_decision, | |
| "gt_category": gt_category, | |
| "gt_severity": gt_severity, | |
| }) | |
| self._recent_decisions.append({ | |
| "content_id": action.content_id, | |
| "decision": action.decision.value, | |
| "category": action.violation_category.value, | |
| }) | |
| if len(self._recent_decisions) > 3: | |
| self._recent_decisions.pop(0) | |
| self._current_idx += 1 | |
| # --- Check done --- | |
| if self._current_idx >= len(self._items) or self._step_number >= self.MAX_STEPS: | |
| self._done = True | |
| reward_obj = ModerationReward( | |
| decision_score = scores["decision_score"] * R_CORRECT_DECISION, | |
| category_score = scores["category_score"] * R_CORRECT_CATEGORY, | |
| severity_score = scores["severity_score"] * R_CORRECT_SEVERITY, | |
| efficiency_bonus = efficiency_bonus, | |
| consistency_bonus = consistency_bonus, | |
| penalty = penalty, | |
| total = total_reward, | |
| explanation = ( | |
| f"Decision: {scores['decision_score']:.2f} | " | |
| f"Category: {scores['category_score']:.2f} | " | |
| f"Severity: {scores['severity_score']:.2f} | " | |
| f"Efficiency: {efficiency_bonus:.2f} | " | |
| f"Consistency: {consistency_bonus:.2f} | " | |
| f"Penalty: {penalty:.2f}" | |
| ), | |
| ) | |
| info: Dict[str, Any] = { | |
| "reward_breakdown": reward_obj.dict(), | |
| "step": self._step_number, | |
| "cumulative_reward": round(self._cumulative_reward, 4), | |
| "items_remaining": len(self._items) - self._current_idx, | |
| } | |
| if self._done: | |
| grader_fn = GRADERS[self._difficulty] | |
| episode_result: EpisodeResult = grader_fn(self._decisions) | |
| info["episode_result"] = episode_result.dict() | |
| obs = self._make_observation() | |
| return obs, total_reward, self._done, info | |
| # ----------------------------------------------------------------------- | |
| # state | |
| # ----------------------------------------------------------------------- | |
| def state(self) -> Dict[str, Any]: | |
| """Return the full current environment state (for debugging / logging).""" | |
| return { | |
| "env_id": self.ENV_ID, | |
| "version": self.VERSION, | |
| "task_id": self._task_id, | |
| "difficulty": self._difficulty, | |
| "step_number": self._step_number, | |
| "current_idx": self._current_idx, | |
| "total_items": len(self._items), | |
| "done": self._done, | |
| "cumulative_reward": round(self._cumulative_reward, 4), | |
| "decisions_made": len(self._decisions), | |
| "violation_decision_history": self._violation_decision_history, | |
| } | |
| # ----------------------------------------------------------------------- | |
| # Internal helpers | |
| # ----------------------------------------------------------------------- | |
| def _make_observation(self) -> ModerationObservation: | |
| """Build the current observation from internal state.""" | |
| if self._done or self._current_idx >= len(self._items): | |
| # Return a terminal observation (episode over) | |
| last = self._items[-1] if self._items else None | |
| content = last["content"] if last else None | |
| user_ctx = last["user_context"] if last else None | |
| else: | |
| item_def = self._items[self._current_idx] | |
| content = item_def["content"] | |
| user_ctx = item_def["user_context"] | |
| task_def = ALL_TASKS.get(self._difficulty, {}) | |
| applicable_rules = task_def.get("applicable_rules", []) | |
| # Backlog pressure: rises as queue fills up | |
| total = len(self._items) | |
| remaining = total - self._current_idx | |
| backlog_pressure = round(remaining / total, 2) if total > 0 else 0.0 | |
| return ModerationObservation( | |
| content = content, | |
| user_context = user_ctx, | |
| applicable_rules = applicable_rules, | |
| queue_position = self._current_idx + 1, | |
| queue_total = len(self._items), | |
| items_processed = self._current_idx, | |
| correct_so_far = sum( | |
| 1 for d in self._decisions | |
| if grade_single_action( | |
| d["action"], d["gt_decision"], d["gt_category"], d["gt_severity"] | |
| )["decision_score"] == 1.0 | |
| ), | |
| step_number = self._step_number, | |
| backlog_pressure = backlog_pressure, | |
| recent_decisions = list(self._recent_decisions), | |
| task_id = self._task_id or "", | |
| task_difficulty = self._difficulty or "", | |
| ) | |