Upload folder using huggingface_hub

client.py

@@ -51,6 +51,9 @@ class AegisEnv(
             rubric=obs_data.get("rubric", ""),
             max_score=float(obs_data.get("max_score", 1.0)),
             student_answer=obs_data.get("student_answer", ""),
+            current_stage=str(obs_data.get("current_stage") or "arbiter"),
+            refinement_loops_taken=int(obs_data.get("refinement_loops_taken") or 0),
+            pipeline_history=str(obs_data.get("pipeline_history") or ""),
             grading_info=obs_data.get("grading_info") or {},
             done=payload.get("done", False),
             reward=payload.get("reward"),

notebooks/grpo_aegis_hf_space_unsloth.ipynb

@@ -0,0 +1,335 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "# GRPO fine-tuning (Unsloth) — AEGIS-Env against your Hugging Face Space\n",
+        "\n",
+        "This notebook trains a **small language model** with **GRPO** using rewards from **[AEGIS-Env](https://huggingface.co/spaces/NishithP2004/aegis-env)**.\n",
+        "\n",
+        "- **Remote environment:** OpenEnv WebSocket API at `…/openenv` on the Space.\n",
+        "- **Three difficulty tiers:** `easy`, `medium`, `hard` → same mapping as the server (`mohler` / `asap-sas` / `ricechem` datasets).\n",
+        "- **Reward signal:** After a **fixed deterministic replay** through arbiter → scrutinizer → validator, the model’s completion is parsed as a **mentor-stage** JSON action (`AegisAction`); the Space returns the **final step reward** in `[0, 1]`.\n",
+        "\n",
+        "**Requirements:** GPU runtime (e.g. T4). The Space must be awake (cold start can take a minute). The clone step pulls `client.py`, `models.py`, and `inference.py` from your Space repo for `build_user_prompt` / `_parse_action`.\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# %%capture\n",
+        "import os, importlib.util\n",
+        "!pip install -qqq uv nest_asyncio\n",
+        "if importlib.util.find_spec(\"torch\") is None or \"COLAB_\" in \"\".join(os.environ.keys()):\n",
+        "    try:\n",
+        "        import numpy\n",
+        "        get_numpy = f\"numpy=={numpy.__version__}\"\n",
+        "    except Exception:\n",
+        "        get_numpy = \"numpy\"\n",
+        "    !uv pip install -qqq \\\n",
+        "        \"torch>=2.8.0\" \"triton>=3.4.0\" {get_numpy} torchvision bitsandbytes \"transformers==4.56.2\" trackio \\\n",
+        "        \"unsloth_zoo[base] @ git+https://github.com/unslothai/unsloth-zoo\" \\\n",
+        "        \"unsloth[base] @ git+https://github.com/unslothai/unsloth\" \\\n",
+        "        git+https://github.com/triton-lang/triton.git@0add68262ab0a2e33b84524346cb27cbb2787356#subdirectory=python/triton_kernels\n",
+        "elif importlib.util.find_spec(\"unsloth\") is None:\n",
+        "    !uv pip install -qqq unsloth trackio\n",
+        "!uv pip install --upgrade --no-deps transformers==4.56.2 tokenizers trl==0.22.2 unsloth unsloth_zoo\n",
+        "!uv pip install -qqq \"openenv-core[core]>=0.2.2\" openai pydantic\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import os\n",
+        "# AEGIS-Env code (client + prompts). Override with your Git repo URL if preferred.\n",
+        "AEGIS_REPO = os.environ.get(\n",
+        "    \"AEGIS_CODE_CLONE_URL\",\n",
+        "    \"https://huggingface.co/spaces/NishithP2004/aegis-env\",\n",
+        ")\n",
+        "!rm -rf /content/aegis-env-space && git clone --depth 1 {AEGIS_REPO} /content/aegis-env-space\n",
+        "\n",
+        "import sys\n",
+        "sys.path.insert(0, \"/content/aegis-env-space\")\n",
+        "\n",
+        "import nest_asyncio\n",
+        "nest_asyncio.apply()\n",
+        "\n",
+        "AEGIS_OPENENV_BASE = os.environ.get(\n",
+        "    \"AEGIS_OPENENV_BASE\",\n",
+        "    \"https://huggingface.co/spaces/NishithP2004/aegis-env/openenv\",\n",
+        ")\n",
+        "print(\"OpenEnv WS/HTTP base:\", AEGIS_OPENENV_BASE)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "id": "7e9a6ebc",
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "import asyncio\n",
+        "from typing import Any, List\n",
+        "\n",
+        "from datasets import Dataset\n",
+        "from transformers import TextStreamer\n",
+        "from trl import GRPOConfig, GRPOTrainer\n",
+        "from unsloth import FastLanguageModel\n",
+        "\n",
+        "from models import AegisAction, AegisObservation\n",
+        "from client import AegisEnv\n",
+        "from inference import build_user_prompt, _parse_action, _strip_markdown_json_fence\n",
+        "\n",
+        "\n",
+        "async def replay_to_mentor_obs(env: AegisEnv, seed: int, task_name: str) -> AegisObservation:\n",
+        "    r = await env.reset(seed=seed, task_name=task_name)\n",
+        "    obs = r.observation\n",
+        "    ms = float(obs.max_score or 1.0)\n",
+        "    mid = max(0.0, min(ms, ms * 0.5))\n",
+        "    a1 = AegisAction(\n",
+        "        proposed_score=mid,\n",
+        "        agent_reasoning=\"Arbiter: preliminary assessment aligned with the rubric.\",\n",
+        "        routing_decision=\"proceed\",\n",
+        "    )\n",
+        "    obs = (await env.step(a1)).observation\n",
+        "    a2 = AegisAction(\n",
+        "        proposed_score=mid,\n",
+        "        agent_reasoning=\"Scrutinizer: refined criterion-by-criterion review.\",\n",
+        "        routing_decision=\"proceed\",\n",
+        "    )\n",
+        "    obs = (await env.step(a2)).observation\n",
+        "    a3 = AegisAction(\n",
+        "        proposed_score=mid,\n",
+        "        agent_reasoning=\"Validator: consistency check passed; proceed to mentor.\",\n",
+        "        routing_decision=\"proceed\",\n",
+        "    )\n",
+        "    obs = (await env.step(a3)).observation\n",
+        "    if str(getattr(obs, \"current_stage\", \"\") or \"\") != \"mentor\":\n",
+        "        raise RuntimeError(f\"Expected mentor stage, got {getattr(obs, 'current_stage', None)}\")\n",
+        "    return obs\n",
+        "\n",
+        "\n",
+        "def _run(coro):\n",
+        "    return asyncio.get_event_loop().run_until_complete(coro)\n",
+        "\n",
+        "\n",
+        "async def build_one_row(task_name: str, seed: int, base_url: str) -> dict[str, Any]:\n",
+        "    env = AegisEnv(base_url=base_url)\n",
+        "    await env.connect()\n",
+        "    try:\n",
+        "        obs = await replay_to_mentor_obs(env, seed, task_name)\n",
+        "        user_text = build_user_prompt(4, None, 0.0, [], obs)\n",
+        "        return {\n",
+        "            \"prompt\": [{\"role\": \"user\", \"content\": user_text}],\n",
+        "            \"answer\": 0,\n",
+        "            \"task_name\": task_name,\n",
+        "            \"episode_seed\": int(seed),\n",
+        "            \"mentor_max_score\": float(obs.max_score or 1.0),\n",
+        "        }\n",
+        "    finally:\n",
+        "        await env.close()\n",
+        "\n",
+        "\n",
+        "def build_dataset(\n",
+        "    base_url: str,\n",
+        "    n_per_tier: int = 40,\n",
+        ") -> Dataset:\n",
+        "    rows: List[dict] = []\n",
+        "    tier_seeds = {\"easy\": 1000, \"medium\": 11000, \"hard\": 21000}\n",
+        "    for task_name in (\"easy\", \"medium\", \"hard\"):\n",
+        "        base_seed = tier_seeds[task_name]\n",
+        "        for i in range(n_per_tier):\n",
+        "            seed = base_seed + i\n",
+        "            row = _run(build_one_row(task_name, seed, base_url))\n",
+        "            rows.append(row)\n",
+        "            print(\"built\", task_name, i, \"max_score\", row[\"mentor_max_score\"])\n",
+        "    return Dataset.from_list(rows)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "# Build prompts against the live Space (sequential; Space allows limited concurrency)\n",
+        "train_ds = build_dataset(AEGIS_OPENENV_BASE, n_per_tier=40)\n",
+        "sample = train_ds[0][\"prompt\"]\n",
+        "print(train_ds)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "MODEL_NAME = \"unsloth/Llama-3.2-3B-Instruct-bnb-4bit\"\n",
+        "max_seq_length = 4096\n",
+        "lora_rank = 8\n",
+        "\n",
+        "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+        "    model_name=MODEL_NAME,\n",
+        "    load_in_4bit=True,\n",
+        "    max_seq_length=max_seq_length,\n",
+        ")\n",
+        "model = FastLanguageModel.get_peft_model(\n",
+        "    model,\n",
+        "    r=lora_rank,\n",
+        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
+        "    lora_alpha=lora_rank * 2,\n",
+        "    use_gradient_checkpointing=\"unsloth\",\n",
+        "    random_state=3407,\n",
+        ")\n",
+        "\n",
+        "maximum_length = len(tokenizer.apply_chat_template(sample, add_generation_prompt=True))\n",
+        "print(\"Prompt tokens (sample):\", maximum_length)\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "async def mentor_env_reward_one(\n",
+        "    completion_text: str,\n",
+        "    seed: int,\n",
+        "    task_name: str,\n",
+        "    mentor_max_score: float,\n",
+        "    base_url: str,\n",
+        ") -> float:\n",
+        "    env = AegisEnv(base_url=base_url)\n",
+        "    await env.connect()\n",
+        "    try:\n",
+        "        obs = await replay_to_mentor_obs(env, seed, task_name)\n",
+        "        ms = float(obs.max_score or mentor_max_score or 1.0)\n",
+        "        action = _parse_action(_strip_markdown_json_fence(completion_text), ms)\n",
+        "        sr = await env.step(action)\n",
+        "        r = float(sr.reward if sr.reward is not None else 0.0)\n",
+        "        return r * 20.0\n",
+        "    except Exception as e:\n",
+        "        print(\"mentor_env_reward error:\", e)\n",
+        "        return -8.0\n",
+        "    finally:\n",
+        "        await env.close()\n",
+        "\n",
+        "\n",
+        "def format_valid(prompts, completions, mentor_max_score, **kwargs):\n",
+        "    scores = []\n",
+        "    for c, ms in zip(completions, mentor_max_score):\n",
+        "        try:\n",
+        "            _parse_action(_strip_markdown_json_fence(c[0][\"content\"]), float(ms))\n",
+        "            scores.append(1.0)\n",
+        "        except Exception:\n",
+        "            scores.append(-1.5)\n",
+        "    return scores\n",
+        "\n",
+        "\n",
+        "def no_markdown_fence(prompts, completions, **kwargs):\n",
+        "    scores = []\n",
+        "    for c in completions:\n",
+        "        t = c[0][\"content\"].strip()\n",
+        "        scores.append(-0.5 if t.startswith(\"```\") else 0.2)\n",
+        "    return scores\n",
+        "\n",
+        "\n",
+        "def mentor_success(prompts, completions, episode_seed, task_name, mentor_max_score, **kwargs):\n",
+        "    scores = []\n",
+        "    for c, seed, task, ms in zip(completions, episode_seed, task_name, mentor_max_score):\n",
+        "        text = c[0][\"content\"]\n",
+        "        coro = mentor_env_reward_one(text, int(seed), str(task), float(ms), AEGIS_OPENENV_BASE)\n",
+        "        scores.append(float(_run(coro)))\n",
+        "    return scores\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "text = tokenizer.apply_chat_template(sample, tokenize=False, add_generation_prompt=True)\n",
+        "_ = model.generate(\n",
+        "    **tokenizer(text, return_tensors=\"pt\").to(\"cuda\"),\n",
+        "    temperature=0.7,\n",
+        "    max_new_tokens=512,\n",
+        "    streamer=TextStreamer(tokenizer, skip_prompt=False),\n",
+        ")\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {},
+      "outputs": [],
+      "source": [
+        "max_prompt_length = min(maximum_length + 8, max_seq_length // 2)\n",
+        "max_completion_length = max_seq_length - max_prompt_length\n",
+        "\n",
+        "training_args = GRPOConfig(\n",
+        "    temperature=0.8,\n",
+        "    learning_rate=2e-4,\n",
+        "    weight_decay=0.001,\n",
+        "    warmup_ratio=0.1,\n",
+        "    lr_scheduler_type=\"linear\",\n",
+        "    optim=\"adamw_8bit\",\n",
+        "    logging_steps=1,\n",
+        "    per_device_train_batch_size=1,\n",
+        "    gradient_accumulation_steps=2,\n",
+        "    num_generations=2,\n",
+        "    max_prompt_length=max_prompt_length,\n",
+        "    max_completion_length=max_completion_length,\n",
+        "    max_steps=200,\n",
+        "    save_steps=50,\n",
+        "    report_to=\"trackio\",\n",
+        "    output_dir=\"outputs_grpo_aegis_hf\",\n",
+        ")\n",
+        "\n",
+        "trainer = GRPOTrainer(\n",
+        "    model=model,\n",
+        "    processing_class=tokenizer,\n",
+        "    reward_funcs=[format_valid, no_markdown_fence, mentor_success],\n",
+        "    args=training_args,\n",
+        "    train_dataset=train_ds,\n",
+        ")\n",
+        "\n",
+        "trainer.train()\n"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": [
+        "## Notes\n",
+        "\n",
+        "- **Latency:** `mentor_success` calls your Space once per completion; keep `n_per_tier`, `max_steps`, and `num_generations` modest while iterating.\n",
+        "- **WebSocket client:** Use an `AegisEnv` / `client.py` that maps `current_stage`, `pipeline_history`, and `refinement_loops_taken` from observations (updated in this repo’s `client.py` — redeploy the Space to pick it up).\n",
+        "- **Overrides:** Set `AEGIS_OPENENV_BASE` or `AEGIS_CODE_CLONE_URL` instead of editing cells.\n"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 5
+}