Blame for scripts/train_sb3_agent.py

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()