scripts/ppo_experiment_v2.py

#!/usr/bin/env python3
"""
PPO Experiment V2 for Symbolic Regression using TRL 0.16+ API

This script implements PPO with a custom RewardModel that computes R² scores
for symbolic expressions. The key insight is that TRL's reward_model parameter
accepts any torch.nn.Module that returns scores.

Key Design:
1. CustomRewardModel wraps R² computation as a neural network module
2. Uses the experimental PPO API from TRL 0.16+
3. JSON format prompts (matches training format)
"""

import os
import sys
import json
import argparse
import logging
import datetime
from pathlib import Path
from typing import Optional, List

import numpy as np
import torch
import torch.nn as nn
from tqdm import tqdm

# Add project root to path
PROJECT_ROOT = Path(__file__).parent.parent
sys.path.insert(0, str(PROJECT_ROOT))
sys.path.insert(0, str(PROJECT_ROOT / "classes"))

from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import Dataset

# Import from experimental PPO (TRL 0.16+)
from trl.experimental.ppo import PPOConfig, PPOTrainer, AutoModelForCausalLMWithValueHead
from peft import PeftModel

from expression import Expression
from dataset import RegressionDataset

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
)
logger = logging.getLogger(__name__)


class SequenceClassifierOutput:
    """Mimics transformers.modeling_outputs.SequenceClassifierOutput"""
    def __init__(self, logits: torch.Tensor):
        self.logits = logits


class SymbolicRegressionRewardModel(nn.Module):
    """
    Custom reward model that computes R² scores for symbolic expressions.

    This wraps the R² computation as a torch.nn.Module that mimics
    AutoModelForSequenceClassification output format, so it can be used
    with TRL's PPOTrainer which expects a reward_model parameter.

    The model doesn't have trainable parameters - it just decodes sequences
    and computes R² scores based on how well the expression fits the data.
    """

    def __init__(self, tokenizer, X: np.ndarray, y: np.ndarray, device: torch.device):
        super().__init__()
        self.tokenizer = tokenizer
        self.X = X
        self.y = y
        self.device = device
        self.n_vars = X.shape[1]

        # Add config attribute to mimic HuggingFace model
        self.config = type('Config', (), {'pad_token_id': tokenizer.pad_token_id})()

        # Dummy parameter so PyTorch recognizes this as a module
        self.dummy = nn.Parameter(torch.zeros(1), requires_grad=False)

        logger.info(f"RewardModel initialized with {len(X)} samples, {self.n_vars} variables")

    def extract_expression(self, generated_text: str) -> str:
        """Extract expression from JSON format output."""
        try:
            # Case 1: Standard JSON with quotes
            if '"expr": "' in generated_text:
                expr_start = generated_text.index('"expr": "') + len('"expr": "')
                remaining = generated_text[expr_start:]
                if '"}' in remaining:
                    return remaining[:remaining.index('"}')].strip()
                if '"' in remaining:
                    return remaining[:remaining.index('"')].strip()
                return remaining.strip()

            # Case 2: Model output without quotes: "expr": value"}
            if '"expr": ' in generated_text:
                expr_start = generated_text.index('"expr": ') + len('"expr": ')
                remaining = generated_text[expr_start:]
                if '"}' in remaining:
                    return remaining[:remaining.index('"}')].strip()
                if '"{' in remaining:
                    return remaining[:remaining.index('"{')].strip().rstrip('}')
                return remaining.strip()

            # Case 3: Compact JSON
            if '"expr":"' in generated_text:
                expr_start = generated_text.index('"expr":"') + len('"expr":"')
                remaining = generated_text[expr_start:]
                if '"}' in remaining:
                    return remaining[:remaining.index('"}')].strip()
                if '"' in remaining:
                    return remaining[:remaining.index('"')].strip()
                return remaining.strip()

        except (ValueError, IndexError):
            pass

        # Fallback
        fallback = generated_text.split('"expr"')[-1].strip(' ":}')
        if '"}' in fallback:
            fallback = fallback[:fallback.index('"}')]
        return fallback.strip()

    def compute_r2(self, expression_str: str) -> float:
        """Compute R² score for an expression."""
        if not expression_str or expression_str.isspace():
            return -1.0

        # Replace constant placeholder C with 1
        if 'C' in expression_str:
            expression_str = expression_str.replace('C', '1')

        try:
            expr = Expression(expression_str, is_prefix=False)

            if not expr.is_valid_on_dataset(self.X):
                return -1.0

            y_pred = expr.evaluate(self.X)

            if not np.all(np.isfinite(y_pred)):
                return -1.0

            ss_res = np.sum((self.y - y_pred) ** 2)
            ss_tot = np.sum((self.y - np.mean(self.y)) ** 2)

            if ss_tot == 0:
                return 0.0

            r2 = 1 - (ss_res / ss_tot)
            return float(np.clip(r2, -1.0, 1.0))
        except Exception:
            return -1.0

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        **kwargs
    ):
        """
        Compute rewards for a batch of sequences.

        Args:
            input_ids: Tensor of shape (batch_size, seq_length)
            attention_mask: Optional attention mask

        Returns:
            SequenceClassifierOutput with logits of shape (batch_size, 1)
        """
        batch_size = input_ids.shape[0]
        rewards = []

        for i in range(batch_size):
            # Decode the sequence
            text = self.tokenizer.decode(input_ids[i], skip_special_tokens=True)

            # Extract expression
            expr_str = self.extract_expression(text)

            # Compute R² score
            r2 = self.compute_r2(expr_str)
            rewards.append(r2)

        # Return in format expected by TRL (SequenceClassifierOutput with logits)
        logits = torch.tensor(rewards, dtype=torch.float32, device=self.device).unsqueeze(-1)
        return SequenceClassifierOutput(logits=logits)


def build_prompt(n_vars: int) -> str:
    """Build JSON format prompt matching training data."""
    vars_list = [f"x_{i+1}" for i in range(n_vars)]
    ops_list = ["+", "-", "*", "sin", "cos"]

    prompt = json.dumps({
        "vars": vars_list,
        "ops": ops_list,
        "cons": None,
        "expr": ""
    })[:-3]  # Remove trailing '"}' for model to complete

    return prompt


def create_ppo_dataset(prompt: str, num_samples: int = 1000) -> Dataset:
    """Create a dataset of prompts for PPO training."""
    return Dataset.from_dict({
        "query": [prompt] * num_samples,
    })


def run_ppo_experiment(
    model_path: str,
    dataset_path: str,
    output_dir: str = "./output/ppo_v2",
    num_episodes: int = 1000,
    batch_size: int = 8,
    learning_rate: float = 1e-5,
):
    """Run PPO experiment with custom R² reward model."""

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # Device setup
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")

    # Load dataset
    logger.info(f"Loading dataset from {dataset_path}")
    dataset_path = Path(dataset_path)
    reg = RegressionDataset(str(dataset_path.parent), dataset_path.name)
    X, y = reg.get_numpy()
    n_vars = X.shape[1]
    logger.info(f"Dataset: {X.shape[0]} samples, {n_vars} variables")

    # Load tokenizer from trained model
    logger.info(f"Loading tokenizer from {model_path}")
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    tokenizer.pad_token = tokenizer.eos_token

    # Load base model and resize embeddings
    logger.info("Loading base GPT-2 model")
    base_model = AutoModelForCausalLM.from_pretrained("gpt2", torch_dtype=torch.float32)

    if len(tokenizer) != base_model.config.vocab_size:
        logger.info(f"Resizing embeddings: {base_model.config.vocab_size} -> {len(tokenizer)}")
        base_model.resize_token_embeddings(len(tokenizer))

    # Load LoRA adapter
    try:
        model_with_lora = PeftModel.from_pretrained(base_model, model_path)
        merged_model = model_with_lora.merge_and_unload()
        logger.info("LoRA adapter loaded and merged")
    except Exception as e:
        logger.warning(f"Could not load as PEFT model: {e}")
        merged_model = AutoModelForCausalLM.from_pretrained(model_path)

    # Wrap with value head for PPO
    policy_model = AutoModelForCausalLMWithValueHead.from_pretrained(merged_model)
    ref_model = AutoModelForCausalLMWithValueHead.from_pretrained(merged_model)
    value_model = AutoModelForCausalLMWithValueHead.from_pretrained(merged_model)

    # Create custom reward model
    reward_model = SymbolicRegressionRewardModel(tokenizer, X, y, device)

    # Build prompt and dataset
    prompt = build_prompt(n_vars)
    logger.info(f"Prompt template: {prompt}...")

    train_dataset = create_ppo_dataset(prompt, num_episodes)

    # PPO Config
    ppo_config = PPOConfig(
        output_dir=str(output_dir),
        learning_rate=learning_rate,
        per_device_train_batch_size=batch_size,
        total_episodes=num_episodes,
        num_ppo_epochs=4,
        gradient_accumulation_steps=1,
        response_length=50,
        temperature=0.7,
        kl_coef=0.05,
        missing_eos_penalty=0.0,  # Don't penalize, expressions don't need EOS
        report_to=None,
    )

    # Create PPO Trainer
    logger.info("Initializing PPO Trainer...")

    try:
        ppo_trainer = PPOTrainer(
            args=ppo_config,
            processing_class=tokenizer,
            model=policy_model,
            ref_model=ref_model,
            reward_model=reward_model,
            value_model=value_model,
            train_dataset=train_dataset,
        )

        logger.info("PPO Trainer initialized successfully!")

        # Run training
        logger.info("Starting PPO training...")
        ppo_trainer.train()

        # Save results
        logger.info(f"Saving model to {output_dir}")
        ppo_trainer.save_model(str(output_dir / "final_model"))

        return {"status": "success", "output_dir": str(output_dir)}

    except Exception as e:
        logger.error(f"PPO training failed: {e}")
        import traceback
        traceback.print_exc()
        return {"status": "error", "error": str(e)}


def main():
    parser = argparse.ArgumentParser(description="PPO Symbolic Regression V2")
    parser.add_argument("--model_path", type=str, default="./output/exp_a_json",
                        help="Path to trained model")
    parser.add_argument("--dataset", type=str, default="./data/ppo_test/sin_x1.csv",
                        help="Path to test dataset CSV")
    parser.add_argument("--output_dir", type=str, default="./output/ppo_v2",
                        help="Output directory")
    parser.add_argument("--num_episodes", type=int, default=1000,
                        help="Number of training episodes")
    parser.add_argument("--batch_size", type=int, default=8,
                        help="Batch size")
    parser.add_argument("--lr", type=float, default=1e-5,
                        help="Learning rate")

    args = parser.parse_args()

    run_ppo_experiment(
        model_path=args.model_path,
        dataset_path=args.dataset,
        output_dir=args.output_dir,
        num_episodes=args.num_episodes,
        batch_size=args.batch_size,
        learning_rate=args.lr,
    )


if __name__ == "__main__":
    main()