Add collaborativeagents module and update gitignore

- Add collaborativeagents subproject with adapters, agents, and evaluation modules
- Update .gitignore to exclude large binary files (.whl, .tar), wandb logs, and results

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

1 files changed, 294 insertions, 0 deletions

diff --git a/collaborativeagents/training/train_grpo.py b/collaborativeagents/training/train_grpo.py
new file mode 100644
index 0000000..42a85a1
--- /dev/null
+++ b/collaborativeagents/training/train_grpo.py
@@ -0,0 +1,294 @@
+#!/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()