"""Main training loop for the Fake Gang Detection LLM Agent.

Learning mechanism (Reflexion + Episodic Memory):

┌─────────────────────────────────────────────────────────────────────┐

│ Episode N │

│ 1. LLM receives: system_prompt + reflections + few_shot + obs │

│ 2. LLM reasons and outputs one action │

│ 3. Action is sent to OpenEnv HTTP server → new obs + reward │

│ 4. Repeat until done │

│ 5. Post-episode: │

│ • If FAIL → Qwen generates a reflection → stored to disk │

│ • If WIN → trajectory saved as few-shot example │

│ 6. Both are injected into Episode N+1's prompt → better policy │

└─────────────────────────────────────────────────────────────────────┘

Curriculum:

Phase 1 (episodes 1-20) : easy — learn basic signal detection

Phase 2 (episodes 21-35): medium — learn to handle evasion

Phase 3 (episodes 36-50): hard — feature-only detection

Usage:

python train.py # defaults: 50 episodes, curriculum

python train.py --task easy --episodes 20

python train.py --env-url http://localhost:8000 --episodes 50 --log-dir runs/

"""

from __future__ import annotations

import argparse

import json

import sys

from collections import defaultdict

from datetime import datetime

from pathlib import Path

from typing import Any, Dict, List, Optional, Tuple

# Resolve imports

_ROOT = Path(__file__).parent

sys.path.insert(0, str(_ROOT))

sys.path.insert(0, str(_ROOT / "server"))

from client import FakeGangEnvClient, StepResult

from models import FakeGangAction, FakeGangObservation, ActionType

from agent.memory import AgentMemory

from agent.hybrid_policy import get_hybrid_action, compute_alpha

from agent.reflection import generate_reflection, generate_success_reflection

# ---------------------------------------------------------------------------

# Curriculum schedule

# ---------------------------------------------------------------------------

def _curriculum_task(episode_num: int, override_task: Optional[str]) -> str:

if override_task:

return override_task

if episode_num < 20:

return "easy"

if episode_num < 35:

return "medium"

return "hard"

def _episode_seed(episode_num: int, task: str) -> int:

"""Rotate through pre-generated seeds so each task sees varied episodes."""

offsets = {"easy": 0, "medium": 50, "hard": 100}

return (episode_num + offsets.get(task, 0)) % 50

# ---------------------------------------------------------------------------

# Episode runner

# ---------------------------------------------------------------------------

def run_episode(

env: FakeGangEnvClient,

task: str,

seed: int,

memory: AgentMemory,

episode_num: int,

alpha: float = 0.20,

temperature: float = 0.4,

verbose: bool = False,

) -> Dict[str, Any]:

"""

Run one full episode using the hybrid policy. Returns a metrics dict.

alpha — current LLM trust weight (0.20 = rules dominate, 1.00 = pure LLM).

"""

result = env.reset(task=task, seed=seed)

obs: FakeGangObservation = result.observation

# Fetch learning context from memory

reflections = memory.get_reflections(task, n=4)

few_shot = memory.get_best_trajectory(task)

action_log: List[str] = []

mode_log: List[str] = []

step_num = 0

while not obs.done:

action, raw_llm, mode = get_hybrid_action(

obs=obs,

reflections=reflections,

few_shot_example=few_shot,

alpha=alpha,

temperature=temperature,

)

# Build a human-readable action string for the log

action_str = action.action_type.value.upper()

if action.account_id:

action_str += f" {action.account_id}"

action_log.append(action_str)

mode_log.append(mode)

if verbose:

mode_tag = mode.split("(")[0] # strip params for brevity

print(f" Step {step_num+1:2d}: {action_str:35s} [{mode_tag}]")

result = env.step(action)

obs = result.observation

step_num += 1

if obs.done:

break

# Parse final message for TP/FP/FN

final_msg = obs.message

won = "[WIN]" in final_msg

reward = result.reward if result.reward is not None else obs.reward or 0.0

steps_used = ({"easy": 30, "medium": 50, "hard": 80}.get(task, 30)) - obs.steps_remaining

# Extract recall/precision from message if present

recall = _extract_float(final_msg, "Recall=")

precision = _extract_float(final_msg, "Precision=")

# Summarise mode distribution for the episode

agree_count = sum(1 for m in mode_log if m == "agree")

rule_count = sum(1 for m in mode_log if m.startswith("rule_override"))

llm_count = sum(1 for m in mode_log if m.startswith("llm"))

return {

"episode": episode_num,

"task": task,

"seed": seed,

"won": won,

"reward": round(reward, 3),

"steps_used": steps_used,

"recall": recall,

"precision": precision,

"action_log": action_log,

"final_message": final_msg,

"n_reflections_used": len(reflections),

"had_few_shot": few_shot is not None,

"alpha_used": round(alpha, 3),

"mode_agree": agree_count,

"mode_rule": rule_count,

"mode_llm": llm_count,

"timestamp": datetime.utcnow().isoformat(),

}

def _extract_float(text: str, key: str) -> float:

import re

m = re.search(re.escape(key) + r"([0-9.]+)", text)

return float(m.group(1)) if m else 0.0

# ---------------------------------------------------------------------------

# Learning step (post-episode)

# ---------------------------------------------------------------------------

def learning_step(

metrics: Dict[str, Any],

memory: AgentMemory,

) -> str:

"""Generate reflection or save trajectory. Returns a short status string."""

task = metrics["task"]

won = metrics["won"]

action_log = metrics["action_log"]

final_msg = metrics["final_message"]

steps_used = metrics["steps_used"]

max_steps = {"easy": 30, "medium": 50, "hard": 80}.get(task, 30)

ep = metrics["episode"]

if won:

# Save trajectory as few-shot example

saved = memory.add_trajectory(

task=task,

action_log=action_log,

final_message=final_msg,

reward=metrics["reward"],

episode_num=ep,

)

# Also generate a success reflection occasionally

if saved or memory.reflection_count(task) == 0:

ref = generate_success_reflection(task, action_log, final_msg, steps_used, max_steps, ep)

memory.add_reflection(task, ref, ep, metrics["reward"])

return f"new best trajectory saved + success reflection generated"

return "trajectory kept (not better than current best)"

else:

# Generate failure reflection

ref = generate_reflection(

task=task,

action_log=action_log,

final_message=final_msg,

won=False,

steps_used=steps_used,

max_steps=max_steps,

episode_num=ep,

)

memory.add_reflection(task, ref, ep, metrics["reward"])

return f"reflection generated: \"{ref[:80]}…\""

# ---------------------------------------------------------------------------

# Progress printer

# ---------------------------------------------------------------------------

class ProgressTracker:

def __init__(self) -> None:

self.by_task: Dict[str, List[Dict]] = defaultdict(list)

self.all: List[Dict] = []

def record(self, m: Dict[str, Any]) -> None:

self.by_task[m["task"]].append(m)

self.all.append(m)

def win_rate(self, task: Optional[str] = None, last_n: int = 10) -> float:

records = self.by_task[task] if task else self.all

if not records:

return 0.0

window = records[-last_n:]

return sum(1 for r in window if r["won"]) / len(window)

def print_episode(self, m: Dict[str, Any], learn_status: str) -> None:

wr = self.win_rate(m["task"], last_n=10)

tag = "WIN " if m["won"] else "LOSS"

alpha = m.get("alpha_used", 0.20)

agree = m.get("mode_agree", 0)

rule = m.get("mode_rule", 0)

llm = m.get("mode_llm", 0)

total_steps = agree + rule + llm

mode_str = (

f"agree={agree} rule={rule} llm={llm}"

if total_steps > 0

else "n/a"

)

print(

f" Ep {m['episode']:3d} | {m['task']:6s} | {tag} | "

f"reward={m['reward']:+7.2f} | "

f"recall={m['recall']:.2f} prec={m['precision']:.2f} | "

f"steps={m['steps_used']:2d} | "

f"wr={wr:.0%} | α={alpha:.2f} | {mode_str}"

)

print(f" └─ learn: {learn_status}")

def print_summary(self, phase: str) -> None:

print(f"\n{'━'*70}")

print(f" {phase} SUMMARY")

print(f"{'━'*70}")

for task in ["easy", "medium", "hard"]:

records = self.by_task[task]

if not records:

continue

wins = sum(1 for r in records if r["won"])

avg_r = sum(r["reward"] for r in records) / len(records)

avg_recall = sum(r["recall"] for r in records) / len(records)

print(

f" {task:6s}: {wins}/{len(records)} wins ({100*wins/len(records):.0f}%) | "

f"avg reward={avg_r:+.2f} | avg recall={avg_recall:.2f}"

)

total = len(self.all)

total_wins = sum(1 for r in self.all if r["won"])

if total:

print(f" {'TOTAL':6s}: {total_wins}/{total} wins ({100*total_wins/total:.0f}%)")

print(f"{'━'*70}\n")

# ---------------------------------------------------------------------------

# Metrics persistence

# ---------------------------------------------------------------------------

def save_metrics(metrics_list: List[Dict], log_dir: Path) -> None:

log_dir.mkdir(parents=True, exist_ok=True)

path = log_dir / "metrics.jsonl"

with open(path, "a") as f:

for m in metrics_list:

f.write(json.dumps(m) + "\n")

# ---------------------------------------------------------------------------

# Main training loop

# ---------------------------------------------------------------------------

def train(

env_url: str = "http://localhost:8000",

task: Optional[str] = None,

n_episodes: int = 50,

temperature: float = 0.4,

verbose: bool = False,

log_dir: Path = Path("runs"),

) -> None:

print(f"\n{'━'*70}")

print(f" Fake Gang Detection — Hybrid Policy Training")

print(f" OpenEnv server : {env_url}")

print(f" Episodes : {n_episodes}")

print(f" Task schedule : {task or 'curriculum (easy→medium→hard)'}")

print(f" LLM : Qwen3 via AWS Bedrock")

print(f" Policy : Hybrid (rules ↔ LLM, dynamic α)")

print(f" Learning : Reflexion (reflections + few-shot in prompt)")

print(f"{'━'*70}\n")

memory = AgentMemory()

tracker = ProgressTracker()

pending_metrics: List[Dict] = []

print(memory.summary())

print()

with FakeGangEnvClient(base_url=env_url) as env:

for ep in range(n_episodes):

current_task = _curriculum_task(ep, task)

seed = _episode_seed(ep, current_task)

# --- Phase announcements ---

if ep == 0:

print("=== Phase 1: Easy — learning basic signal detection ===")

elif ep == 20 and task is None:

print("\n=== Phase 2: Medium — handling evasion ===")

elif ep == 35 and task is None:

print("\n=== Phase 3: Hard — feature-only detection ===")

# --- Compute α for this episode ---

# Load persisted α and update with latest win rate + reflection count

n_refs = memory.reflection_count(current_task)

wr = memory.recent_win_rate(current_task, n=10)

alpha = compute_alpha(recent_win_rate=wr, n_reflections=n_refs, task=current_task)

# --- Run episode ---

metrics = run_episode(

env=env,

task=current_task,

seed=seed,

memory=memory,

episode_num=ep + 1,

alpha=alpha,

temperature=temperature,

verbose=verbose,

)

# --- Post-episode: record win, update + persist α ---

memory.record_win(current_task, metrics["won"], ep + 1)

# Recompute with the just-added result for next episode's alpha

new_wr = memory.recent_win_rate(current_task, n=10)

new_alpha = compute_alpha(recent_win_rate=new_wr, n_reflections=n_refs, task=current_task)

memory.save_alpha(current_task, new_alpha)

# --- Learning step ---

learn_status = learning_step(metrics, memory)

# --- Log ---

tracker.record(metrics)

pending_metrics.append(metrics)

tracker.print_episode(metrics, learn_status)

# Flush metrics every 5 episodes

if len(pending_metrics) >= 5:

save_metrics(pending_metrics, log_dir)

pending_metrics.clear()

# Phase summary every 10 episodes

if (ep + 1) % 10 == 0:

tracker.print_summary(f"Episodes 1–{ep+1}")

# Final flush and summary

if pending_metrics:

save_metrics(pending_metrics, log_dir)

tracker.print_summary("FINAL")

print(memory.summary())

print(f"\nMetrics saved to: {log_dir / 'metrics.jsonl'}")

print(f"Memory saved to: {memory.memory_dir}/")

# ---------------------------------------------------------------------------

# CLI

# ---------------------------------------------------------------------------

if __name__ == "__main__":

parser = argparse.ArgumentParser(

description="Train the Fake Gang Detection LLM agent (Reflexion + Episodic Memory)."

)

parser.add_argument("--env-url", default="http://localhost:8000",

help="OpenEnv server URL (default: http://localhost:8000)")

parser.add_argument("--task", choices=["easy", "medium", "hard"], default=None,

help="Fix task (default: curriculum easy→medium→hard)")

parser.add_argument("--episodes", type=int, default=50,

help="Total training episodes (default: 50)")

parser.add_argument("--temperature", type=float, default=0.4,

help="LLM temperature (default: 0.4)")

parser.add_argument("--verbose", action="store_true",

help="Print each action step")

parser.add_argument("--log-dir", default="runs",

help="Directory for metrics output (default: runs/)")

args = parser.parse_args()

train(

env_url=args.env_url,

task=args.task,

n_episodes=args.episodes,

temperature=args.temperature,

verbose=args.verbose,

log_dir=Path(args.log_dir),

)