discovery plot

.gitignore

@@ -2,7 +2,6 @@
 __pycache__/
 *.py[cod]
 .pytest_cache/
-# artifacts/
 *.egg-info/
 test_results
 .claude/

Alien_Physics_Discovery_OpenEnv.ipynb

@@ -1,787 +1 @@
-{
- "nbformat": 4,
- "nbformat_minor": 5,
- "metadata": {
-  "kernelspec": {
-   "display_name": "Python 3",
-   "language": "python",
-   "name": "python3"
-  },
-  "language_info": {
-   "name": "python",
-   "version": "3.10.0"
-  },
-  "colab": {
-   "provenance": [],
-   "gpuType": "T4"
-  },
-  "accelerator": "GPU"
- },
- "cells": [
-  {
-   "cell_type": "markdown",
-   "id": "md-title",
-   "metadata": {},
-   "source": [
-    "# Alien Physics Discovery \u2014 OpenEnv Hackathon 2026\n",
-    "### Train an LLM to do unsupervised scientific discovery from raw sensor data\n",
-    "\n",
-    "**Theme:** Wild Card (Theme 5) + World Modeling (Theme 3.1) + Self-Improvement (Theme 4)\n",
-    "\n",
-    "**The Problem:** LLMs can write code and explain physics, but have never been trained to\n",
-    "*discover* physical structure from raw, ambiguous time-series data through iterative theory revision.\n",
-    "\n",
-    "**The Environment:** An agent observes noisy sensor logs from a simulated alien universe\n",
-    "governed by hidden ODEs. It must write a Python `Theory` class that models the dynamics.\n",
-    "The environment executes the code against future observations. Reward = log-likelihood \u2212 MDL penalty.\n",
-    "\n",
-    "---\n",
-    "> \u26a1 **Runtime:** T4 GPU (free Colab tier works)  \n",
-    "> \ud83d\udce6 **Stack:** Unsloth + HF TRL (GRPO) + project's sandboxed execution"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s1",
-   "metadata": {},
-   "source": [
-    "## Step 1 \u2014 Install Dependencies"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "install",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%%capture\n",
-    "# Unsloth must come first \u2014 it patches transformers in-place\n",
-    "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
-    "!pip install \"trl>=0.9.0\" peft accelerate bitsandbytes datasets wandb huggingface_hub\n",
-    "!pip install matplotlib numpy scipy\n",
-    "\n",
-    "# Install the Ontology Architect environment from the repo\n",
-    "# Replace with your actual repo URL before submitting\n",
-    "!pip install \"git+https://github.com/rohan-27p/ontology_architect.git\"\n",
-    "\n",
-    "print('All dependencies installed')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s2",
-   "metadata": {},
-   "source": [
-    "## Step 2 \u2014 Imports & Config"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "imports",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import json, os, re, warnings\n",
-    "from pathlib import Path\n",
-    "\n",
-    "import numpy as np\n",
-    "import torch\n",
-    "from datasets import Dataset\n",
-    "from unsloth import FastLanguageModel\n",
-    "from trl import GRPOConfig, GRPOTrainer\n",
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "warnings.filterwarnings('ignore')\n",
-    "\n",
-    "# \u2500\u2500 Project imports \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
-    "from ontology_architect.baselines import get_baseline\n",
-    "from ontology_architect.config import (\n",
-    "    ExperimentConfig, FeedbackConfig, RewardConfig, SandboxConfig, UniverseConfig\n",
-    ")\n",
-    "from ontology_architect.models import OntologyArchitectAction\n",
-    "from ontology_architect.sandbox import TheorySandbox\n",
-    "from ontology_architect.reward import score_theory\n",
-    "from ontology_architect.server.ontology_architect_environment import OntologyArchitectEnvironment\n",
-    "from ontology_architect.universe import ProceduralAlienUniverse\n",
-    "\n",
-    "# \u2500\u2500 Reproducibility \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
-    "SEED = 42\n",
-    "np.random.seed(SEED)\n",
-    "torch.manual_seed(SEED)\n",
-    "\n",
-    "# \u2500\u2500 Model \u2014 Qwen2.5-Coder-1.5B fits T4 (16 GB) comfortably \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
-    "MODEL_NAME  = 'unsloth/Qwen2.5-Coder-1.5B-Instruct-bnb-4bit'\n",
-    "MAX_SEQ_LEN = 2048\n",
-    "\n",
-    "# \u2500\u2500 Training knobs \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n",
-    "NUM_EPISODES   = 60    # oracle curriculum episodes\n",
-    "TRAIN_EPOCHS   = 2\n",
-    "NUM_GENS       = 2     # GRPO group size \u2014 2 keeps T4 safe\n",
-    "MAX_COMPLETION = 512   # theory code tokens\n",
-    "BATCH_SIZE     = 1\n",
-    "GRAD_ACCUM     = 8\n",
-    "LR             = 2e-5\n",
-    "USE_WANDB      = False\n",
-    "HF_TOKEN       = ''    # set to push model to Hub\n",
-    "\n",
-    "# \u2500\u2500 Environment config \u2014 matches tiny_smoke but longer for richer signal \u2500\u2500\n",
-    "ENV_CONFIG = ExperimentConfig(\n",
-    "    universe=UniverseConfig(\n",
-    "        seed=SEED,\n",
-    "        family='dual_fluid',\n",
-    "        observation_window=12,\n",
-    "        future_window=4,\n",
-    "        max_steps=4,\n",
-    "        anomaly_rate=0.06,\n",
-    "        drift_interval=6,\n",
-    "    ),\n",
-    "    sandbox=SandboxConfig(mode='subprocess', timeout_seconds=4.0),\n",
-    "    reward=RewardConfig(\n",
-    "        mdl_lambda=0.0005,\n",
-    "        prediction_sigma=0.12,\n",
-    "        anomaly_bonus=2.0,\n",
-    "        drift_bonus=1.0,\n",
-    "        execution_error_penalty=-25.0,\n",
-    "    ),\n",
-    "    feedback=FeedbackConfig(peer_review_window=4, lineage_window=4),\n",
-    ")\n",
-    "\n",
-    "print(f'Config ready')\n",
-    "print(f'  Model   : {MODEL_NAME}')\n",
-    "print(f'  Device  : {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')\n",
-    "print(f'  VRAM    : {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB' if torch.cuda.is_available() else '')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s3",
-   "metadata": {},
-   "source": [
-    "## Step 3 \u2014 Generate Oracle Curriculum\n\nThe teacher baseline rolls out `NUM_EPISODES` episodes. Each step saves `(prompt, completion, reward, future_records)`. The `future_records` are the ground-truth sensor windows \u2014 we store them so the GRPO reward function can score any new theory against the exact same future without re-running the environment."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "curriculum",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "CURRICULUM_PATH = Path('artifacts/curriculum/grpo_curriculum.jsonl')\n",
-    "CURRICULUM_PATH.parent.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "FAMILIES = ['dual_fluid', 'coupled_oscillator']\n",
-    "teacher  = get_baseline('teacher')\n",
-    "sandbox  = TheorySandbox(ENV_CONFIG.sandbox)\n",
-    "\n",
-    "examples = []\n",
-    "\n",
-    "with CURRICULUM_PATH.open('w') as fh:\n",
-    "    for ep in range(NUM_EPISODES):\n",
-    "        # Rotate through families so the model sees both universes\n",
-    "        family = FAMILIES[ep % len(FAMILIES)]\n",
-    "        cfg = ExperimentConfig(\n",
-    "            universe=UniverseConfig(\n",
-    "                seed=SEED + ep * 997,\n",
-    "                family=family,\n",
-    "                observation_window=ENV_CONFIG.universe.observation_window,\n",
-    "                future_window=ENV_CONFIG.universe.future_window,\n",
-    "                max_steps=ENV_CONFIG.universe.max_steps,\n",
-    "                anomaly_rate=ENV_CONFIG.universe.anomaly_rate,\n",
-    "                drift_interval=ENV_CONFIG.universe.drift_interval,\n",
-    "            ),\n",
-    "            sandbox=ENV_CONFIG.sandbox,\n",
-    "            reward=ENV_CONFIG.reward,\n",
-    "            feedback=ENV_CONFIG.feedback,\n",
-    "        )\n",
-    "        env = OntologyArchitectEnvironment(cfg)\n",
-    "        obs = env.reset()\n",
-    "\n",
-    "        for step in range(cfg.universe.max_steps):\n",
-    "            # Capture the future window that this step will be scored against\n",
-    "            cursor = env._cursor\n",
-    "            future_records = env._records[cursor : cursor + cfg.universe.future_window]\n",
-    "            sensor_names   = list(env._universe.sensor_names)\n",
-    "\n",
-    "            row = {\n",
-    "                'prompt'        : obs.text,\n",
-    "                'completion'    : teacher,\n",
-    "                'episode'       : ep,\n",
-    "                'step'          : step,\n",
-    "                'family'        : family,\n",
-    "                'sensor_names'  : sensor_names,\n",
-    "                # Serialise future so reward_fn can use it without re-running env\n",
-    "                'future_json'   : json.dumps([\n",
-    "                    {'t': r.t, 'sensors': r.sensors, 'anomaly': r.anomaly, 'drift': r.drift}\n",
-    "                    for r in future_records\n",
-    "                ]),\n",
-    "            }\n",
-    "\n",
-    "            obs = env.step(OntologyArchitectAction(\n",
-    "                theory_module=teacher,\n",
-    "                revision_note=f'curriculum ep={ep} step={step}',\n",
-    "            ))\n",
-    "            row['reward'] = obs.reward\n",
-    "            row['metrics'] = json.dumps(obs.metadata.get('last_metrics', {}))\n",
-    "\n",
-    "            fh.write(json.dumps(row) + '\\n')\n",
-    "            examples.append(row)\n",
-    "\n",
-    "            if obs.done:\n",
-    "                break\n",
-    "\n",
-    "print(f'Generated {len(examples)} curriculum examples across {NUM_EPISODES} episodes')\n",
-    "print(f'  Families: {set(e[\"family\"] for e in examples)}')\n",
-    "print(f'  Reward range: [{min(e[\"reward\"] for e in examples):.3f}, {max(e[\"reward\"] for e in examples):.3f}]')\n",
-    "teacher_mean = float(np.mean([e['reward'] for e in examples]))\n",
-    "print(f'  Teacher mean reward: {teacher_mean:.4f}  (this is the bar GRPO must beat)')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s4",
-   "metadata": {},
-   "source": [
-    "## Step 4 \u2014 Build HuggingFace Dataset"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "dataset",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "SYSTEM_PROMPT = \"\"\"You are a scientific discovery agent. An alien universe is governed by hidden physical laws.\n",
-    "You observe noisy sensor logs and must write a compact Python Theory class that predicts future readings.\n",
-    "\n",
-    "Your class MUST implement:\n",
-    "  fit(history)            \u2014 history is a list of dicts: [{'sensors': {'sigma': float, ...}}, ...]\n",
-    "  predict(window)         \u2014 window is {'steps': int, 'history': [...], 'sensor_names': [...]}\n",
-    "                            return a list of {'sensors': {...}, 'anomaly_prob': float}\n",
-    "  log_prob(observations)  \u2014 return a float\n",
-    "\n",
-    "Optional but rewarded:\n",
-    "  detect_drift(history) -> bool\n",
-    "  describe()            -> str\n",
-    "\n",
-    "Allowed imports: numpy, scipy, math, statistics, collections, typing, random, functools, itertools\n",
-    "DO NOT import anything else. DO NOT use file I/O, network, or system calls.\n",
-    "Shorter correct theories score higher than longer ones (MDL principle).\n",
-    "Reply with ONLY the Python code \u2014 no markdown fences, no explanation.\"\"\"\n",
-    "\n",
-    "\n",
-    "def make_chat_prompt(obs_text: str, tokenizer) -> str:\n",
-    "    messages = [\n",
-    "        {'role': 'system', 'content': SYSTEM_PROMPT},\n",
-    "        {'role': 'user',   'content': obs_text},\n",
-    "    ]\n",
-    "    return tokenizer.apply_chat_template(\n",
-    "        messages, tokenize=False, add_generation_prompt=True\n",
-    "    )\n",
-    "\n",
-    "\n",
-    "# We load raw rows first; prompts are formatted after the tokenizer is loaded\n",
-    "raw_rows = [json.loads(l) for l in CURRICULUM_PATH.read_text().splitlines() if l.strip()]\n",
-    "print(f'Loaded {len(raw_rows)} rows from curriculum')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s5",
-   "metadata": {},
-   "source": [
-    "## Step 5 \u2014 Load Model with Unsloth"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "load-model",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "model, tokenizer = FastLanguageModel.from_pretrained(\n",
-    "    model_name     = MODEL_NAME,\n",
-    "    max_seq_length = MAX_SEQ_LEN,\n",
-    "    dtype          = None,       # auto: bf16 on Ampere+, fp16 otherwise\n",
-    "    load_in_4bit   = True,\n",
-    ")\n",
-    "\n",
-    "model = FastLanguageModel.get_peft_model(\n",
-    "    model,\n",
-    "    r                          = 16,\n",
-    "    target_modules             = ['q_proj', 'k_proj', 'v_proj', 'o_proj',\n",
-    "                                   'gate_proj', 'up_proj', 'down_proj'],\n",
-    "    lora_alpha                 = 16,\n",
-    "    lora_dropout               = 0.05,\n",
-    "    bias                       = 'none',\n",
-    "    use_gradient_checkpointing = 'unsloth',\n",
-    "    random_state               = SEED,\n",
-    ")\n",
-    "\n",
-    "total     = sum(p.numel() for p in model.parameters())\n",
-    "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
-    "print(f'Model loaded: {total:,} total params, {trainable:,} trainable ({100*trainable/total:.2f}%)')\n",
-    "\n",
-    "# Ensure pad token is set (needed for left-padding in GRPO)\n",
-    "if tokenizer.pad_token is None:\n",
-    "    tokenizer.pad_token = tokenizer.eos_token\n",
-    "tokenizer.padding_side = 'left'"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s6",
-   "metadata": {},
-   "source": [
-    "## Step 6 \u2014 Format Dataset with Chat Template"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "format-dataset",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "dataset_rows = []\n",
-    "for row in raw_rows:\n",
-    "    dataset_rows.append({\n",
-    "        'prompt'      : make_chat_prompt(row['prompt'], tokenizer),\n",
-    "        'future_json' : row['future_json'],\n",
-    "        'history_json' : row.get('history_json', '[]'),\n",
-    "        'sensor_names': json.dumps(row['sensor_names']),\n",
-    "    })\n",
-    "\n",
-    "dataset = Dataset.from_list(dataset_rows)\n",
-    "print(f'Dataset: {len(dataset)} rows')\n",
-    "print(f'Columns: {dataset.column_names}')\n",
-    "print(f'Sample prompt length: {len(dataset[0][\"prompt\"])} chars')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s7",
-   "metadata": {},
-   "source": [
-    "## Step 7 \u2014 GRPO Reward Function\n\n",
-    "Each completion is scored by running it through the project's real sandbox and scorer against\n",
-    "the **pre-captured future window** stored in the dataset row. This avoids the stale-trajectory\n",
-    "bug in the original notebook where the reward used a separately-seeded universe."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "reward-fn",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from ontology_architect.universe import SensorRecord\n",
-    "\n",
-    "_sandbox = TheorySandbox(ENV_CONFIG.sandbox)\n",
-    "_reward_cfg = ENV_CONFIG.reward\n",
-    "\n",
-    "\n",
-    "def _deserialise_future(future_json: str) -> list[SensorRecord]:\n",
-    "    return [\n",
-    "        SensorRecord(t=r['t'], sensors=r['sensors'], anomaly=r['anomaly'], drift=r['drift'])\n",
-    "        for r in json.loads(future_json)\n",
-    "    ]\n",
-    "\n",
-    "\n",
-    "def physics_reward_fn(\n",
-    "    prompts     : list[str],\n",
-    "    completions : list[str],\n",
-    "    future_json : list[str],\n",
-    "    history_json: list[str],\n",
-    "    sensor_names: list[str],\n",
-    "    **kwargs,\n",
-    ") -> list[float]:\n",
-    "    \"\"\"\n",
-    "    GRPO reward function.  Columns `future_json` and `sensor_names` are\n",
-    "    passed automatically by GRPOTrainer from the dataset.\n",
-    "    \"\"\"\n",
-    "    rewards = []\n",
-    "    for theory_code, fj, hj_str, sn_json in zip(completions, future_json, history_json, sensor_names):\n",
-    "        try:\n",
-    "            # Strip markdown fences the model sometimes emits\n",
-    "            theory_code = re.sub(r'^```(?:python)?\\s*', '', theory_code.strip())\n",
-    "            theory_code = re.sub(r'```\\s*$', '', theory_code)\n",
-    "\n",
-    "            future       = _deserialise_future(fj)\n",
-    "            names        = tuple(json.loads(sn_json))\n",
-    "            history_data = json.loads(hj_str) if hj_str else []\n",
-    "            sandbox_res  = _sandbox.execute(theory_code, history_data, future, names)\n",
-    "            breakdown    = score_theory(sandbox_res, future, theory_code, _reward_cfg)\n",
-    "            reward       = float(np.clip(breakdown.reward, -30.0, 5.0))\n",
-    "        except Exception:\n",
-    "            reward = float(_reward_cfg.execution_error_penalty)\n",
-    "        rewards.append(reward)\n",
-    "    return rewards\n",
-    "\n",
-    "\n",
-    "# Sanity check \u2014 teacher should score positively\n",
-    "row0 = dataset_rows[0]\n",
-    "test_rewards = physics_reward_fn(\n",
-    "    prompts      = [row0['prompt']],\n",
-    "    completions  = [get_baseline('teacher')],\n",
-    "    future_json  = [row0['future_json']],\n",
-    "    sensor_names = [row0['sensor_names']],\n",
-    ")\n",
-    "print(f'Sanity check \u2014 teacher reward on row 0: {test_rewards[0]:.4f}')\n",
-    "print(f'  (Expected roughly > {teacher_mean:.2f} for a good theory)')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s8",
-   "metadata": {},
-   "source": [
-    "## Step 8 \u2014 GRPO Training"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "train",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "if USE_WANDB:\n",
-    "    import wandb\n",
-    "    wandb.init(project='alien-physics-discovery', name='qwen2.5-coder-1.5b-grpo')\n",
-    "\n",
-    "grpo_config = GRPOConfig(\n",
-    "    output_dir                  = 'artifacts/checkpoints/grpo',\n",
-    "    num_train_epochs            = TRAIN_EPOCHS,\n",
-    "    per_device_train_batch_size = BATCH_SIZE,\n",
-    "    gradient_accumulation_steps = GRAD_ACCUM,\n",
-    "    learning_rate               = LR,\n",
-    "    warmup_ratio                = 0.05,\n",
-    "    lr_scheduler_type           = 'cosine',\n",
-    "    num_generations             = NUM_GENS,\n",
-    "    max_completion_length       = MAX_COMPLETION,   # correct TRL >=0.9 field\n",
-    "    logging_steps               = 5,\n",
-    "    save_steps                  = 50,\n",
-    "    save_total_limit            = 2,\n",
-    "    fp16                        = not torch.cuda.is_bf16_supported(),\n",
-    "    bf16                        = torch.cuda.is_bf16_supported(),\n",
-    "    report_to                   = 'wandb' if USE_WANDB else 'none',\n",
-    "    beta                        = 0.001,\n",
-    "    max_grad_norm               = 0.3,\n",
-    ")\n",
-    "\n",
-    "trainer = GRPOTrainer(\n",
-    "    model            = model,\n",
-    "    processing_class = tokenizer,\n",
-    "    args             = grpo_config,          # 'args', not 'config'\n",
-    "    train_dataset    = dataset,\n",
-    "    reward_funcs     = [physics_reward_fn],\n",
-    ")\n",
-    "\n",
-    "print('Starting GRPO training')\n",
-    "print(f'  Model       : {MODEL_NAME}')\n",
-    "print(f'  Samples     : {len(dataset)}')\n",
-    "print(f'  Epochs      : {TRAIN_EPOCHS}')\n",
-    "print(f'  Group size  : {NUM_GENS}')\n",
-    "print(f'  Batch\u00d7accum : {BATCH_SIZE}\u00d7{GRAD_ACCUM} = {BATCH_SIZE*GRAD_ACCUM} effective')\n",
-    "\n",
-    "train_result = trainer.train()\n",
-    "print(f'\\nTraining complete \u2014 final loss: {train_result.training_loss:.4f}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s9",
-   "metadata": {},
-   "source": [
-    "## Step 9 \u2014 Plot Training Curves"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "plot",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "logs    = trainer.state.log_history\n",
-    "steps   = [l['step']           for l in logs if 'loss'   in l]\n",
-    "losses  = [l['loss']           for l in logs if 'loss'   in l]\n",
-    "r_pairs = [(l['step'], l['reward']) for l in logs if 'reward' in l]\n",
-    "\n",
-    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\n",
-    "fig.suptitle('Alien Physics Discovery \u2014 GRPO Training (Qwen2.5-Coder-1.5B)', fontweight='bold')\n",
-    "\n",
-    "ax1.plot(steps, losses, color='#4f86c6', linewidth=2)\n",
-    "ax1.set_xlabel('Training Step')\n",
-    "ax1.set_ylabel('Loss')\n",
-    "ax1.set_title('GRPO Loss')\n",
-    "ax1.grid(alpha=0.3)\n",
-    "\n",
-    "if r_pairs:\n",
-    "    rs, rv = zip(*r_pairs)\n",
-    "    ax2.plot(rs, rv, color='#e67e22', linewidth=2, label='Mean group reward')\n",
-    "    ax2.axhline(teacher_mean, color='gray', linestyle='--', alpha=0.7,\n",
-    "                label=f'Teacher baseline ({teacher_mean:.2f})')\n",
-    "    ax2.set_xlabel('Training Step')\n",
-    "    ax2.set_ylabel('MDL Reward')\n",
-    "    ax2.set_title('Reward vs Teacher Baseline')\n",
-    "    ax2.legend()\n",
-    "    ax2.grid(alpha=0.3)\n",
-    "else:\n",
-    "    ax2.text(0.5, 0.5, 'Enable USE_WANDB for live\\nreward curves, or check\\nlog_history manually',\n",
-    "             ha='center', va='center', transform=ax2.transAxes, color='gray')\n",
-    "\n",
-    "plt.tight_layout()\n",
-    "plt.savefig('artifacts/training_curves.png', dpi=150, bbox_inches='tight')\n",
-    "plt.show()\n",
-    "print('Saved artifacts/training_curves.png')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s10",
-   "metadata": {},
-   "source": [
-    "## Step 10 \u2014 Before / After Evaluation"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "eval",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "FastLanguageModel.for_inference(model)\n",
-    "\n",
-    "def generate_theory(obs_text: str, max_new: int = 480) -> str:\n",
-    "    prompt  = make_chat_prompt(obs_text, tokenizer)\n",
-    "    inputs  = tokenizer(prompt, return_tensors='pt',\n",
-    "                        truncation=True, max_length=MAX_SEQ_LEN - max_new).to(model.device)\n",
-    "    with torch.no_grad():\n",
-    "        out = model.generate(\n",
-    "            **inputs,\n",
-    "            max_new_tokens  = max_new,\n",
-    "            temperature     = 0.6,\n",
-    "            do_sample       = True,\n",
-    "            pad_token_id    = tokenizer.eos_token_id,\n",
-    "        )\n",
-    "    theory = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n",
-    "    theory = re.sub(r'^```(?:python)?\\s*', '', theory.strip())\n",
-    "    theory = re.sub(r'```\\s*$', '', theory)\n",
-    "    return theory\n",
-    "\n",
-    "\n",
-    "def eval_full_episode(family: str, seed: int) -> dict:\n",
-    "    from dataclasses import replace as dc_replace\n",
-    "    cfg = ExperimentConfig(\n",
-    "        universe=dc_replace(ENV_CONFIG.universe, family=family, seed=seed),\n",
-    "        sandbox=ENV_CONFIG.sandbox,\n",
-    "        reward=ENV_CONFIG.reward,\n",
-    "        feedback=ENV_CONFIG.feedback,\n",
-    "    )\n",
-    "    env  = OntologyArchitectEnvironment(cfg)\n",
-    "    obs  = env.reset()\n",
-    "    rewards, done = [], False\n",
-    "    final_theory  = ''\n",
-    "    while not done:\n",
-    "        theory = generate_theory(obs.text)\n",
-    "        obs    = env.step(OntologyArchitectAction(\n",
-    "            theory_module=theory, revision_note='eval'\n",
-    "        ))\n",
-    "        rewards.append(obs.reward)\n",
-    "        final_theory = theory\n",
-    "        done = obs.done\n",
-    "    return {'family': family, 'rewards': rewards,\n",
-    "            'mean_reward': float(np.mean(rewards)), 'final_theory': final_theory}\n",
-    "\n",
-    "\n",
-    "# Teacher baseline (untrained)\n",
-    "def eval_teacher_episode(family: str, seed: int) -> dict:\n",
-    "    from dataclasses import replace as dc_replace\n",
-    "    cfg = ExperimentConfig(\n",
-    "        universe=dc_replace(ENV_CONFIG.universe, family=family, seed=seed),\n",
-    "        sandbox=ENV_CONFIG.sandbox,\n",
-    "        reward=ENV_CONFIG.reward,\n",
-    "        feedback=ENV_CONFIG.feedback,\n",
-    "    )\n",
-    "    env     = OntologyArchitectEnvironment(cfg)\n",
-    "    obs     = env.reset()\n",
-    "    teacher = get_baseline('teacher')\n",
-    "    rewards, done = [], False\n",
-    "    while not done:\n",
-    "        obs  = env.step(OntologyArchitectAction(theory_module=teacher, revision_note='baseline'))\n",
-    "        rewards.append(obs.reward)\n",
-    "        done = obs.done\n",
-    "    return {'family': family, 'rewards': rewards, 'mean_reward': float(np.mean(rewards))}\n",
-    "\n",
-    "\n",
-    "print('Running eval episodes (trained model vs teacher baseline)...')\n",
-    "eval_seeds  = [9001, 9002]\n",
-    "eval_families = ['thermal_split', 'coupled_oscillator']\n",
-    "\n",
-    "trained_results  = [eval_full_episode(f, s)    for f in eval_families for s in eval_seeds]\n",
-    "baseline_results = [eval_teacher_episode(f, s) for f in eval_families for s in eval_seeds]\n",
-    "\n",
-    "trained_mean  = float(np.mean([r['mean_reward'] for r in trained_results]))\n",
-    "baseline_mean = float(np.mean([r['mean_reward'] for r in baseline_results]))\n",
-    "\n",
-    "print(f'\\nResults:')\n",
-    "print(f'  Trained model mean reward  : {trained_mean:.4f}')\n",
-    "print(f'  Teacher baseline mean reward: {baseline_mean:.4f}')\n",
-    "print(f'  Delta                       : {trained_mean - baseline_mean:+.4f}')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s11",
-   "metadata": {},
-   "source": [
-    "## Step 11 \u2014 Per-Turn Improvement Plot"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "per-turn-plot",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# Aggregate rewards by turn index across all eval episodes\n",
-    "max_turns = max(len(r['rewards']) for r in trained_results)\n",
-    "\n",
-    "def pad_rewards(results, length):\n",
-    "    padded = [r['rewards'] + [r['rewards'][-1]] * (length - len(r['rewards'])) for r in results]\n",
-    "    return np.array(padded)\n",
-    "\n",
-    "trained_mat  = pad_rewards(trained_results,  max_turns)\n",
-    "baseline_mat = pad_rewards(baseline_results, max_turns)\n",
-    "turns = list(range(1, max_turns + 1))\n",
-    "\n",
-    "fig, ax = plt.subplots(figsize=(8, 4))\n",
-    "ax.plot(turns, trained_mat.mean(0),  color='#3498db', marker='o', lw=2, label='Trained model')\n",
-    "ax.fill_between(turns,\n",
-    "                trained_mat.mean(0) - trained_mat.std(0),\n",
-    "                trained_mat.mean(0) + trained_mat.std(0),\n",
-    "                alpha=0.15, color='#3498db')\n",
-    "ax.plot(turns, baseline_mat.mean(0), color='gray',    marker='s', lw=2, linestyle='--', label='Teacher baseline')\n",
-    "ax.set_xlabel('Refinement Turn')\n",
-    "ax.set_ylabel('MDL Reward')\n",
-    "ax.set_title('Per-Turn Theory Improvement')\n",
-    "ax.set_xticks(turns)\n",
-    "ax.legend()\n",
-    "ax.grid(alpha=0.3)\n",
-    "plt.tight_layout()\n",
-    "plt.savefig('artifacts/per_turn_improvement.png', dpi=150, bbox_inches='tight')\n",
-    "plt.show()\n",
-    "print('Saved artifacts/per_turn_improvement.png')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s12",
-   "metadata": {},
-   "source": [
-    "## Step 12 \u2014 Show a Discovered Theory"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "show-theory",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "best = max(trained_results, key=lambda r: r['mean_reward'])\n",
-    "print(f\"Best result \u2014 family={best['family']}  mean_reward={best['mean_reward']:.4f}\")\n",
-    "print('=' * 70)\n",
-    "theory = best['final_theory']\n",
-    "print(theory[:2000])\n",
-    "if len(theory) > 2000:\n",
-    "    print(f'\\n... [{len(theory)-2000} more chars] ...')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s13",
-   "metadata": {},
-   "source": [
-    "## Step 13 \u2014 Save Model & Push to Hub"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "save",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "Path('artifacts').mkdir(exist_ok=True)\n",
-    "model.save_pretrained('artifacts/model')\n",
-    "tokenizer.save_pretrained('artifacts/model')\n",
-    "print('Model saved to artifacts/model')\n",
-    "\n",
-    "if HF_TOKEN:\n",
-    "    from huggingface_hub import login\n",
-    "    login(token=HF_TOKEN)\n",
-    "    model.push_to_hub('lostdecimal27/alien-physics-discovery-qwen2.5-1.5b')\n",
-    "    tokenizer.push_to_hub('lostdecimal27/alien-physics-discovery-qwen2.5-1.5b')\n",
-    "    print('Pushed to HuggingFace Hub')\n",
-    "else:\n",
-    "    print('Set HF_TOKEN above to push to Hub')"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "md-s14",
-   "metadata": {},
-   "source": [
-    "## Step 14 \u2014 Submission Checklist"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "checklist",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "checks = {\n",
-    "    'artifacts/training_curves.png'    : Path('artifacts/training_curves.png').exists(),\n",
-    "    'artifacts/per_turn_improvement.png': Path('artifacts/per_turn_improvement.png').exists(),\n",
-    "    'artifacts/model (local save)'     : Path('artifacts/model').exists(),\n",
-    "    'openenv.yaml in repo'             : Path('openenv.yaml').exists(),\n",
-    "    'README.md in repo'                : Path('README.md').exists(),\n",
-    "    'Colab notebook (this file)'       : True,\n",
-    "}\n",
-    "\n",
-    "print('OpenEnv Hackathon \u2014 Submission Checklist')\n",
-    "print('=' * 50)\n",
-    "for item, ok in checks.items():\n",
-    "    print(f\"  {'\u2705' if ok else '\u274c'}  {item}\")\n",
-    "\n",
-    "print('\\nRemaining manual steps:')\n",
-    "for todo in [\n",
-    "    '\u25a1  Push model to HuggingFace Hub (set HF_TOKEN)',\n",
-    "    '\u25a1  Deploy environment to HuggingFace Space',\n",
-    "    '\u25a1  Record < 2 min YouTube demo or write HF blog post',\n",
-    "    '\u25a1  Update lostdecimal27 placeholders above',\n",
-    "    '\u25a1  Update README with video/blog URL and HF Space URL',\n",
-    "    '\u25a1  Submit HF Space URL to hackathon portal',\n",
-    "]:\n",
-    "    print(f'  {todo}')\n",
-    "\n",
-    "print(f'\\nFinal score summary:')\n",
-    "print(f'  Teacher baseline mean reward : {baseline_mean:.4f}')\n",
-    "print(f'  Trained model mean reward    : {trained_mean:.4f}')\n",
-    "print(f'  Improvement                  : {trained_mean - baseline_mean:+.4f}')"
-   ]
-  }
- ]
-}
+{"nbformat": 4, "nbformat_minor": 5, "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.10.0"}, "colab": {"provenance": [], "gpuType": "T4"}, "accelerator": "GPU"}, "cells": [{"cell_type": "markdown", "id": "md-title", "metadata": {}, "source": ["# Alien Physics Discovery \u2014 OpenEnv Hackathon 2026\n", "### Train an LLM to do unsupervised scientific discovery from raw sensor data\n", "\n", "**Theme:** Wild Card (Theme 5) + World Modeling (Theme 3.1) + Self-Improvement (Theme 4)\n", "\n", "**The Problem:** LLMs can write code and explain physics, but have never been trained to\n", "*discover* physical structure from raw, ambiguous time-series data through iterative theory revision.\n", "\n", "**The Environment:** An agent observes noisy sensor logs from a simulated alien universe\n", "governed by hidden ODEs. It must write a Python `Theory` class that models the dynamics.\n", "The environment executes the code against future observations. Reward = log-likelihood \u2212 MDL penalty.\n", "\n", "---\n", "> \u26a1 **Runtime:** T4 GPU (free Colab tier works)  \n", "> \ud83d\udce6 **Stack:** Unsloth + HF TRL (GRPO) + project's sandboxed execution"]}, {"cell_type": "markdown", "id": "md-s1", "metadata": {}, "source": ["## Step 1 \u2014 Install Dependencies"]}, {"cell_type": "code", "execution_count": null, "id": "install", "metadata": {}, "outputs": [], "source": ["%%capture\n", "# Unsloth must come first \u2014 it patches transformers in-place\n", "!pip install \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n", "!pip install \"trl>=0.9.0\" peft accelerate bitsandbytes datasets wandb huggingface_hub\n", "!pip install matplotlib numpy scipy\n", "\n", "# Install the Ontology Architect environment from the repo\n", "# Replace with your actual repo URL before submitting\n", "!pip install \"git+https://github.com/rohan-27p/ontology_architect.git\"\n", "\n", "print('All dependencies installed')"]}, {"cell_type": "markdown", "id": "md-s2", "metadata": {}, "source": ["## Step 2 \u2014 Imports & Config"]}, {"cell_type": "code", "execution_count": null, "id": "imports", "metadata": {}, "outputs": [], "source": ["import json, os, re, warnings\n", "from pathlib import Path\n", "\n", "import numpy as np\n", "import torch\n", "from datasets import Dataset\n", "from unsloth import FastLanguageModel\n", "from trl import GRPOConfig, GRPOTrainer\n", "import matplotlib.pyplot as plt\n", "\n", "warnings.filterwarnings('ignore')\n", "\n", "# \u2500\u2500 Project imports \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", "from ontology_architect.baselines import get_baseline\n", "from ontology_architect.config import (\n", "    ExperimentConfig, FeedbackConfig, RewardConfig, SandboxConfig, UniverseConfig\n", ")\n", "from ontology_architect.models import OntologyArchitectAction\n", "from ontology_architect.sandbox import TheorySandbox\n", "from ontology_architect.reward import score_theory\n", "from ontology_architect.server.ontology_architect_environment import OntologyArchitectEnvironment\n", "from ontology_architect.universe import ProceduralAlienUniverse\n", "\n", "# \u2500\u2500 Reproducibility \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", "SEED = 42\n", "np.random.seed(SEED)\n", "torch.manual_seed(SEED)\n", "\n", "# \u2500\u2500 Model \u2014 Qwen2.5-Coder-1.5B fits T4 (16 GB) comfortably \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", "MODEL_NAME  = 'unsloth/Qwen2.5-Coder-7B-Instruct-bnb-4bit'\n", "MAX_SEQ_LEN = 2048\n", "\n", "# \u2500\u2500 Training knobs \u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\u2500\n", "NUM_EPISODES   = 60    # oracle curriculum episodes\n", "TRAIN_EPOCHS   = 2\n", "NUM_GENS       = 4     # GRPO group size \u2014 2 keeps T4 safe\n", "MAX_COMPLETION = 768   # theory code tokens\n", "BATCH_SIZE     = 1\n", "GRAD_ACCUM     = 8\n", "LR             = 2e-5\n", "USE_WANDB      = False\n", "HF_TOKEN       = ''    # set to push model to Hub\n", "\n", "# \u2500\u2500 Environment config \u2014 matches tiny_smoke but longer for richer signal \u2500\u2500\n", "ENV_CONFIG = ExperimentConfig(\n", "    universe=UniverseConfig(\n", "        seed=SEED,\n", "        family='dual_fluid',\n", "        observation_window=12,\n", "        future_window=4,\n", "        max_steps=4,\n", "        anomaly_rate=0.06,\n", "        drift_interval=6,\n", "    ),\n", "    sandbox=SandboxConfig(mode='subprocess', timeout_seconds=4.0),\n", "    reward=RewardConfig(\n", "        mdl_lambda=0.0005,\n", "        prediction_sigma=0.12,\n", "        anomaly_bonus=2.0,\n", "        drift_bonus=1.0,\n", "        execution_error_penalty=-25.0,\n", "    ),\n", "    feedback=FeedbackConfig(peer_review_window=4, lineage_window=4),\n", ")\n", "\n", "print(f'Config ready')\n", "print(f'  Model   : {MODEL_NAME}')\n", "print(f'  Device  : {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"CPU\"}')\n", "print(f'  VRAM    : {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB' if torch.cuda.is_available() else '')"]}, {"cell_type": "markdown", "id": "md-s3", "metadata": {}, "source": ["## Step 3 \u2014 Generate Oracle Curriculum\n\nThe teacher baseline rolls out `NUM_EPISODES` episodes. Each step saves `(prompt, completion, reward, future_records)`. The `future_records` are the ground-truth sensor windows \u2014 we store them so the GRPO reward function can score any new theory against the exact same future without re-running the environment."]}, {"cell_type": "code", "execution_count": null, "id": "curriculum", "metadata": {}, "outputs": [], "source": ["CURRICULUM_PATH = Path('artifacts/curriculum/grpo_curriculum.jsonl')\n", "CURRICULUM_PATH.parent.mkdir(parents=True, exist_ok=True)\n", "\n", "FAMILIES = ['dual_fluid', 'coupled_oscillator']\n", "teacher  = get_baseline('teacher')\n", "sandbox  = TheorySandbox(ENV_CONFIG.sandbox)\n", "\n", "examples = []\n", "\n", "with CURRICULUM_PATH.open('w') as fh:\n", "    for ep in range(NUM_EPISODES):\n", "        # Rotate through families so the model sees both universes\n", "        family = FAMILIES[ep % len(FAMILIES)]\n", "        cfg = ExperimentConfig(\n", "            universe=UniverseConfig(\n", "                seed=SEED + ep * 997,\n", "                family=family,\n", "                observation_window=ENV_CONFIG.universe.observation_window,\n", "                future_window=ENV_CONFIG.universe.future_window,\n", "                max_steps=ENV_CONFIG.universe.max_steps,\n", "                anomaly_rate=ENV_CONFIG.universe.anomaly_rate,\n", "                drift_interval=ENV_CONFIG.universe.drift_interval,\n", "            ),\n", "            sandbox=ENV_CONFIG.sandbox,\n", "            reward=ENV_CONFIG.reward,\n", "            feedback=ENV_CONFIG.feedback,\n", "        )\n", "        env = OntologyArchitectEnvironment(cfg)\n", "        obs = env.reset()\n", "\n", "        for step in range(cfg.universe.max_steps):\n", "            # Capture the future window that this step will be scored against\n", "            cursor = env._cursor\n", "            future_records = env._records[cursor : cursor + cfg.universe.future_window]\n", "            sensor_names   = list(env._universe.sensor_names)\n", "\n", "            row = {\n", "                'prompt'        : obs.text,\n", "                'completion'    : teacher,\n", "                'episode'       : ep,\n", "                'step'          : step,\n", "                'family'        : family,\n", "                'sensor_names'  : sensor_names,\n", "                # Serialise future so reward_fn can use it without re-running env\n", "                'future_json'   : json.dumps([\n", "                    {'t': r.t, 'sensors': r.sensors, 'anomaly': r.anomaly, 'drift': r.drift}\n", "                    for r in future_records\n", "                ]),\n", "            }\n", "\n", "            obs = env.step(OntologyArchitectAction(\n", "                theory_module=teacher,\n", "                revision_note=f'curriculum ep={ep} step={step}',\n", "            ))\n", "            row['reward'] = obs.reward\n", "            row['metrics'] = json.dumps(obs.metadata.get('last_metrics', {}))\n", "\n", "            fh.write(json.dumps(row) + '\\n')\n", "            examples.append(row)\n", "\n", "            if obs.done:\n", "                break\n", "\n", "print(f'Generated {len(examples)} curriculum examples across {NUM_EPISODES} episodes')\n", "print(f'  Families: {set(e[\"family\"] for e in examples)}')\n", "print(f'  Reward range: [{min(e[\"reward\"] for e in examples):.3f}, {max(e[\"reward\"] for e in examples):.3f}]')\n", "teacher_mean = float(np.mean([e['reward'] for e in examples]))\n", "print(f'  Teacher mean reward: {teacher_mean:.4f}  (this is the bar GRPO must beat)')"]}, {"cell_type": "markdown", "id": "md-s4", "metadata": {}, "source": ["## Step 4 \u2014 Build HuggingFace Dataset"]}, {"cell_type": "code", "execution_count": null, "id": "dataset", "metadata": {}, "outputs": [], "source": ["SYSTEM_PROMPT = \"\"\"You are a scientific discovery agent. An alien universe is governed by hidden physical laws.\n", "You observe noisy sensor logs and must write a compact Python Theory class that predicts future readings.\n", "\n", "Your class MUST implement:\n", "  fit(history)            \u2014 history is a list of dicts: [{'sensors': {'sigma': float, ...}}, ...]\n", "  predict(window)         \u2014 window is {'steps': int, 'history': [...], 'sensor_names': [...]}\n", "                            return a list of {'sensors': {...}, 'anomaly_prob': float}\n", "  log_prob(observations)  \u2014 return a float\n", "\n", "Optional but rewarded:\n", "  detect_drift(history) -> bool\n", "  describe()            -> str\n", "\n", "Allowed imports: numpy, scipy, math, statistics, collections, typing, random, functools, itertools\n", "DO NOT import anything else. DO NOT use file I/O, network, or system calls.\n", "Shorter correct theories score higher than longer ones (MDL principle).\n", "Reply with ONLY the Python code \u2014 no markdown fences, no explanation.\"\"\"\n", "\n", "\n", "def make_chat_prompt(obs_text: str, tokenizer) -> str:\n", "    messages = [\n", "        {'role': 'system', 'content': SYSTEM_PROMPT},\n", "        {'role': 'user',   'content': obs_text},\n", "    ]\n", "    return tokenizer.apply_chat_template(\n", "        messages, tokenize=False, add_generation_prompt=True\n", "    )\n", "\n", "\n", "# We load raw rows first; prompts are formatted after the tokenizer is loaded\n", "raw_rows = [json.loads(l) for l in CURRICULUM_PATH.read_text().splitlines() if l.strip()]\n", "print(f'Loaded {len(raw_rows)} rows from curriculum')"]}, {"cell_type": "markdown", "id": "md-s5", "metadata": {}, "source": ["## Step 5 \u2014 Load Model with Unsloth"]}, {"cell_type": "code", "execution_count": null, "id": "load-model", "metadata": {}, "outputs": [], "source": ["model, tokenizer = FastLanguageModel.from_pretrained(\n", "    model_name     = MODEL_NAME,\n", "    max_seq_length = MAX_SEQ_LEN,\n", "    dtype          = None,       # auto: bf16 on Ampere+, fp16 otherwise\n", "    load_in_4bit   = True,\n", ")\n", "\n", "model = FastLanguageModel.get_peft_model(\n", "    model,\n", "    r                          = 16,\n", "    target_modules             = ['q_proj', 'k_proj', 'v_proj', 'o_proj',\n", "                                   'gate_proj', 'up_proj', 'down_proj'],\n", "    lora_alpha                 = 16,\n", "    lora_dropout               = 0.05,\n", "    bias                       = 'none',\n", "    use_gradient_checkpointing = 'unsloth',\n", "    random_state               = SEED,\n", ")\n", "\n", "total     = sum(p.numel() for p in model.parameters())\n", "trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n", "print(f'Model loaded: {total:,} total params, {trainable:,} trainable ({100*trainable/total:.2f}%)')\n", "\n", "# Ensure pad token is set (needed for left-padding in GRPO)\n", "if tokenizer.pad_token is None:\n", "    tokenizer.pad_token = tokenizer.eos_token\n", "tokenizer.padding_side = 'left'"]}, {"cell_type": "markdown", "id": "md-s6", "metadata": {}, "source": ["## Step 6 \u2014 Format Dataset with Chat Template"]}, {"cell_type": "code", "execution_count": null, "id": "format-dataset", "metadata": {}, "outputs": [], "source": ["dataset_rows = []\n", "for row in raw_rows:\n", "    dataset_rows.append({\n", "        'prompt'      : make_chat_prompt(row['prompt'], tokenizer),\n", "        'future_json' : row['future_json'],\n", "        'history_json' : row.get('history_json', '[]'),\n", "        'sensor_names': json.dumps(row['sensor_names']),\n", "    })\n", "\n", "dataset = Dataset.from_list(dataset_rows)\n", "print(f'Dataset: {len(dataset)} rows')\n", "print(f'Columns: {dataset.column_names}')\n", "print(f'Sample prompt length: {len(dataset[0][\"prompt\"])} chars')"]}, {"cell_type": "markdown", "id": "md-s7", "metadata": {}, "source": ["## Step 7 \u2014 GRPO Reward Function\n\n", "Each completion is scored by running it through the project's real sandbox and scorer against\n", "the **pre-captured future window** stored in the dataset row. This avoids the stale-trajectory\n", "bug in the original notebook where the reward used a separately-seeded universe."]}, {"cell_type": "code", "execution_count": null, "id": "reward-fn", "metadata": {}, "outputs": [], "source": ["from ontology_architect.universe import SensorRecord\n", "\n", "_sandbox = TheorySandbox(ENV_CONFIG.sandbox)\n", "_reward_cfg = ENV_CONFIG.reward\n", "\n", "\n", "def _deserialise_future(future_json: str) -> list[SensorRecord]:\n", "    return [\n", "        SensorRecord(t=r['t'], sensors=r['sensors'], anomaly=r['anomaly'], drift=r['drift'])\n", "        for r in json.loads(future_json)\n", "    ]\n", "\n", "\n", "def physics_reward_fn(\n", "    prompts     : list[str],\n", "    completions : list[str],\n", "    future_json : list[str],\n", "    history_json: list[str],\n", "    sensor_names: list[str],\n", "    **kwargs,\n", ") -> list[float]:\n", "    \"\"\"\n", "    GRPO reward function.  Columns `future_json` and `sensor_names` are\n", "    passed automatically by GRPOTrainer from the dataset.\n", "    \"\"\"\n", "    rewards = []\n", "    for theory_code, fj, hj_str, sn_json in zip(completions, future_json, history_json, sensor_names):\n", "        try:\n", "            # Strip markdown fences the model sometimes emits\n", "            theory_code = re.sub(r'^```(?:python)?\\s*', '', theory_code.strip())\n", "            theory_code = re.sub(r'```\\s*$', '', theory_code)\n", "\n", "            future       = _deserialise_future(fj)\n", "            names        = tuple(json.loads(sn_json))\n", "            history_data = json.loads(hj_str) if hj_str else []\n", "            sandbox_res  = _sandbox.execute(theory_code, history_data, future, names)\n", "            breakdown    = score_theory(sandbox_res, future, theory_code, _reward_cfg)\n", "            reward       = float(np.clip(breakdown.reward, -30.0, 5.0))\n", "        except Exception:\n", "            reward = float(_reward_cfg.execution_error_penalty)\n", "        rewards.append(reward)\n", "    return rewards\n", "\n", "\n", "# Sanity check \u2014 teacher should score positively\n", "row0 = dataset_rows[0]\n", "test_rewards = physics_reward_fn(\n", "    prompts      = [row0['prompt']],\n", "    completions  = [get_baseline('teacher')],\n", "    future_json  = [row0['future_json']],\n", "    sensor_names = [row0['sensor_names']],\n", ")\n", "print(f'Sanity check \u2014 teacher reward on row 0: {test_rewards[0]:.4f}')\n", "print(f'  (Expected roughly > {teacher_mean:.2f} for a good theory)')"]}, {"cell_type": "markdown", "id": "md-s8", "metadata": {}, "source": ["## Step 8 \u2014 GRPO Training"]}, {"cell_type": "code", "execution_count": null, "id": "train", "metadata": {}, "outputs": [], "source": ["if USE_WANDB:\n", "    import wandb\n", "    wandb.init(project='alien-physics-discovery', name='qwen2.5-coder-1.5b-grpo')\n", "\n", "grpo_config = GRPOConfig(\n", "    output_dir                  = 'artifacts/checkpoints/grpo',\n", "    num_train_epochs            = TRAIN_EPOCHS,\n", "    per_device_train_batch_size = BATCH_SIZE,\n", "    gradient_accumulation_steps = GRAD_ACCUM,\n", "    learning_rate               = LR,\n", "    warmup_ratio                = 0.05,\n", "    lr_scheduler_type           = 'cosine',\n", "    num_generations             = NUM_GENS,\n", "    max_completion_length       = MAX_COMPLETION,   # correct TRL >=0.9 field\n", "    logging_steps               = 5,\n", "    save_steps                  = 50,\n", "    save_total_limit            = 2,\n", "    fp16                        = not torch.cuda.is_bf16_supported(),\n", "    bf16                        = torch.cuda.is_bf16_supported(),\n", "    report_to                   = 'wandb' if USE_WANDB else 'none',\n", "    beta                        = 0.001,\n", "    max_grad_norm               = 0.3,\n", ")\n", "\n", "trainer = GRPOTrainer(\n", "    model            = model,\n", "    processing_class = tokenizer,\n", "    args             = grpo_config,          # 'args', not 'config'\n", "    train_dataset    = dataset,\n", "    reward_funcs     = [physics_reward_fn],\n", ")\n", "\n", "print('Starting GRPO training')\n", "print(f'  Model       : {MODEL_NAME}')\n", "print(f'  Samples     : {len(dataset)}')\n", "print(f'  Epochs      : {TRAIN_EPOCHS}')\n", "print(f'  Group size  : {NUM_GENS}')\n", "print(f'  Batch\u00d7accum : {BATCH_SIZE}\u00d7{GRAD_ACCUM} = {BATCH_SIZE*GRAD_ACCUM} effective')\n", "\n", "train_result = trainer.train()\n", "print(f'\\nTraining complete \u2014 final loss: {train_result.training_loss:.4f}')"]}, {"cell_type": "markdown", "id": "md-s9", "metadata": {}, "source": ["## Step 9 \u2014 Plot Training Curves"]}, {"cell_type": "code", "execution_count": null, "id": "plot", "metadata": {}, "outputs": [], "source": ["logs    = trainer.state.log_history\n", "steps   = [l['step']           for l in logs if 'loss'   in l]\n", "losses  = [l['loss']           for l in logs if 'loss'   in l]\n", "r_pairs = [(l['step'], l['reward']) for l in logs if 'reward' in l]\n", "\n", "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(13, 4))\n", "fig.suptitle('Alien Physics Discovery \u2014 GRPO Training (Qwen2.5-Coder-1.5B)', fontweight='bold')\n", "\n", "ax1.plot(steps, losses, color='#4f86c6', linewidth=2)\n", "ax1.set_xlabel('Training Step')\n", "ax1.set_ylabel('Loss')\n", "ax1.set_title('GRPO Loss')\n", "ax1.grid(alpha=0.3)\n", "\n", "if r_pairs:\n", "    rs, rv = zip(*r_pairs)\n", "    ax2.plot(rs, rv, color='#e67e22', linewidth=2, label='Mean group reward')\n", "    ax2.axhline(teacher_mean, color='gray', linestyle='--', alpha=0.7,\n", "                label=f'Teacher baseline ({teacher_mean:.2f})')\n", "    ax2.set_xlabel('Training Step')\n", "    ax2.set_ylabel('MDL Reward')\n", "    ax2.set_title('Reward vs Teacher Baseline')\n", "    ax2.legend()\n", "    ax2.grid(alpha=0.3)\n", "else:\n", "    ax2.text(0.5, 0.5, 'Enable USE_WANDB for live\\nreward curves, or check\\nlog_history manually',\n", "             ha='center', va='center', transform=ax2.transAxes, color='gray')\n", "\n", "plt.tight_layout()\n", "plt.savefig('artifacts/training_curves.png', dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print('Saved artifacts/training_curves.png')"]}, {"cell_type": "markdown", "id": "md-s10", "metadata": {}, "source": ["## Step 10 \u2014 Before / After Evaluation"]}, {"cell_type": "code", "execution_count": null, "id": "eval", "metadata": {}, "outputs": [], "source": ["FastLanguageModel.for_inference(model)\n", "\n", "def generate_theory(obs_text: str, max_new: int = 480) -> str:\n", "    prompt  = make_chat_prompt(obs_text, tokenizer)\n", "    inputs  = tokenizer(prompt, return_tensors='pt',\n", "                        truncation=True, max_length=MAX_SEQ_LEN - max_new).to(model.device)\n", "    with torch.no_grad():\n", "        out = model.generate(\n", "            **inputs,\n", "            max_new_tokens  = max_new,\n", "            temperature     = 0.6,\n", "            do_sample       = True,\n", "            pad_token_id    = tokenizer.eos_token_id,\n", "        )\n", "    theory = tokenizer.decode(out[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)\n", "    theory = re.sub(r'^```(?:python)?\\s*', '', theory.strip())\n", "    theory = re.sub(r'```\\s*$', '', theory)\n", "    return theory\n", "\n", "\n", "def eval_full_episode(family: str, seed: int) -> dict:\n", "    from dataclasses import replace as dc_replace\n", "    cfg = ExperimentConfig(\n", "        universe=dc_replace(ENV_CONFIG.universe, family=family, seed=seed),\n", "        sandbox=ENV_CONFIG.sandbox,\n", "        reward=ENV_CONFIG.reward,\n", "        feedback=ENV_CONFIG.feedback,\n", "    )\n", "    env  = OntologyArchitectEnvironment(cfg)\n", "    obs  = env.reset()\n", "    rewards, done = [], False\n", "    final_theory  = ''\n", "    while not done:\n", "        theory = generate_theory(obs.text)\n", "        obs    = env.step(OntologyArchitectAction(\n", "            theory_module=theory, revision_note='eval'\n", "        ))\n", "        rewards.append(obs.reward)\n", "        final_theory = theory\n", "        done = obs.done\n", "    return {'family': family, 'rewards': rewards,\n", "            'mean_reward': float(np.mean(rewards)), 'final_theory': final_theory}\n", "\n", "\n", "# Teacher baseline (untrained)\n", "def eval_teacher_episode(family: str, seed: int) -> dict:\n", "    from dataclasses import replace as dc_replace\n", "    cfg = ExperimentConfig(\n", "        universe=dc_replace(ENV_CONFIG.universe, family=family, seed=seed),\n", "        sandbox=ENV_CONFIG.sandbox,\n", "        reward=ENV_CONFIG.reward,\n", "        feedback=ENV_CONFIG.feedback,\n", "    )\n", "    env     = OntologyArchitectEnvironment(cfg)\n", "    obs     = env.reset()\n", "    teacher = get_baseline('teacher')\n", "    rewards, done = [], False\n", "    while not done:\n", "        obs  = env.step(OntologyArchitectAction(theory_module=teacher, revision_note='baseline'))\n", "        rewards.append(obs.reward)\n", "        done = obs.done\n", "    return {'family': family, 'rewards': rewards, 'mean_reward': float(np.mean(rewards))}\n", "\n", "\n", "print('Running eval episodes (trained model vs teacher baseline)...')\n", "eval_seeds  = [9001, 9002]\n", "eval_families = ['thermal_split', 'coupled_oscillator']\n", "\n", "trained_results  = [eval_full_episode(f, s)    for f in eval_families for s in eval_seeds]\n", "baseline_results = [eval_teacher_episode(f, s) for f in eval_families for s in eval_seeds]\n", "\n", "trained_mean  = float(np.mean([r['mean_reward'] for r in trained_results]))\n", "baseline_mean = float(np.mean([r['mean_reward'] for r in baseline_results]))\n", "\n", "print(f'\\nResults:')\n", "print(f'  Trained model mean reward  : {trained_mean:.4f}')\n", "print(f'  Teacher baseline mean reward: {baseline_mean:.4f}')\n", "print(f'  Delta                       : {trained_mean - baseline_mean:+.4f}')"]}, {"cell_type": "markdown", "id": "md-s11", "metadata": {}, "source": ["## Step 11 \u2014 Per-Turn Improvement Plot"]}, {"cell_type": "code", "execution_count": null, "id": "per-turn-plot", "metadata": {}, "outputs": [], "source": ["# Aggregate rewards by turn index across all eval episodes\n", "max_turns = max(len(r['rewards']) for r in trained_results)\n", "\n", "def pad_rewards(results, length):\n", "    padded = [r['rewards'] + [r['rewards'][-1]] * (length - len(r['rewards'])) for r in results]\n", "    return np.array(padded)\n", "\n", "trained_mat  = pad_rewards(trained_results,  max_turns)\n", "baseline_mat = pad_rewards(baseline_results, max_turns)\n", "turns = list(range(1, max_turns + 1))\n", "\n", "fig, ax = plt.subplots(figsize=(8, 4))\n", "ax.plot(turns, trained_mat.mean(0),  color='#3498db', marker='o', lw=2, label='Trained model')\n", "ax.fill_between(turns,\n", "                trained_mat.mean(0) - trained_mat.std(0),\n", "                trained_mat.mean(0) + trained_mat.std(0),\n", "                alpha=0.15, color='#3498db')\n", "ax.plot(turns, baseline_mat.mean(0), color='gray',    marker='s', lw=2, linestyle='--', label='Teacher baseline')\n", "ax.set_xlabel('Refinement Turn')\n", "ax.set_ylabel('MDL Reward')\n", "ax.set_title('Per-Turn Theory Improvement')\n", "ax.set_xticks(turns)\n", "ax.legend()\n", "ax.grid(alpha=0.3)\n", "plt.tight_layout()\n", "plt.savefig('artifacts/per_turn_improvement.png', dpi=150, bbox_inches='tight')\n", "plt.show()\n", "print('Saved artifacts/per_turn_improvement.png')"]}, {"cell_type": "markdown", "id": "md-s12", "metadata": {}, "source": ["## Step 12 \u2014 Show a Discovered Theory"]}, {"cell_type": "code", "execution_count": null, "id": "show-theory", "metadata": {}, "outputs": [], "source": ["best = max(trained_results, key=lambda r: r['mean_reward'])\n", "print(f\"Best result \u2014 family={best['family']}  mean_reward={best['mean_reward']:.4f}\")\n", "print('=' * 70)\n", "theory = best['final_theory']\n", "print(theory[:2000])\n", "if len(theory) > 2000:\n", "    print(f'\\n... [{len(theory)-2000} more chars] ...')"]}, {"cell_type": "markdown", "id": "md-s13", "metadata": {}, "source": ["## Step 13 \u2014 Save Model & Push to Hub"]}, {"cell_type": "code", "execution_count": null, "id": "save", "metadata": {}, "outputs": [], "source": ["Path('artifacts').mkdir(exist_ok=True)\n", "model.save_pretrained('artifacts/model')\n", "tokenizer.save_pretrained('artifacts/model')\n", "print('Model saved to artifacts/model')\n", "\n", "if HF_TOKEN:\n", "    from huggingface_hub import login\n", "    login(token=HF_TOKEN)\n", "    model.push_to_hub('lostdecimal27/alien-physics-discovery-qwen2.5-1.5b')\n", "    tokenizer.push_to_hub('lostdecimal27/alien-physics-discovery-qwen2.5-1.5b')\n", "    print('Pushed to HuggingFace Hub')\n", "else:\n", "    print('Set HF_TOKEN above to push to Hub')"]}, {"cell_type": "markdown", "id": "md-s14", "metadata": {}, "source": ["## Step 14 \u2014 Submission Checklist"]}, {"cell_type": "code", "execution_count": null, "id": "checklist", "metadata": {}, "outputs": [], "source": ["checks = {\n", "    'artifacts/training_curves.png'    : Path('artifacts/training_curves.png').exists(),\n", "    'artifacts/per_turn_improvement.png': Path('artifacts/per_turn_improvement.png').exists(),\n", "    'artifacts/model (local save)'     : Path('artifacts/model').exists(),\n", "    'openenv.yaml in repo'             : Path('openenv.yaml').exists(),\n", "    'README.md in repo'                : Path('README.md').exists(),\n", "    'Colab notebook (this file)'       : True,\n", "}\n", "\n", "print('OpenEnv Hackathon \u2014 Submission Checklist')\n", "print('=' * 50)\n", "for item, ok in checks.items():\n", "    print(f\"  {'\u2705' if ok else '\u274c'}  {item}\")\n", "\n", "print('\\nRemaining manual steps:')\n", "for todo in [\n", "    '\u25a1  Push model to HuggingFace Hub (set HF_TOKEN)',\n", "    '\u25a1  Deploy environment to HuggingFace Space',\n", "    '\u25a1  Record < 2 min YouTube demo or write HF blog post',\n", "    '\u25a1  Update lostdecimal27 placeholders above',\n", "    '\u25a1  Update README with video/blog URL and HF Space URL',\n", "    '\u25a1  Submit HF Space URL to hackathon portal',\n", "]:\n", "    print(f'  {todo}')\n", "\n", "print(f'\\nFinal score summary:')\n", "print(f'  Teacher baseline mean reward : {baseline_mean:.4f}')\n", "print(f'  Trained model mean reward    : {trained_mean:.4f}')\n", "print(f'  Improvement                  : {trained_mean - baseline_mean:+.4f}')"]}]}

baselines.py

@@ -143,9 +143,9 @@ DUAL_FLUID_DSL_ORACLE = _json.dumps({
         "C": 0.0,
     },
     "dynamics": {
-        "A": {"linear": {"terms": {"A": -0.3, "B": 0.25}, "bias": 0.0}},
-        "B": {"add": [{"sin": "A"}, {"linear": {"terms": {"A": -0.08}}}]},
-        "C": {"linear": {"terms": {"A": 0.06, "C": -0.12}}},
+        "A": {"add": [{"linear": {"terms": {"A": -0.3}}}, {"mul": [0.5, {"var": "A"}, {"var": "B"}]}]},
+        "B": {"add": [{"mul": [0.4, {"sin": {"var": "t"}}]}, {"mul": [-0.2, {"pow": [{"var": "A"}, 2]}]}]},
+        "C": {"add": [{"mul": [0.15, {"sin": {"mul": [0.4, {"var": "A"}]}}]}, {"linear": {"terms": {"C": -0.12}}}]},
     },
     "observations": {
         "pressure": {"var": "A"},

scripts/visualize_discovery.py

@@ -73,8 +73,8 @@ def main() -> None:
 
     # Run theory through sandbox to get predictions
     sandbox = TheorySandbox(config.sandbox)
-    # Use first half as history, predict second half
-    split = len(records) // 2
+    # Use a small initial window as history so the oracle can initialize correctly
+    split = 8
     history = records[:split]
     future = records[split:]
 

theory_dsl.py

@@ -72,6 +72,7 @@ def normalize_theory_spec(spec: dict[str, Any]) -> dict[str, Any]:
         raise TheoryDSLValidationError(f"DSL observations exceed {MAX_SENSOR_OUTPUTS} outputs")
 
     state_set = set(state)
+    state_set.add("t")
     for name, expr in dynamics.items():
         if str(name) not in state_set:
             raise TheoryDSLValidationError(f"Dynamic variable '{name}' is not in state")
@@ -284,7 +285,9 @@ class Theory:
         recent = history[-lookback:] if history else []
         self.trend = self._fit_trend(recent)
         self.residual_scale = self._fit_residual_scale(recent)
+        self.last_t = float(history[-1].get("t", 0.0)) if history else 0.0
         self.state = self._initial_state()
+        self.state["t"] = self.last_t
         return {
             "dsl_version": self.spec["dsl_version"],
             "state": list(self.spec["state"]),
@@ -387,6 +390,7 @@ class Theory:
         for _ in range(max(substeps, 1)):
             deriv = self._derivative(state)
             state = {name: self._clip(value + step_dt * deriv.get(name, 0.0)) for name, value in state.items()}
+            state["t"] = state.get("t", 0.0) + step_dt
         return state
 
     def _derivative(self, state):