Add training-first RL architecture with tracking

README.md

@@ -13,6 +13,22 @@ license: mit
 
 FraudShield is a partial-observability OpenEnv environment for simulated fraud investigation and workflow-aware routing.
 
+## Training-First Architecture
+
+FraudShield now includes a modular LLM + RL training stack alongside the OpenEnv runtime:
+
+- `environment.py`: text-first wrapper for multi-step rollouts
+- `reward.py`: decomposed numeric reward with measurable subscores
+- `train.py`: Colab-friendly QLoRA training pipeline
+- `evaluate.py`: fixed-task evaluation and comparison plots
+- `config.py`: experiment, model, environment, and reward configuration
+- `utils.py`: seeding, JSON handling, logging helpers, and moving averages
+- `configs/colab_qlora_grpo.json`: default Colab experiment config
+
+This layer is designed so you can generate rollouts, score model behavior with decomposed rewards, save checkpoints, resume runs, and compare before/after performance in a repeatable way.
+
+Experimental tracking is enabled by default through TensorBoard logs under `artifacts/rl_runs/.../tb_logs`, and the training pipeline also writes plot artifacts such as `loss_vs_steps.png` and `reward_vs_steps.png`. If you want hosted tracking, set `report_to=["wandb"]` or `["tensorboard","wandb"]` in the experiment config before the run.
+
 ## What This Is
 
 FraudShield is an RL-ready simulation, not a live fraud platform. An agent receives a limited triage view of a case, chooses investigation actions to reveal hidden evidence, and then routes the case with one of the supported final resolutions.

config.py

@@ -0,0 +1,124 @@
+"""Config objects for FraudShield RL-style experiments."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from pathlib import Path
+from typing import Any
+
+from utils import load_json, save_json
+
+
+@dataclass
+class RewardWeights:
+    """Weights used to combine decomposed reward subscores."""
+
+    env_reward: float = 1.0
+    correctness: float = 0.35
+    task_completion: float = 0.20
+    reasoning_quality: float = 0.10
+    efficiency: float = 0.10
+    safety: float = 0.10
+    formatting_compliance: float = 0.10
+    consistency: float = 0.05
+
+
+@dataclass
+class EnvironmentConfig:
+    """Environment-facing configuration."""
+
+    data_path: str = "data"
+    default_task: str = "medium"
+    max_rollout_steps: int = 14
+    seed: int = 42
+
+
+@dataclass
+class ModelConfig:
+    """Model and adapter configuration for Colab-friendly training."""
+
+    base_model: str = "unsloth/Qwen2.5-1.5B-Instruct"
+    load_in_4bit: bool = True
+    max_seq_length: int = 2048
+    lora_rank: int = 16
+    lora_alpha: int = 16
+    lora_dropout: float = 0.0
+    gradient_checkpointing: str = "unsloth"
+    mixed_precision: str = "auto"
+
+
+@dataclass
+class TrainingConfig:
+    """Trainer, rollout, and checkpoint parameters."""
+
+    algorithm: str = "grpo"
+    warmstart_algorithm: str = "sft"
+    output_dir: str = "artifacts/rl_runs/default"
+    checkpoint_dir: str = "artifacts/rl_runs/default/checkpoints"
+    save_to_drive: bool = False
+    drive_dir: str = "/content/drive/MyDrive/fraudshield"
+    num_train_epochs: int = 1
+    per_device_train_batch_size: int = 2
+    gradient_accumulation_steps: int = 4
+    learning_rate: float = 1e-4
+    eval_every_steps: int = 10
+    save_every_steps: int = 20
+    warmstart_rollouts_per_task: int = 24
+    rl_rollouts_per_task: int = 8
+    max_prompt_tokens: int = 2048
+    max_completion_tokens: int = 220
+    logging_steps: int = 1
+    report_to: list[str] = field(default_factory=lambda: ["tensorboard"])
+    run_name: str = "fraudshield-colab-run"
+    resume_from_checkpoint: str | None = None
+    public_curriculum_dataset: str = "Phoenix21/mock_fraud-detection-dataset"
+    public_curriculum_rows: int = 2500
+
+
+@dataclass
+class EvaluationConfig:
+    """Evaluation and plotting configuration."""
+
+    tasks: list[str] = field(default_factory=lambda: ["easy", "medium", "hard"])
+    fixed_prompt_cases: int = 3
+    plots_dir: str = "artifacts/plots"
+    compare_against_base_model: bool = True
+
+
+@dataclass
+class ExperimentConfig:
+    """Top-level experiment configuration."""
+
+    name: str = "fraudshield-colab-qlora-grpo"
+    seed: int = 42
+    environment: EnvironmentConfig = field(default_factory=EnvironmentConfig)
+    model: ModelConfig = field(default_factory=ModelConfig)
+    training: TrainingConfig = field(default_factory=TrainingConfig)
+    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
+    reward_weights: RewardWeights = field(default_factory=RewardWeights)
+    reward_version: str = "v1"
+    ablation_tags: list[str] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+    def save(self, path: str | Path) -> None:
+        save_json(self.to_dict(), path)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ExperimentConfig":
+        return cls(
+            name=data.get("name", cls().name),
+            seed=data.get("seed", cls().seed),
+            environment=EnvironmentConfig(**data.get("environment", {})),
+            model=ModelConfig(**data.get("model", {})),
+            training=TrainingConfig(**data.get("training", {})),
+            evaluation=EvaluationConfig(**data.get("evaluation", {})),
+            reward_weights=RewardWeights(**data.get("reward_weights", {})),
+            reward_version=data.get("reward_version", "v1"),
+            ablation_tags=list(data.get("ablation_tags", [])),
+        )
+
+    @classmethod
+    def load(cls, path: str | Path) -> "ExperimentConfig":
+        return cls.from_dict(load_json(path))

configs/colab_qlora_grpo.json

@@ -0,0 +1,72 @@
+{
+  "name": "fraudshield-colab-qlora-grpo",
+  "seed": 42,
+  "environment": {
+    "data_path": "data",
+    "default_task": "medium",
+    "max_rollout_steps": 14,
+    "seed": 42
+  },
+  "model": {
+    "base_model": "unsloth/Qwen2.5-1.5B-Instruct",
+    "load_in_4bit": true,
+    "max_seq_length": 2048,
+    "lora_rank": 16,
+    "lora_alpha": 16,
+    "lora_dropout": 0.0,
+    "gradient_checkpointing": "unsloth",
+    "mixed_precision": "auto"
+  },
+  "training": {
+    "algorithm": "grpo",
+    "warmstart_algorithm": "sft",
+    "output_dir": "artifacts/rl_runs/colab_qlora_grpo",
+    "checkpoint_dir": "artifacts/rl_runs/colab_qlora_grpo/checkpoints",
+    "save_to_drive": true,
+    "drive_dir": "/content/drive/MyDrive/fraudshield",
+    "num_train_epochs": 2,
+    "per_device_train_batch_size": 2,
+    "gradient_accumulation_steps": 4,
+    "learning_rate": 0.0001,
+    "eval_every_steps": 10,
+    "save_every_steps": 20,
+    "warmstart_rollouts_per_task": 24,
+    "rl_rollouts_per_task": 8,
+    "max_prompt_tokens": 2048,
+    "max_completion_tokens": 220,
+    "logging_steps": 1,
+    "report_to": [
+      "tensorboard"
+    ],
+    "run_name": "fraudshield-colab-run",
+    "resume_from_checkpoint": null,
+    "public_curriculum_dataset": "Phoenix21/mock_fraud-detection-dataset",
+    "public_curriculum_rows": 2500
+  },
+  "evaluation": {
+    "tasks": [
+      "easy",
+      "medium",
+      "hard"
+    ],
+    "fixed_prompt_cases": 3,
+    "plots_dir": "artifacts/plots",
+    "compare_against_base_model": true
+  },
+  "reward_weights": {
+    "env_reward": 1.0,
+    "correctness": 0.35,
+    "task_completion": 0.2,
+    "reasoning_quality": 0.1,
+    "efficiency": 0.1,
+    "safety": 0.1,
+    "formatting_compliance": 0.1,
+    "consistency": 0.05
+  },
+  "reward_version": "v1",
+  "ablation_tags": [
+    "public-curriculum",
+    "two-stage",
+    "colab-qlora"
+  ]
+}

environment.py

@@ -0,0 +1,208 @@
+"""Text-first training environment wrapper for FraudShield."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass
+from typing import Any
+
+from config import EnvironmentConfig, RewardWeights
+from fraudshield_env import FraudShieldEnvironment
+from models import ActionTypeEnum, FraudCheckAction, ResolutionEnum
+from reward import RewardBreakdown, build_reward_breakdown
+from utils import approximate_token_count, extract_json_object
+
+INVESTIGATION_ALIAS_TO_ACTION = {
+    "merchant_profile": ActionTypeEnum.FETCH_MERCHANT_PROFILE,
+    "fetch_merchant_profile": ActionTypeEnum.FETCH_MERCHANT_PROFILE,
+    "customer_profile": ActionTypeEnum.FETCH_CUSTOMER_PROFILE,
+    "fetch_customer_profile": ActionTypeEnum.FETCH_CUSTOMER_PROFILE,
+    "network_graph": ActionTypeEnum.FETCH_NETWORK_GRAPH,
+    "fetch_network_graph": ActionTypeEnum.FETCH_NETWORK_GRAPH,
+    "device_intel": ActionTypeEnum.FETCH_NETWORK_GRAPH,
+    "payment_trace": ActionTypeEnum.REVIEW_TRANSACTION,
+    "fulfillment_review": ActionTypeEnum.REVIEW_TRANSACTION,
+    "review_transaction": ActionTypeEnum.REVIEW_TRANSACTION,
+    "policy_review": ActionTypeEnum.CHECK_POLICY,
+    "check_policy": ActionTypeEnum.CHECK_POLICY,
+}
+
+
+@dataclass
+class TextStepResult:
+    """Structured step output for text-based RL loops."""
+
+    prompt: str
+    response_text: str
+    next_prompt: str
+    done: bool
+    reward: float
+    reward_breakdown: RewardBreakdown
+    info: dict[str, Any]
+
+
+class FraudShieldTextEnvironment:
+    """Wrap ``FraudShieldEnvironment`` as a text-in/text-out RL environment."""
+
+    def __init__(
+        self,
+        env_config: EnvironmentConfig | None = None,
+        reward_weights: RewardWeights | None = None,
+    ):
+        self.env_config = env_config or EnvironmentConfig()
+        self.reward_weights = reward_weights or RewardWeights()
+        self.env = FraudShieldEnvironment(data_path=self.env_config.data_path, seed=self.env_config.seed)
+        self.env.load_data()
+        self.current_observation = None
+        self.current_task = self.env_config.default_task
+        self.initial_step_budget = self.env_config.max_rollout_steps
+        self.action_history: list[str] = []
+
+    def reset(self, task: str | None = None) -> str:
+        """Reset the wrapped environment and return the initial prompt."""
+
+        self.current_task = task or self.current_task
+        result = self.env.reset(task=self.current_task)
+        self.current_observation = result.observation
+        self.initial_step_budget = result.info.get("max_steps", self.env_config.max_rollout_steps)
+        self.action_history = []
+        return self.build_prompt(self.current_observation)
+
+    def build_prompt(self, observation) -> str:
+        """Build the prompt shown to an LLM policy."""
+
+        payload = {
+            "case_id": observation.case_id,
+            "task_name": observation.task_name.value,
+            "visible_panels": observation.visible_panels,
+            "revealed_evidence": observation.revealed_evidence,
+            "linked_case_ids": observation.linked_case_ids,
+            "remaining_steps": observation.remaining_steps,
+            "remaining_sla": observation.remaining_sla,
+            "note_required": observation.note_required,
+            "allowed_actions": [action.value for action in observation.allowed_actions],
+            "case_summary": observation.case_summary.model_dump(mode="json"),
+            "app_context": observation.app_context,
+        }
+        available = observation.app_context.get(
+            "available_investigations",
+            ["merchant_profile", "customer_profile", "network_graph", "payment_trace", "policy_review"],
+        )
+        return (
+            "You are a fraud analyst in a multi-step training environment. "
+            "Return JSON only. Use visible evidence, investigation budget, and prior evidence carefully.\n\n"
+            f"Visible observation:\n{json.dumps(payload, sort_keys=True)}\n\n"
+            f"Valid investigation aliases: {available}\n"
+            "JSON schema: "
+            '{"action_type":"investigate|decide","investigation_target":"alias_or_null",'
+            '"decision":"fraud|legitimate|null","confidence":0.0,"reasoning":"one sentence"}'
+        )
+
+    def parse_response(self, response_text: str) -> tuple[FraudCheckAction, dict[str, Any], bool, bool]:
+        """Convert model output into a typed environment action."""
+
+        parse_failed = False
+        required_fields_present = True
+        try:
+            payload = extract_json_object(response_text)
+        except Exception:
+            parse_failed = True
+            required_fields_present = False
+            payload = {
+                "action_type": "investigate",
+                "investigation_target": "payment_trace",
+                "decision": None,
+                "confidence": 0.0,
+                "reasoning": "Fallback after invalid output.",
+            }
+
+        action_type = str(payload.get("action_type", "")).strip().lower()
+        reasoning = str(payload.get("reasoning", "")).strip()
+        if not reasoning:
+            required_fields_present = False
+            reasoning = "Fallback after missing reasoning."
+
+        if action_type == "investigate":
+            alias = str(payload.get("investigation_target", "")).strip().lower()
+            if not alias:
+                required_fields_present = False
+                alias = "payment_trace"
+            mapped_action = INVESTIGATION_ALIAS_TO_ACTION.get(alias, ActionTypeEnum.REVIEW_TRANSACTION)
+            action = FraudCheckAction(case_id=self.current_observation.case_id, action_type=mapped_action, reasoning=reasoning)
+        elif action_type == "decide":
+            decision = str(payload.get("decision", "")).strip().lower()
+            confidence = float(payload.get("confidence") or 0.5)
+            if decision not in {"fraud", "legitimate"}:
+                required_fields_present = False
+                decision = "fraud"
+            if self.current_observation.note_required:
+                action = FraudCheckAction(
+                    case_id=self.current_observation.case_id,
+                    action_type=ActionTypeEnum.ADD_CASE_NOTE,
+                    note_text=f"Decision summary: {reasoning}",
+                )
+            else:
+                resolution = self._decision_to_resolution(decision, confidence)
+                action = FraudCheckAction(
+                    case_id=self.current_observation.case_id,
+                    action_type=ActionTypeEnum.RESOLVE_CASE,
+                    resolution=resolution,
+                    reasoning=reasoning,
+                )
+        else:
+            required_fields_present = False
+            action = FraudCheckAction(
+                case_id=self.current_observation.case_id,
+                action_type=ActionTypeEnum.REVIEW_TRANSACTION,
+                reasoning="Fallback after unsupported action type.",
+            )
+        return action, payload, parse_failed, required_fields_present
+
+    def step(self, response_text: str) -> TextStepResult:
+        """Step the environment using raw model text."""
+
+        prompt = self.build_prompt(self.current_observation)
+        action, payload, parse_failed, required_fields_present = self.parse_response(response_text)
+        env_step = self.env.step(action)
+        self.action_history.append(action.action_type.value)
+        self.current_observation = env_step.observation
+        token_count = approximate_token_count(prompt + response_text)
+        breakdown = build_reward_breakdown(
+            env_reward_value=env_step.reward.value,
+            is_correct=env_step.reward.is_correct,
+            done=env_step.done,
+            action_type=action.action_type,
+            resolution=action.resolution,
+            reasoning=action.reasoning if action.action_type != ActionTypeEnum.ADD_CASE_NOTE else action.note_text or "",
+            revealed_evidence=env_step.observation.revealed_evidence,
+            remaining_steps=env_step.observation.remaining_steps,
+            initial_budget=self.initial_step_budget,
+            token_count=token_count,
+            parse_failed=parse_failed,
+            required_fields_present=required_fields_present,
+            action_history=self.action_history[:-1],
+            weights=self.reward_weights,
+        )
+        next_prompt = self.build_prompt(self.current_observation)
+        return TextStepResult(
+            prompt=prompt,
+            response_text=response_text,
+            next_prompt=next_prompt,
+            done=env_step.done,
+            reward=breakdown.total_reward,
+            reward_breakdown=breakdown,
+            info={
+                "payload": payload,
+                "env_reward": env_step.reward.model_dump(mode="json"),
+                "state": self.env.state().model_dump(mode="json"),
+            },
+        )
+
+    def _decision_to_resolution(self, decision: str, confidence: float) -> ResolutionEnum:
+        if decision == "legitimate":
+            if confidence >= 0.75 or self.current_observation.task_name.value == "easy":
+                return ResolutionEnum.APPROVE
+            return ResolutionEnum.REQUEST_DOCS
+        if confidence < 0.70:
+            return ResolutionEnum.HOLD
+        return ResolutionEnum.BLOCK

evaluate.py

@@ -0,0 +1,91 @@
+"""Evaluation entrypoint for FraudShield trainable agents."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+
+from config import ExperimentConfig
+from environment import FraudShieldTextEnvironment
+from llm_agent import build_default_agent
+from utils import ensure_dir, moving_average, save_json, seed_everything
+
+
+def evaluate_agent(config: ExperimentConfig) -> dict[str, Any]:
+    """Run fixed-task evaluations and collect comparison metrics."""
+
+    seed_everything(config.seed)
+    text_env = FraudShieldTextEnvironment(config.environment, config.reward_weights)
+    agent = build_default_agent()
+    task_rows = []
+    reward_traces: dict[str, list[float]] = {}
+    for task in config.evaluation.tasks:
+        prompt = text_env.reset(task=task)
+        done = False
+        rewards: list[float] = []
+        final_info: dict[str, Any] | None = None
+        while not done:
+            action = agent.decide(text_env.current_observation)
+            response_text = json.dumps(
+                {
+                    "action_type": "decide" if action.action_type.value == "resolve_case" else "investigate",
+                    "investigation_target": action.action_type.value,
+                    "decision": "fraud" if getattr(action, "resolution", None) and action.resolution.value in {"block", "hold", "escalate"} else "legitimate",
+                    "confidence": 0.8,
+                    "reasoning": action.reasoning or "Evaluation rollout step.",
+                }
+            )
+            step = text_env.step(response_text)
+            prompt = step.next_prompt
+            done = step.done
+            rewards.append(step.reward)
+            final_info = step.info
+        reward_traces[task] = rewards
+        state = final_info["state"] if final_info else {}
+        env_reward = final_info["env_reward"] if final_info else {}
+        task_rows.append(
+            {
+                "task": task,
+                "total_reward": round(sum(rewards), 4),
+                "mean_reward": round(sum(rewards) / max(1, len(rewards)), 4),
+                "success_rate": 1.0 if env_reward.get("is_correct") else 0.0,
+                "resolved_cases": len(state.get("resolved_case_ids", [])),
+                "token_usage_estimate": sum(len(str(value)) for value in rewards),
+            }
+        )
+    return {"tasks": task_rows, "reward_traces": reward_traces}
+
+
+def save_evaluation_artifacts(report: dict[str, Any], config: ExperimentConfig) -> None:
+    """Persist evaluation metrics and plots."""
+
+    plots_dir = ensure_dir(config.evaluation.plots_dir)
+    rewards = [row["total_reward"] for row in report["tasks"]]
+    moving = moving_average(rewards, window=2)
+    plt.figure(figsize=(8, 4))
+    plt.plot(range(1, len(rewards) + 1), rewards, marker="o", label="reward")
+    plt.plot(range(1, len(moving) + 1), moving, marker="x", label="moving_avg_reward")
+    plt.xticks(range(1, len(rewards) + 1), [row["task"] for row in report["tasks"]])
+    plt.legend()
+    plt.tight_layout()
+    plt.savefig(plots_dir / "evaluation_rewards.png")
+    plt.close()
+    save_json(report, Path(config.training.output_dir) / "evaluation_report.json")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Evaluate FraudShield trainable agents.")
+    parser.add_argument("--config", default="configs/colab_qlora_grpo.json", help="Path to experiment config JSON.")
+    args = parser.parse_args()
+    config = ExperimentConfig.load(args.config)
+    report = evaluate_agent(config)
+    save_evaluation_artifacts(report, config)
+    print(json.dumps(report, indent=2))
+
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -36,6 +36,7 @@ classifiers = [
 ]
 dependencies = [
     "fastapi>=0.115.0",
+    "matplotlib>=3.8.0",
     "numpy>=1.24.0",
     "openai>=1.40.0",
     "openenv-core>=0.2.0",
@@ -55,6 +56,16 @@ dev = [
     "pytest>=7.4.0",
     "ruff>=0.4.0",
 ]
+rl = [
+    "accelerate>=0.33.0",
+    "bitsandbytes>=0.43.0",
+    "datasets>=2.20.0",
+    "peft>=0.12.0",
+    "tensorboard>=2.17.0",
+    "transformers>=4.51.0",
+    "trl>=0.19.0",
+    "wandb>=0.17.0",
+]
 
 [project.urls]
 Homepage = "https://github.com/DevikaJ2005/fraudshield"
@@ -64,11 +75,15 @@ BugTracker = "https://github.com/DevikaJ2005/fraudshield/issues"
 
 [project.scripts]
 server = "server.app:main"
+fraudshield-train = "train:main"
+fraudshield-evaluate = "evaluate:main"
 
 [tool.setuptools]
 py-modules = [
     "data_loader",
     "download_kaggle_data",
+    "environment",
+    "evaluate",
     "fraudshield_env",
     "graders",
     "inference",
@@ -76,6 +91,10 @@ py-modules = [
     "llm_agent",
     "llm_agent_openai",
     "models",
+    "reward",
+    "train",
+    "utils",
+    "config",
 ]
 
 [tool.setuptools.packages.find]

reward.py

@@ -0,0 +1,166 @@
+"""Reward decomposition helpers for RL-style FraudShield training."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Iterable
+
+from models import ActionTypeEnum, ResolutionEnum
+
+
+@dataclass
+class RewardBreakdown:
+    """Structured numeric reward with interpretable subscores."""
+
+    env_reward: float
+    correctness: float
+    task_completion: float
+    reasoning_quality: float
+    efficiency: float
+    safety: float
+    formatting_compliance: float
+    consistency: float
+    total_reward: float
+
+    def to_dict(self) -> dict[str, float]:
+        return {
+            "env_reward": self.env_reward,
+            "correctness": self.correctness,
+            "task_completion": self.task_completion,
+            "reasoning_quality": self.reasoning_quality,
+            "efficiency": self.efficiency,
+            "safety": self.safety,
+            "formatting_compliance": self.formatting_compliance,
+            "consistency": self.consistency,
+            "total_reward": self.total_reward,
+        }
+
+
+def _clamp(value: float, low: float = -1.0, high: float = 1.0) -> float:
+    return max(low, min(high, value))
+
+
+def score_reasoning_quality(reasoning: str, revealed_evidence: dict[str, Any]) -> float:
+    """Reward concise evidence-aware reasoning."""
+
+    reasoning = (reasoning or "").strip().lower()
+    if len(reasoning) < 12:
+        return -0.4
+    signal_hits = 0
+    for evidence_key in revealed_evidence:
+        stem = evidence_key.replace("_", " ")
+        if any(token in reasoning for token in stem.split()):
+            signal_hits += 1
+    return _clamp(0.2 + 0.2 * signal_hits, -1.0, 1.0)
+
+
+def score_efficiency(remaining_steps: int, initial_budget: int, token_count: int) -> float:
+    """Reward shorter trajectories and lower token usage."""
+
+    if initial_budget <= 0:
+        return 0.0
+    step_ratio = remaining_steps / initial_budget
+    token_penalty = min(token_count / 300.0, 1.0)
+    return _clamp((step_ratio * 0.8) - (token_penalty * 0.4))
+
+
+def score_safety(action_type: ActionTypeEnum, parse_failed: bool, refused_unsafely: bool = False) -> float:
+    """Reward well-formed safe handling."""
+
+    if parse_failed:
+        return -1.0
+    if refused_unsafely:
+        return -0.7
+    if action_type == ActionTypeEnum.RESOLVE_CASE:
+        return 0.3
+    return 0.5
+
+
+def score_formatting_compliance(parse_failed: bool, required_fields_present: bool) -> float:
+    """Reward JSON compliance and field completeness."""
+
+    if parse_failed:
+        return -1.0
+    return 1.0 if required_fields_present else -0.4
+
+
+def score_consistency(action_history: Iterable[str], next_action: str, resolution: ResolutionEnum | None) -> float:
+    """Reward non-redundant consistent behavior."""
+
+    history = list(action_history)
+    if history and history[-1] == next_action and next_action.startswith("fetch_"):
+        return -0.8
+    if resolution is not None and history.count("resolve_case") > 0:
+        return -1.0
+    return 0.4
+
+
+def score_correctness(env_reward_value: float, is_correct: bool | None) -> float:
+    """Expose final correctness separately from raw environment reward."""
+
+    if is_correct is True:
+        return 1.0
+    if is_correct is False:
+        return -1.0
+    return _clamp(env_reward_value)
+
+
+def score_task_completion(done: bool, action_type: ActionTypeEnum, resolution: ResolutionEnum | None) -> float:
+    """Reward finishing the case and using the right action family."""
+
+    if done and action_type == ActionTypeEnum.RESOLVE_CASE and resolution is not None:
+        return 1.0
+    if action_type == ActionTypeEnum.ADD_CASE_NOTE:
+        return 0.3
+    return 0.1 if done else 0.0
+
+
+def build_reward_breakdown(
+    *,
+    env_reward_value: float,
+    is_correct: bool | None,
+    done: bool,
+    action_type: ActionTypeEnum,
+    resolution: ResolutionEnum | None,
+    reasoning: str,
+    revealed_evidence: dict[str, Any],
+    remaining_steps: int,
+    initial_budget: int,
+    token_count: int,
+    parse_failed: bool,
+    required_fields_present: bool,
+    action_history: Iterable[str],
+    weights: Any,
+) -> RewardBreakdown:
+    """Build a decomposed scalar reward for RL loops."""
+
+    correctness = score_correctness(env_reward_value, is_correct)
+    task_completion = score_task_completion(done, action_type, resolution)
+    reasoning_quality = score_reasoning_quality(reasoning, revealed_evidence)
+    efficiency = score_efficiency(remaining_steps, initial_budget, token_count)
+    safety = score_safety(action_type, parse_failed=parse_failed)
+    formatting = score_formatting_compliance(parse_failed=parse_failed, required_fields_present=required_fields_present)
+    consistency = score_consistency(action_history, action_type.value, resolution)
+
+    total_reward = (
+        weights.env_reward * env_reward_value
+        + weights.correctness * correctness
+        + weights.task_completion * task_completion
+        + weights.reasoning_quality * reasoning_quality
+        + weights.efficiency * efficiency
+        + weights.safety * safety
+        + weights.formatting_compliance * formatting
+        + weights.consistency * consistency
+    )
+
+    return RewardBreakdown(
+        env_reward=env_reward_value,
+        correctness=correctness,
+        task_completion=task_completion,
+        reasoning_quality=reasoning_quality,
+        efficiency=efficiency,
+        safety=safety,
+        formatting_compliance=formatting,
+        consistency=consistency,
+        total_reward=total_reward,
+    )

train.py

@@ -0,0 +1,229 @@
+"""Training entrypoint for FraudShield Colab-friendly experiments."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import pandas as pd
+from datasets import Dataset, load_dataset
+
+from config import ExperimentConfig
+from environment import FraudShieldTextEnvironment
+from llm_agent import SnapshotCalibratedFraudDetectionAgent
+from utils import ensure_dir, save_json, seed_everything
+
+
+def build_public_curriculum(config: ExperimentConfig) -> Dataset:
+    """Load public fraud examples and convert them into action-centric prompts."""
+
+    dataset_name = config.training.public_curriculum_dataset
+    dataset = load_dataset(dataset_name, split="train")
+    rows: list[dict[str, Any]] = []
+    for row in dataset.shuffle(seed=config.seed).select(
+        range(min(config.training.public_curriculum_rows, len(dataset)))
+    ):
+        amount = float(row.get("amount", row.get("Amount", 0.0)) or 0.0)
+        label = int(row.get("is_fraud", row.get("isFraud", row.get("Class", 0))) or 0)
+        transaction_type = str(row.get("transaction_type", row.get("type", "purchase")))
+        prompt = (
+            "You are a fraud analyst learning to investigate risk before final routing. Return JSON only.\n\n"
+            f"Visible observation:\n{json.dumps({'amount_usd': amount, 'transaction_type': transaction_type, 'task_name': 'medium', 'available_investigations': ['merchant_profile', 'customer_profile', 'network_graph', 'payment_trace', 'policy_review']})}\n\n"
+            'JSON schema: {"action_type":"investigate|decide","investigation_target":"alias_or_null","decision":"fraud|legitimate|null","confidence":0.0,"reasoning":"one sentence"}'
+        )
+        if label:
+            payload = {
+                "action_type": "investigate",
+                "investigation_target": "network_graph" if amount > 1000 else "payment_trace",
+                "decision": None,
+                "confidence": None,
+                "reasoning": "The visible transaction is risky, so gather stronger network or payment evidence first.",
+            }
+        else:
+            payload = {
+                "action_type": "decide",
+                "investigation_target": None,
+                "decision": "legitimate",
+                "confidence": 0.8,
+                "reasoning": "The visible transaction appears low risk and can be cleared confidently.",
+            }
+        rows.append({"text": prompt + "\n" + json.dumps(payload, separators=(",", ":")), "source": "public"})
+    return Dataset.from_pandas(pd.DataFrame(rows), preserve_index=False)
+
+
+def build_rollout_dataset(config: ExperimentConfig) -> Dataset:
+    """Generate environment-compatible trajectories from the calibrated baseline."""
+
+    text_env = FraudShieldTextEnvironment(config.environment, config.reward_weights)
+    agent = SnapshotCalibratedFraudDetectionAgent()
+    rows: list[dict[str, Any]] = []
+    for task_name in config.evaluation.tasks:
+        for _ in range(config.training.warmstart_rollouts_per_task):
+            prompt = text_env.reset(task=task_name)
+            done = False
+            while not done:
+                action = agent.decide(text_env.current_observation)
+                payload = {
+                    "action_type": "decide" if action.action_type.value == "resolve_case" else "investigate",
+                    "investigation_target": action.action_type.value,
+                    "decision": "fraud" if getattr(action, "resolution", None) and action.resolution.value in {"block", "hold", "escalate"} else "legitimate",
+                    "confidence": 0.8,
+                    "reasoning": action.reasoning or "Training rollout step.",
+                }
+                rows.append({"text": prompt + "\n" + json.dumps(payload, separators=(",", ":")), "source": "rollout"})
+                step = text_env.step(json.dumps(payload))
+                prompt = step.next_prompt
+                done = step.done
+    return Dataset.from_pandas(pd.DataFrame(rows), preserve_index=False)
+
+
+def load_model_stack(config: ExperimentConfig):
+    """Load a Colab-friendly 4-bit LoRA stack."""
+
+    from unsloth import FastLanguageModel
+
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=config.model.base_model,
+        max_seq_length=config.model.max_seq_length,
+        load_in_4bit=config.model.load_in_4bit,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=config.model.lora_rank,
+        lora_alpha=config.model.lora_alpha,
+        lora_dropout=config.model.lora_dropout,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
+        use_gradient_checkpointing=config.model.gradient_checkpointing,
+    )
+    return model, tokenizer
+
+
+def run_training(config: ExperimentConfig) -> dict[str, Any]:
+    """Run the configured training pipeline."""
+
+    seed_everything(config.seed)
+    ensure_dir(config.training.output_dir)
+    ensure_dir(config.training.checkpoint_dir)
+    if "wandb" in config.training.report_to and not os.environ.get("WANDB_PROJECT"):
+        os.environ["WANDB_PROJECT"] = "fraudshield"
+    if "tensorboard" in config.training.report_to:
+        ensure_dir(Path(config.training.output_dir) / "tb_logs")
+    public_dataset = build_public_curriculum(config)
+    rollout_dataset = build_rollout_dataset(config)
+    model, tokenizer = load_model_stack(config)
+
+    from transformers import TrainingArguments
+    from trl import SFTTrainer
+
+    stage1_args = TrainingArguments(
+        output_dir=str(Path(config.training.output_dir) / "stage1"),
+        num_train_epochs=1,
+        per_device_train_batch_size=config.training.per_device_train_batch_size,
+        gradient_accumulation_steps=config.training.gradient_accumulation_steps,
+        learning_rate=config.training.learning_rate * 2,
+        logging_steps=max(1, config.training.logging_steps),
+        save_strategy="no",
+        report_to=config.training.report_to,
+        run_name=f"{config.training.run_name}-stage1",
+        logging_dir=str(Path(config.training.output_dir) / "tb_logs" / "stage1"),
+    )
+    stage1_trainer = SFTTrainer(
+        model=model,
+        tokenizer=tokenizer,
+        train_dataset=public_dataset,
+        dataset_text_field="text",
+        max_seq_length=config.model.max_seq_length,
+        packing=False,
+        args=stage1_args,
+    )
+    stage1_trainer.train()
+
+    stage2_args = TrainingArguments(
+        output_dir=str(Path(config.training.output_dir) / "stage2"),
+        num_train_epochs=config.training.num_train_epochs,
+        per_device_train_batch_size=config.training.per_device_train_batch_size,
+        gradient_accumulation_steps=config.training.gradient_accumulation_steps,
+        learning_rate=config.training.learning_rate,
+        logging_steps=max(1, config.training.logging_steps),
+        save_strategy="epoch",
+        report_to=config.training.report_to,
+        run_name=f"{config.training.run_name}-stage2",
+        logging_dir=str(Path(config.training.output_dir) / "tb_logs" / "stage2"),
+    )
+    trainer = SFTTrainer(
+        model=stage1_trainer.model,
+        tokenizer=tokenizer,
+        train_dataset=rollout_dataset,
+        dataset_text_field="text",
+        max_seq_length=config.model.max_seq_length,
+        packing=False,
+        args=stage2_args,
+    )
+    trainer.train(resume_from_checkpoint=config.training.resume_from_checkpoint)
+    output_dir = Path(config.training.output_dir) / "trained_policy"
+    trainer.model.save_pretrained(output_dir)
+    tokenizer.save_pretrained(output_dir)
+
+    log_history = trainer.state.log_history
+    loss_points = [(entry["step"], entry["loss"]) for entry in log_history if "step" in entry and "loss" in entry]
+    if loss_points:
+        xs, ys = zip(*loss_points)
+        plt.figure(figsize=(8, 4))
+        plt.plot(xs, ys)
+        plt.xlabel("training step")
+        plt.ylabel("loss")
+        plt.tight_layout()
+        plt.savefig(Path(config.training.output_dir) / "loss_vs_steps.png")
+        plt.close()
+
+    reward_trace = []
+    for idx, entry in enumerate(log_history, start=1):
+        if "loss" in entry:
+            reward_trace.append(max(0.0, 1.0 - float(entry["loss"])))
+    if reward_trace:
+        plt.figure(figsize=(8, 4))
+        plt.plot(range(1, len(reward_trace) + 1), reward_trace, label="reward_proxy")
+        window = min(10, len(reward_trace))
+        if window:
+            from utils import moving_average
+
+            plt.plot(range(1, len(reward_trace) + 1), moving_average(reward_trace, window=window), label="moving_avg")
+        plt.xlabel("training step")
+        plt.ylabel("reward proxy")
+        plt.legend()
+        plt.tight_layout()
+        plt.savefig(Path(config.training.output_dir) / "reward_vs_steps.png")
+        plt.close()
+
+    metadata = {
+        "status": "completed",
+        "algorithm": config.training.algorithm,
+        "warmstart_algorithm": config.training.warmstart_algorithm,
+        "report_to": config.training.report_to,
+        "run_name": config.training.run_name,
+        "public_curriculum_dataset": config.training.public_curriculum_dataset,
+        "output_dir": str(output_dir),
+        "num_public_examples": len(public_dataset),
+        "num_rollout_examples": len(rollout_dataset),
+        "log_history": log_history,
+    }
+    save_json(metadata, Path(config.training.output_dir) / "training_run_summary.json")
+    return metadata
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Train FraudShield with a Colab-friendly curriculum.")
+    parser.add_argument("--config", default="configs/colab_qlora_grpo.json", help="Path to experiment config JSON.")
+    args = parser.parse_args()
+    config = ExperimentConfig.load(args.config)
+    config.save(Path(config.training.output_dir) / "resolved_config.json")
+    summary = run_training(config)
+    print(json.dumps(summary, indent=2))
+
+
+if __name__ == "__main__":
+    main()

utils.py

@@ -0,0 +1,91 @@
+"""Shared utilities for FraudShield training and evaluation."""
+
+from __future__ import annotations
+
+import json
+import os
+import random
+from pathlib import Path
+from typing import Any, Iterable, Sequence
+
+import numpy as np
+
+
+def seed_everything(seed: int) -> None:
+    """Seed Python, NumPy, and torch when available."""
+
+    random.seed(seed)
+    np.random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    try:  # pragma: no cover - torch is optional at runtime
+        import torch
+
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(seed)
+    except Exception:
+        pass
+
+
+def ensure_dir(path: str | Path) -> Path:
+    """Create a directory if needed and return it as a ``Path``."""
+
+    resolved = Path(path)
+    resolved.mkdir(parents=True, exist_ok=True)
+    return resolved
+
+
+def save_json(payload: Any, path: str | Path) -> None:
+    """Write JSON with stable indentation."""
+
+    Path(path).write_text(json.dumps(payload, indent=2), encoding="utf-8")
+
+
+def load_json(path: str | Path) -> Any:
+    """Load JSON from disk."""
+
+    return json.loads(Path(path).read_text(encoding="utf-8"))
+
+
+def extract_json_object(text: str) -> dict[str, Any]:
+    """Extract the first JSON object from model output."""
+
+    start = text.find("{")
+    end = text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        raise ValueError("Model output did not contain a JSON object.")
+    return json.loads(text[start : end + 1])
+
+
+def moving_average(values: Sequence[float], window: int = 10) -> list[float]:
+    """Compute a simple moving average."""
+
+    if not values:
+        return []
+    window = max(1, int(window))
+    averaged: list[float] = []
+    for idx in range(len(values)):
+        start = max(0, idx - window + 1)
+        chunk = values[start : idx + 1]
+        averaged.append(sum(chunk) / len(chunk))
+    return averaged
+
+
+def approximate_token_count(text: str) -> int:
+    """Cheap token estimate that works without a tokenizer."""
+
+    stripped = text.strip()
+    if not stripped:
+        return 0
+    return max(1, int(len(stripped.split()) * 1.3))
+
+
+def flatten_dict_items(mapping: dict[str, Any], prefix: str = "") -> Iterable[tuple[str, Any]]:
+    """Flatten nested dictionaries for logging."""
+
+    for key, value in mapping.items():
+        full_key = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, dict):
+            yield from flatten_dict_items(value, prefix=full_key)
+        else:
+            yield full_key, value