feat: baseline eval and GRPO training script

eval/baseline_eval.py

@@ -0,0 +1,61 @@
+import json
+import random
+from pathlib import Path
+
+from environment.env import ContextCorruptionEnv
+from environment.actions import ContextCorruptionAction, ActionType
+
+NUM_EPISODES = 100
+RESULTS_PATH = Path(__file__).parent / "baseline_results.json"
+
+
+def run_baseline():
+    env = ContextCorruptionEnv()
+    rewards = []
+
+    for ep in range(NUM_EPISODES):
+        obs = env.reset()
+        done = False
+
+        while not done:
+            # Randomly flag 0-4 docs
+            if random.random() < 0.4 and obs.budget_remaining > 1:
+                doc_id = random.randint(0, len(obs.documents) - 1)
+                action = ContextCorruptionAction(
+                    action_type=ActionType.flag_suspicious,
+                    doc_id=doc_id,
+                )
+            else:
+                action = ContextCorruptionAction(
+                    action_type=ActionType.submit_answer,
+                    answer="unknown",
+                    confidence=0.5,
+                )
+
+            obs = env.step(action)
+            done = obs.episode_done
+
+        rewards.append(obs.reward)
+        if (ep + 1) % 10 == 0:
+            print(f"Episode {ep + 1}/{NUM_EPISODES} | reward: {obs.reward:.4f} | avg so far: {sum(rewards)/len(rewards):.4f}")
+
+    avg = round(sum(rewards) / len(rewards), 4)
+    minimum = round(min(rewards), 4)
+    maximum = round(max(rewards), 4)
+
+    results = {
+        "num_episodes": NUM_EPISODES,
+        "avg_reward": avg,
+        "min_reward": minimum,
+        "max_reward": maximum,
+        "all_rewards": rewards,
+    }
+
+    RESULTS_PATH.write_text(json.dumps(results, indent=2))
+    print(f"\nBaseline results — avg: {avg} | min: {minimum} | max: {maximum}")
+    print(f"Saved to {RESULTS_PATH}")
+    return results
+
+
+if __name__ == "__main__":
+    run_baseline()

eval/baseline_results.json

@@ -0,0 +1,108 @@
+{
+  "num_episodes": 100,
+  "avg_reward": 0.1314,
+  "min_reward": -0.0792,
+  "max_reward": 0.3292,
+  "all_rewards": [
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.0417,
+    0.1875,
+    0.1458,
+    0.2917,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.0417,
+    -0.0792,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.2917,
+    0.1458,
+    0.0417,
+    0.1458,
+    0.2917,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.1458,
+    0.1458,
+    0.3292,
+    0.0417,
+    0.0417,
+    0.0833,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.1458,
+    0.1125,
+    0.0417,
+    0.2417,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.2417,
+    0.1458,
+    0.0417,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    -0.0667,
+    0.1458,
+    0.1458,
+    0.1125,
+    0.0417,
+    0.1125,
+    0.0417,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.2167,
+    0.1458,
+    0.1125,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.0417,
+    0.1292,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.2333,
+    0.2167,
+    0.1458,
+    0.0417,
+    0.0417,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458,
+    0.1458
+  ]
+}

training/train_grpo.py

@@ -0,0 +1,97 @@
+import os
+import wandb
+
+# Fill in after HF Space is live
+ENV_URL = os.getenv("ENV_URL", "http://localhost:8000")
+
+SYSTEM_PROMPT = """You are an epistemic agent. You will be given a question and a set of documents.
+Your goal is to identify the correct answer and flag any documents that contain corrupted or false information.
+
+You have three tools:
+- read_doc(doc_id): Spend 1 budget unit to focus on a specific document (content already visible).
+- flag_suspicious(doc_id): Mark a document as potentially corrupted.
+- submit_answer(answer, confidence): End the episode with your final answer and a confidence score (0.0–1.0).
+
+Strategy:
+- Cross-reference claims across documents. Corrupted documents will contradict the majority or your parametric knowledge.
+- When documents conflict, trust your parametric knowledge and flag the outlier.
+- Use your budget wisely — you have 12 actions total. Flag only when confident, as false positives are penalised.
+- Submit as soon as you are confident; unused budget gives a small efficiency bonus."""
+
+
+def build_prompt(obs_dict: dict) -> str:
+    docs_text = "\n\n".join(
+        f"[Doc {d['id']}] {d['title']}\n{d['content']}"
+        for d in obs_dict["documents"]
+    )
+    flagged = obs_dict.get("flagged_ids", [])
+    return (
+        f"Question: {obs_dict['question']}\n\n"
+        f"Documents:\n{docs_text}\n\n"
+        f"Flagged so far: {flagged}\n"
+        f"Budget remaining: {obs_dict['budget_remaining']}\n"
+        f"Turn: {obs_dict['turn']}\n\n"
+        "What is your next action? Respond with a JSON action object:\n"
+        '{"action_type": "submit_answer", "answer": "...", "confidence": 0.9}\n'
+        "or\n"
+        '{"action_type": "flag_suspicious", "doc_id": 2}'
+    )
+
+
+def main():
+    from unsloth import FastLanguageModel
+    from trl import GRPOTrainer, GRPOConfig
+
+    wandb.init(project="context-corruption-env", name="qwen-1.5b-grpo-run1")
+
+    # Load model with Unsloth
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name="unsloth/Qwen2-1.5B-Instruct",
+        max_seq_length=2048,
+        load_in_4bit=True,
+    )
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=16,
+        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
+        lora_dropout=0.0,
+        use_gradient_checkpointing="unsloth",
+    )
+
+    # Environment factory for GRPO — one env instance per parallel generation
+    from environment.env import ContextCorruptionEnv
+
+    def env_factory():
+        return ContextCorruptionEnv()
+
+    config = GRPOConfig(
+        num_train_epochs=3,
+        per_device_train_batch_size=4,
+        gradient_accumulation_steps=4,
+        learning_rate=5e-5,
+        max_completion_length=512,
+        num_generations=8,
+        report_to="wandb",
+        logging_steps=10,
+        save_steps=50,
+        output_dir="checkpoints/grpo-qwen-1.5b",
+    )
+
+    trainer = GRPOTrainer(
+        model=model,
+        args=config,
+        processing_class=tokenizer,
+        environment_factory=env_factory,
+    )
+
+    trainer.train()
+    wandb.finish()
+
+    # Save model
+    model.save_pretrained("checkpoints/grpo-qwen-1.5b-final")
+    tokenizer.save_pretrained("checkpoints/grpo-qwen-1.5b-final")
+    print("Training complete. Model saved to checkpoints/grpo-qwen-1.5b-final")
+
+
+if __name__ == "__main__":
+    main()