train_adversarial.py

import logging
import pickle
from pathlib import Path

from src.adversarial.agent import FraudsterAgent

logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


def adversarial_training(
    initial_model_path: str = "models/fraud_model.pkl",
    feature_engineer_path: str = "models/feature_engineer.pkl",
    n_rounds: int = 5,
    n_episodes_per_round: int = 3,
):
    """Run adversarial training rounds against the current fraud model."""

    logger.info("Loading initial model and feature engineer")
    with open(initial_model_path, "rb") as f:
        model = pickle.load(f)

    with open(feature_engineer_path, "rb") as f:
        feature_engineer = pickle.load(f)

    agent = FraudsterAgent()

    history: dict[str, list] = {"round": [], "evasion_rate": [], "q_table_size": []}

    for round_num in range(n_rounds):
        logger.info("Round %s/%s", round_num + 1, n_rounds)

        round_evasion_rates = []

        for episode in range(n_episodes_per_round):
            metrics = agent.train_episode(model, feature_engineer, n_transactions=100)
            round_evasion_rates.append(metrics["evasion_rate"])
            logger.info(
                f"Episode {episode + 1}: "
                f"Evasion={metrics['evasion_rate']:.1%}, "
                f"Q-table size={metrics['q_table_size']}"
            )

        avg_evasion = sum(round_evasion_rates) / len(round_evasion_rates)

        history["round"].append(round_num)
        history["evasion_rate"].append(avg_evasion)
        history["q_table_size"].append(len(agent.q_table))

        logger.info(f"Round {round_num + 1} avg evasion rate: {avg_evasion:.1%}")

    logger.info("Saving adversarial agent")
    Path("models").mkdir(exist_ok=True)
    agent.save("models/adversarial_agent.pkl")

    logger.info("Adversarial training complete")
    logger.info(f"Final evasion rate: {history['evasion_rate'][-1]:.1%}")
    logger.info(f"Q-table size: {history['q_table_size'][-1]}")

    return history


if __name__ == "__main__":
    history = adversarial_training()