train.py

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path

from bdo_ai_env.baseline_agent import greedy_policy, stingy_policy
from bdo_ai_env.training import build_prompt, rollout_episode, save_episode_traces
from server.bdo_environment import BDOEnvironment


def scripted_training_policy(observation: dict) -> dict:
    """
    Lightweight stand-in policy for local development.

    This gives us a deterministic training-like path before Unsloth/TRL is
    installed. It tries to be more coherent than the baselines so we can test
    reward plumbing and trace export.
    """

    highest = max(observation["nodes"], key=lambda node: node["reported_demand"])
    weakest_hardware = min(observation["nodes"], key=lambda node: node["biometric_signal"])
    actions = []

    if weakest_hardware["biometric_signal"] < 0.55:
        actions.append(
            {"name": "dispatch_repair", "params": {"village": weakest_hardware["village"]}}
        )
    elif highest["report_lag_days"] > 0 and highest["reported_demand"] > 2600:
        actions.append(
            {"name": "trigger_field_audit", "params": {"village": highest["village"]}}
        )

    remaining_budget = observation["treasury"]["district_budget"]
    spend = min(max(2000, int(highest["reported_demand"] * 0.8)), remaining_budget)
    actions.append({"name": "allocate_funds", "params": {"village": highest["village"], "amount": spend}})
    actions.append({"name": "approve_batch", "params": {"village": highest["village"], "mode": "conservative"}})

    predicted_fraud = min(
        0.9,
        max(0.1, 1 - sum(node["biometric_signal"] for node in observation["nodes"]) / len(observation["nodes"])),
    )
    thought = (
        f"{highest['village']} shows the highest reported demand, while "
        f"{weakest_hardware['village']} looks most fragile on biometrics. "
        f"Use conservative controls and targeted intervention."
    )
    return {
        "thought_process": thought,
        "predicted_fraud_level": round(predicted_fraud, 3),
        "actions": actions,
    }


def try_unsloth_training() -> str:
    """
    Detect whether the actual training stack is available.

    We don't execute a full GRPO pipeline in this environment yet because the
    dependencies are not installed here, but this hook makes the next step very
    small once they are.
    """

    try:
        import unsloth  # noqa: F401
        import trl  # noqa: F401
        import transformers  # noqa: F401
    except Exception:
        return "unavailable"
    return "available"


def main() -> None:
    parser = argparse.ArgumentParser(description="Phase 3 training scaffold for BDO.ai")
    parser.add_argument("--scenario", default="black_swan")
    parser.add_argument("--episodes", type=int, default=12)
    parser.add_argument("--output", default="artifacts/training_traces.json")
    args = parser.parse_args()

    stack_status = try_unsloth_training()
    env = BDOEnvironment(scenario=args.scenario)

    traces = []
    policies = [scripted_training_policy, greedy_policy, stingy_policy]
    for episode in range(args.episodes):
        policy_fn = policies[episode % len(policies)]
        trace = rollout_episode(env, policy_fn=policy_fn, episode_index=episode)
        traces.append(trace)

    save_episode_traces(traces, args.output)

    prompt_preview = build_prompt(env.reset().model_dump(mode="json", exclude_none=True))
    report = {
        "scenario": args.scenario,
        "episodes": args.episodes,
        "output": args.output,
        "training_stack": stack_status,
        "prompt_preview": prompt_preview[:700],
        "best_total_training_reward": max(trace.total_training_reward for trace in traces),
        "best_avg_belief_accuracy": max(trace.avg_belief_accuracy for trace in traces),
    }
    Path("artifacts").mkdir(exist_ok=True)
    Path("artifacts/train_report.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    random.seed(7)
    main()