rlm_forge_training.py

import torch
import json
import re
import random
from typing import Optional
from dataclasses import dataclass
from datasets import Dataset
from trl import GRPOConfig, GRPOTrainer

print(f"GPU: {torch.cuda.get_device_name(0)}")
# print(f"VRAM: {torch.cuda.get_device_properties(0). / 1e9:.1f} GB")
print(f"PyTorch: {torch.__version__}")

from rlm_forge.server.environment import RLMForgeEnvironment
from rlm_forge.models import RLMForgeAction

env = RLMForgeEnvironment()

# Run a quick episode
obs = env.reset(seed=1)
print(f"Task: {obs.task_description[:200]}...")
print(f"Available tools: {obs.available_functions}")

# Take a step — list files
obs2 = env.step(RLMForgeAction(code="print(list_dir())"))
print(f"\nStep 1 stdout: {obs2.stdout[:200]}")

# Finalize and get reward
obs3 = env.step(RLMForgeAction(code="FINAL()"))
print(f"\nBaseline reward (no implementation): {obs3.reward:.4f}")
print(f"Test results: {obs3.test_results}")

env.cleanup()


# Model config — adjust based on available VRAM
# MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"  # 32B for H100
MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"  # Fallback for smaller GPUs
HF_TOKEN = ''
MAX_STEPS_PER_EPISODE = 6  # Max REPL interactions per episode
NUM_EPISODES_PER_PROMPT = 2  # GRPO group size (completions per prompt)
NUM_TRAINING_PROMPTS = 8  # 16 # Total unique prompts (episodes) for training
GRPO_EPOCHS = 2  # Training epochs over collected data
BATCH_SIZE = 2
GRAD_ACCUM = 4

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import LoraConfig, get_peft_model

# 4-bit quantization for 32B model on H100
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True,token=HF_TOKEN)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token



model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    # attn_implementation="flash_attention_2",
    token=HF_TOKEN
)

# LoRA config for efficient training
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    lora_dropout=0.05,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    task_type="CAUSAL_LM",
)

model = get_peft_model(model, lora_config)
model.print_trainable_parameters()

SYSTEM_PROMPT = """You are an expert Python developer. You are given a repository where a source file has been replaced with a broken stub. Your task is to explore the repository, understand the expected behavior from the tests, and rewrite the source file so all tests pass.

You interact via a Python REPL. Available functions:
- read_file(path) — Read a file from the repo
- list_dir(path='.') — List directory contents
- search(pattern, path='.') — Grep for a pattern
- write_file(path, content) — Write/create a file
- run_tests(test_path=None) — Run pytest on a test file
- FINAL() — Signal that your implementation is complete

Strategy:
1. Read the failing test file to understand expected behavior
2. Read other source files for context (imports, dependencies)
3. Write the implementation
4. Run tests to verify
5. Fix any failures
6. Call FINAL() when done

Output ONLY valid Python code. No markdown, no explanations — just code to execute."""


def build_prompt(task_description: str, failing_tests: list[str]) -> list[dict]:
    """Build the chat prompt for the initial observation."""
    user_msg = f"{task_description}\n\nFailing tests:\n" + "\n".join(failing_tests)
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_msg},
    ]


def extract_code_from_response(response: str) -> str:
    """Extract executable Python code from model response."""
    # Try to find code blocks first
    code_blocks = re.findall(r"```(?:python)?\n(.*?)```", response, re.DOTALL)
    if code_blocks:
        return "\n".join(code_blocks)
    # Otherwise treat the whole response as code
    lines = response.strip().split("\n")
    code_lines = []
    for line in lines:
        stripped = line.strip()
        if stripped and not stripped.startswith("#") and any(c in stripped for c in "=()[]{}:"):
            code_lines.append(line)
        elif stripped.startswith("#") or stripped.startswith("import") or stripped.startswith("from"):
            code_lines.append(line)
        elif not stripped:
            code_lines.append(line)
        else:
            code_lines.append(f"# {line}")
    return "\n".join(code_lines)


print("Prompt builder ready.")

@dataclass
class Trajectory:
    """A full multi-step episode trajectory for GRPO training."""
    prompt_text: str        # Tokenized prompt (system + task)
    completion_text: str    # All model outputs concatenated
    reward: float           # Final episode reward
    steps: int              # Number of steps taken
    seed: int               # Environment seed (for reproducibility)
    tests_passed: int
    tests_total: int


def run_episode(
    model,
    tokenizer,
    env: RLMForgeEnvironment,
    seed: int,
    max_steps: int = MAX_STEPS_PER_EPISODE,
    temperature: float = 0.7,
    max_new_tokens: int = 2048,
) -> Trajectory:
    """Run a single episode: generate code actions, execute them, collect trajectory."""
    obs = env.reset(seed=seed)

    messages = build_prompt(obs.task_description, obs.failing_tests or [])
    prompt_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)

    all_completions = []  # All model outputs for this episode

    for step_i in range(max_steps):
        # Build the full conversation so far for the model
        if step_i > 0:
            # Add the observation as assistant feedback
            messages.append({"role": "user", "content": f"REPL output:\n{obs.stdout}\n{obs.stderr}"})

        # Generate next action
        full_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(full_text, return_tensors="pt", truncation=True, max_length=8192).to(model.device)

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_new_tokens,
                temperature=temperature,
                top_p=0.95,
                do_sample=True,
                pad_token_id=tokenizer.pad_token_id,
            )

        # Decode only the new tokens
        new_tokens = outputs[0][inputs["input_ids"].shape[1]:]
        response = tokenizer.decode(new_tokens, skip_special_tokens=True)
        all_completions.append(response)

        # Add to conversation history
        messages.append({"role": "assistant", "content": response})

        # Extract and execute code
        code = extract_code_from_response(response)

        # Check if model wants to finalize
        if "FINAL()" in code:
            obs = env.step(RLMForgeAction(code=code))
            break
        else:
            obs = env.step(RLMForgeAction(code=code))

        if obs.done:
            break

    # If we exhausted steps without FINAL, force finalize
    if not obs.done:
        obs = env.step(RLMForgeAction(code="FINAL()"))

    # Build the full completion text (all model outputs joined)
    completion_text = "\n<|step|>\n".join(all_completions)

    reward = obs.reward or 0.0
    test_results = obs.test_results or {}

    return Trajectory(
        prompt_text=prompt_text,
        completion_text=completion_text,
        reward=reward,
        steps=step_i + 1,
        seed=seed,
        tests_passed=test_results.get("tests_passed", 0),
        tests_total=test_results.get("tests_total", 0),
    )


print("Episode runner ready.")


def collect_trajectories(
    model,
    tokenizer,
    num_prompts: int = NUM_TRAINING_PROMPTS,
    episodes_per_prompt: int = NUM_EPISODES_PER_PROMPT,
    temperature: float = 0.7,
) -> list[list[Trajectory]]:
    """Collect GRPO groups: multiple trajectories per unique prompt/seed."""
    env = RLMForgeEnvironment()
    all_groups = []

    for prompt_idx in range(num_prompts):
        seed = prompt_idx * 100  # Deterministic seeds
        group = []

        for ep_idx in range(episodes_per_prompt):
            print(f"  Prompt {prompt_idx+1}/{num_prompts}, Episode {ep_idx+1}/{episodes_per_prompt}...", end=" ")
            traj = run_episode(
                model, tokenizer, env,
                seed=seed,  # Same seed = same task for GRPO group
                temperature=temperature + 0.1 * ep_idx,  # Vary temperature for diversity
            )
            group.append(traj)
            print(f"reward={traj.reward:.3f}, steps={traj.steps}, "
                  f"tests={traj.tests_passed}/{traj.tests_total}")

        all_groups.append(group)

    env.cleanup()
    return all_groups

# GRPO Training configuration
grpo_config = GRPOConfig(
    output_dir="./rlm_forge_grpo_output",
    num_train_epochs=GRPO_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM,
    learning_rate=1e-5,
    warmup_ratio=0.1,
    max_completion_length=4096,
    # max_prompt_length=4096,
    num_generations=NUM_EPISODES_PER_PROMPT,  # GRPO group size
    logging_steps=1,
    save_strategy="epoch",
    bf16=True,
    gradient_checkpointing=True,
    # GRPO-specific
    beta=0.1,  # KL penalty coefficient
    report_to="none",
)

# Collect pre-training baseline
print("=" * 60)
print("COLLECTING BASELINE TRAJECTORIES")
print("=" * 60)
baseline_groups = collect_trajectories(model, tokenizer)

# Summary stats
all_rewards = [t.reward for g in baseline_groups for t in g]
print(f"\nBaseline: mean_reward={sum(all_rewards)/len(all_rewards):.4f}, "
      f"min={min(all_rewards):.4f}, max={max(all_rewards):.4f}")





def trajectories_to_dataset(groups: list[list[Trajectory]]) -> Dataset:
    """Convert trajectory groups into a HuggingFace Dataset for GRPO training."""
    records = []
    for group in groups:
        prompt = group[0].prompt_text
        for traj in group:
            records.append({
                "prompt": prompt,
                "completion": traj.completion_text,
                "reward": traj.reward,
            })
    return Dataset.from_list(records)


def build_reward_fn(groups: list[list[Trajectory]]):
    """Build a reward function from pre-collected trajectories."""
    reward_map = {}
    for group in groups:
        for traj in group:
            key = traj.completion_text[:200]
            reward_map[key] = traj.reward

    def reward_fn(completions: list[str], **kwargs) -> list[float]:
        rewards = []
        for c in completions:
            key = c[:200]
            rewards.append(reward_map.get(key, 0.0))
        return rewards

    return reward_fn


# Build dataset from baseline trajectories
train_dataset = trajectories_to_dataset(baseline_groups)
print(f"Training dataset: {len(train_dataset)} examples")
print(f"Sample prompt length: {len(train_dataset[0]['prompt'])} chars")
print(f"Sample completion length: {len(train_dataset[0]['completion'])} chars")
print(f"Sample reward: {train_dataset[0]['reward']:.4f}")




# Build reward function from collected trajectories
reward_fn = build_reward_fn(baseline_groups)

# Prepare prompts dataset (unique prompts only, GRPO generates completions)
prompt_dataset = Dataset.from_list([
    {"prompt": group[0].prompt_text}
    for group in baseline_groups
])

# Initialize GRPO trainer
trainer = GRPOTrainer(
    model=model,
    args=grpo_config,
    train_dataset=prompt_dataset,
    reward_funcs=reward_fn,
    processing_class=tokenizer,
)

print("GRPO Trainer initialized. Starting training...")
trainer.train()
print("Training complete!")


# Collect post-training trajectories with the same seeds
print("=" * 60)
print("COLLECTING POST-TRAINING TRAJECTORIES")
print("=" * 60)
post_groups = collect_trajectories(model, tokenizer, temperature=0.5)

post_rewards = [t.reward for g in post_groups for t in g]
baseline_rewards = [t.reward for g in baseline_groups for t in g]

print(f"\n{'='*60}")
print(f"RESULTS COMPARISON")
print(f"{'='*60}")
print(f"Baseline: mean={sum(baseline_rewards)/len(baseline_rewards):.4f}, "
      f"max={max(baseline_rewards):.4f}")
print(f"Trained:  mean={sum(post_rewards)/len(post_rewards):.4f}, "
      f"max={max(post_rewards):.4f}")
print(f"Improvement: {(sum(post_rewards)/len(post_rewards) - sum(baseline_rewards)/len(baseline_rewards)):.4f}")

# Per-task comparison
print(f"\nPer-task breakdown:")
for i, (bg, pg) in enumerate(zip(baseline_groups, post_groups)):
    b_mean = sum(t.reward for t in bg) / len(bg)
    p_mean = sum(t.reward for t in pg) / len(pg)
    delta = p_mean - b_mean
    arrow = "\u2191" if delta > 0 else "\u2193" if delta < 0 else "\u2192"
    print(f"  Task {i}: baseline={b_mean:.3f} \u2192 trained={p_mean:.3f} ({arrow} {abs(delta):.3f})")

import matplotlib.pyplot as plt
import numpy as np

fig, axes = plt.subplots(1, 3, figsize=(16, 5))

# 1. Reward distribution: baseline vs trained
ax1 = axes[0]
ax1.hist(baseline_rewards, bins=20, alpha=0.6, label="Baseline", color="steelblue")
ax1.hist(post_rewards, bins=20, alpha=0.6, label="After GRPO", color="coral")
ax1.set_xlabel("Episode Reward")
ax1.set_ylabel("Count")
ax1.set_title("Reward Distribution")
ax1.legend()
ax1.axvline(np.mean(baseline_rewards), color="steelblue", linestyle="--", alpha=0.8)
ax1.axvline(np.mean(post_rewards), color="coral", linestyle="--", alpha=0.8)

# 2. Per-task mean reward comparison
ax2 = axes[1]
task_ids = list(range(len(baseline_groups)))
b_means = [np.mean([t.reward for t in g]) for g in baseline_groups]
p_means = [np.mean([t.reward for t in g]) for g in post_groups]
x = np.arange(len(task_ids))
width = 0.35
ax2.bar(x - width/2, b_means, width, label="Baseline", color="steelblue", alpha=0.8)
ax2.bar(x + width/2, p_means, width, label="After GRPO", color="coral", alpha=0.8)
ax2.set_xlabel("Task ID")
ax2.set_ylabel("Mean Reward")
ax2.set_title("Per-Task Reward Improvement")
ax2.legend()
ax2.set_xticks(x)

# 3. Test pass rate improvement
ax3 = axes[2]
b_pass_rates = [np.mean([t.tests_passed / max(t.tests_total, 1) for t in g]) for g in baseline_groups]
p_pass_rates = [np.mean([t.tests_passed / max(t.tests_total, 1) for t in g]) for g in post_groups]
ax3.bar(x - width/2, b_pass_rates, width, label="Baseline", color="steelblue", alpha=0.8)
ax3.bar(x + width/2, p_pass_rates, width, label="After GRPO", color="coral", alpha=0.8)
ax3.set_xlabel("Task ID")
ax3.set_ylabel("Test Pass Rate")
ax3.set_title("Test Pass Rate Improvement")
ax3.legend()
ax3.set_xticks(x)

plt.tight_layout()
plt.savefig("rlm_forge_results.png", dpi=150, bbox_inches="tight")
plt.show()

print(f"\nOverall test pass rate:")
print(f"  Baseline: {np.mean(b_pass_rates):.1%}")
print(f"  Trained:  {np.mean(p_pass_rates):.1%}")


# Save the trained LoRA adapter
model.save_pretrained("./rlm_forge_lora_adapter")
tokenizer.save_pretrained("./rlm_forge_lora_adapter")

# Save training log
training_log = {
    "model_id": MODEL_ID,
    "num_prompts": NUM_TRAINING_PROMPTS,
    "episodes_per_prompt": NUM_EPISODES_PER_PROMPT,
    "max_steps_per_episode": MAX_STEPS_PER_EPISODE,
    "grpo_epochs": GRPO_EPOCHS,
    "baseline_mean_reward": float(np.mean(baseline_rewards)),
    "baseline_max_reward": float(max(baseline_rewards)),
    "trained_mean_reward": float(np.mean(post_rewards)),
    "trained_max_reward": float(max(post_rewards)),
    "improvement": float(np.mean(post_rewards) - np.mean(baseline_rewards)),
    "baseline_test_pass_rate": float(np.mean(b_pass_rates)),
    "trained_test_pass_rate": float(np.mean(p_pass_rates)),
}

with open("training_log.json", "w") as f:
    json.dump(training_log, f, indent=2)

print("Saved LoRA adapter to ./rlm_forge_lora_adapter")
print("Saved training log to training_log.json")
print(f"\nFinal summary:")
print(json.dumps(training_log, indent=2))