from __future__ import annotations import argparse import json import random import sys from pathlib import Path from typing import Any, Dict, List import gymnasium as gym import matplotlib.pyplot as plt import numpy as np from stable_baselines3 import PPO from stable_baselines3.common.callbacks import BaseCallback from stable_baselines3.common.vec_env import DummyVecEnv PROJECT_ROOT = Path(__file__).resolve().parents[1] if str(PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(PROJECT_ROOT)) from models import DataRecord, EraseOpsAction, EraseOpsObservation from server.environment import EraseOpsEnvironment from server.graders import grade ALL_TASKS = ["clean_sweep", "tangled_lake", "legal_minefield", "cascade_collapse"] EVAL_SEEDS = [42, 43, 44, 45, 46] class GdprEraseOpsGymEnv(gym.Env[np.ndarray, int]): metadata = {"render_modes": []} def __init__(self, tasks: List[str], max_records: int = 32): super().__init__() self.tasks = tasks self.max_records = max_records self._rng = random.Random(42) self._env = EraseOpsEnvironment() self._obs: EraseOpsObservation | None = None self._current_task = tasks[0] self._escalated_ids: set[str] = set() self._task_weights: Dict[str, float] = {task: 1.0 for task in tasks} self.action_space = gym.spaces.Discrete(1 + 4 * max_records) self.observation_space = gym.spaces.Box( low=-10.0, high=10.0, shape=(10 + max_records * 10,), dtype=np.float32, ) def reset(self, *, seed: int | None = None, options: Dict[str, Any] | None = None): if seed is not None: self._rng.seed(seed) opts = options or {} task_name = str(opts.get("task_name") or self._sample_task()) task_seed = int(opts.get("task_seed") or self._rng.randint(1, 10_000)) self._current_task = task_name self._obs = self._env.reset(task_name=task_name, seed=task_seed) self._escalated_ids = set() return self._encode_observation(self._obs), {"task_name": task_name, "task_seed": task_seed} def step(self, action: int): assert self._obs is not None, "Environment must be reset before step" action_obj = self._decode_action(int(action), self._obs) action_obj = self._guardrail_action(action_obj, self._obs) obs, reward, done, info = self._env.step(action_obj) if action_obj.action_type == "escalate" and action_obj.record_id is not None: self._escalated_ids.add(action_obj.record_id) self._obs = obs info = dict(info) info["task_name"] = self._current_task return self._encode_observation(obs), float(reward.total), bool(done), False, info def _guardrail_action(self, action: EraseOpsAction, obs: EraseOpsObservation) -> EraseOpsAction: # Keep a valid policy surface for PPO by replacing obviously invalid actions with safe defaults. by_id = {record.record_id: record for record in obs.records} record = by_id.get(action.record_id) if action.record_id else None target_user = obs.request.user_id if action.action_type in {"inspect", "delete", "anonymize", "escalate"} and record is None: return self._heuristic_action(obs) if record is not None and record.status != "active": return self._heuristic_action(obs) if action.action_type in {"delete", "anonymize"} and record is not None and record.user_id != target_user: return self._heuristic_action(obs) if action.action_type == "delete" and record is not None: if record.legal_hold or record.data_type != "PII": return self._heuristic_action(obs) unresolved = self._unresolved_dependencies(record, by_id) if unresolved: return self._heuristic_action(obs) if action.action_type == "escalate" and record is not None and not record.legal_hold: return self._heuristic_action(obs) return action def _heuristic_action(self, obs: EraseOpsObservation) -> EraseOpsAction: by_id = {record.record_id: record for record in obs.records} target_user = obs.request.user_id # 1) Escalate unresolved legal-hold records for target user. for record in obs.records: if ( record.user_id == target_user and record.legal_hold and record.status == "active" and record.record_id not in self._escalated_ids ): return EraseOpsAction(action_type="escalate", record_id=record.record_id, justification="guardrail-escalate") # 2) Delete safe, in-scope, dependency-resolved PII records. deletable = [] for record in obs.records: if record.user_id != target_user or record.status != "active" or record.data_type != "PII" or record.legal_hold: continue if not self._unresolved_dependencies(record, by_id): deletable.append(record) if deletable: deletable.sort(key=lambda rec: (-rec.age_days, rec.record_id)) return EraseOpsAction( action_type="delete", record_id=deletable[0].record_id, justification="guardrail-delete", ) # 3) Inspect unresolved target-user PII for additional signal. for record in obs.records: if record.user_id == target_user and record.status == "active" and record.data_type == "PII": return EraseOpsAction(action_type="inspect", record_id=record.record_id, justification="guardrail-inspect") return EraseOpsAction(action_type="noop", record_id=None, justification="guardrail-noop") def _unresolved_dependencies(self, record: DataRecord, by_id: Dict[str, DataRecord]) -> List[str]: unresolved: List[str] = [] for dep_id in record.dependencies: dep = by_id.get(dep_id) if dep is None: unresolved.append(dep_id) continue if dep.status == "deleted": continue if dep.legal_hold and dep.record_id in self._escalated_ids: continue unresolved.append(dep_id) return unresolved def set_task_weights(self, weights: Dict[str, float]) -> None: filtered = {task: max(0.0, float(weights.get(task, 0.0))) for task in self.tasks} total = sum(filtered.values()) if total <= 0.0: self._task_weights = {task: 1.0 for task in self.tasks} return self._task_weights = filtered def _sample_task(self) -> str: total = sum(self._task_weights.values()) if total <= 0.0: return self._rng.choice(self.tasks) draw = self._rng.uniform(0.0, total) running = 0.0 for task in self.tasks: running += self._task_weights.get(task, 0.0) if draw <= running: return task return self.tasks[-1] def _decode_action(self, action_idx: int, obs: EraseOpsObservation) -> EraseOpsAction: if action_idx <= 0: return EraseOpsAction(action_type="noop", record_id=None, justification="ppo noop") record_slot = (action_idx - 1) % self.max_records action_block = (action_idx - 1) // self.max_records action_types = ["inspect", "delete", "anonymize", "escalate"] action_type = action_types[min(action_block, len(action_types) - 1)] record_id = None if 0 <= record_slot < len(obs.records): record_id = obs.records[record_slot].record_id return EraseOpsAction( action_type=action_type, record_id=record_id, justification=f"ppo:{action_type}:{record_slot}", ) def _encode_observation(self, obs: EraseOpsObservation) -> np.ndarray: req_user = obs.request.user_id global_features = np.array( [ float(obs.step) / max(float(obs.step_budget), 1.0), float(obs.pii_remaining) / 16.0, float(obs.legal_hold_violations) / 8.0, float(obs.dependency_violations) / 8.0, float(obs.contamination_violations) / 8.0, float(obs.requests_completed) / 16.0, float(obs.cumulative_penalty) / 20.0, float(obs.deadline_pressure), float(obs.escalation_queue_size) / 16.0, float(obs.blocked_dependency_count) / 16.0, ], dtype=np.float32, ) per_record: List[float] = [] for idx in range(self.max_records): if idx >= len(obs.records): per_record.extend([0.0] * 10) continue record = obs.records[idx] status_active = 1.0 if record.status == "active" else 0.0 status_deleted = 1.0 if record.status == "deleted" else 0.0 status_anonymized = 1.0 if record.status == "anonymized" else 0.0 is_target = 1.0 if record.user_id == req_user else 0.0 is_pii = 1.0 if record.data_type == "PII" else 0.0 is_legal_hold = 1.0 if record.legal_hold else 0.0 dep_norm = min(1.0, float(len(record.dependencies)) / 4.0) age_norm = min(1.0, float(record.age_days) / 730.0) cost_norm = min(1.0, float(record.deletion_cost)) blocked = 1.0 if is_target and is_pii and status_active and len(record.dependencies) > 0 else 0.0 per_record.extend( [ is_target, is_pii, is_legal_hold, status_active, status_deleted, status_anonymized, dep_norm, age_norm, cost_norm, blocked, ] ) vector = np.concatenate([global_features, np.array(per_record, dtype=np.float32)], axis=0) return vector.astype(np.float32) class GradeEvalCallback(BaseCallback): def __init__(self, eval_task: str, eval_freq: int, episodes: int = 5, max_records: int = 32): super().__init__() self.eval_task = eval_task self.eval_freq = max(1, eval_freq) self.episodes = max(1, episodes) self.max_records = max_records self.timesteps: List[int] = [] self.mean_scores: List[float] = [] def _on_step(self) -> bool: if self.num_timesteps % self.eval_freq != 0: return True mean_score = evaluate_model( self.model, [self.eval_task], seeds=EVAL_SEEDS[: self.episodes], max_records=self.max_records, )[self.eval_task]["avg_score"] self.timesteps.append(self.num_timesteps) self.mean_scores.append(mean_score) return True def run_episode_for_grade(model: PPO, task_name: str, seed: int, max_records: int) -> float: env = GdprEraseOpsGymEnv(tasks=[task_name], max_records=max_records) obs_vec, _ = env.reset(options={"task_name": task_name, "task_seed": seed}) done = False while not done: action, _ = model.predict(obs_vec, deterministic=True) obs_vec, _, done, _, _ = env.step(int(action)) final_state = env._env.state() return float(grade(task_name, final_state)) def evaluate_model(model: PPO, tasks: List[str], seeds: List[int], max_records: int = 32) -> Dict[str, Dict[str, Any]]: output: Dict[str, Dict[str, Any]] = {} for task in tasks: scores = [run_episode_for_grade(model, task_name=task, seed=seed, max_records=max_records) for seed in seeds] output[task] = { "scores": [round(value, 4) for value in scores], "avg_score": round(float(np.mean(scores)), 4), "min_score": round(float(np.min(scores)), 4), "max_score": round(float(np.max(scores)), 4), } return output def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Train PPO agent on GDPR-EraseOps and export learning artifacts") parser.add_argument("--timesteps", type=int, default=50_000) parser.add_argument("--eval-freq", type=int, default=2_500) parser.add_argument("--n-envs", type=int, default=4) parser.add_argument("--max-records", type=int, default=32) parser.add_argument("--tasks", type=str, default=",".join(ALL_TASKS)) parser.add_argument("--eval-task", type=str, default="legal_minefield") parser.add_argument("--output-dir", type=str, default="artifacts/ppo") parser.add_argument("--stage-steps", type=str, default="16000,17000,17000") parser.add_argument("--entropy-schedule", type=str, default="0.02,0.01,0.003") parser.add_argument( "--curriculum", type=str, default=( "clean_sweep:0.55,tangled_lake:0.35,legal_minefield:0.10,cascade_collapse:0.00;" "clean_sweep:0.20,tangled_lake:0.30,legal_minefield:0.40,cascade_collapse:0.10;" "clean_sweep:0.05,tangled_lake:0.15,legal_minefield:0.50,cascade_collapse:0.30" ), ) return parser.parse_args() def _parse_float_list(csv_values: str) -> List[float]: return [float(item.strip()) for item in csv_values.split(",") if item.strip()] def _parse_int_list(csv_values: str) -> List[int]: return [int(item.strip()) for item in csv_values.split(",") if item.strip()] def _parse_curriculum(curriculum_text: str, tasks: List[str]) -> List[Dict[str, float]]: stages: List[Dict[str, float]] = [] for chunk in [stage.strip() for stage in curriculum_text.split(";") if stage.strip()]: weights = {task: 0.0 for task in tasks} for pair in [item.strip() for item in chunk.split(",") if item.strip()]: if ":" not in pair: continue key, value = pair.split(":", 1) task_name = key.strip() if task_name in weights: weights[task_name] = max(0.0, float(value.strip())) stages.append(weights) return stages def main() -> None: args = parse_args() tasks = [item.strip() for item in args.tasks.split(",") if item.strip()] for task in tasks: if task not in ALL_TASKS: raise ValueError(f"Unknown task '{task}'. Allowed values: {ALL_TASKS}") output_dir = PROJECT_ROOT / args.output_dir output_dir.mkdir(parents=True, exist_ok=True) stage_steps = _parse_int_list(args.stage_steps) entropy_schedule = _parse_float_list(args.entropy_schedule) curriculum = _parse_curriculum(args.curriculum, tasks) if not stage_steps: stage_steps = [int(args.timesteps)] if not entropy_schedule: entropy_schedule = [0.01] if not curriculum: curriculum = [{task: 1.0 for task in tasks}] stage_count = max(len(stage_steps), len(entropy_schedule), len(curriculum)) def _expand_list(values: List[Any]) -> List[Any]: if len(values) >= stage_count: return values[:stage_count] return values + [values[-1]] * (stage_count - len(values)) stage_steps = [int(v) for v in _expand_list(stage_steps)] entropy_schedule = [float(v) for v in _expand_list(entropy_schedule)] curriculum = [dict(v) for v in _expand_list(curriculum)] def make_env() -> GdprEraseOpsGymEnv: return GdprEraseOpsGymEnv(tasks=tasks, max_records=args.max_records) vec_env = DummyVecEnv([make_env for _ in range(args.n_envs)]) callback = GradeEvalCallback( eval_task=args.eval_task, eval_freq=args.eval_freq, episodes=5, max_records=args.max_records, ) model = PPO( policy="MlpPolicy", env=vec_env, learning_rate=3e-4, n_steps=512, batch_size=256, n_epochs=10, gamma=0.99, gae_lambda=0.95, ent_coef=entropy_schedule[0], verbose=1, seed=42, ) stage_plan: List[Dict[str, Any]] = [] for idx in range(stage_count): weights = curriculum[idx] entropy = entropy_schedule[idx] steps = stage_steps[idx] vec_env.env_method("set_task_weights", weights) model.ent_coef = float(entropy) model.learn(total_timesteps=steps, callback=callback, reset_num_timesteps=False) stage_plan.append( { "stage": idx + 1, "steps": steps, "ent_coef": round(float(entropy), 6), "weights": {task: round(float(weights.get(task, 0.0)), 4) for task in tasks}, "cum_timesteps": int(model.num_timesteps), } ) model_path = output_dir / "ppo_gdpr_eraseops.zip" model.save(model_path) eval_summary = evaluate_model(model, tasks=ALL_TASKS, seeds=EVAL_SEEDS, max_records=args.max_records) curve_path = output_dir / "learning_curve.png" if callback.timesteps: plt.figure(figsize=(8, 4.5)) plt.plot(callback.timesteps, callback.mean_scores, marker="o", linewidth=2) plt.title(f"PPO learning curve on {args.eval_task}") plt.xlabel("Training timesteps") plt.ylabel("Mean grader score") plt.ylim(0.0, 1.0) plt.grid(True, alpha=0.25) plt.tight_layout() plt.savefig(curve_path, dpi=140) plt.close() curve_json_path = output_dir / "learning_curve.json" curve_json_path.write_text( json.dumps( { "eval_task": args.eval_task, "timesteps": callback.timesteps, "scores": [round(value, 4) for value in callback.mean_scores], }, indent=2, ), encoding="utf-8", ) summary_path = output_dir / "evaluation_summary.json" summary_path.write_text(json.dumps(eval_summary, indent=2), encoding="utf-8") print( json.dumps( { "model": str(model_path), "curve": str(curve_path), "stages": stage_plan, "summary": eval_summary, }, indent=2, ) ) if __name__ == "__main__": main()