train.py

"""
Example RL training loop for the BA Agent environment.

Wraps the OpenEnv HTTP env into a Gymnasium-style adapter and runs a tiny
PPO-flavoured loop over an LLM policy. Intended as a starting reference
for users plugging the env into TRL / verl / SkyRL / custom training stacks.

Two policy backends supported:
    - "stub"        — fixed plan (constant payloads per stage)
    - "openrouter"  — calls a chat model to produce each stage's payload

Usage:
    python train.py --policy stub --episodes 5
    python train.py --policy openrouter --model openai/gpt-4o-mini --episodes 5
"""
from __future__ import annotations

import argparse
import json
import os
import statistics
import time
from typing import Any, Dict, List

from client import BAAgentClient


_SYS_PROMPT_PER_STAGE = {
    "EXTRACT": "You are extracting a structured summary of one feature. Return a short paragraph listing scope, entities, and constraints.",
    "INTERVIEW": "Return a JSON list of 3 stakeholder Q&A pairs ([{q,a}]) covering primary actor, trigger, and exit criteria.",
    "GRAPH": "Return a JSON object {nodes: [..], edges: [[a,b], ...]} for the entity / workflow graph.",
    "STORY_GEN": "Return a JSON list of user stories. Each story MUST have title, description ('As a ... I want ... so that ...'), and acceptance_criteria (Given/When/Then). Cover every workflow stage in the input docs.",
    "FINISH": "",
}


def _stub_policy(stage: str, obs: Dict[str, Any]) -> str:
    return {
        "EXTRACT": f"Feature scope summary for {obs.get('title','')}. Entities and constraints derived from input docs.",
        "INTERVIEW": json.dumps([{"q": "Primary actor?", "a": "Operator"}, {"q": "Trigger?", "a": "Inbound work item"}]),
        "GRAPH": json.dumps({"nodes": ["Actor", "Document"], "edges": [["Actor", "Document"]]}),
        "STORY_GEN": json.dumps([
            {"title": "Receive item", "description": "As an operator, I want to receive a work item, so that I can process it.",
             "acceptance_criteria": "Given inbound, When received, Then a record is created."},
            {"title": "Validate item", "description": "As an operator, I want validation, so that errors are caught.",
             "acceptance_criteria": "Given a record, When validated, Then errors are flagged."},
            {"title": "Finalise item", "description": "As an operator, I want to finalise, so that the record is auditable.",
             "acceptance_criteria": "Given a validated record, When finalised, Then status=Finalised."},
        ]),
        "FINISH": "",
    }[stage]


def _openrouter_policy(stage: str, obs: Dict[str, Any], model: str) -> str:
    api_key = os.environ.get("OPENROUTER_API_KEY", "").strip()
    if not api_key:
        return _stub_policy(stage, obs)
    from openai import OpenAI
    client = OpenAI(api_key=api_key, base_url="https://openrouter.ai/api/v1")
    docs = "\n\n".join(f"=== {d.get('filename','')} ===\n{(d.get('content','') or '')[:1500]}" for d in obs.get("input_documents", []) or [])
    prompt = (
        f"Feature: {obs.get('title','')}\n"
        f"Description: {obs.get('description','')[:1000]}\n\n"
        f"Input docs:\n{docs[:6000]}\n\n"
        f"Task: {_SYS_PROMPT_PER_STAGE[stage]}\n\n"
        "Return ONLY the requested payload, no commentary."
    )
    resp = client.chat.completions.create(
        model=model,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,
        max_tokens=1500,
    )
    return resp.choices[0].message.content.strip()


def run_episode(env: BAAgentClient, policy, model: str) -> Dict[str, Any]:
    obs = env.reset()
    total = 0.0
    step_rewards: List[float] = []
    terminal: Dict[str, Any] = {}
    for stage in ["EXTRACT", "INTERVIEW", "GRAPH", "STORY_GEN", "FINISH"]:
        if stage == "FINISH":
            payload = ""
        else:
            payload = policy(stage, obs) if model is None else _openrouter_policy(stage, obs, model)
        obs, r, done, meta = env.step(stage, payload)
        total += r
        step_rewards.append(r)
        if done:
            terminal = meta
            break
    return {
        "task_id": obs.get("task_id"),
        "step_rewards": [round(x, 3) for x in step_rewards],
        "episode_reward": round(total, 4),
        "composite_0_to_100": terminal.get("composite_0_to_100"),
        "engine": terminal.get("engine"),
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--server", default="http://localhost:8000")
    parser.add_argument("--policy", choices=["stub", "openrouter"], default="stub")
    parser.add_argument("--model", default="openai/gpt-4o-mini")
    parser.add_argument("--episodes", type=int, default=5)
    args = parser.parse_args()

    env = BAAgentClient(args.server)
    if args.policy == "stub":
        pol = _stub_policy
        model_id = None
    else:
        pol = _stub_policy  # fallback inside
        model_id = args.model

    print(f"Server: {args.server}  Policy: {args.policy}  Episodes: {args.episodes}")
    composites: List[float] = []
    ep_rewards: List[float] = []
    for i in range(args.episodes):
        t0 = time.time()
        out = run_episode(env, pol, model_id)
        dt = time.time() - t0
        c = out.get("composite_0_to_100")
        composites.append(c if c is not None else 0.0)
        ep_rewards.append(out["episode_reward"])
        print(f"ep {i+1:>2}/{args.episodes}  task={out['task_id']:<10}  total_r={out['episode_reward']:+.3f}  composite={c}/100  {dt:.1f}s")

    print()
    print(f"Mean episode reward : {statistics.mean(ep_rewards):+.3f} ± {statistics.pstdev(ep_rewards):.3f}")
    print(f"Mean composite       : {statistics.mean(composites):.2f}/100")


if __name__ == "__main__":
    main()