sync: push from tools/sync_space_to_hub.py (no artifacts/)

README.md

@@ -112,6 +112,18 @@ REASON: max 20 words
 CONFIDENCE: 0.0-1.0
 ```
 
+### Road-graph RL (connected-edge actions)
+
+This repo also includes a road-graph RL environment mounted under `POST /road/reset` and `POST /road/step`.
+Its action schema is:
+
+```text
+CURRENT_NODE: <int>
+NEXT_NODE: <int>
+REASON: max 20 words
+CONFIDENCE: 0.0-1.0
+```
+
 ### Reward (verifiable + anti‑hack)
 
 Total reward is the sum of components (each logged) in `ev_grid_oracle/reward.py`:

tools/road_reward_smoke.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import re
+
+from ev_grid_oracle.road_env import RoadCore
+from ev_grid_oracle.road_models import RoadAction, RoadState
+
+
+def main() -> int:
+    core = RoadCore(g=None, nodes=[])  # type: ignore[arg-type]
+    obs = core.reset(seed=0)
+    st = obs.state
+    nb = list(core.g.neighbors(st.node))[0]
+
+    ok = f"CURRENT_NODE: {st.node}\nNEXT_NODE: {int(nb)}\nREASON: go\nCONFIDENCE: 0.7\n"
+    bad = f"CURRENT_NODE: {st.node + 999}\nNEXT_NODE: {int(nb)}\nREASON: hack\nCONFIDENCE: 0.7\n"
+
+    r = re.compile(r"CURRENT_NODE:\s*(\d+)\s*\nNEXT_NODE:\s*(\d+)\s*\n", re.I)
+
+    def parse(t: str) -> RoadAction | None:
+        m = r.search(t.strip())
+        if not m:
+            return None
+        return RoadAction(current_node=int(m.group(1)), next_node=int(m.group(2)))
+
+    def reward(comp: str) -> float:
+        st2 = RoadState.model_validate(st.model_dump(mode="json"))
+        a = parse(comp)
+        if a is None or int(a.current_node) != int(st2.node):
+            return -1.0
+        local = RoadCore(g=core.g, nodes=core.nodes)
+        local.node = int(st2.node)
+        local.battery_pct = float(st2.battery_pct_0_100)
+        local.target_station_id = str(st2.target_station_id)
+        local.steps_remaining = int(st2.steps_remaining)
+        ob = local.step(a)
+        base = float(ob.reward_breakdown.get("total", 0.0))
+        cheat = -1.0 if ob.anti_cheat_flags else 0.0
+        return base + cheat
+
+    print("ok", reward(ok))
+    print("bad", reward(bad))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
+

training/train_grpo.ipynb

@@ -6,24 +6,25 @@
       "source": [
         "# EV Grid Oracle — GRPO Training (Colab T4)\n",
         "\n",
-        "This notebook trains a small LLM (Qwen 2.5 3B Instruct) with **verifier-based GRPO** to route EVs in the `EVGridCore` simulation.\n",
+        "This notebook trains a small LLM (Qwen 2.5 3B Instruct) with **verifier-based GRPO** on the **real Bangalore road graph** (connected-edge actions only).\n",
+        "\n",
+        "- **Environment**: OpenEnv-compatible `EVGridRoadEnvironment` mounted at `/road/` in the Space.\n",
+        "- **Key constraint**: the policy can only choose a **connected neighbor** in the OSM-derived graph (no teleporting).\n",
         "\n",
         "## Run order (Colab — once per runtime)\n",
         "\n",
         "1. **Runtime → Change runtime type → GPU** (T4 matches the defaults below).\n",
         "2. Run the **next code cell** first. It clones this repo, moves into it, and runs `pip install -e .` so `import ev_grid_oracle` works.\n",
-        "3. Run cells **top to bottom**. After changing only hyperparameters, you can restart from the dataset cell once setup has run.\n",
+        "3. Run cells **top to bottom**.\n",
         "4. After training, use the save cell and upload `ev_oracle_lora/` to the Hub (or copy to Drive).\n",
         "\n",
-        "**Links:** [Open in Colab](https://colab.research.google.com/github/NITISH-R-G/ev-grid-oracle/blob/main/training/train_grpo.ipynb) · [Notebook on GitHub](https://github.com/NITISH-R-G/ev-grid-oracle/blob/main/training/train_grpo.ipynb) · [HF mini-blog (markdown in repo)](https://github.com/NITISH-R-G/ev-grid-oracle/blob/main/docs/hf-mini-blog-ev-grid-oracle.md)\n",
+        "**Links:** [Open in Colab](https://colab.research.google.com/github/NITISH-R-G/ev-grid-oracle/blob/main/training/train_grpo.ipynb) · [Notebook on GitHub](https://github.com/NITISH-R-G/ev-grid-oracle/blob/main/training/train_grpo.ipynb)\n",
         "\n",
         "**Action schema (strict):**\n",
         "\n",
-        "```\n",
-        "ACTION: route|defer|load_shift\n",
-        "STATION: BLR-01..BLR-25 or NONE\n",
-        "CHARGE_RATE: slow|fast|ultra_fast\n",
-        "DEFER_MINUTES: integer\n",
+        "```text\n",
+        "CURRENT_NODE: <int>\n",
+        "NEXT_NODE: <int>\n",
         "REASON: max 20 words\n",
         "CONFIDENCE: 0.0-1.0\n",
         "```\n",
@@ -85,13 +86,12 @@
         "\n",
         "from datasets import Dataset\n",
         "\n",
-        "from ev_grid_oracle.city_graph import build_city_graph\n",
-        "from ev_grid_oracle.env import EVGridCore\n",
-        "from ev_grid_oracle.models import ActionType, ChargeRate, EVGridAction, GridState\n",
+        "from ev_grid_oracle.road_env import RoadCore\n",
+        "from ev_grid_oracle.road_models import RoadAction, RoadState\n",
         "\n",
         "\n",
-        "graph = build_city_graph()\n",
-        "core = EVGridCore(city_graph=graph)\n"
+        "core = RoadCore(g=None, nodes=[])  # graph is loaded inside reset()\n",
+        ""
       ],
       "execution_count": null,
       "outputs": [],
@@ -102,65 +102,21 @@
       "metadata": {},
       "source": [
         "ACTION_RE = re.compile(\n",
-        "    r\"ACTION:\\s*(?P<action>route|defer|load_shift)\\s*\\n\"\n",
-        "    r\"STATION:\\s*(?P<station>BLR-\\d\\d|NONE)\\s*\\n\"\n",
-        "    r\"CHARGE_RATE:\\s*(?P<rate>slow|fast|ultra_fast)\\s*\\n\"\n",
-        "    r\"DEFER_MINUTES:\\s*(?P<defer>\\d+)\\s*\\n\",\n",
+        "    r\"CURRENT_NODE:\\s*(?P<cur>\\d+)\\s*\\n\"\n",
+        "    r\"NEXT_NODE:\\s*(?P<nxt>\\d+)\\s*\\n\",\n",
         "    re.IGNORECASE,\n",
         ")\n",
         "\n",
-        "SIM_RE = re.compile(\n",
-        "    r\"<SIMULATE>\\s*\\n\"\n",
-        "    r\"T\\+5_GRID_LOAD_PCT:\\s*(?P<grid>[01](?:\\.\\d+)?)\\s*\\n\"\n",
-        "    r\"T\\+5_RENEWABLE_PCT:\\s*(?P<ren>[01](?:\\.\\d+)?)\\s*\\n\"\n",
-        "    r\"T\\+5_TOP_STATIONS:\\s*(?P<tops>.+?)\\s*\\n\"\n",
-        "    r\"</SIMULATE>\",\n",
-        "    re.IGNORECASE | re.DOTALL,\n",
-        ")\n",
         "\n",
-        "\n",
-        "def parse_action(text: str, *, ev_id: str) -> Optional[EVGridAction]:\n",
+        "def parse_action(text: str) -> Optional[RoadAction]:\n",
         "    m = ACTION_RE.search(text.strip())\n",
         "    if not m:\n",
         "        return None\n",
-        "\n",
-        "    action_type = ActionType(m.group(\"action\").lower())\n",
-        "    station = m.group(\"station\").upper()\n",
-        "    rate = ChargeRate(m.group(\"rate\").lower())\n",
-        "    defer = int(m.group(\"defer\"))\n",
-        "\n",
-        "    station_id = None if station == \"NONE\" else station\n",
-        "\n",
         "    try:\n",
-        "        return EVGridAction(\n",
-        "            action_type=action_type,\n",
-        "            ev_id=ev_id,\n",
-        "            station_id=station_id,\n",
-        "            charge_rate=rate,\n",
-        "            defer_minutes=defer,\n",
-        "        )\n",
+        "        return RoadAction(current_node=int(m.group(\"cur\")), next_node=int(m.group(\"nxt\")))\n",
         "    except Exception:\n",
         "        return None\n",
-        "\n",
-        "\n",
-        "def parse_sim(text: str):\n",
-        "    m = SIM_RE.search(text)\n",
-        "    if not m:\n",
-        "        return None\n",
-        "    try:\n",
-        "        grid = float(m.group(\"grid\"))\n",
-        "        ren = float(m.group(\"ren\"))\n",
-        "        tops_raw = m.group(\"tops\").strip()\n",
-        "        parts = [p.strip() for p in tops_raw.split(\"|\") if p.strip()]\n",
-        "        tops = []\n",
-        "        for p in parts[:3]:\n",
-        "            sid, load_s, q_s = [x.strip() for x in p.split(\":\")]\n",
-        "            tops.append((sid.upper(), float(load_s), int(q_s)))\n",
-        "        if not tops:\n",
-        "            return None\n",
-        "        return {\"grid\": grid, \"ren\": ren, \"tops\": tops}\n",
-        "    except Exception:\n",
-        "        return None\n"
+        ""
       ],
       "execution_count": null,
       "outputs": [],
@@ -170,22 +126,40 @@
       "cell_type": "code",
       "metadata": {},
       "source": [
-        "def generate_episode_dataset(n: int = 500, *, seed: int = 123) -> Dataset:\n",
+        "def _format_neighbors(st: RoadState, *, max_k: int = 12) -> str:\n",
+        "    # Expose valid actions (neighbors) so the LLM can't claim it didn't know.\n",
+        "    g = core.g\n",
+        "    neigh = list(g.neighbors(int(st.node)))[:max_k]\n",
+        "    return \", \".join(str(int(x)) for x in neigh)\n",
+        "\n",
+        "\n",
+        "def generate_episode_dataset(n: int = 800, *, seed: int = 123) -> Dataset:\n",
         "    rows = []\n",
         "    for i in range(n):\n",
         "        obs = core.reset(seed=seed + i)\n",
-        "        # Keep verifier honest: reward_fn uses state_json, not prompt parsing.\n",
-        "        rows.append(\n",
-        "            {\n",
-        "                \"prompt\": obs.prompt,\n",
-        "                \"state_json\": obs.state.model_dump(mode=\"json\"),\n",
-        "            }\n",
+        "        st = obs.state\n",
+        "        neigh = _format_neighbors(st)\n",
+        "        prompt = (\n",
+        "            \"You are routing an EV on Bangalore's real road graph. You must pick NEXT_NODE as a connected neighbor only.\\n\\n\"\n",
+        "            f\"CURRENT_NODE: {st.node}\\n\"\n",
+        "            f\"BATTERY_PCT: {st.battery_pct_0_100:.1f}\\n\"\n",
+        "            f\"TARGET_STATION_ID: {st.target_station_id}\\n\"\n",
+        "            f\"TARGET_LATLNG: {st.target_lat:.6f},{st.target_lng:.6f}\\n\"\n",
+        "            f\"STEPS_REMAINING: {st.steps_remaining}\\n\"\n",
+        "            f\"VALID_NEXT_NODES: {neigh}\\n\\n\"\n",
+        "            \"Respond in this exact schema:\\n\"\n",
+        "            \"CURRENT_NODE: <int>\\n\"\n",
+        "            \"NEXT_NODE: <int>\\n\"\n",
+        "            \"REASON: max 20 words\\n\"\n",
+        "            \"CONFIDENCE: 0.0-1.0\\n\"\n",
         "        )\n",
+        "        rows.append({\"prompt\": prompt, \"state_json\": st.model_dump(mode=\"json\")})\n",
         "    return Dataset.from_list(rows)\n",
         "\n",
         "\n",
-        "train_ds = generate_episode_dataset(n=500)\n",
-        "train_ds[0][\"prompt\"][:400]\n"
+        "train_ds = generate_episode_dataset(n=800)\n",
+        "train_ds[0][\"prompt\"][:450]\n",
+        ""
       ],
       "execution_count": null,
       "outputs": [],
@@ -221,55 +195,45 @@
       "cell_type": "code",
       "metadata": {},
       "source": [
-        "from ev_grid_oracle.models import SimulationPrediction, SimTopStation\n",
-        "from ev_grid_oracle.world_model_verifier import score_prediction\n",
+        "from ev_grid_oracle.road_models import RoadState\n",
         "\n",
         "\n",
         "def reward_fn(prompts, completions, **kwargs):\n",
         "    rewards = []\n",
         "\n",
-        "    # TRL passes prompt strings; we recover the matching state via dataset column.\n",
-        "    # We rely on GRPOTrainer passing `kwargs[\"batch\"]` with original examples.\n",
         "    batch = kwargs.get(\"batch\")\n",
         "    state_jsons = batch[\"state_json\"] if batch is not None and \"state_json\" in batch else None\n",
         "\n",
-        "    for prompt, completion, state_json in zip(prompts, completions, state_jsons or [None] * len(prompts)):\n",
+        "    for completion, state_json in zip(completions, state_jsons or [None] * len(completions)):\n",
         "        if state_json is None:\n",
         "            rewards.append(0.0)\n",
         "            continue\n",
         "\n",
-        "        # pick target ev_id = first pending EV in state (matches prompt_builder v0)\n",
-        "        state = GridState.model_validate(state_json)\n",
-        "        ev_id = state.pending_evs[0].ev_id if state.pending_evs else \"EV-000\"\n",
-        "\n",
-        "        action = parse_action(completion, ev_id=ev_id)\n",
+        "        st = RoadState.model_validate(state_json)\n",
+        "        action = parse_action(completion)\n",
         "        if action is None:\n",
         "            rewards.append(-1.0)\n",
         "            continue\n",
         "\n",
-        "        # Base env reward\n",
-        "        local = EVGridCore(city_graph=graph)\n",
-        "        local._grid_state = state\n",
+        "        # Hard anti-cheat: must match the provided current node.\n",
+        "        if int(action.current_node) != int(st.node):\n",
+        "            rewards.append(-1.0)\n",
+        "            continue\n",
+        "\n",
+        "        # Step local env from the same state.\n",
+        "        local = RoadCore(g=core.g, nodes=core.nodes)\n",
+        "        local.node = int(st.node)\n",
+        "        local.battery_pct = float(st.battery_pct_0_100)\n",
+        "        local.target_station_id = str(st.target_station_id)\n",
+        "        local.steps_remaining = int(st.steps_remaining)\n",
+        "\n",
         "        obs = local.step(action)\n",
         "        base_r = float(obs.reward_breakdown.get(\"total\", 0.0))\n",
         "\n",
-        "        # Dream-state prediction reward (aggregate-only, verifiable)\n",
-        "        sim = parse_sim(completion)\n",
-        "        if sim is None:\n",
-        "            pred_r = -1.0\n",
-        "        else:\n",
-        "            pred = SimulationPrediction(\n",
-        "                t5_grid_load_pct=sim[\"grid\"],\n",
-        "                t5_renewable_pct=sim[\"ren\"],\n",
-        "                t5_top_stations=[\n",
-        "                    SimTopStation(station_id=sid, load_pct=load, queue=q) for sid, load, q in sim[\"tops\"]\n",
-        "                ],\n",
-        "            )\n",
-        "            sc = score_prediction(state, action, pred)\n",
-        "            # map [0,1] -> [-1,+1]\n",
-        "            pred_r = (sc.score_0_1 * 2.0) - 1.0\n",
-        "\n",
-        "        rewards.append(base_r + 2.0 * pred_r)\n",
+        "        # Penalize any anti-cheat flags from the verifier.\n",
+        "        cheat_pen = -1.0 if obs.anti_cheat_flags else 0.0\n",
+        "\n",
+        "        rewards.append(base_r + cheat_pen)\n",
         "\n",
         "    return rewards\n",
         ""
@@ -285,25 +249,52 @@
         "from trl import GRPOConfig, GRPOTrainer\n",
         "\n",
         "config = GRPOConfig(\n",
-        "    output_dir=\"ev_oracle_grpo\",\n",
+        "    output_dir=\"ev_oracle_grpo_road\",\n",
         "    num_train_epochs=1,\n",
         "    per_device_train_batch_size=2,\n",
         "    gradient_accumulation_steps=8,\n",
         "    learning_rate=5e-5,\n",
         "    num_generations=4,\n",
-        "    max_completion_length=160,\n",
+        "    max_completion_length=120,\n",
         "    report_to=[],\n",
+        "    logging_steps=1,\n",
         ")\n",
         "\n",
+        "# Minimal guardrail sampling: print a few raw generations early.\n",
+        "class SampleCallback:\n",
+        "    def __init__(self, every_steps: int = 10, n: int = 3):\n",
+        "        self.every_steps = every_steps\n",
+        "        self.n = n\n",
+        "\n",
+        "    def on_step_end(self, args, state, control, **kwargs):\n",
+        "        step = int(getattr(state, \"global_step\", 0) or 0)\n",
+        "        if step == 1 or (self.every_steps and step % self.every_steps == 0):\n",
+        "            ex = train_ds.select(range(min(self.n, len(train_ds))))\n",
+        "            for i, p in enumerate(ex[\"prompt\"]):\n",
+        "                out = tokenizer.decode(\n",
+        "                    model.generate(\n",
+        "                        **tokenizer(p, return_tensors=\"pt\").to(model.device),\n",
+        "                        max_new_tokens=80,\n",
+        "                        do_sample=True,\n",
+        "                        temperature=0.7,\n",
+        "                    )[0],\n",
+        "                    skip_special_tokens=True,\n",
+        "                )\n",
+        "                print(f\"\\n--- sample step={step} i={i} ---\\n\", out[-400:])\n",
+        "        return control\n",
+        "\n",
+        "\n",
         "trainer = GRPOTrainer(\n",
         "    model=model,\n",
         "    processing_class=tokenizer,\n",
         "    reward_funcs=reward_fn,\n",
         "    args=config,\n",
         "    train_dataset=train_ds,\n",
+        "    callbacks=[SampleCallback(every_steps=25, n=2)],\n",
         ")\n",
         "\n",
-        "trainer.train()\n"
+        "trainer.train()\n",
+        ""
       ],
       "execution_count": null,
       "outputs": [],