Spaces:
Running
Running
| 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() | |