Add MAX_CONCURRENT_ENVS, sync latest changes

Dockerfile

@@ -9,6 +9,8 @@ RUN pip install --no-cache-dir -r requirements.txt \
 
 COPY . /app
 
+ENV MAX_CONCURRENT_ENVS=16
+
 EXPOSE 8000
 
 CMD ["uvicorn", "origami_server.app:app", "--host", "0.0.0.0", "--port", "8000"]

origami_server/app.py

@@ -16,6 +16,7 @@ app = create_app(
     OrigamiAction,
     OrigamiObservation,
     env_name="origami_env",
+    max_concurrent_envs=int(os.environ.get("MAX_CONCURRENT_ENVS", 1)),
 )
 
 from .tasks import TASKS

tests/test_origami.py

@@ -9,7 +9,7 @@ from origami_server.engine.simulate import simulate
 from origami_server.environment import OrigamiEnvironment
 from origami_server.models import OrigamiAction
 from origami_server.tasks import TASKS, get_task, list_tasks
-from training.reward import extract_fold_json, shape_match, valid_fold
+from training.reward import extract_fold_json, valid_fold
 
 # --- Fixtures ---
 
@@ -221,14 +221,20 @@ class TestRewards:
         assert scores[0] == 1.0
         assert scores[1] == -2.0
 
-    def test_shape_match_reward(self):
-        import json
+    def test_shape_match_via_server(self):
+        """shape_match reward now goes through the server (WebSocket).
+        Test the same flow via TestClient's websocket to verify end-to-end."""
+        from fastapi.testclient import TestClient
 
-        good = [[{"content": json.dumps(TRIANGLE_FOLD)}]]
-        bad = [[{"content": "nope"}]]
-        scores = shape_match(good + bad, task_name="triangle")
-        assert scores[0] == 20.0
-        assert scores[1] == -2.0
+        from origami_server.app import app
+
+        client = TestClient(app)
+        with client.websocket_connect("/ws") as ws:
+            ws.send_json({"type": "reset", "data": {"task_name": "triangle"}})
+            ws.receive_json()
+            ws.send_json({"type": "step", "data": {"fold_data": TRIANGLE_FOLD}})
+            resp = ws.receive_json()
+            assert resp["data"]["reward"] == 20.0
 
 
 # --- API ---

training/reward.py

@@ -1,21 +1,17 @@
 """GRPO reward functions for origami RL training.
 
-Two reward functions (matching the 2048 pattern):
-1. valid_fold: Does the LLM output parse as valid FOLD JSON?
-2. shape_match: Simulate and compare to target shape.
+Follows the OpenEnv 2048 pattern exactly:
+- launch_openenv() spawns/reuses the origami server
+- Reward functions call the server via EnvClient
+- Server computes simulation + shape matching, returns reward
+
+These functions are also importable for use in notebooks.
 """
 
 import json
 import re
 from typing import Any
 
-import numpy as np
-
-from origami_server.engine.fold_parser import validate_fold
-from origami_server.engine.shape_match import compute_shape_match
-from origami_server.engine.simulate import simulate
-from origami_server.tasks import get_task
-
 
 def extract_fold_json(response: str) -> dict | None:
     """Extract FOLD JSON from LLM response text.
@@ -55,6 +51,8 @@ def valid_fold(completions: list, **kwargs: Any) -> list[float]:
     +1.0  valid FOLD JSON with correct structure
     -0.5  parseable JSON but invalid FOLD structure
     -2.0  not parseable as JSON at all
+
+    Local check — no server needed.
     """
     scores = []
     for completion in completions:
@@ -65,58 +63,35 @@ def valid_fold(completions: list, **kwargs: Any) -> list[float]:
             scores.append(-2.0)
             continue
 
-        is_valid, error = validate_fold(fold_data)
-        if is_valid:
-            scores.append(1.0)
-        else:
+        # Basic structural validation
+        required = {"vertices_coords", "edges_vertices", "edges_assignment"}
+        if not required.issubset(fold_data.keys()):
             scores.append(-0.5)
+            continue
 
-    return scores
-
-
-def shape_match(
-    completions: list,
-    task_name: str = "triangle",
-    **kwargs: Any,
-) -> list[float]:
-    """Reward 2: Simulate the fold and compare to target shape.
-
-    Score = similarity × 20.0 (range: 0 to 20)
-    -1.0  if simulation fails/diverges
-    -2.0  if FOLD data is invalid
-
-    This is the main reward signal — AlphaFold-style shape comparison.
-    """
-    task = get_task(task_name)
-    target_fold = task["target_fold"]
-
-    # Pre-compute target positions
-    try:
-        target_result = simulate(target_fold, crease_percent=1.0)
-        target_positions = target_result.positions
-    except Exception:
-        # Target itself fails — all scores 0
-        return [0.0] * len(completions)
+        verts = fold_data.get("vertices_coords", [])
+        edges = fold_data.get("edges_vertices", [])
+        assigns = fold_data.get("edges_assignment", [])
 
-    scores = []
-    for completion in completions:
-        response = completion[0]["content"]
-        fold_data = extract_fold_json(response)
+        if len(edges) != len(assigns):
+            scores.append(-0.5)
+            continue
 
-        if fold_data is None:
-            scores.append(-2.0)
+        has_fold = any(a in ("M", "V") for a in assigns)
+        has_boundary = any(a == "B" for a in assigns)
+        if not has_fold or not has_boundary:
+            scores.append(-0.5)
             continue
 
-        is_valid, error = validate_fold(fold_data)
-        if not is_valid:
-            scores.append(-1.0)
+        n = len(verts)
+        valid_indices = all(
+            0 <= e[0] < n and 0 <= e[1] < n and e[0] != e[1]
+            for e in edges
+        )
+        if not valid_indices:
+            scores.append(-0.5)
             continue
 
-        try:
-            result = simulate(fold_data, crease_percent=1.0)
-            similarity = compute_shape_match(result.positions, target_positions)
-            scores.append(similarity * 20.0)
-        except Exception:
-            scores.append(-1.0)
+        scores.append(1.0)
 
     return scores

training/train_grpo.py

@@ -1,20 +1,28 @@
 """GRPO training script for origami RL.
 
-Follows the 2048 OpenEnv + Unsloth pattern:
-- LLM generates FOLD JSON crease patterns
-- Two reward functions: valid_fold + shape_match
+Follows the OpenEnv 2048 pattern exactly:
+- Environment runs as a FastAPI server (origami_server.app)
+- Training connects via WebSocket client (OrigamiEnv)
+- Reward functions call the server, never import engine code
 - GRPOTrainer from TRL handles the RL loop
 
-Usage (local/Colab):
+Usage:
+    # 1. Start the environment server first:
+    uvicorn origami_server.app:app --host 0.0.0.0 --port 8000
+
+    # 2. Run training (connects to server):
     python -m training.train_grpo --task triangle --max_steps 600
 
-Usage (Northflank — env vars set in Dockerfile.train):
-    python -m training.train_grpo --task $TASK --model $MODEL --max_steps $MAX_STEPS
+    # Or specify server URL:
+    python -m training.train_grpo --server http://gpu-host:8000
 """
 
 import argparse
+import functools
 import os
 
+import requests
+
 PROMPT_TEMPLATE = """You are an origami designer. Generate a FOLD-format crease pattern
 that, when folded, produces the target shape described below.
 
@@ -49,60 +57,109 @@ def main():
     parser = argparse.ArgumentParser(description="GRPO training for origami RL")
     parser.add_argument("--task", default="triangle", help="Task name")
     parser.add_argument("--max_steps", type=int, default=600)
-    parser.add_argument("--num_generations", type=int, default=4)
-    parser.add_argument("--model", default="Qwen/Qwen2.5-3B-Instruct")
+    parser.add_argument("--num_generations", type=int, default=2)
+    parser.add_argument("--model", default="unsloth/Qwen3-14B")
     parser.add_argument("--lr", type=float, default=2e-4)
+    parser.add_argument("--lora_rank", type=int, default=4)
+    parser.add_argument(
+        "--server", default="http://localhost:8000",
+        help="URL of the origami environment server",
+    )
     args = parser.parse_args()
 
-    # --- These imports are heavy, only load when actually training ---
+    # --- Verify server is running ---
+    print(f"Connecting to environment server at {args.server}...")
+    try:
+        r = requests.get(f"{args.server}/health", timeout=5)
+        assert r.status_code == 200
+        print("Server is healthy.")
+    except Exception as e:
+        print(f"ERROR: Cannot connect to server at {args.server}")
+        print(f"Start it first: uvicorn origami_server.app:app --port 8000")
+        raise SystemExit(1)
+
+    # --- Get task info from server ---
+    task = requests.get(f"{args.server}/tasks/{args.task}").json()
+    prompt_text = build_prompt(task)
+    print(f"Task: {task['name']} — {task['description']}")
+
+    # --- Configure reward functions (OpenEnv pattern) ---
+    from client import OrigamiEnv
+    from origami_server.models import OrigamiAction
+    from training.reward import extract_fold_json, valid_fold
+    from unsloth import is_port_open, launch_openenv
+
+    global port, openenv_process
+    port = int(args.server.split(":")[-1]) if ":" in args.server else 8000
+    openenv_process = None
+
+    launch_openenv = functools.partial(
+        launch_openenv,
+        working_directory=os.getcwd(),
+        server="origami_server.app:app",
+        environment={**os.environ, "PYTHONPATH": os.getcwd()},
+        openenv_class=OrigamiEnv,
+    )
+
+    def shape_match_reward(completions, **kwargs):
+        global port, openenv_process
+        scores = []
+        for completion in completions:
+            response = completion[0]["content"]
+            fold_data = extract_fold_json(response)
+            if fold_data is None:
+                scores.append(0.0)
+                continue
+            try:
+                port, openenv_process = launch_openenv(port, openenv_process)
+                openenv_process.reset(task_name=args.task)
+                result = openenv_process.step(OrigamiAction(fold_data=fold_data))
+                scores.append(result.reward if result.reward is not None else 0.0)
+            except TimeoutError:
+                scores.append(-1.0)
+            except Exception:
+                scores.append(-3.0)
+        return scores
+
+    # --- Build dataset (same prompt repeated, like 2048) ---
     from datasets import Dataset
-    from trl import GRPOConfig, GRPOTrainer
 
-    from origami_server.tasks import get_task
-    from training.reward import shape_match, valid_fold
+    dataset = Dataset.from_list(
+        [{"prompt": [{"role": "user", "content": prompt_text}]}] * 1000
+    )
 
-    # Try Unsloth first (CUDA), fall back to HF+PEFT
+    # --- Load model with QLoRA ---
     try:
         from unsloth import FastLanguageModel
         USE_UNSLOTH = True
     except ImportError:
         USE_UNSLOTH = False
 
-    task = get_task(args.task)
-    prompt_text = build_prompt(task)
-
-    # Build dataset (1000 copies of same prompt, like 2048)
-    dataset = Dataset.from_list(
-        [
-            {
-                "prompt": [{"role": "user", "content": prompt_text}],
-                "answer": 0,
-            }
-        ]
-        * 1000
-    )
+    max_seq_length = 768  # FOLD JSON is compact
 
-    # Load model with LoRA
     if USE_UNSLOTH:
+        print(f"Loading {args.model} with Unsloth QLoRA (rank={args.lora_rank})...")
         model, tokenizer = FastLanguageModel.from_pretrained(
             model_name=args.model,
             load_in_4bit=True,
-            max_seq_length=2048,
+            max_seq_length=max_seq_length,
+            offload_embedding=True,  # Needed for 14B on limited VRAM
         )
         model = FastLanguageModel.get_peft_model(
             model,
-            r=8,
+            r=args.lora_rank,
             target_modules=[
                 "q_proj", "k_proj", "v_proj", "o_proj",
                 "gate_proj", "up_proj", "down_proj",
             ],
-            lora_alpha=16,
+            lora_alpha=args.lora_rank * 2,
             use_gradient_checkpointing="unsloth",
+            random_state=3407,
         )
     else:
         import torch
-        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
         from peft import LoraConfig, get_peft_model
+        from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
         bnb_config = BitsAndBytesConfig(
             load_in_4bit=True,
@@ -118,19 +175,23 @@ def main():
             torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
         )
         model = get_peft_model(model, LoraConfig(
-            r=8, lora_alpha=16, task_type="CAUSAL_LM",
-            target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
-                            "gate_proj", "up_proj", "down_proj"],
+            r=args.lora_rank,
+            lora_alpha=args.lora_rank * 2,
+            task_type="CAUSAL_LM",
+            target_modules=[
+                "q_proj", "k_proj", "v_proj", "o_proj",
+                "gate_proj", "up_proj", "down_proj",
+            ],
         ))
 
     if tokenizer.pad_token is None:
         tokenizer.pad_token = tokenizer.eos_token
 
-    # Wrap shape_match to inject task_name
-    def shape_match_reward(completions, **kwargs):
-        return shape_match(completions, task_name=args.task, **kwargs)
+    model.print_trainable_parameters()
+
+    # --- GRPO config (matches 2048 pattern) ---
+    from trl import GRPOConfig, GRPOTrainer
 
-    # GRPO config
     training_args = GRPOConfig(
         temperature=1.0,
         learning_rate=args.lr,
@@ -142,8 +203,8 @@ def main():
         per_device_train_batch_size=1,
         gradient_accumulation_steps=1,
         num_generations=args.num_generations,
-        max_prompt_length=1024,
-        max_completion_length=1024,
+        max_prompt_length=512,
+        max_completion_length=max_seq_length - 512,
         max_steps=args.max_steps,
         save_steps=100,
         output_dir=os.environ.get("OUTPUT_DIR", "outputs"),
@@ -157,9 +218,10 @@ def main():
         train_dataset=dataset,
     )
 
+    print(f"Training: {args.max_steps} steps, {args.num_generations} generations/step")
     trainer.train()
 
-    # Save the LoRA adapter
+    # Save LoRA adapter
     save_path = os.path.join(
         os.environ.get("OUTPUT_DIR", "outputs"),
         f"origami-{args.task}-lora-final",

training/train_origami.ipynb

@@ -3,7 +3,7 @@
   {
    "cell_type": "markdown",
    "id": "p8uwc5bkc4n",
-   "source": "# Origami RL — GRPO Training Notebook\n\nTrain an LLM to generate valid FOLD-format crease patterns that fold into target shapes.\n\n**Pipeline:**\n1. LLM receives a prompt describing a target shape (e.g. \"fold diagonally into a triangle\")\n2. LLM generates a FOLD JSON crease pattern\n3. Physics simulator folds the paper analytically\n4. Reward = shape similarity (chamfer distance) to target × 20\n\n**Reward functions:**\n- `valid_fold`: +1.0 valid FOLD JSON, −0.5 parseable but invalid, −2.0 unparseable\n- `shape_match`: similarity × 20.0 (0–20), −1.0 sim fails, −2.0 invalid FOLD\n\n**Algorithm:** GRPO (Group Relative Policy Optimization) via TRL + Unsloth LoRA",
+   "source": "# Origami RL — GRPO Training\n\nTrain an LLM to generate FOLD crease patterns using OpenEnv + Unsloth + TRL.\n\nFollows the [2048 OpenEnv notebook](https://colab.research.google.com/github/unslothai/notebooks/blob/main/nb/OpenEnv_gpt_oss_(20B)_Reinforcement_Learning_2048_Game.ipynb) pattern exactly:\n1. `launch_openenv()` spawns the origami environment server\n2. LLM generates FOLD JSON crease patterns\n3. Reward functions call the server via OpenEnv client\n4. GRPO updates policy based on relative rewards",
    "metadata": {}
   },
   {
@@ -15,7 +15,7 @@
   {
    "cell_type": "code",
    "id": "ulhu8a5p5ti",
-   "source": "# Run this cell once to install all dependencies\n# For Colab: unsloth has a specific install process\nimport sys\nIN_COLAB = \"google.colab\" in sys.modules\n\nif IN_COLAB:\n    # Unsloth's recommended Colab install\n    !pip install --no-deps \"unsloth[colab-new]\"\n    !pip install --no-deps trl datasets peft accelerate bitsandbytes xformers\nelse:\n    !pip install -q \"trl>=0.7\" \"datasets>=2.14\" unsloth torch transformers accelerate bitsandbytes\n\n# Core origami env deps (numpy, scipy, pydantic)\n!pip install -q numpy scipy pydantic",
+   "source": "%%capture\nimport os, importlib.util\n!pip install --upgrade -qqq uv\nif importlib.util.find_spec(\"torch\") is None or \"COLAB_\" in \"\".join(os.environ.keys()):\n    try: import numpy; get_numpy = f\"numpy=={numpy.__version__}\"\n    except: 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\"\nelif 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!pip install -qqq fastapi uvicorn requests numpy scipy pydantic",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -23,13 +23,13 @@
   {
    "cell_type": "markdown",
    "id": "qcetkmcq1hf",
-   "source": "## 2. Setup Python Path & Imports",
+   "source": "## 2. Clone Origami Env + Setup Paths",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "3hr273dhqiv",
-   "source": "import os\nimport sys\nimport json\n\n# Add the repo root to Python path so origami_server and training modules are importable\nREPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(\"__file__\"), \"..\"))\nif REPO_ROOT not in sys.path:\n    sys.path.insert(0, REPO_ROOT)\n\nprint(f\"Repo root: {REPO_ROOT}\")\nprint(f\"Python: {sys.version}\")",
+   "source": "%%capture\n# Clone the origami env repo (skip if running locally)\nimport subprocess, sys, os\nfrom pathlib import Path\n\nREPO_URL = \"https://github.com/YOUR_USERNAME/origami_env.git\"  # TODO: update with your repo\nLOCAL_DIR = \"origami_env\"\n\nif not Path(LOCAL_DIR).exists():\n    # Running on Colab — clone the repo\n    !git clone {REPO_URL} {LOCAL_DIR} > /dev/null 2>&1\n    !pip install -e {LOCAL_DIR} > /dev/null 2>&1\n\n# Add repo to Python path\nworking_directory = str(Path(LOCAL_DIR).absolute()) if Path(LOCAL_DIR).exists() else str(Path.cwd().parent.absolute())\nsys.path.insert(0, working_directory)\nprint(f\"Working directory: {working_directory}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -37,7 +37,7 @@
   {
    "cell_type": "code",
    "id": "bnm2w57r3lc",
-   "source": "import numpy as np\n\n# Verify origami env modules load correctly\nfrom origami_server.tasks import TASKS, get_task, list_tasks\nfrom origami_server.engine.fold_parser import validate_fold, parse_fold\nfrom origami_server.engine.simulate import simulate\nfrom origami_server.engine.shape_match import compute_shape_match\nfrom training.reward import valid_fold, shape_match, extract_fold_json\n\nprint(f\"Available tasks: {list_tasks()}\")\nprint(\"All origami modules loaded successfully.\")",
+   "source": "# Import OpenEnv client + models (same pattern as 2048 notebook)\nfrom client import OrigamiEnv\nfrom origami_server.models import OrigamiAction, OrigamiObservation, OrigamiState\nprint(\"Origami OpenEnv modules loaded.\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -45,43 +45,13 @@
   {
    "cell_type": "markdown",
    "id": "lcaus7mtuj",
-   "source": "## 3. Explore the Environment\n\nSanity-check the simulator and reward functions before training.",
+   "source": "## 3. Load Model + QLoRA",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "hlqp4y30m87",
-   "source": "# Print all tasks with their details\nfor name, task in TASKS.items():\n    print(f\"\\n{'='*50}\")\n    print(f\"Task: {task['name']}\")\n    print(f\"Description: {task['description']}\")\n    print(f\"Difficulty: {task['difficulty']}\")\n    print(f\"Paper: {task['paper']}\")\n    fold = task[\"target_fold\"]\n    n_verts = len(fold[\"vertices_coords\"])\n    n_edges = len(fold[\"edges_vertices\"])\n    n_folds = sum(1 for a in fold[\"edges_assignment\"] if a in (\"M\", \"V\"))\n    print(f\"Vertices: {n_verts}, Edges: {n_edges}, Fold creases: {n_folds}\")",
-   "metadata": {},
-   "execution_count": null,
-   "outputs": []
-  },
-  {
-   "cell_type": "code",
-   "id": "dwqqus8mhlj",
-   "source": "# Test the simulator on each task\nfor name in list_tasks():\n    task = get_task(name)\n    target_fold = task[\"target_fold\"]\n    \n    # Simulate flat (0%), half (50%), and fully folded (100%)\n    r_flat = simulate(target_fold, crease_percent=0.0)\n    r_half = simulate(target_fold, crease_percent=0.5)\n    r_full = simulate(target_fold, crease_percent=1.0)\n    \n    z_half = r_half.positions[:, 2].max() - r_half.positions[:, 2].min()\n    \n    # Shape match: target vs itself should be 1.0\n    self_sim = compute_shape_match(r_full.positions, r_full.positions)\n    \n    print(f\"{name:15s} | converged={r_full.converged} | strain={r_full.max_strain:.6f} | \"\n          f\"z_range@50%={z_half:.3f} | self_similarity={self_sim:.3f}\")",
-   "metadata": {},
-   "execution_count": null,
-   "outputs": []
-  },
-  {
-   "cell_type": "code",
-   "id": "p1weq9kv5q",
-   "source": "# Test reward functions with mock LLM outputs\ntriangle_fold = TASKS[\"triangle\"][\"target_fold\"]\n\n# Simulate what the reward functions see during training:\n# completions = list of [{\"content\": \"...LLM response...\"}]\ngood_response = json.dumps(triangle_fold)\nbad_json = \"I think we should fold it like this...\"\ninvalid_fold = json.dumps({\"vertices_coords\": [[0, 0]], \"edges_vertices\": [], \"edges_assignment\": []})\n\ncompletions = [\n    [{\"content\": f\"```json\\n{good_response}\\n```\"}],   # correct answer in fenced block\n    [{\"content\": bad_json}],                              # garbage\n    [{\"content\": invalid_fold}],                          # parseable but invalid FOLD\n]\n\nprint(\"valid_fold rewards:\", valid_fold(completions))\nprint(\"shape_match rewards:\", shape_match(completions, task_name=\"triangle\"))\nprint()\nprint(\"Expected: valid_fold  = [1.0, -2.0, -0.5]\")\nprint(\"Expected: shape_match = [20.0, -2.0, -1.0]\")",
-   "metadata": {},
-   "execution_count": null,
-   "outputs": []
-  },
-  {
-   "cell_type": "markdown",
-   "id": "45l0n1hgvr",
-   "source": "## 4. Visualize Tasks\n\n2D crease patterns for each task (matplotlib).",
-   "metadata": {}
-  },
-  {
-   "cell_type": "code",
-   "id": "fkopb9lgg7i",
-   "source": "import matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nfrom mpl_toolkits.mplot3d.art3d import Poly3DCollection\n\nEDGE_COLORS = {\"M\": \"red\", \"V\": \"blue\", \"B\": \"black\"}\nEDGE_STYLES = {\"M\": \"--\", \"V\": \":\", \"B\": \"-\"}\n\nfig, axes = plt.subplots(2, 4, figsize=(16, 8))\n\nfor idx, (name, task) in enumerate(TASKS.items()):\n    fold = task[\"target_fold\"]\n    verts = np.array(fold[\"vertices_coords\"])\n    \n    # Row 1: 2D crease pattern\n    ax = axes[0, idx]\n    ax.set_title(f\"{name}\\n{task['description']}\", fontsize=9)\n    ax.set_aspect(\"equal\")\n    ax.set_xlim(-0.1, 1.1)\n    ax.set_ylim(-0.1, 1.1)\n    ax.grid(True, alpha=0.2)\n    \n    for i, (e, a) in enumerate(zip(fold[\"edges_vertices\"], fold[\"edges_assignment\"])):\n        v1, v2 = verts[e[0]], verts[e[1]]\n        color = EDGE_COLORS.get(a, \"gray\")\n        style = EDGE_STYLES.get(a, \"-\")\n        lw = 2.5 if a == \"B\" else 1.8\n        ax.plot([v1[0], v2[0]], [v1[1], v2[1]], color=color, linestyle=style, linewidth=lw)\n    \n    ax.scatter(verts[:, 0], verts[:, 1], c=\"black\", s=15, zorder=5)\n    \n    # Row 2: 3D folded shape\n    ax3 = fig.add_subplot(2, 4, idx + 5, projection=\"3d\")\n    result = simulate(fold, crease_percent=1.0)\n    pos = result.positions\n    \n    if \"faces_vertices\" in fold:\n        for face in fold[\"faces_vertices\"]:\n            tri_verts = [pos[vi] for vi in face]\n            poly = Poly3DCollection([tri_verts], alpha=0.3, facecolor=\"lightskyblue\", edgecolor=\"steelblue\")\n            ax3.add_collection3d(poly)\n    \n    for i, (e, a) in enumerate(zip(fold[\"edges_vertices\"], fold[\"edges_assignment\"])):\n        p1, p2 = pos[e[0]], pos[e[1]]\n        color = EDGE_COLORS.get(a, \"gray\")\n        ax3.plot([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]], color=color, linewidth=1.2)\n    \n    ax3.scatter(pos[:, 0], pos[:, 1], pos[:, 2], c=\"black\", s=10, zorder=5)\n    ax3.set_title(f\"Folded (3D)\", fontsize=9)\n    ax3.set_xlim(-0.2, 1.2)\n    ax3.set_ylim(-0.2, 1.2)\n    ax3.set_zlim(-0.6, 0.6)\n    \n    # Remove the empty 2D subplot that was in row 2\n    axes[1, idx].remove()\n\nplt.tight_layout()\nplt.show()",
+   "source": "from unsloth import FastLanguageModel\nimport torch\n\nmax_seq_length = 768\nlora_rank = 4\n\nmodel, tokenizer = FastLanguageModel.from_pretrained(\n    model_name = \"unsloth/Qwen3-14B\",\n    load_in_4bit = True,\n    max_seq_length = max_seq_length,\n    offload_embedding = True,\n)",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -89,13 +59,13 @@
   {
    "cell_type": "markdown",
    "id": "a14w2fkoewq",
-   "source": "## 5. Training Configuration",
+   "source": "## 4. LoRA Adapter",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "2phdejbobq3",
-   "source": "# ============================================================\n# Training hyperparameters — edit these before launching\n# ============================================================\n\nTASK_NAME = \"triangle\"          # \"triangle\", \"half_fold\", \"quarter_fold\", \"letter_fold\"\nMODEL_NAME = \"Qwen/Qwen2.5-3B-Instruct\"  # Change to your preferred model\nMAX_STEPS = 600                 # Total GRPO training steps\nNUM_GENERATIONS = 4             # Completions per prompt per step\nLEARNING_RATE = 2e-4\nLORA_R = 8                     # LoRA rank\nLORA_ALPHA = 16                # LoRA alpha\nMAX_PROMPT_LENGTH = 1024\nMAX_COMPLETION_LENGTH = 1024\nDATASET_SIZE = 1000             # Number of prompt copies (same prompt repeated)\nOUTPUT_DIR = \"outputs\"\nSAVE_STEPS = 100",
+   "source": "model = FastLanguageModel.get_peft_model(\n    model,\n    r = lora_rank,\n    target_modules = [\n        \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n        \"gate_proj\", \"up_proj\", \"down_proj\",\n    ],\n    lora_alpha = lora_rank * 2,\n    use_gradient_checkpointing = \"unsloth\",\n    random_state = 3407,\n)",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -103,13 +73,13 @@
   {
    "cell_type": "markdown",
    "id": "feal20fr8j5",
-   "source": "## 6. Build the Prompt & Dataset",
+   "source": "## 5. Launch OpenEnv Server",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "uo7zh1dwp6r",
-   "source": "from training.train_grpo import PROMPT_TEMPLATE, build_prompt\n\ntask = get_task(TASK_NAME)\nprompt_text = build_prompt(task)\n\nprint(\"=\"*60)\nprint(\"PROMPT THAT THE LLM WILL SEE:\")\nprint(\"=\"*60)\nprint(prompt_text)",
+   "source": "# Launch origami environment server (same pattern as 2048 notebook)\nglobal port\nglobal openenv_process\nport = 8000\nopenenv_process = None\nserver = \"origami_server.app:app\"\nenvironment = {\n    **os.environ,\n    \"PYTHONPATH\": working_directory,\n}\n\n# Augment Unsloth's launch_openenv with our config\nimport functools\nfrom unsloth import is_port_open, launch_openenv\nlaunch_openenv = functools.partial(\n    launch_openenv,\n    working_directory = working_directory,\n    server = server,\n    environment = environment,\n    openenv_class = OrigamiEnv,\n)",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -117,7 +87,7 @@
   {
    "cell_type": "code",
    "id": "900vyqwb8g",
-   "source": "from datasets import Dataset\n\n# GRPO pattern: same prompt repeated many times, the RL loop generates\n# multiple completions per prompt and uses relative rewards to update policy\ndataset = Dataset.from_list(\n    [\n        {\n            \"prompt\": [{\"role\": \"user\", \"content\": prompt_text}],\n            \"answer\": 0,  # placeholder, not used by GRPO\n        }\n    ]\n    * DATASET_SIZE\n)\n\nprint(f\"Dataset size: {len(dataset)}\")\nprint(f\"Sample prompt (first 100 chars): {dataset[0]['prompt'][0]['content'][:100]}...\")",
+   "source": "# Test the connection — reset and inspect\nport, openenv_process = launch_openenv(port, openenv_process)\nresult = openenv_process.reset(task_name=\"triangle\")\nprint(f\"Server running on port {port}\")\nprint(f\"Observation: done={result.done}, reward={result.reward}\")\nprint(f\"Task: {result.observation.task}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -125,48 +95,54 @@
   {
    "cell_type": "markdown",
    "id": "xn6n1hpx2aa",
-   "source": "## 7. Load Model + LoRA\n\nUses Unsloth for fast 4-bit LoRA fine-tuning. Falls back to standard HuggingFace if Unsloth isn't available.",
+   "source": "## 6. Prompt + Dataset",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "vkfaeuu9dq",
-   "source": "import torch\nprint(f\"CUDA available: {torch.cuda.is_available()}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n    print(f\"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")\nelif hasattr(torch.backends, \"mps\") and torch.backends.mps.is_available():\n    print(\"Apple MPS (Metal) available — note: Unsloth requires CUDA, will use HF fallback\")\nelse:\n    print(\"No GPU detected — training will be very slow\")",
+   "source": "import requests\n\nTASK_NAME = \"triangle\"  # \"triangle\", \"half_fold\", \"quarter_fold\", \"letter_fold\"\n\n# Fetch task params from the server (paper size, description, etc.)\ntask_info = requests.get(f\"http://localhost:{port}/tasks/{TASK_NAME}\").json()\n\nPROMPT_TEMPLATE = \"\"\"You are an origami designer. Generate a FOLD-format crease pattern\nthat, when folded, produces the target shape described below.\n\nTarget: {description}\nPaper size: {width} x {height}\n\nOutput a JSON object with these exact fields:\n- vertices_coords: [[x, y], ...] — 2D positions on the flat paper (0 to {width} for x, 0 to {height} for y)\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 (V: 180, M: -180, B: 0)\n\nRules:\n- Boundary edges (B) must outline the paper rectangle\n- At least one fold crease (M or V) must exist\n- All vertex indices must be valid (0 to N-1)\n\nOutput ONLY the JSON object wrapped in ```json ... ``` markers.\"\"\"\n\nprompt = PROMPT_TEMPLATE.format(\n    description=task_info[\"description\"],\n    width=task_info[\"paper\"][\"width\"],\n    height=task_info[\"paper\"][\"height\"],\n).strip()\n\n# Build dataset — same prompt repeated 1000x (identical to 2048 pattern)\nfrom datasets import Dataset\ndataset = Dataset.from_list([{\n    \"prompt\": [{\"role\": \"user\", \"content\": prompt}],\n}] * 1000)\n\nprint(f\"Task: {task_info['name']} — {task_info['description']}\")\nprint(f\"Paper: {task_info['paper']['width']} x {task_info['paper']['height']}\")\nprint(f\"Difficulty: {task_info['difficulty']}\")\nprint(f\"Dataset: {len(dataset)} rows\")\nprint(f\"\\nPrompt:\\n{prompt[:200]}...\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
   },
+  {
+   "cell_type": "markdown",
+   "id": "3f7ritml396",
+   "source": "## 7. Reward Functions\n\nTwo reward functions (same pattern as 2048 notebook):\n- `valid_fold` — local JSON structure check (fast, no server call)\n- `shape_match` — calls the origami server via `launch_openenv`, submits the fold, returns similarity × 20",
+   "metadata": {}
+  },
   {
    "cell_type": "code",
-   "id": "xwlkfw3xxoo",
-   "source": "USE_UNSLOTH = False\n\ntry:\n    from unsloth import FastLanguageModel\n    USE_UNSLOTH = True\n    print(\"Using Unsloth for fast LoRA loading\")\nexcept ImportError:\n    print(\"Unsloth not available, using standard HuggingFace + PEFT\")\n\nif USE_UNSLOTH:\n    model, tokenizer = FastLanguageModel.from_pretrained(\n        model_name=MODEL_NAME,\n        load_in_4bit=True,\n        max_seq_length=MAX_PROMPT_LENGTH + MAX_COMPLETION_LENGTH,\n    )\n    model = FastLanguageModel.get_peft_model(\n        model,\n        r=LORA_R,\n        target_modules=[\n            \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n            \"gate_proj\", \"up_proj\", \"down_proj\",\n        ],\n        lora_alpha=LORA_ALPHA,\n        use_gradient_checkpointing=\"unsloth\",\n    )\nelse:\n    from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig\n    from peft import LoraConfig, get_peft_model\n\n    bnb_config = BitsAndBytesConfig(\n        load_in_4bit=True,\n        bnb_4bit_quant_type=\"nf4\",\n        bnb_4bit_compute_dtype=torch.bfloat16,\n    ) if torch.cuda.is_available() else None\n\n    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\n    model = AutoModelForCausalLM.from_pretrained(\n        MODEL_NAME,\n        quantization_config=bnb_config,\n        device_map=\"auto\" if torch.cuda.is_available() else \"cpu\",\n        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,\n    )\n\n    lora_config = LoraConfig(\n        r=LORA_R,\n        lora_alpha=LORA_ALPHA,\n        target_modules=[\n            \"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n            \"gate_proj\", \"up_proj\", \"down_proj\",\n        ],\n        task_type=\"CAUSAL_LM\",\n    )\n    model = get_peft_model(model, lora_config)\n\nif tokenizer.pad_token is None:\n    tokenizer.pad_token = tokenizer.eos_token\n\nmodel.print_trainable_parameters()",
+   "id": "4dqsw30e9nq",
+   "source": "import json, re\n\n# --- Reward 1: valid_fold (local check, no server needed) ---\n\ndef extract_fold_json(response):\n    \"\"\"Extract FOLD JSON from LLM response text.\"\"\"\n    # Try fenced code block\n    match = re.search(r\"```(?:json)?\\s*(\\{.*?\\})\\s*```\", response, re.DOTALL)\n    if match:\n        try: return json.loads(match.group(1))\n        except json.JSONDecodeError: pass\n    # Try raw JSON with vertices_coords\n    match = re.search(r\"\\{[^{}]*\\\"vertices_coords\\\"[^{}]*\\}\", response, re.DOTALL)\n    if match:\n        try: return json.loads(match.group(0))\n        except json.JSONDecodeError: pass\n    # Try whole response\n    try:\n        data = json.loads(response.strip())\n        if isinstance(data, dict) and \"vertices_coords\" in data:\n            return data\n    except (json.JSONDecodeError, ValueError): pass\n    return None\n\ndef valid_fold(completions, **kwargs):\n    \"\"\"Does the LLM output parse as valid FOLD JSON?\n    +1.0 valid, -0.5 parseable but invalid, -2.0 unparseable.\"\"\"\n    scores = []\n    for completion in completions:\n        response = completion[0][\"content\"]\n        fold_data = extract_fold_json(response)\n        if fold_data is None:\n            scores.append(-2.0); continue\n        required = {\"vertices_coords\", \"edges_vertices\", \"edges_assignment\"}\n        if not required.issubset(fold_data.keys()):\n            scores.append(-0.5); continue\n        verts = fold_data.get(\"vertices_coords\", [])\n        edges = fold_data.get(\"edges_vertices\", [])\n        assigns = fold_data.get(\"edges_assignment\", [])\n        if len(edges) != len(assigns):\n            scores.append(-0.5); continue\n        if not any(a in (\"M\", \"V\") for a in assigns) or not any(a == \"B\" for a in assigns):\n            scores.append(-0.5); continue\n        n = len(verts)\n        if not all(0 <= e[0] < n and 0 <= e[1] < n and e[0] != e[1] for e in edges):\n            scores.append(-0.5); continue\n        scores.append(1.0)\n    return scores\n\n# --- Reward 2: shape_match (calls server via launch_openenv) ---\n\ndef shape_match(completions, **kwargs):\n    \"\"\"Submit fold to origami server, get shape similarity reward.\n    Calls launch_openenv to ensure server is running, then reset + step.\"\"\"\n    global port, openenv_process\n    scores = []\n    for completion in completions:\n        response = completion[0][\"content\"]\n        fold_data = extract_fold_json(response)\n        if fold_data is None:\n            scores.append(0.0)\n            continue\n        try:\n            port, openenv_process = launch_openenv(port, openenv_process)\n            openenv_process.reset(task_name=TASK_NAME)\n            result = openenv_process.step(OrigamiAction(fold_data=fold_data))\n            reward = result.reward if result.reward is not None else 0.0\n            scores.append(reward)\n        except TimeoutError:\n            scores.append(-1.0)\n        except Exception as e:\n            scores.append(-3.0)\n    return scores\n\n# Quick test\ntest_good = [[{\"content\": json.dumps({\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})}]]\ntest_bad = [[{\"content\": \"not json\"}]]\nprint(f\"valid_fold — good: {valid_fold(test_good)}, bad: {valid_fold(test_bad)}\")\nprint(f\"shape_match — good: {shape_match(test_good)}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
   },
   {
    "cell_type": "markdown",
-   "id": "3f7ritml396",
-   "source": "## 8. Setup GRPO Trainer",
+   "id": "62lvkfoyu1p",
+   "source": "## 8. GRPO Trainer",
    "metadata": {}
   },
   {
    "cell_type": "code",
-   "id": "4dqsw30e9nq",
-   "source": "from trl import GRPOConfig, GRPOTrainer\n\n# Wrap shape_match to inject the task name\ndef shape_match_reward(completions, **kwargs):\n    return shape_match(completions, task_name=TASK_NAME, **kwargs)\n\ntraining_args = GRPOConfig(\n    temperature=1.0,\n    learning_rate=LEARNING_RATE,\n    weight_decay=0.001,\n    warmup_ratio=0.1,\n    lr_scheduler_type=\"linear\",\n    optim=\"adamw_8bit\" if torch.cuda.is_available() else \"adamw_torch\",\n    logging_steps=1,\n    per_device_train_batch_size=1,\n    gradient_accumulation_steps=1,\n    num_generations=NUM_GENERATIONS,\n    max_prompt_length=MAX_PROMPT_LENGTH,\n    max_completion_length=MAX_COMPLETION_LENGTH,\n    max_steps=MAX_STEPS,\n    save_steps=SAVE_STEPS,\n    output_dir=OUTPUT_DIR,\n    report_to=\"none\",  # Set to \"wandb\" if you want W&B logging\n)\n\ntrainer = GRPOTrainer(\n    model=model,\n    processing_class=tokenizer,\n    reward_funcs=[valid_fold, shape_match_reward],\n    args=training_args,\n    train_dataset=dataset,\n)\n\nprint(f\"Trainer ready. Task: {TASK_NAME}, Model: {MODEL_NAME}\")\nprint(f\"Max steps: {MAX_STEPS}, Generations per step: {NUM_GENERATIONS}\")\nprint(f\"Reward functions: valid_fold + shape_match\")",
+   "id": "eohisxhna96",
+   "source": "from trl import GRPOConfig, GRPOTrainer\n\ntraining_args = GRPOConfig(\n    temperature = 1.0,\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 = 1,\n    num_generations = 2,\n    max_prompt_length = 512,\n    max_completion_length = max_seq_length - 512,\n    max_steps = 600,\n    save_steps = 100,\n    output_dir = \"outputs\",\n    report_to = \"none\",\n)\n\ntrainer = GRPOTrainer(\n    model = model,\n    processing_class = tokenizer,\n    reward_funcs = [valid_fold, shape_match],\n    args = training_args,\n    train_dataset = dataset,\n)\n\nprint(f\"Trainer ready: {training_args.max_steps} steps, {training_args.num_generations} generations/step\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
   },
   {
    "cell_type": "markdown",
-   "id": "62lvkfoyu1p",
+   "id": "ve98mq6rgot",
    "source": "## 9. Train!",
    "metadata": {}
   },
   {
    "cell_type": "code",
-   "id": "eohisxhna96",
+   "id": "8il1yknetfg",
    "source": "trainer.train()",
    "metadata": {},
    "execution_count": null,
@@ -181,7 +157,7 @@
   {
    "cell_type": "code",
    "id": "t3d4tu6o5mc",
-   "source": "SAVE_PATH = f\"origami-{TASK_NAME}-lora\"\n\n# Save LoRA adapter\nmodel.save_pretrained(SAVE_PATH)\ntokenizer.save_pretrained(SAVE_PATH)\nprint(f\"LoRA adapter saved to {SAVE_PATH}/\")\n\n# Optional: merge LoRA into base model and save full model\n# merged_path = f\"origami-{TASK_NAME}-merged\"\n# if USE_UNSLOTH:\n#     model.save_pretrained_merged(merged_path, tokenizer)\n# else:\n#     merged_model = model.merge_and_unload()\n#     merged_model.save_pretrained(merged_path)\n#     tokenizer.save_pretrained(merged_path)\n# print(f\"Merged model saved to {merged_path}/\")",
+   "source": "save_path = f\"origami-{TASK_NAME}-lora\"\nmodel.save_pretrained(save_path)\ntokenizer.save_pretrained(save_path)\nprint(f\"LoRA adapter saved to {save_path}/\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -189,13 +165,13 @@
   {
    "cell_type": "markdown",
    "id": "q18eizy1ok",
-   "source": "## 11. Evaluate — Generate & Score Completions\n\nTest the trained model by generating crease patterns and scoring them.",
+   "source": "## 11. Evaluate — Generate & Score",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "on56augj41",
-   "source": "# Put model in inference mode\nif USE_UNSLOTH:\n    FastLanguageModel.for_inference(model)\n\nNUM_EVAL_SAMPLES = 8\n\n# Build chat messages\nmessages = [{\"role\": \"user\", \"content\": prompt_text}]\ninput_ids = tokenizer.apply_chat_template(\n    messages, tokenize=True, add_generation_prompt=True, return_tensors=\"pt\"\n).to(model.device)\n\nprint(f\"Generating {NUM_EVAL_SAMPLES} completions...\")\nprint(f\"Input length: {input_ids.shape[1]} tokens\\n\")\n\neval_completions = []\nfor i in range(NUM_EVAL_SAMPLES):\n    with torch.no_grad():\n        output = model.generate(\n            input_ids,\n            max_new_tokens=MAX_COMPLETION_LENGTH,\n            temperature=0.7,\n            top_p=0.9,\n            do_sample=True,\n            pad_token_id=tokenizer.pad_token_id,\n        )\n    response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)\n    eval_completions.append([{\"content\": response}])\n    \n    # Quick score\n    fold_data = extract_fold_json(response)\n    if fold_data is None:\n        status = \"UNPARSEABLE\"\n    else:\n        is_valid, err = validate_fold(fold_data)\n        if not is_valid:\n            status = f\"INVALID: {err}\"\n        else:\n            try:\n                result = simulate(fold_data, crease_percent=1.0)\n                target_result = simulate(task[\"target_fold\"], crease_percent=1.0)\n                sim = compute_shape_match(result.positions, target_result.positions)\n                status = f\"similarity={sim:.3f} (reward={sim * 20:.1f})\"\n            except Exception as e:\n                status = f\"SIM ERROR: {e}\"\n    \n    print(f\"  Sample {i+1}: {status}\")\n\n# Compute aggregate reward scores\nprint(f\"\\nAggregate rewards:\")\nvf_scores = valid_fold(eval_completions)\nsm_scores = shape_match(eval_completions, task_name=TASK_NAME)\nprint(f\"  valid_fold:  mean={np.mean(vf_scores):.2f}, scores={vf_scores}\")\nprint(f\"  shape_match: mean={np.mean(sm_scores):.2f}, scores={sm_scores}\")",
+   "source": "import numpy as np\nFastLanguageModel.for_inference(model)\n\nNUM_EVAL = 8\nmessages = [{\"role\": \"user\", \"content\": prompt}]\ninput_ids = tokenizer.apply_chat_template(\n    messages, tokenize=True, add_generation_prompt=True, return_tensors=\"pt\"\n).to(model.device)\n\nprint(f\"Generating {NUM_EVAL} completions (input: {input_ids.shape[1]} tokens)...\\n\")\n\neval_completions = []\nfor i in range(NUM_EVAL):\n    with torch.no_grad():\n        output = model.generate(\n            input_ids,\n            max_new_tokens=max_seq_length - 512,\n            temperature=0.7, top_p=0.9, do_sample=True,\n            pad_token_id=tokenizer.pad_token_id,\n        )\n    response = tokenizer.decode(output[0][input_ids.shape[1]:], skip_special_tokens=True)\n    eval_completions.append([{\"content\": response}])\n    fold = extract_fold_json(response)\n    status = f\"parsed ({len(fold.get('vertices_coords', []))} verts)\" if fold else \"UNPARSEABLE\"\n    print(f\"  Sample {i+1}: {status}\")\n\nprint(f\"\\nScoring via server...\")\nvf_scores = valid_fold(eval_completions)\nsm_scores = shape_match(eval_completions)\nprint(f\"  valid_fold:  mean={np.mean(vf_scores):.2f}, scores={vf_scores}\")\nprint(f\"  shape_match: mean={np.mean(sm_scores):.2f}, scores={sm_scores}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -203,13 +179,13 @@
   {
    "cell_type": "markdown",
    "id": "tb1y8hszrk",
-   "source": "## 12. Visualize a Generated Fold\n\nPick the best completion and visualize its crease pattern + 3D fold vs the target.",
+   "source": "## 12. Visualize Best Result",
    "metadata": {}
   },
   {
    "cell_type": "code",
    "id": "0zo3krbkiqej",
-   "source": "# Find the best valid completion\nbest_idx = int(np.argmax(sm_scores))\nbest_response = eval_completions[best_idx][0][\"content\"]\nbest_fold = extract_fold_json(best_response)\n\nif best_fold is None or sm_scores[best_idx] <= 0:\n    print(\"No valid completions to visualize.\")\nelse:\n    is_valid, _ = validate_fold(best_fold)\n    if not is_valid:\n        print(\"Best completion has invalid FOLD structure.\")\n    else:\n        pred_result = simulate(best_fold, crease_percent=1.0)\n        target_result = simulate(task[\"target_fold\"], crease_percent=1.0)\n        \n        fig = plt.figure(figsize=(14, 5))\n        \n        # 1) Generated 2D crease pattern\n        ax1 = fig.add_subplot(131)\n        ax1.set_title(f\"Generated Crease Pattern\\n(sample {best_idx+1})\", fontsize=10)\n        ax1.set_aspect(\"equal\")\n        verts = np.array(best_fold[\"vertices_coords\"])\n        for i, (e, a) in enumerate(zip(best_fold[\"edges_vertices\"], best_fold[\"edges_assignment\"])):\n            v1, v2 = verts[e[0]], verts[e[1]]\n            color = EDGE_COLORS.get(a, \"gray\")\n            style = EDGE_STYLES.get(a, \"-\")\n            ax1.plot([v1[0], v2[0]], [v1[1], v2[1]], color=color, linestyle=style, linewidth=2)\n        ax1.scatter(verts[:, 0], verts[:, 1], c=\"black\", s=20, zorder=5)\n        ax1.grid(True, alpha=0.2)\n        \n        # 2) Generated 3D fold\n        ax2 = fig.add_subplot(132, projection=\"3d\")\n        ax2.set_title(f\"Generated 3D Fold\\nsimilarity={sm_scores[best_idx]/20:.3f}\", fontsize=10)\n        pos = pred_result.positions\n        for i, (e, a) in enumerate(zip(best_fold[\"edges_vertices\"], best_fold[\"edges_assignment\"])):\n            p1, p2 = pos[e[0]], pos[e[1]]\n            color = EDGE_COLORS.get(a, \"gray\")\n            ax2.plot([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]], color=color, linewidth=1.5)\n        ax2.scatter(pos[:, 0], pos[:, 1], pos[:, 2], c=\"black\", s=15, zorder=5)\n        \n        # 3) Target 3D fold\n        ax3 = fig.add_subplot(133, projection=\"3d\")\n        ax3.set_title(\"Target 3D Fold\", fontsize=10)\n        tpos = target_result.positions\n        tfold = task[\"target_fold\"]\n        for i, (e, a) in enumerate(zip(tfold[\"edges_vertices\"], tfold[\"edges_assignment\"])):\n            p1, p2 = tpos[e[0]], tpos[e[1]]\n            color = EDGE_COLORS.get(a, \"gray\")\n            ax3.plot([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]], color=color, linewidth=1.5)\n        ax3.scatter(tpos[:, 0], tpos[:, 1], tpos[:, 2], c=\"black\", s=15, zorder=5)\n        \n        plt.tight_layout()\n        plt.show()\n        \n        print(f\"\\nBest generated FOLD JSON:\")\n        print(json.dumps(best_fold, indent=2))",
+   "source": "import matplotlib.pyplot as plt\nimport requests\n\nEDGE_COLORS = {\"M\": \"red\", \"V\": \"blue\", \"B\": \"black\"}\nEDGE_STYLES = {\"M\": \"--\", \"V\": \":\", \"B\": \"-\"}\n\nbest_idx = int(np.argmax(sm_scores))\nbest_fold = extract_fold_json(eval_completions[best_idx][0][\"content\"])\n\nif best_fold is None or sm_scores[best_idx] <= 0:\n    print(\"No valid completions to visualize.\")\nelse:\n    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))\n\n    # Generated crease pattern\n    ax1.set_title(f\"Generated (sample {best_idx+1})\\nreward={sm_scores[best_idx]:.1f}\", fontsize=10)\n    ax1.set_aspect(\"equal\")\n    verts = np.array(best_fold[\"vertices_coords\"])\n    for e, a in zip(best_fold[\"edges_vertices\"], best_fold[\"edges_assignment\"]):\n        v1, v2 = verts[e[0]], verts[e[1]]\n        ax1.plot([v1[0], v2[0]], [v1[1], v2[1]],\n                 color=EDGE_COLORS.get(a, \"gray\"),\n                 linestyle=EDGE_STYLES.get(a, \"-\"), linewidth=2)\n    ax1.scatter(verts[:, 0], verts[:, 1], c=\"black\", s=20, zorder=5)\n    ax1.grid(True, alpha=0.2)\n\n    # Target crease pattern (from server)\n    ax2.set_title(\"Target\", fontsize=10)\n    ax2.set_aspect(\"equal\")\n    port, openenv_process = launch_openenv(port, openenv_process)\n    # Get target from server via HTTP\n    target_resp = requests.get(f\"http://localhost:{port}/tasks/{TASK_NAME}\")\n    target = target_resp.json()[\"target_fold\"]\n    tverts = np.array(target[\"vertices_coords\"])\n    for e, a in zip(target[\"edges_vertices\"], target[\"edges_assignment\"]):\n        v1, v2 = tverts[e[0]], tverts[e[1]]\n        ax2.plot([v1[0], v2[0]], [v1[1], v2[1]],\n                 color=EDGE_COLORS.get(a, \"gray\"),\n                 linestyle=EDGE_STYLES.get(a, \"-\"), linewidth=2)\n    ax2.scatter(tverts[:, 0], tverts[:, 1], c=\"black\", s=20, zorder=5)\n    ax2.grid(True, alpha=0.2)\n\n    plt.tight_layout()\n    plt.show()\n    print(f\"\\nBest FOLD JSON:\\n{json.dumps(best_fold, indent=2)}\")",
    "metadata": {},
    "execution_count": null,
    "outputs": []
@@ -217,7 +193,7 @@
   {
    "cell_type": "markdown",
    "id": "qlakksqmoe",
-   "source": "## 13. Plot Training Logs",
+   "source": "## 13. Training Logs",
    "metadata": {}
   },
   {