train.py

#!/usr/bin/env python3
"""
ESCTR Training Script — GRPO with TRL + vLLM
=============================================

Train an LLM to be an autonomous financial controller using
Group Relative Policy Optimization (GRPO) against the ESCTR environment.

Usage (Colab / HF Jobs):
    pip install -Uq "trl[vllm]" trackio datasets
    pip install -e .        # install esctr-environment package
    python train.py

The environment runs in-process (no HTTP server needed during training).
The HF Space deployment is only for judges to test the environment interactively.
"""

import random
import sys
import os

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer

# ---------------------------------------------------------------------------
# Import ESCTR environment (runs in-process, no server needed)
# ---------------------------------------------------------------------------
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from server.environment import ESCTREnvironment
from server.models import ESCTRAction


# ---------------------------------------------------------------------------
# System prompt — tells the model what it is and what tools are available
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """You are an autonomous Financial Controller AI operating within an enterprise ERP system.

Your job is to investigate financial discrepancies in procurement records by using the available tools, then submit a precise monetary adjustment.

INVESTIGATION WORKFLOW:
1. Query databases to discover what records exist (purchase_orders, invoices, shipping_logs, sla_contracts, warehouse_logs)
2. Read specific documents to get full details
3. Compare line items, delivery dates, and contract terms
4. Calculate the exact adjustment amount
5. Submit your financial decision with the calculated amount and reasoning

CRITICAL RULES:
- Always query AND read documents before submitting. Never guess.
- Your adjustment_amount must be the EXACT monetary difference you calculated.
- Show your arithmetic in the adjustment_reason.
- If a vendor offers a settlement, verify their claims against internal records before accepting.

You have access to the following tools. Call them to interact with the ERP system."""


# ---------------------------------------------------------------------------
# ESCTR Environment wrapper for TRL environment_factory
# ---------------------------------------------------------------------------
# TRL discovers public methods (with docstrings) as callable tools.
# The model generates tool calls; TRL executes them and feeds results back.
# ---------------------------------------------------------------------------

# Task to train on — start with the easiest task for stable training
TRAIN_TASKS = [
    t.strip()
    for t in os.environ.get("ESCTR_TASKS", os.environ.get("ESCTR_TASK", "procurement_reconciliation")).split(",")
    if t.strip()
]


class ESCTRToolEnv:
    """TRL-compatible wrapper around the ESCTR environment.

    Public methods with docstrings are auto-discovered as tools by TRL's
    environment_factory. The trainer handles the multi-turn loop automatically.
    """

    def __init__(self):
        self.env = ESCTREnvironment()
        self.reward = 0.0
        self.done = False
        self._tasks = TRAIN_TASKS or ["procurement_reconciliation"]
        self._task = self._tasks[0]

    def reset(self, **kwargs) -> str | None:
        """Reset the environment and return the initial briefing."""
        seed = random.randint(0, 100_000)
        # Simple curriculum: sample tasks from configured task pool.
        self._task = random.choice(self._tasks)
        obs = self.env.reset(
            task_name=self._task,
            seed=seed,
        )
        self.reward = 0.0
        self.done = False
        return obs.system_response

    def query_database(self, table: str) -> str:
        """
        Query a corporate database table to discover available records.

        Args:
            table: The database table to query. One of: 'purchase_orders', 'invoices', 'shipping_logs', 'sla_contracts', 'warehouse_logs'

        Returns:
            A summary of records found in the specified table.
        """
        if self.done:
            raise ValueError("Episode is over. No more actions allowed.")

        action = ESCTRAction(
            action_type="query_database",
            query_parameters={"table": table},
        )
        obs = self.env.step(action)
        self.reward = obs.reward
        self.done = obs.done
        return obs.system_response

    def read_document(self, document_id: str) -> str:
        """
        Read a specific document by its unique identifier to see full details.

        Args:
            document_id: The document ID to read, e.g. 'PO-2024-0055' or 'INV-2024-0055'

        Returns:
            The full contents of the requested document.
        """
        if self.done:
            raise ValueError("Episode is over. No more actions allowed.")

        action = ESCTRAction(
            action_type="read_document",
            document_id=document_id,
        )
        obs = self.env.step(action)
        self.reward = obs.reward
        self.done = obs.done
        return obs.system_response

    def communicate_vendor(self, message_content: str) -> str:
        """
        Send a message to the vendor during a dispute negotiation.

        Args:
            message_content: The message to send to the vendor, such as requesting clarification or rejecting a settlement offer.

        Returns:
            The vendor's response to your message.
        """
        if self.done:
            raise ValueError("Episode is over. No more actions allowed.")

        action = ESCTRAction(
            action_type="communicate_vendor",
            message_content=message_content,
        )
        obs = self.env.step(action)
        self.reward = obs.reward
        self.done = obs.done
        return obs.system_response

    def submit_financial_decision(self, adjustment_amount: float, adjustment_reason: str) -> str:
        """
        Submit the final financial adjustment. This is the terminal action that ends the episode.

        Args:
            adjustment_amount: The exact monetary adjustment amount as a float (e.g. 450.00). Must be calculated from the documents.
            adjustment_reason: A brief explanation of why this adjustment is correct, including your arithmetic.

        Returns:
            The grading result with your score and feedback.
        """
        if self.done:
            raise ValueError("Episode is over. No more actions allowed.")

        action = ESCTRAction(
            action_type="submit_financial_decision",
            adjustment_amount=adjustment_amount,
            adjustment_reason=adjustment_reason,
        )
        obs = self.env.step(action)
        self.reward = obs.reward
        self.done = obs.done
        return obs.system_response


# ---------------------------------------------------------------------------
# Reward function — reads from env instances after each episode
# ---------------------------------------------------------------------------

def reward_func(environments, **kwargs) -> list[float]:
    """Extract reward from each environment instance after episode completion."""
    return [env.reward for env in environments]


# ---------------------------------------------------------------------------
# Training configuration
# ---------------------------------------------------------------------------

def main():
    # Model selection — Qwen3-1.7B is efficient on T4 GPU
    model_name = os.environ.get("ESCTR_MODEL", "Qwen/Qwen3-1.7B")
    output_dir = os.environ.get("ESCTR_OUTPUT", "esctr-grpo-trained")
    num_episodes = int(os.environ.get("ESCTR_EPISODES", "1000"))

    # Create dataset — each entry triggers one rollout episode
    dataset = Dataset.from_dict({
        "prompt": [[{"role": "user", "content": SYSTEM_PROMPT}]] * num_episodes
    })

    # GRPO configuration
    grpo_config = GRPOConfig(
        # Training schedule
        num_train_epochs=1,
        learning_rate=1e-6,
        gradient_accumulation_steps=4,
        per_device_train_batch_size=1,
        warmup_steps=10,
        optim="adamw_torch",
        max_grad_norm=1.0,

        # GRPO settings
        num_generations=2,
        max_completion_length=768,
        log_completions=True,
        num_completions_to_print=2,
        chat_template_kwargs={"enable_thinking": False},

        # Logging
        output_dir=output_dir,
        report_to="trackio",
        trackio_space_id=output_dir,
        logging_steps=1,
        save_steps=25,
        save_total_limit=2,

        # Memory optimization
        gradient_checkpointing=True,
        bf16=False,
        fp16=True,

        # Hub integration
        push_to_hub=True,
    )

    # Create trainer
    trainer = GRPOTrainer(
        model=model_name,
        reward_funcs=reward_func,
        train_dataset=dataset,
        args=grpo_config,
        environment_factory=ESCTRToolEnv,
    )

    # Show GPU stats before training
    import torch
    if torch.cuda.is_available():
        gpu_stats = torch.cuda.get_device_properties(0)
        start_gpu_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
        max_memory = round(gpu_stats.total_memory / 1024 / 1024 / 1024, 3)
        print(f"GPU = {gpu_stats.name}. Max memory = {max_memory} GB.")
        print(f"{start_gpu_memory} GB of memory reserved.")

    print(f"\n{'='*60}")
    print(f"ESCTR Training — {model_name}")
    print(f"Tasks: {', '.join(TRAIN_TASKS)}")
    print(f"Episodes: {num_episodes}")
    print(f"Output: {output_dir}")
    print(f"{'='*60}\n")

    # Train!
    trainer_stats = trainer.train()

    # Show training stats
    if torch.cuda.is_available():
        used_memory = round(torch.cuda.max_memory_reserved() / 1024 / 1024 / 1024, 3)
        print(f"\nTraining completed in {trainer_stats.metrics['train_runtime']:.0f} seconds")
        print(f"Peak GPU memory: {used_memory} GB / {max_memory} GB")

    # Save and push
    trainer.save_model(output_dir)
    trainer.push_to_hub()
    print(f"\nModel saved to {output_dir} and pushed to Hub!")


if __name__ == "__main__":
    main()