Route rewards through OpenEnv API instead of local computation

- Add Next.js rewrite proxy: /api/env/* → localhost:8000 (OpenEnv backend)
- Replace local shape_match_reward with openenv_reward that calls:
POST /api/env/reset (start episode for task)
POST /api/env/step (submit FOLD, get reward from environment)
- Add OptigamiEnvClient class for clean API interaction
- Fetch task list from environment instead of hardcoding
- Test cell verifies full reset→step→reward loop before training
- valid_fold_reward stays local (just JSON parsing, no simulation)

The training loop now properly uses the OpenEnv environment for
reward computation, satisfying the submission requirement.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

next.config.ts

@@ -20,6 +20,14 @@ const nextConfig: NextConfig = {
     ],
   },
   output: 'standalone',
+  async rewrites() {
+    return [
+      {
+        source: '/api/env/:path*',
+        destination: 'http://localhost:8000/:path*',
+      },
+    ];
+  },
   transpilePackages: ['motion'],
   webpack: (config, {dev}) => {
     // HMR is disabled in AI Studio via DISABLE_HMR env var.

training/train_grpo.ipynb

@@ -4,12 +4,12 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# Optigami — GRPO Training for Origami Crease Pattern Generation\n",
+    "# Optigami — GRPO Training with OpenEnv\n",
     "\n",
     "Train an LLM to generate FOLD-format crease patterns using **GRPO** with **Unsloth** + **TRL**.\n",
     "\n",
-    "The model receives diverse origami descriptions and must learn to generate valid crease patterns\n",
-    "that fold into the target shapes — scored by structural validity + shape similarity.\n",
+    "**Rewards come from the OpenEnv environment** deployed on HF Spaces — the training loop\n",
+    "calls `POST /reset` and `POST /step` on the live environment to get simulation results.\n",
     "\n",
     "**Environment**: [openenv-community/optigami_](https://huggingface.co/spaces/openenv-community/optigami_) (OpenEnv 0.2.1)\n",
     "\n",
@@ -45,159 +45,19 @@
     "elif importlib.util.find_spec(\"unsloth\") is None:\n",
     "    !pip install -qqq unsloth trackio\n",
     "!pip install --upgrade --no-deps transformers==4.56.2 tokenizers trl==0.22.2 unsloth unsloth_zoo\n",
-    "!pip install -qqq scipy datasets"
+    "!pip install -qqq scipy datasets requests"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 2. Origami Engine (Inlined)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import json, re, random\n",
-    "from collections import defaultdict\n",
-    "from dataclasses import dataclass\n",
-    "from typing import Any\n",
-    "\n",
-    "import numpy as np\n",
-    "from scipy.spatial.distance import cdist\n",
-    "from scipy.spatial.transform import Rotation\n",
-    "\n",
-    "\n",
-    "def validate_fold(fold_data):\n",
-    "    for key in (\"vertices_coords\", \"edges_vertices\", \"edges_assignment\"):\n",
-    "        if key not in fold_data: return False, f\"Missing: {key}\"\n",
-    "    verts, edges, asgn = fold_data[\"vertices_coords\"], fold_data[\"edges_vertices\"], fold_data[\"edges_assignment\"]\n",
-    "    if len(verts) < 3: return False, \"<3 verts\"\n",
-    "    if len(edges) < 3: return False, \"<3 edges\"\n",
-    "    if len(edges) != len(asgn): return False, \"len mismatch\"\n",
-    "    if \"edges_foldAngle\" in fold_data and len(fold_data[\"edges_foldAngle\"]) != len(edges): return False, \"angle len\"\n",
-    "    nv = len(verts)\n",
-    "    for i, v in enumerate(verts):\n",
-    "        if not isinstance(v, (list, tuple)) or len(v) < 2: return False, f\"vert {i}\"\n",
-    "    for i, e in enumerate(edges):\n",
-    "        if not isinstance(e, (list, tuple)) or len(e) != 2: return False, f\"edge {i}\"\n",
-    "        if e[0]<0 or e[0]>=nv or e[1]<0 or e[1]>=nv: return False, f\"edge {i} range\"\n",
-    "        if e[0]==e[1]: return False, f\"edge {i} degen\"\n",
-    "    for i, a in enumerate(asgn):\n",
-    "        if a not in {\"M\",\"V\",\"B\",\"F\",\"U\",\"C\"}: return False, f\"asgn {i}\"\n",
-    "    if not any(a in (\"M\",\"V\") for a in asgn): return False, \"no folds\"\n",
-    "    if not any(a==\"B\" for a in asgn): return False, \"no boundary\"\n",
-    "    return True, \"\"\n",
-    "\n",
-    "def parse_fold(fd):\n",
-    "    verts = fd[\"vertices_coords\"]\n",
-    "    vertices = np.zeros((len(verts),3), dtype=np.float64)\n",
-    "    for i,v in enumerate(verts): vertices[i,0]=v[0]; vertices[i,1]=v[1]; (len(v)>2 and setattr(vertices, '__setitem__', None) is None) or (vertices.__setitem__((i,2), v[2]) if len(v)>2 else None)\n",
-    "    # Fix the above — simple approach\n",
-    "    vertices = np.zeros((len(verts),3), dtype=np.float64)\n",
-    "    for i,v in enumerate(verts):\n",
-    "        vertices[i,0]=v[0]; vertices[i,1]=v[1]\n",
-    "        if len(v)>2: vertices[i,2]=v[2]\n",
-    "    edges = np.array(fd[\"edges_vertices\"], dtype=np.int32)\n",
-    "    asgn = list(fd[\"edges_assignment\"])\n",
-    "    if \"edges_foldAngle\" in fd: fa = np.array(fd[\"edges_foldAngle\"], dtype=np.float64)\n",
-    "    else:\n",
-    "        fa = np.zeros(len(edges), dtype=np.float64)\n",
-    "        for i,a in enumerate(asgn):\n",
-    "            if a==\"V\": fa[i]=180.0\n",
-    "            elif a==\"M\": fa[i]=-180.0\n",
-    "    fa_rad = np.radians(fa)\n",
-    "    if \"faces_vertices\" in fd:\n",
-    "        tris = []\n",
-    "        for face in fd[\"faces_vertices\"]:\n",
-    "            for i in range(1,len(face)-1): tris.append([face[0],face[i],face[i+1]])\n",
-    "        faces = np.array(tris, dtype=np.int32) if tris else np.zeros((0,3), dtype=np.int32)\n",
-    "    else:\n",
-    "        adj=defaultdict(set)\n",
-    "        for v1,v2 in edges: adj[v1].add(v2); adj[v2].add(v1)\n",
-    "        ts=set()\n",
-    "        for v1,v2 in edges:\n",
-    "            for v3 in adj[v1]&adj[v2]: ts.add(tuple(sorted([v1,v2,v3])))\n",
-    "        faces = np.array(list(ts), dtype=np.int32) if ts else np.zeros((0,3), dtype=np.int32)\n",
-    "    return {\"vertices\":vertices,\"edges\":edges,\"assignments\":asgn,\"fold_angles\":fa_rad,\"faces\":faces}\n",
-    "\n",
-    "@dataclass\n",
-    "class SimResult:\n",
-    "    positions: np.ndarray; converged: bool; steps_taken: int; max_strain: float; total_energy: float\n",
-    "\n",
-    "def simulate(fold_data, crease_percent=1.0):\n",
-    "    p = parse_fold(fold_data)\n",
-    "    flat=p[\"vertices\"].copy(); edges=p[\"edges\"]; asgn=p[\"assignments\"]; fa=p[\"fold_angles\"]; faces=p[\"faces\"]; pos=flat.copy()\n",
-    "    if len(faces)==0: return SimResult(pos,True,0,0.,0.)\n",
-    "    fadj={}\n",
-    "    for fi,face in enumerate(faces):\n",
-    "        for j in range(len(face)):\n",
-    "            v1,v2=int(face[j]),int(face[(j+1)%len(face)]); k=(min(v1,v2),max(v1,v2)); fadj.setdefault(k,[]).append(fi)\n",
-    "    cm={}\n",
-    "    for i,(v1,v2) in enumerate(edges):\n",
-    "        k=(min(int(v1),int(v2)),max(int(v1),int(v2)))\n",
-    "        if asgn[i] in (\"M\",\"V\"): cm[k]=fa[i]*crease_percent\n",
-    "    nf=len(faces); fR=[None]*nf; ft=[None]*nf; fR[0]=np.eye(3); ft[0]=np.zeros(3)\n",
-    "    vis=[False]*nf; vis[0]=True; placed=set(int(vi) for vi in faces[0])\n",
-    "    q=[0]\n",
-    "    while q:\n",
-    "        fi=q.pop(0); face=faces[fi]\n",
-    "        for j in range(len(face)):\n",
-    "            v1,v2=int(face[j]),int(face[(j+1)%len(face)]); ek=(min(v1,v2),max(v1,v2))\n",
-    "            for fj in fadj.get(ek,[]):\n",
-    "                if vis[fj]: continue\n",
-    "                vis[fj]=True; q.append(fj); ang=cm.get(ek,0.0)\n",
-    "                if abs(ang)>1e-10:\n",
-    "                    p1=pos[v1].copy(); ax=pos[v2]-p1; al=np.linalg.norm(ax)\n",
-    "                    fr=Rotation.from_rotvec(ang*ax/al).as_matrix() if al>1e-12 else np.eye(3)\n",
-    "                    fR[fj]=fr@fR[fi]; ft[fj]=fr@(ft[fi]-p1)+p1\n",
-    "                else: fR[fj]=fR[fi].copy(); ft[fj]=ft[fi].copy()\n",
-    "                for vi in faces[fj]:\n",
-    "                    vi2=int(vi)\n",
-    "                    if vi2 not in placed: pos[vi2]=fR[fj]@flat[vi2]+ft[fj]; placed.add(vi2)\n",
-    "    ms=0.0\n",
-    "    for v1,v2 in edges:\n",
-    "        r=np.linalg.norm(flat[v2]-flat[v1]); c=np.linalg.norm(pos[v2]-pos[v1])\n",
-    "        if r>1e-12: ms=max(ms,abs(c-r)/r)\n",
-    "    return SimResult(pos,True,1,ms,0.)\n",
-    "\n",
-    "def _rotations():\n",
-    "    rs=[np.eye(3)]\n",
-    "    for k in range(1,4):\n",
-    "        a=k*np.pi/2; c,s=np.cos(a),np.sin(a)\n",
-    "        rs.append(np.array([[c,-s,0],[s,c,0],[0,0,1]]))\n",
-    "    for k in range(1,4):\n",
-    "        a=k*np.pi/2; c,s=np.cos(a),np.sin(a)\n",
-    "        rs.append(np.array([[1,0,0],[0,c,-s],[0,s,c]]))\n",
-    "    for k in range(1,4):\n",
-    "        a=k*np.pi/2; c,s=np.cos(a),np.sin(a)\n",
-    "        rs.append(np.array([[c,0,s],[0,1,0],[-s,0,c]]))\n",
-    "    rs+=[np.diag([-1.,1.,1.]),np.diag([1.,-1.,1.]),np.diag([1.,1.,-1.])]\n",
-    "    return rs\n",
-    "\n",
-    "def compute_shape_match(pred, tgt):\n",
-    "    if len(pred)==0 or len(tgt)==0: return 0.0\n",
-    "    pc=pred-pred.mean(0); tc=tgt-tgt.mean(0); best=0.0\n",
-    "    for r in _rotations():\n",
-    "        rot=pc@r.T; d=cdist(rot,tc); ch=(d.min(1).mean()+d.min(0).mean())/2\n",
-    "        ap=np.vstack([rot,tc]); dg=np.linalg.norm(ap.max(0)-ap.min(0))\n",
-    "        sc=max(0.,1.-ch/dg) if dg>1e-12 else (1. if ch<1e-12 else 0.)\n",
-    "        best=max(best,sc)\n",
-    "    return best\n",
-    "\n",
-    "print(\"Engine loaded.\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## 3. Task Definitions\n",
+    "## 2. OpenEnv Client\n",
     "\n",
-    "Diverse origami tasks — the model must generalize across different shapes, not memorize one."
+    "Connect to the Optigami OpenEnv environment deployed on HF Spaces.\n",
+    "The environment handles fold simulation and reward computation via its API:\n",
+    "- `POST /api/env/reset` — start a new episode for a task\n",
+    "- `POST /api/env/step` — submit a FOLD crease pattern, get back reward + shape similarity"
    ]
   },
   {
@@ -206,104 +66,75 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "TASKS = {\n",
-    "    \"triangle\": {\n",
-    "        \"name\": \"triangle\",\n",
-    "        \"description\": \"Fold the paper in half diagonally to make a triangle\",\n",
-    "        \"difficulty\": 1, \"paper\": {\"width\": 1.0, \"height\": 1.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[0,0],[1,0],[1,1],[0,1]],\n",
-    "            \"edges_vertices\": [[0,1],[1,2],[2,3],[3,0],[0,2]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"V\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,180],\n",
-    "            \"faces_vertices\": [[0,1,2],[0,2,3]],\n",
-    "        },\n",
-    "    },\n",
-    "    \"half_fold\": {\n",
-    "        \"name\": \"half_fold\",\n",
-    "        \"description\": \"Fold the paper in half horizontally along the middle\",\n",
-    "        \"difficulty\": 1, \"paper\": {\"width\": 1.0, \"height\": 1.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[0,0],[1,0],[1,1],[0,1],[0,0.5],[1,0.5]],\n",
-    "            \"edges_vertices\": [[0,1],[1,5],[5,2],[2,3],[3,4],[4,0],[4,5]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"V\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,0,0,180],\n",
-    "            \"faces_vertices\": [[0,1,5,4],[4,5,2,3]],\n",
-    "        },\n",
-    "    },\n",
-    "    \"quarter_fold\": {\n",
-    "        \"name\": \"quarter_fold\",\n",
-    "        \"description\": \"Fold the paper into quarters with two perpendicular creases through the center\",\n",
-    "        \"difficulty\": 2, \"paper\": {\"width\": 1.0, \"height\": 1.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[0,0],[0.5,0],[1,0],[0,0.5],[0.5,0.5],[1,0.5],[0,1],[0.5,1],[1,1]],\n",
-    "            \"edges_vertices\": [[0,1],[1,2],[2,5],[5,8],[8,7],[7,6],[6,3],[3,0],[1,4],[4,7],[3,4],[4,5]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"V\",\"V\",\"V\",\"V\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,0,0,0,0,180,180,180,180],\n",
-    "            \"faces_vertices\": [[0,1,4,3],[1,2,5,4],[3,4,7,6],[4,5,8,7]],\n",
-    "        },\n",
-    "    },\n",
-    "    \"accordion\": {\n",
-    "        \"name\": \"accordion\",\n",
-    "        \"description\": \"Make a zig-zag accordion fold with alternating mountain and valley creases like a paper fan\",\n",
-    "        \"difficulty\": 2, \"paper\": {\"width\": 2.0, \"height\": 2.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[-1,1],[-0.5,1],[0,1],[0.5,1],[1,1],[-1,-1],[-0.5,-1],[0,-1],[0.5,-1],[1,-1]],\n",
-    "            \"edges_vertices\": [[0,1],[1,2],[2,3],[3,4],[5,6],[6,7],[7,8],[8,9],[0,5],[4,9],[1,6],[2,7],[3,8]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"V\",\"M\",\"V\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,0,0,0,0,0,0,180,-180,180],\n",
-    "            \"faces_vertices\": [[0,1,6,5],[1,2,7,6],[2,3,8,7],[3,4,9,8]],\n",
-    "        },\n",
-    "    },\n",
-    "    \"waterbomb\": {\n",
-    "        \"name\": \"waterbomb\",\n",
-    "        \"description\": \"Create a waterbomb base with valley folds on diagonals and mountain folds on midlines of a square\",\n",
-    "        \"difficulty\": 3, \"paper\": {\"width\": 2.0, \"height\": 2.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[-1,1],[0,1],[1,1],[-1,0],[0,0],[1,0],[-1,-1],[0,-1],[1,-1]],\n",
-    "            \"edges_vertices\": [[0,1],[1,2],[2,5],[5,8],[8,7],[7,6],[6,3],[3,0],[0,4],[2,4],[6,4],[8,4],[1,4],[3,4],[5,4],[7,4]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"V\",\"V\",\"V\",\"V\",\"M\",\"M\",\"M\",\"M\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,0,0,0,0,180,180,180,180,-180,-180,-180,-180],\n",
-    "            \"faces_vertices\": [[0,1,4],[1,2,4],[2,5,4],[5,8,4],[8,7,4],[7,6,4],[6,3,4],[3,0,4]],\n",
-    "        },\n",
-    "    },\n",
-    "    \"miura_ori\": {\n",
-    "        \"name\": \"miura_ori\",\n",
-    "        \"description\": \"Create a Miura-ori tessellation with offset zigzag vertices that folds flat in one motion\",\n",
-    "        \"difficulty\": 3, \"paper\": {\"width\": 2.0, \"height\": 2.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[-1,1],[0,1.2],[1,1],[-1,0],[0,0.2],[1,0],[-1,-1],[0,-0.8],[1,-1]],\n",
-    "            \"edges_vertices\": [[0,1],[1,2],[2,5],[5,8],[8,7],[7,6],[6,3],[3,0],[1,4],[4,7],[3,4],[4,5]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"M\",\"M\",\"V\",\"M\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,0,0,0,0,-180,-180,180,-180],\n",
-    "            \"faces_vertices\": [[0,1,4,3],[1,2,5,4],[3,4,7,6],[4,5,8,7]],\n",
-    "        },\n",
-    "    },\n",
-    "    \"letter_fold\": {\n",
-    "        \"name\": \"letter_fold\",\n",
-    "        \"description\": \"Tri-fold the paper into thirds like a letter envelope with two parallel horizontal creases\",\n",
-    "        \"difficulty\": 2, \"paper\": {\"width\": 1.0, \"height\": 1.0},\n",
-    "        \"target_fold\": {\n",
-    "            \"vertices_coords\": [[0,0],[1,0],[0,0.333],[1,0.333],[0,0.667],[1,0.667],[0,1],[1,1]],\n",
-    "            \"edges_vertices\": [[0,1],[1,3],[3,5],[5,7],[7,6],[6,4],[4,2],[2,0],[2,3],[4,5]],\n",
-    "            \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"B\",\"V\",\"M\"],\n",
-    "            \"edges_foldAngle\": [0,0,0,0,0,0,0,0,180,-180],\n",
-    "            \"faces_vertices\": [[0,1,3,2],[2,3,5,4],[4,5,7,6]],\n",
-    "        },\n",
-    "    },\n",
-    "}\n",
+    "import requests\n",
+    "import json\n",
+    "import re\n",
+    "\n",
+    "# OpenEnv environment URL — deployed on HF Spaces\n",
+    "OPENENV_URL = \"https://openenv-community-optigami-c92c300.hf.space/api/env\"\n",
+    "\n",
     "\n",
-    "def get_task(name): return TASKS[name]\n",
-    "print(f\"Tasks ({len(TASKS)}): {list(TASKS.keys())}\")"
+    "class OptigamiEnvClient:\n",
+    "    \"\"\"Client for the Optigami OpenEnv environment on HF Spaces.\"\"\"\n",
+    "\n",
+    "    def __init__(self, base_url=OPENENV_URL):\n",
+    "        self.base_url = base_url\n",
+    "        self.session = requests.Session()\n",
+    "        self.session.headers.update({\"Content-Type\": \"application/json\"})\n",
+    "\n",
+    "    def health(self):\n",
+    "        \"\"\"Check if the environment is reachable.\"\"\"\n",
+    "        try:\n",
+    "            r = self.session.get(f\"{self.base_url}/health\", timeout=10)\n",
+    "            return r.status_code == 200\n",
+    "        except:\n",
+    "            return False\n",
+    "\n",
+    "    def reset(self, task_name=\"triangle\"):\n",
+    "        \"\"\"Reset environment for a task. Returns observation dict.\"\"\"\n",
+    "        r = self.session.post(\n",
+    "            f\"{self.base_url}/reset\",\n",
+    "            json={\"task_name\": task_name},\n",
+    "            timeout=30,\n",
+    "        )\n",
+    "        r.raise_for_status()\n",
+    "        return r.json()\n",
+    "\n",
+    "    def step(self, fold_data):\n",
+    "        \"\"\"Submit a FOLD crease pattern. Returns observation with reward.\"\"\"\n",
+    "        r = self.session.post(\n",
+    "            f\"{self.base_url}/step\",\n",
+    "            json={\"action\": {\"fold_data\": fold_data}},\n",
+    "            timeout=30,\n",
+    "        )\n",
+    "        r.raise_for_status()\n",
+    "        return r.json()\n",
+    "\n",
+    "    def get_tasks(self):\n",
+    "        \"\"\"Fetch available tasks from the environment.\"\"\"\n",
+    "        r = self.session.get(f\"{self.base_url.replace('/api/env', '')}/tasks\", timeout=10)\n",
+    "        r.raise_for_status()\n",
+    "        return r.json()\n",
+    "\n",
+    "\n",
+    "# Initialize client and verify connection\n",
+    "env = OptigamiEnvClient()\n",
+    "if env.health():\n",
+    "    print(f\"Connected to OpenEnv at {OPENENV_URL}\")\n",
+    "    tasks = env.get_tasks()\n",
+    "    print(f\"Available tasks: {list(tasks.keys())}\")\n",
+    "else:\n",
+    "    print(f\"WARNING: Cannot reach OpenEnv at {OPENENV_URL}\")\n",
+    "    print(\"Make sure the HF Space is running!\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 4. Prompt Template & Multi-Task Dataset\n",
+    "## 3. Test the OpenEnv Environment\n",
     "\n",
-    "Each row in the dataset is a **different task** — the model must generalize across shapes."
+    "Verify the full loop: reset → step → get reward."
    ]
   },
   {
@@ -312,62 +143,37 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from datasets import Dataset\n",
-    "\n",
-    "PROMPT_TEMPLATE = \"\"\"You are an origami designer. Generate a FOLD-format crease pattern\n",
-    "that, when folded, produces the target shape described below.\n",
-    "\n",
-    "Target: {description}\n",
-    "Paper size: {width} x {height}\n",
-    "\n",
-    "Output a JSON object with these exact fields:\n",
-    "- vertices_coords: [[x, y], ...] — 2D positions on the flat paper\n",
-    "- edges_vertices: [[v1, v2], ...] — pairs of vertex indices forming edges\n",
-    "- edges_assignment: [\"B\"|\"M\"|\"V\", ...] — B=boundary, M=mountain fold, V=valley fold\n",
-    "- edges_foldAngle: [angle, ...] — fold angles in degrees (M: negative, V: positive, B: 0)\n",
-    "\n",
-    "Rules:\n",
-    "- Boundary edges (B) must outline the paper rectangle\n",
-    "- At least one fold crease (M or V) must exist\n",
-    "- Mountain fold angles are negative (-180 to 0)\n",
-    "- Valley fold angles are positive (0 to 180)\n",
-    "- All vertex indices in edges must be valid (0 to N-1)\n",
-    "\n",
-    "Output ONLY the JSON object wrapped in ```json ... ``` markers.\"\"\"\n",
-    "\n",
-    "\n",
-    "def build_prompt(task):\n",
-    "    return PROMPT_TEMPLATE.format(\n",
-    "        description=task[\"description\"],\n",
-    "        width=task[\"paper\"][\"width\"],\n",
-    "        height=task[\"paper\"][\"height\"],\n",
-    "    )\n",
-    "\n",
-    "\n",
-    "# Build multi-task dataset — rotate through all tasks\n",
-    "task_names = list(TASKS.keys())\n",
-    "rows = []\n",
-    "for i in range(1000):\n",
-    "    tn = task_names[i % len(task_names)]\n",
-    "    t = TASKS[tn]\n",
-    "    rows.append({\n",
-    "        \"prompt\": [{\"role\": \"user\", \"content\": build_prompt(t)}],\n",
-    "        \"task_name\": tn,\n",
-    "        \"answer\": 0,\n",
-    "    })\n",
-    "\n",
-    "grpo_dataset = Dataset.from_list(rows)\n",
-    "print(f\"Dataset: {len(grpo_dataset)} rows across {len(task_names)} tasks\")\n",
-    "print(f\"Task distribution: {dict((tn, sum(1 for r in rows if r['task_name']==tn)) for tn in task_names)}\")"
+    "# Reset for triangle task\n",
+    "obs = env.reset(task_name=\"triangle\")\n",
+    "print(\"Reset observation:\")\n",
+    "print(f\"  Task: {obs['observation']['task']}\")\n",
+    "print(f\"  Done: {obs['observation']['done']}\")\n",
+    "print(f\"  Target positions: {len(obs['observation']['target_positions'])} vertices\")\n",
+    "\n",
+    "# Submit the known-correct triangle fold\n",
+    "correct_fold = {\n",
+    "    \"vertices_coords\": [[0,0],[1,0],[1,1],[0,1]],\n",
+    "    \"edges_vertices\": [[0,1],[1,2],[2,3],[3,0],[0,2]],\n",
+    "    \"edges_assignment\": [\"B\",\"B\",\"B\",\"B\",\"V\"],\n",
+    "    \"edges_foldAngle\": [0,0,0,0,180],\n",
+    "    \"faces_vertices\": [[0,1,2],[0,2,3]],\n",
+    "}\n",
+    "result = env.step(correct_fold)\n",
+    "print(f\"\\nStep result:\")\n",
+    "print(f\"  Reward: {result['reward']}\")\n",
+    "print(f\"  Shape similarity: {result['observation']['shape_similarity']}\")\n",
+    "print(f\"  Done: {result['observation']['done']}\")\n",
+    "print(f\"  Is stable: {result['observation']['is_stable']}\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 5. Reward Functions\n",
+    "## 4. Reward Functions (via OpenEnv API)\n",
     "\n",
-    "The reward functions read `task_name` from the dataset row to know which target to compare against."
+    "- **`valid_fold_reward`**: Local JSON validation (fast, no API call needed)\n",
+    "- **`openenv_reward`**: Calls the OpenEnv environment to simulate the fold and get the real reward"
    ]
   },
   {
@@ -379,6 +185,7 @@
     "PRINTER = 0\n",
     "\n",
     "def extract_fold_json(response):\n",
+    "    \"\"\"Extract FOLD JSON from LLM response text.\"\"\"\n",
     "    m = re.search(r\"```(?:json)?\\s*(\\{.*?\\})\\s*```\", response, re.DOTALL)\n",
     "    if m:\n",
     "        try: return json.loads(m.group(1))\n",
@@ -395,52 +202,95 @@
     "\n",
     "\n",
     "def valid_fold_reward(completions, **kwargs):\n",
-    "    \"\"\"Reward 1: +1.0 valid FOLD, -0.5 bad structure, -2.0 unparseable.\"\"\"\n",
+    "    \"\"\"Reward 1 (local): +1.0 valid FOLD structure, -0.5 bad structure, -2.0 unparseable.\"\"\"\n",
+    "    REQUIRED = {\"vertices_coords\", \"edges_vertices\", \"edges_assignment\"}\n",
     "    scores = []\n",
     "    for c in completions:\n",
     "        fold = extract_fold_json(c[0][\"content\"])\n",
-    "        if fold is None: scores.append(-2.0); continue\n",
-    "        ok, _ = validate_fold(fold)\n",
-    "        scores.append(1.0 if ok else -0.5)\n",
+    "        if fold is None:\n",
+    "            scores.append(-2.0)\n",
+    "            continue\n",
+    "        # Basic structural checks\n",
+    "        if not REQUIRED.issubset(fold.keys()):\n",
+    "            scores.append(-0.5); continue\n",
+    "        verts = fold[\"vertices_coords\"]\n",
+    "        edges = fold[\"edges_vertices\"]\n",
+    "        asgn = fold[\"edges_assignment\"]\n",
+    "        if len(verts) < 3 or len(edges) < 3 or len(edges) != len(asgn):\n",
+    "            scores.append(-0.5); continue\n",
+    "        if not any(a in (\"M\",\"V\") for a in asgn):\n",
+    "            scores.append(-0.5); continue\n",
+    "        if not any(a == \"B\" for a in asgn):\n",
+    "            scores.append(-0.5); continue\n",
+    "        scores.append(1.0)\n",
     "    return scores\n",
     "\n",
     "\n",
-    "def shape_match_reward(completions, task_name, **kwargs):\n",
-    "    \"\"\"Reward 2: similarity*20 for matching target shape. Uses task_name from dataset.\"\"\"\n",
+    "def openenv_reward(completions, task_name, **kwargs):\n",
+    "    \"\"\"Reward 2 (OpenEnv API): Submit fold to environment, get simulation reward.\n",
+    "\n",
+    "    Calls POST /reset and POST /step on the HF Space OpenEnv environment.\n",
+    "    The environment runs the fold simulation and computes shape similarity.\n",
+    "    \"\"\"\n",
     "    global PRINTER\n",
-    "    # task_name comes as a list (one per completion in the batch)\n",
-    "    if isinstance(task_name, list):\n",
-    "        tn = task_name[0]  # All completions in a group share the same prompt/task\n",
-    "    else:\n",
-    "        tn = task_name\n",
-    "    task = get_task(tn)\n",
-    "    try: tgt = simulate(task[\"target_fold\"]).positions\n",
-    "    except: return [0.0]*len(completions)\n",
+    "    # task_name comes as a list from the dataset\n",
+    "    tn = task_name[0] if isinstance(task_name, list) else task_name\n",
     "\n",
     "    scores = []\n",
     "    for c in completions:\n",
     "        resp = c[0][\"content\"]\n",
+    "\n",
+    "        # Periodic logging\n",
     "        if PRINTER % 10 == 0:\n",
     "            print(f\"\\n--- [{tn}] Sample {PRINTER} ---\\n{resp[:300]}\")\n",
     "        PRINTER += 1\n",
+    "\n",
+    "        # Parse the FOLD JSON from the LLM response\n",
     "        fold = extract_fold_json(resp)\n",
-    "        if fold is None: scores.append(-2.0); continue\n",
-    "        ok, _ = validate_fold(fold)\n",
-    "        if not ok: scores.append(-1.0); continue\n",
+    "        if fold is None:\n",
+    "            scores.append(-2.0)\n",
+    "            continue\n",
+    "\n",
     "        try:\n",
-    "            sim = compute_shape_match(simulate(fold).positions, tgt)\n",
-    "            scores.append(sim * 20.0)\n",
-    "        except: scores.append(-1.0)\n",
+    "            # Reset environment for this task\n",
+    "            env.reset(task_name=tn)\n",
+    "\n",
+    "            # Submit the fold to OpenEnv — environment simulates and scores it\n",
+    "            result = env.step(fold)\n",
+    "\n",
+    "            # Get reward from the environment\n",
+    "            reward = result.get(\"reward\", None)\n",
+    "            if reward is not None:\n",
+    "                scores.append(float(reward))\n",
+    "            else:\n",
+    "                # Fallback: extract from observation\n",
+    "                obs = result.get(\"observation\", {})\n",
+    "                if obs.get(\"error\"):\n",
+    "                    scores.append(-2.0)\n",
+    "                else:\n",
+    "                    sim = obs.get(\"shape_similarity\", 0.0)\n",
+    "                    scores.append(float(sim) * 20.0)\n",
+    "\n",
+    "        except requests.exceptions.RequestException as e:\n",
+    "            print(f\"OpenEnv API error: {e}\")\n",
+    "            scores.append(-1.0)\n",
+    "        except Exception as e:\n",
+    "            print(f\"Reward error: {e}\")\n",
+    "            scores.append(-1.0)\n",
+    "\n",
     "    return scores\n",
     "\n",
-    "print(\"Reward functions ready.\")"
+    "\n",
+    "print(\"Reward functions ready (valid_fold=local, openenv_reward=API).\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 6. Quick Engine Test"
+    "## 5. Prompt Template & Multi-Task Dataset\n",
+    "\n",
+    "Each row is a different task — the model must generalize across shapes."
    ]
   },
   {
@@ -449,18 +299,63 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "for tn in TASKS:\n",
-    "    t = TASKS[tn]\n",
-    "    r = simulate(t[\"target_fold\"])\n",
-    "    ss = compute_shape_match(r.positions, r.positions)\n",
-    "    print(f\"{tn:15s} | verts={len(t['target_fold']['vertices_coords']):2d} | self-sim={ss:.3f} | strain={r.max_strain:.4f}\")"
+    "from datasets import Dataset\n",
+    "\n",
+    "PROMPT_TEMPLATE = \"\"\"You are an origami designer. Generate a FOLD-format crease pattern\n",
+    "that, when folded, produces the target shape described below.\n",
+    "\n",
+    "Target: {description}\n",
+    "Paper size: {width} x {height}\n",
+    "\n",
+    "Output a JSON object with these exact fields:\n",
+    "- vertices_coords: [[x, y], ...] — 2D positions on the flat paper\n",
+    "- edges_vertices: [[v1, v2], ...] — pairs of vertex indices forming edges\n",
+    "- edges_assignment: [\"B\"|\"M\"|\"V\", ...] — B=boundary, M=mountain fold, V=valley fold\n",
+    "- edges_foldAngle: [angle, ...] — fold angles in degrees (M: negative, V: positive, B: 0)\n",
+    "\n",
+    "Rules:\n",
+    "- Boundary edges (B) must outline the paper rectangle\n",
+    "- At least one fold crease (M or V) must exist\n",
+    "- Mountain fold angles are negative (-180 to 0)\n",
+    "- Valley fold angles are positive (0 to 180)\n",
+    "- All vertex indices in edges must be valid (0 to N-1)\n",
+    "\n",
+    "Output ONLY the JSON object wrapped in ```json ... ``` markers.\"\"\"\n",
+    "\n",
+    "\n",
+    "def build_prompt(task):\n",
+    "    return PROMPT_TEMPLATE.format(\n",
+    "        description=task[\"description\"],\n",
+    "        width=task[\"paper\"][\"width\"],\n",
+    "        height=task[\"paper\"][\"height\"],\n",
+    "    )\n",
+    "\n",
+    "\n",
+    "# Fetch tasks from the OpenEnv environment\n",
+    "env_tasks = env.get_tasks()\n",
+    "task_names = list(env_tasks.keys())\n",
+    "print(f\"Tasks from OpenEnv: {task_names}\")\n",
+    "\n",
+    "# Build multi-task dataset\n",
+    "rows = []\n",
+    "for i in range(1000):\n",
+    "    tn = task_names[i % len(task_names)]\n",
+    "    t = env_tasks[tn]\n",
+    "    rows.append({\n",
+    "        \"prompt\": [{\"role\": \"user\", \"content\": build_prompt(t)}],\n",
+    "        \"task_name\": tn,\n",
+    "        \"answer\": 0,\n",
+    "    })\n",
+    "\n",
+    "grpo_dataset = Dataset.from_list(rows)\n",
+    "print(f\"Dataset: {len(grpo_dataset)} rows across {len(task_names)} tasks\")"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 7. Configuration"
+    "## 6. Configuration"
    ]
   },
   {
@@ -494,7 +389,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 8. Load Model + LoRA"
+    "## 7. Load Model + LoRA"
    ]
   },
   {
@@ -523,13 +418,15 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 9. GRPO Training\n",
+    "## 8. GRPO Training\n",
+    "\n",
+    "Each training step:\n",
+    "1. Pick a task prompt (rotates through all 7 origami tasks)\n",
+    "2. Generate N completions (FOLD JSON attempts)\n",
+    "3. Score each: `valid_fold_reward` (local) + `openenv_reward` (calls HF Space API)\n",
+    "4. GRPO computes group-relative advantages and updates the policy\n",
     "\n",
-    "Each step:\n",
-    "1. Pick a task prompt (rotates through triangle, half_fold, quarter_fold, accordion, waterbomb, miura_ori, letter_fold)\n",
-    "2. Generate N completions\n",
-    "3. Score each: valid_fold (+1/-0.5/-2) + shape_match (0-20)\n",
-    "4. GRPO uses within-group variance to update the policy"
+    "The **OpenEnv environment** on HF Spaces handles the fold simulation and shape comparison."
    ]
   },
   {
@@ -562,12 +459,12 @@
     "trainer = GRPOTrainer(\n",
     "    model=model,\n",
     "    processing_class=tokenizer,\n",
-    "    reward_funcs=[valid_fold_reward, shape_match_reward],\n",
+    "    reward_funcs=[valid_fold_reward, openenv_reward],\n",
     "    args=training_args,\n",
     "    train_dataset=grpo_dataset,\n",
     ")\n",
     "\n",
-    "print(\"Starting GRPO training...\")\n",
+    "print(\"Starting GRPO training (rewards from OpenEnv API)...\")\n",
     "trainer.train()\n",
     "print(\"Done!\")"
    ]
@@ -576,7 +473,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 10. Test Across Tasks"
+    "## 9. Test Across Tasks"
    ]
   },
   {
@@ -587,8 +484,8 @@
    "source": [
     "FastLanguageModel.for_inference(model)\n",
     "\n",
-    "for tn in [\"triangle\", \"half_fold\", \"accordion\", \"waterbomb\"]:\n",
-    "    task = get_task(tn)\n",
+    "for tn in list(env_tasks.keys())[:4]:\n",
+    "    task = env_tasks[tn]\n",
     "    prompt = build_prompt(task)\n",
     "    inputs = tokenizer.apply_chat_template(\n",
     "        [{\"role\": \"user\", \"content\": prompt}],\n",
@@ -599,12 +496,14 @@
     "\n",
     "    fold = extract_fold_json(response)\n",
     "    if fold:\n",
-    "        ok, err = validate_fold(fold)\n",
-    "        if ok:\n",
-    "            sim = compute_shape_match(simulate(fold).positions, simulate(task[\"target_fold\"]).positions)\n",
-    "            print(f\"{tn:15s} | valid | similarity={sim:.3f} | reward={sim*20:.1f}\")\n",
-    "        else:\n",
-    "            print(f\"{tn:15s} | invalid: {err}\")\n",
+    "        try:\n",
+    "            env.reset(task_name=tn)\n",
+    "            result = env.step(fold)\n",
+    "            reward = result.get(\"reward\", 0)\n",
+    "            sim = result[\"observation\"].get(\"shape_similarity\", 0)\n",
+    "            print(f\"{tn:15s} | reward={reward:6.2f} | similarity={sim:.3f}\")\n",
+    "        except Exception as e:\n",
+    "            print(f\"{tn:15s} | env error: {e}\")\n",
     "    else:\n",
     "        print(f\"{tn:15s} | no JSON extracted\")"
    ]
@@ -613,7 +512,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 11. Save Model"
+    "## 10. Save Model"
    ]
   },
   {