path: root/collaborativeagents/training/train_grpo.py

#!/usr/bin/env python3
"""
GRPO Training for Session-Level Reflection.

Based on the MULTISESSIONCOLLAB paper:
- Uses TRL's GRPOTrainer
- Reward = Coverage (LLM judge) + Format bonus
- Judge evaluates reflection quality against user preference enforcements
"""

import os
import json
import argparse
from pathlib import Path
import concurrent.futures

os.environ["WANDB_PROJECT"] = "collaborative-agent-reflection-grpo"

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
from trl import GRPOConfig, GRPOTrainer
from json_repair import repair_json
from tenacity import retry, stop_after_attempt, wait_exponential
import openai


# ===== Configuration =====
DEFAULT_BASE_MODEL = "/projects/bfqt/users/yurenh2/ml-projects/personalization-user-model/collaborativeagents/training/outputs/sft_reflection/checkpoint-best"
MAX_SEQ_LENGTH = 3072
MAX_PROMPT_LENGTH = 2048
MAX_COMPLETION_LENGTH = MAX_SEQ_LENGTH - MAX_PROMPT_LENGTH


# ===== Global tracking =====
reflection_scores_tracker = {
    "scores": [],
    "batch_count": 0
}


def extract_json_answer(text: str) -> str:
    """Extract agent_notes from JSON response."""
    try:
        answer = repair_json(text, return_objects=True)
        if isinstance(answer, dict) and "agent_notes" in answer:
            return answer["agent_notes"]
        return str(answer)
    except Exception as e:
        print(f"Error extracting JSON answer: {e}")
        return ""


def load_grpo_data(data_path: str, tokenizer, max_prompt_length: int):
    """Load and filter GRPO training data."""
    with open(data_path) as f:
        data = json.load(f)

    grpo_data = []
    for elem in data:
        prompt = elem['messages'][0]['content']
        gold_response = extract_json_answer(elem['messages'][1]['content'])
        responses_that_enforce_preferences = elem.get('responses_that_enforce_preferences', [])

        # Filter by token length
        tokens = tokenizer.tokenize(prompt)
        if len(tokens) > max_prompt_length:
            continue

        grpo_data.append({
            'prompt': prompt,
            'gold_response': gold_response,
            'responses_that_enforce_preferences': responses_that_enforce_preferences
        })

    return grpo_data


# ===== Reward Functions =====

def create_judge_client(judge_url: str):
    """Create OpenAI client for judge model."""
    return openai.OpenAI(base_url=judge_url, api_key="EMPTY")


@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=4, max=10))
def ask_judge(client, prompt: str, model_name: str = "meta-llama/Llama-3.3-70B-Instruct"):
    """Query judge model for evaluation."""
    messages = [{"role": "user", "content": prompt}]
    response = client.chat.completions.create(
        model=model_name,
        messages=messages,
        max_tokens=1024,
        temperature=0.0,
    )
    return response.choices[0].message.content.strip()


def evaluate_single_reflection(args):
    """Evaluate a single reflection using LLM judge."""
    prompt, completion, gold_response, responses_that_enforce_preferences, judge_url, judge_model = args

    if not completion:
        print(f"Empty completion, score=0")
        return 0

    # Format user messages that enforce preferences
    user_messages_str = ""
    for i, response in enumerate(responses_that_enforce_preferences):
        user_messages_str += f"User message #{i+1}: {response}\n"

    if not user_messages_str:
        user_messages_str = "No explicit preference enforcement found in conversation."

    evaluation_prompt = f"""You are an expert evaluator analyzing a conversational agent's reflection of a conversation, where they analyze the conversation to identify the user's preferences and create actionable notes to help them satisfy these preferences in future conversations.

Throughout the conversation, the user explicitly enforces their preferences whenever necessary. The agent analyzes the conversation to identify the user's preferences and create actionable notes to help them satisfy these preferences in future conversations.

# Your Task:
Evaluate whether the agent's reflection successfully captures the user's preferences and provides actionable notes to help them satisfy these preferences in future conversations.

# Agent's Reflection:
{completion}

# User Messages Where They Enforce Their Preferences:
{user_messages_str}

# Gold Reflection:
Here is a gold reflection for the same conversation. Use this as a reference to evaluate the agent's reflection.
{gold_response}

# Evaluation Criteria:
Assess the reflection on four dimensions:
- **Coverage (Completeness):** Does the agent's reflection capture all of the user's preferences?
- **Actionability (Quality):** Does the agent's reflection provide actionable notes and details that help the agent satisfy these preferences in future conversations?
- **Accuracy (No Hallucination):** Are all points grounded in actual user statements? Does the reflection avoid inventing preferences or misrepresenting user statements?
- **Clarity:** Is the reflection well-organized and clearly formatted? Does the reflection avoid redundancy, with each preference stated once without repetitive or overlapping notes?

You will output a score from 0-3, where:
- 0: Does not effectively capture user preferences: gaps in coverage, or significant hallucinations
- 1: Captures some preferences with limited actionable notes, may hallucinate some preferences
- 2: Captures most preferences with actionable notes, may have some slight hallucinations
- 3: Comprehensively captures all preferences with highly actionable notes and no hallucinations

# Output Format:
{{
    "reasoning": "Brief explanation of your decision",
    "reflection_score": <0-3>
}}

Output a properly formatted JSON response, as specified by the Output Format."""

    try:
        client = create_judge_client(judge_url)
        result = ask_judge(client, evaluation_prompt, judge_model)
        parsed = repair_json(result, return_objects=True)
        score = parsed.get("reflection_score", 0)
        print(f"Reflection score: {score}")
        return score
    except Exception as e:
        print(f"Error evaluating reflection: {e}")
        return 0


def reflection_reward_func(prompts, completions, gold_response, responses_that_enforce_preferences,
                          judge_url, judge_model, **kwargs) -> list[float]:
    """Compute reflection quality reward using LLM judge."""
    args_list = []
    for i, completion in enumerate(completions):
        completion_notes = extract_json_answer(completion)
        args_list.append((
            prompts[i],
            completion_notes,
            gold_response[i],
            responses_that_enforce_preferences[i],
            judge_url,
            judge_model
        ))

    # Parallel evaluation
    with concurrent.futures.ProcessPoolExecutor(max_workers=8) as executor:
        rewards = list(executor.map(evaluate_single_reflection, args_list))

    # Track scores
    reflection_scores_tracker["scores"].extend(rewards)
    reflection_scores_tracker["batch_count"] += 1
    avg_score = sum(reflection_scores_tracker["scores"]) / len(reflection_scores_tracker["scores"])
    print(f"Batch {reflection_scores_tracker['batch_count']}: rewards={rewards}, running_avg={avg_score:.2f}")

    return rewards


def format_reward_func(prompts, completions, **kwargs) -> list[float]:
    """Reward for proper JSON format with required fields."""
    rewards = []
    for completion in completions:
        reward = 0.0
        try:
            parsed = repair_json(completion, return_objects=True)
            if isinstance(parsed, dict):
                if "agent_notes" in parsed:
                    reward += 0.25
                if "user_preferences_reasoning" in parsed:
                    reward += 0.25
        except:
            pass
        rewards.append(reward)

    print(f"Format rewards: {rewards}")
    return rewards


def main():
    parser = argparse.ArgumentParser(description="Train GRPO for session-level reflection")
    parser.add_argument("--model-path", type=str, default=DEFAULT_BASE_MODEL,
                       help="Path to SFT checkpoint to start from")
    parser.add_argument("--data-path", type=str, required=True,
                       help="Path to GRPO training data JSON")
    parser.add_argument("--output-dir", type=str, default="./outputs_grpo_reflection",
                       help="Output directory for checkpoints")
    parser.add_argument("--judge-url", type=str, default="http://localhost:8000/v1",
                       help="URL for judge model vLLM server")
    parser.add_argument("--judge-model", type=str, default="meta-llama/Llama-3.3-70B-Instruct",
                       help="Model name for judge")
    parser.add_argument("--max-steps", type=int, default=200,
                       help="Maximum training steps")
    parser.add_argument("--learning-rate", type=float, default=1e-6,
                       help="Learning rate")
    parser.add_argument("--num-generations", type=int, default=8,
                       help="Number of rollout generations per prompt")
    args = parser.parse_args()

    print(f"Loading model from {args.model_path}...")
    tokenizer = AutoTokenizer.from_pretrained(args.model_path)

    generation_config = GenerationConfig.from_pretrained(args.model_path)
    generation_config.max_length = MAX_SEQ_LENGTH
    generation_config.do_sample = True
    generation_config.top_p = 0.9

    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        generation_config=generation_config
    )

    print(f"Loading training data from {args.data_path}...")
    dataset = load_grpo_data(args.data_path, tokenizer, MAX_PROMPT_LENGTH)
    print(f"Loaded {len(dataset)} examples")

    # Training config (from paper Table 4)
    training_args = GRPOConfig(
        learning_rate=args.learning_rate,
        gradient_accumulation_steps=4,
        num_train_epochs=1,
        bf16=True,
        max_prompt_length=MAX_PROMPT_LENGTH,
        max_completion_length=MAX_COMPLETION_LENGTH,
        report_to="wandb",
        logging_steps=1,
        num_generations=args.num_generations,
        max_steps=args.max_steps,
        save_steps=50,
        output_dir=args.output_dir,
        kl_coef=0.003,  # From paper Table 4
    )

    # Create reward functions with judge config
    def reflection_reward_with_judge(prompts, completions, gold_response, responses_that_enforce_preferences, **kwargs):
        return reflection_reward_func(
            prompts, completions, gold_response, responses_that_enforce_preferences,
            args.judge_url, args.judge_model, **kwargs
        )

    trainer = GRPOTrainer(
        model=model,
        processing_class=tokenizer,
        reward_funcs=[format_reward_func, reflection_reward_with_judge],
        args=training_args,
        train_dataset=dataset,
    )

    print("Starting GRPO training...")
    trainer.train()

    # Save final model
    final_path = Path(args.output_dir) / "final"
    trainer.save_model(str(final_path))
    tokenizer.save_pretrained(str(final_path))
    print(f"Saved final model to {final_path}")


if __name__ == "__main__":
    main()