+
+ WAITING FOR TRAINING...
+
+
+React build not found. Run npm run build in the frontend directory.
Training viewer: /viewer/training.html
" + ) + + +def run(host: str = "0.0.0.0", port: int = 9001) -> None: + """Start the training server. Call from Colab notebook.""" + uvicorn.run(app, host=host, port=port) + + +if __name__ == "__main__": + run() diff --git a/server/models.py b/server/models.py new file mode 100644 index 0000000000000000000000000000000000000000..6d468ca73921b45db9d77ed28f0ba951a2787449 --- /dev/null +++ b/server/models.py @@ -0,0 +1,72 @@ +""" +OpenEnv Pydantic models for the origami RL environment. + +OrigamiAction — one fold per step +OrigamiObservation — everything the LLM and Three.js viewer need +OrigamiState — server-side episode tracking +""" +from __future__ import annotations + +from typing import Any, Optional + +from pydantic import BaseModel, Field + +# openenv base classes — use them if available, fall back to plain Pydantic +try: + from openenv.core.env_server.types import Action, Observation, State +except ImportError: + Action = BaseModel + class State(BaseModel): + """Minimal stand-in for openenv State base class.""" + episode_id: Optional[str] = None + step_count: int = 0 + + class Observation(BaseModel): + """Minimal stand-in for openenv Observation base class.""" + done: bool = False + reward: Optional[float] = None + + +class OrigamiAction(Action): + """One fold operation sent by the client each step.""" + + fold_type: str = Field( + default="valley", + description="'valley' | 'mountain' | 'pleat' | 'crimp' | 'stop'", + ) + fold_line: dict[str, list[float]] = Field( + default_factory=lambda: {"start": [0.0, 0.5], "end": [1.0, 0.5]}, + description="{'start': [x, y], 'end': [x, y]} normalized 0-1", + ) + fold_angle: float = Field( + default=180.0, + description="Fold angle in degrees, 0-180", + ) + layer_select: str = Field( + default="all", + description="'all' | 'top' | 'bottom'", + ) + + +class OrigamiObservation(Observation): + """Everything the LLM and Three.js viewer need. + + paper_state contains FOLD-compatible geometry + physics data. + metrics contains all computed quality metrics. + No render_urls — the browser renders from paper_state directly. + """ + + task: dict[str, Any] = Field(default_factory=dict) + paper_state: dict[str, Any] = Field(default_factory=dict) + metrics: dict[str, Any] = Field(default_factory=dict) + fold_history: list[dict[str, Any]] = Field(default_factory=list) + error: Optional[str] = Field(default=None) + + +class OrigamiState(State): + """Server-side episode tracking.""" + + task_name: str = Field(default="") + num_folds_applied: int = Field(default=0) + is_valid: bool = Field(default=True) + total_reward: float = Field(default=0.0) diff --git a/server/origami_environment.py b/server/origami_environment.py new file mode 100644 index 0000000000000000000000000000000000000000..886e7c97b2b6379d09fa0b8caeb86fa80f70a99c --- /dev/null +++ b/server/origami_environment.py @@ -0,0 +1,221 @@ +""" +OrigamiEnvironment — OpenEnv environment wrapping the origami physics engine. + +Implements reset() / step() / state following the OpenEnv interface. +Engine (physics, fold, validation, metrics) lives in engine/. +No server-side image rendering — paper_state contains all geometry data. +""" +from __future__ import annotations + +import json +import os +import uuid +from typing import Any, Optional + +# openenv base class — fall back to plain object if not installed +try: + from openenv.core.env_server.interfaces import Environment +except ImportError: + from typing import Generic, TypeVar + A = TypeVar("A") + O = TypeVar("O") + S = TypeVar("S") + class Environment(Generic[A, O, S]): + """Minimal stand-in for openenv.core.env_server.interfaces.Environment.""" + def __init__(self, **kwargs): pass + +from engine.paper import Paper +from engine.fold_engine import apply_fold +from engine.physics import simulate +from engine.validation import validate_state +from engine.metrics import compute_all_metrics +from server.models import OrigamiAction, OrigamiObservation, OrigamiState +from server.tasks import get_task_by_name, sample_task + + +def _get_material(name: str): + """Get material by name, falling back to paper.""" + try: + from engine.materials import get_material + return get_material(name) + except Exception: + from engine.materials import get_material + return get_material("paper") + + +class OrigamiEnvironment(Environment[OrigamiAction, OrigamiObservation, OrigamiState]): + """Origami folding RL environment. + + Each episode: agent receives paper_state + task, applies folds one at a + time via step(), receives metrics + reward, ends with 'stop' action or + when max_folds is reached. + """ + + SUPPORTS_CONCURRENT_SESSIONS = False + + def __init__(self, **kwargs): + super().__init__(**kwargs) + self._paper: Optional[Paper] = None + self._task: Optional[dict] = None + self._fold_history: list[dict] = [] + self._metrics: dict = {} + self._validation: dict = {} + self._error: Optional[str] = None + self._episode_id: Optional[str] = None + self._step_count: int = 0 + self._total_reward: float = 0.0 + + # ── reset ───────────────────────────────────────────────────────── + + def reset( + self, + seed: Optional[int] = None, + episode_id: Optional[str] = None, + **kwargs: Any, + ) -> OrigamiObservation: + self._episode_id = episode_id or str(uuid.uuid4()) + self._step_count = 0 + self._fold_history = [] + self._error = None + self._total_reward = 0.0 + + # Select task + task_name = kwargs.get("task_name") + if task_name: + self._task = get_task_by_name(task_name) + if not self._task: + self._task = sample_task(seed=seed) + + # Create flat sheet + mat = _get_material(self._task["material"]) + self._paper = Paper.create_flat_sheet( + width=self._task["width"], + height=self._task["height"], + material=mat, + ) + + # Initial validation + metrics (no physics needed for flat sheet) + self._validation = validate_state(self._paper) + self._metrics = compute_all_metrics(self._paper, self._task, self._validation) + + return self._make_observation(done=False, reward=None) + + # ── step ────────────────────────────────────────────────────────── + + def step( + self, + action: OrigamiAction, + timeout_s: Optional[float] = None, + **kwargs: Any, + ) -> OrigamiObservation: + if self._paper is None or self._task is None: + return self._make_observation(done=True, reward=-5.0) + + self._step_count += 1 + self._error = None + + # ── Stop action ─────────────────────────────────────────────── + if action.fold_type == "stop": + return self._finalize_episode() + + # ── Build fold dict ─────────────────────────────────────────── + fold_dict = { + "type": action.fold_type, + "line": action.fold_line, + "angle": action.fold_angle, + } + + # ── Apply fold ──────────────────────────────────────────────── + new_paper, err = apply_fold(self._paper, fold_dict) + if err: + self._error = err + return self._make_observation(done=True, reward=-5.0) + + self._paper = new_paper + self._fold_history.append({**fold_dict, "step": self._step_count}) + + # ── Physics relaxation ──────────────────────────────────────── + try: + self._paper = simulate(self._paper, fold_percent=1.0) + except Exception as exc: + self._error = f"Physics failed: {exc}" + # Continue — don't abort episode on physics failure + + # ── Validate ────────────────────────────────────────────────── + self._validation = validate_state(self._paper) + + # ── Metrics ─────────────────────────────────────────────────── + self._metrics = compute_all_metrics(self._paper, self._task, self._validation) + + # ── Check termination ───────────────────────────────────────── + max_folds = self._task.get("max_folds", 50) + if self._step_count >= max_folds: + return self._finalize_episode() + + if self._validation.get("self_intersections", 0) > 0: + self._error = "Self-intersection detected" + return self._finalize_episode() + + return self._make_observation(done=False, reward=None) + + # ── state ───────────────────────────────────────────────────────── + + @property + def state(self) -> OrigamiState: + return OrigamiState( + episode_id=self._episode_id, + step_count=self._step_count, + task_name=self._task.get("name", "") if self._task else "", + num_folds_applied=len(self._fold_history), + is_valid=self._metrics.get("is_valid", True), + total_reward=self._total_reward, + ) + + # ── internals ───────────────────────────────────────────────────── + + def _finalize_episode(self) -> OrigamiObservation: + reward = self._compute_reward() + self._total_reward = reward + return self._make_observation(done=True, reward=reward) + + def _make_observation(self, done: bool, reward: Optional[float]) -> OrigamiObservation: + return OrigamiObservation( + done=done, + reward=reward, + task=self._task or {}, + paper_state=self._paper.to_observation_dict() if self._paper else {}, + metrics=self._metrics, + fold_history=self._fold_history, + error=self._error, + ) + + def _compute_reward(self) -> float: + m = self._metrics + reward = 0.0 + + # Compactness is the main signal + reward += m.get("compactness", 0.0) * 20.0 + + # Bonus for fitting in target box + if m.get("fits_target_box", False): + reward += 10.0 + + # Bonus for deployability (if task requires it) + if m.get("is_deployable", False): + reward += 5.0 + + # Penalties for violations + reward -= m.get("kawasaki_violations", 0) * 2.0 + reward -= m.get("maekawa_violations", 0) * 2.0 + reward -= m.get("self_intersections", 0) * 5.0 + + # Penalty for too many folds (encourage efficiency) + reward -= m.get("fold_count", 0) * 0.5 + + # Penalty for exceeding material strain limit + max_strain = m.get("max_strain", 0.0) + strain_limit = self._paper.material.max_strain if self._paper else 0.05 + if max_strain > strain_limit: + reward -= 3.0 * (max_strain / strain_limit) + + return float(reward) diff --git a/server/tasks.py b/server/tasks.py new file mode 100644 index 0000000000000000000000000000000000000000..bbcaa8add7105bd162c0d2bc00b135c303a45a9c --- /dev/null +++ b/server/tasks.py @@ -0,0 +1,123 @@ +""" +Task pool and curriculum for the origami RL environment. + +7 tasks across 4 difficulty levels. +""" +from __future__ import annotations + +import random +from typing import Optional + + +TASKS: dict[str, dict] = { + "half_fold": { + "name": "half_fold", + "description": "Fold a 1x1 paper sheet in half along the horizontal midline.", + "width": 1.0, + "height": 1.0, + "material": "paper", + "target_ratio": 0.50, + "max_folds": 3, + "target_box": [1.0, 0.5, 0.02], + "must_deploy": False, + "difficulty": 1, + }, + "quarter_fold": { + "name": "quarter_fold", + "description": "Fold a 1x1 paper sheet into quarters using two perpendicular folds.", + "width": 1.0, + "height": 1.0, + "material": "paper", + "target_ratio": 0.25, + "max_folds": 5, + "target_box": [0.5, 0.5, 0.04], + "must_deploy": False, + "difficulty": 1, + }, + "letter_fold": { + "name": "letter_fold", + "description": "Fold a 1x1 paper into thirds (letter fold) using two parallel folds.", + "width": 1.0, + "height": 1.0, + "material": "paper", + "target_ratio": 0.33, + "max_folds": 5, + "target_box": [1.0, 0.34, 0.03], + "must_deploy": False, + "difficulty": 2, + }, + "map_fold": { + "name": "map_fold", + "description": "Fold a 1x1 paper into eighths using a grid fold pattern. Must be re-deployable.", + "width": 1.0, + "height": 1.0, + "material": "paper", + "target_ratio": 0.125, + "max_folds": 8, + "target_box": [0.5, 0.25, 0.08], + "must_deploy": True, + "difficulty": 2, + }, + "solar_panel": { + "name": "solar_panel", + "description": "Pack a 1x1 Mylar solar panel into a compact configuration using a Miura-ori style fold. Must deploy.", + "width": 1.0, + "height": 1.0, + "material": "mylar", + "target_ratio": 0.05, + "max_folds": 20, + "target_box": [0.25, 0.25, 0.05], + "must_deploy": True, + "difficulty": 3, + }, + "shelter_wall": { + "name": "shelter_wall", + "description": "Fold a 1x1 aluminum sheet into a compact structural panel within strain limits.", + "width": 1.0, + "height": 1.0, + "material": "aluminum", + "target_ratio": 0.10, + "max_folds": 15, + "target_box": [0.5, 0.25, 0.1], + "must_deploy": False, + "difficulty": 3, + }, + "stent": { + "name": "stent", + "description": "Fold a 0.5x1.5 nitinol sheet into a compact tube configuration for a medical stent. Superelastic material.", + "width": 0.5, + "height": 1.5, + "material": "nitinol", + "target_ratio": 0.09, + "max_folds": 25, + "target_box": [0.1, 0.1, 0.15], + "must_deploy": True, + "difficulty": 4, + }, +} + + +def get_task_by_name(name: str) -> Optional[dict]: + """Return task dict by name, or None if not found.""" + return TASKS.get(name) + + +def sample_task(seed: Optional[int] = None, difficulty: Optional[int] = None) -> dict: + """Sample a random task, optionally filtered by difficulty level.""" + rng = random.Random(seed) + pool = list(TASKS.values()) + if difficulty is not None: + pool = [t for t in pool if t["difficulty"] == difficulty] + if not pool: + pool = list(TASKS.values()) + return dict(rng.choice(pool)) + + +def get_tasks_by_difficulty(level: int) -> list[dict]: + """Return all tasks at a given difficulty level.""" + return [dict(t) for t in TASKS.values() if t["difficulty"] == level] + + +def available_task_names() -> list[str]: + """Return sorted list of all task names.""" + return sorted(TASKS.keys()) diff --git a/server/training_broadcast.py b/server/training_broadcast.py new file mode 100644 index 0000000000000000000000000000000000000000..1d1a2518a0425fa41a41fb450a225fc6e2bc970e --- /dev/null +++ b/server/training_broadcast.py @@ -0,0 +1,216 @@ +""" +TrainingBroadcastServer — fire-and-forget broadcast hub for live training viewer. + +The RL training process calls publish() after each env.step(). +Spectator browsers connect via /ws/training WebSocket. +Broadcast is async and non-blocking: if no viewers are connected, observations are dropped. +""" +from __future__ import annotations + +import asyncio +import json +import logging +from dataclasses import dataclass, field +from typing import Any, Optional + +from fastapi import WebSocket, WebSocketDisconnect + +logger = logging.getLogger(__name__) + + +@dataclass +class EpisodeInfo: + episode_id: str + task_name: str + status: str = "running" # "running" | "done" | "timeout" | "error" + step: int = 0 + observation: dict = field(default_factory=dict) + metrics: dict = field(default_factory=dict) + fold_history: list = field(default_factory=list) + steps: list = field(default_factory=list) # full step history for replay + score: Optional[float] = None + final_metrics: Optional[dict] = None + + +class TrainingBroadcastServer: + """Central hub for broadcasting RL training observations to spectator WebSockets. + + Thread-safe: publish() can be called from training threads (ThreadPoolExecutor). + WebSocket handlers run in the asyncio event loop. + """ + + def __init__(self) -> None: + self._spectators: list[WebSocket] = [] + self._registry: dict[str, EpisodeInfo] = {} + self._batch_id: int = 0 + self._loop: Optional[asyncio.AbstractEventLoop] = None + self._lock = asyncio.Lock() + + # ── Episode publishing (called from training thread / async context) ── + + def publish(self, episode_id: str, data: dict) -> None: + """Fire-and-forget: push an update from the training process. + + Safe to call from any thread. Schedules onto the stored event loop + (set by the FastAPI startup handler). No-op if no loop is available. + """ + loop = self._loop + if loop is None or loop.is_closed(): + return + asyncio.run_coroutine_threadsafe(self._async_publish(episode_id, data), loop) + + async def _async_publish(self, episode_id: str, data: dict) -> None: + msg_type = data.get("type", "episode_update") + + async with self._lock: + if msg_type == "batch_start": + self._batch_id = data.get("batch_id", self._batch_id + 1) + self._registry.clear() + await self._broadcast(data) + return + + if msg_type == "batch_done": + await self._broadcast(data) + return + + if msg_type == "training_done": + await self._broadcast(data) + return + + # episode_update or episode_done + ep = self._registry.setdefault( + episode_id, + EpisodeInfo(episode_id=episode_id, task_name=data.get("task_name", "")), + ) + + if msg_type == "episode_done": + ep.status = data.get("status", "done") + ep.score = data.get("score") + ep.final_metrics = data.get("final_metrics") + else: + step_num = data.get("step", ep.step) + ep.step = step_num + ep.status = "running" + obs = data.get("observation", {}) + ep.observation = obs + ep.metrics = obs.get("metrics", {}) + ep.fold_history = obs.get("fold_history", ep.fold_history) + # Accumulate full step history for /episode/replay + if step_num > 0: + fold_hist = obs.get("fold_history", []) + latest_fold = fold_hist[-1] if fold_hist else {} + ep.steps.append({ + "step": step_num, + "fold": latest_fold, + "paper_state": obs.get("paper_state", {}), + "metrics": obs.get("metrics", {}), + "done": obs.get("done", False), + }) + + await self._broadcast({"episode_id": episode_id, **data}) + + # ── Spectator management ── + + async def connect_spectator(self, websocket: WebSocket) -> None: + """Accept a new viewer WebSocket and serve it until disconnect.""" + await websocket.accept() + + async with self._lock: + self._spectators.append(websocket) + + # Send current registry snapshot immediately + await self._send_registry(websocket) + + try: + while True: + # Viewers are read-only; drain any incoming messages (pings etc) + await asyncio.wait_for(websocket.receive_text(), timeout=30.0) + except (WebSocketDisconnect, asyncio.TimeoutError, Exception): + pass + finally: + await self.disconnect_spectator(websocket) + + async def disconnect_spectator(self, websocket: WebSocket) -> None: + async with self._lock: + self._spectators = [s for s in self._spectators if s is not websocket] + + # ── Batch control ── + + async def start_batch(self, batch_id: int, num_episodes: int, prompt_index: int = 0) -> None: + """Call before starting a new training batch.""" + data = { + "type": "batch_start", + "batch_id": batch_id, + "num_episodes": num_episodes, + "prompt_index": prompt_index, + } + await self._async_publish("__batch__", data) + + async def finish_batch( + self, + batch_id: int, + scores: list[float], + best_episode_id: str = "", + ) -> None: + """Call after all episodes in a batch complete.""" + data = { + "type": "batch_done", + "batch_id": batch_id, + "scores": scores, + "best_episode_id": best_episode_id, + "avg_score": sum(scores) / len(scores) if scores else 0.0, + } + await self._async_publish("__batch__", data) + + async def clear_batch(self) -> None: + """Reset episode registry for next batch.""" + async with self._lock: + self._registry.clear() + + # ── Internals ── + + async def _broadcast(self, message: dict) -> None: + """Send message to all spectators, removing dead connections.""" + if not self._spectators: + return + payload = json.dumps(message, default=str) + dead: list[WebSocket] = [] + for ws in list(self._spectators): + try: + await ws.send_text(payload) + except Exception: + dead.append(ws) + for ws in dead: + self._spectators = [s for s in self._spectators if s is not ws] + + async def _send_registry(self, websocket: WebSocket) -> None: + """Send the full episode registry to a newly connected viewer.""" + async with self._lock: + episodes = { + ep_id: { + "status": ep.status, + "task": ep.task_name, + "step": ep.step, + "observation": ep.observation, + "metrics": ep.metrics, + "score": ep.score, + } + for ep_id, ep in self._registry.items() + } + payload = { + "type": "registry", + "batch_id": self._batch_id, + "episodes": episodes, + } + try: + await websocket.send_text(json.dumps(payload, default=str)) + except Exception: + pass + + @property + def spectator_count(self) -> int: + return len(self._spectators) + + @property + def active_episodes(self) -> int: + return sum(1 for ep in self._registry.values() if ep.status == "running") diff --git a/src/App.css b/src/App.css index 74b5e053450a48a6bdb4d71aad648e7af821975c..1631274e73a4fbc2f4fde20e54234f4eadc07b91 100644 --- a/src/App.css +++ b/src/App.css @@ -1,38 +1,572 @@ -.App { - text-align: center; +:root { + --bg: #0d0d14; + --surface: #13131d; + --surface-2: #1a1a2e; + --paper-white: #fafaf5; + --paper-edge: #2a2a3a; + --mountain: #f59e0b; + --valley: #38bdf8; + --target-ghost: rgba(124, 58, 237, 0.20); + --target-ghost-stroke: rgba(124, 58, 237, 0.45); + --validity: #22d3ee; + --progress: #22c55e; + --economy: #a78bfa; + --text-primary: #f8fafc; + --text-dim: #64748b; + --border: #2a2a3a; + --border-bright: #3a3a5a; + --font-display: 'JetBrains Mono', monospace; + --font-mono: 'IBM Plex Mono', monospace; +} + +.app { + display: flex; + flex-direction: column; + height: 100vh; + background: var(--bg); + overflow: hidden; +} + +/* ─── HEADER ─── */ +.app-header { + display: flex; + align-items: center; + gap: 24px; + padding: 0 20px; + height: 48px; + border-bottom: 1px solid var(--border); + background: var(--surface); + flex-shrink: 0; + z-index: 10; +} + +.app-title { + font-family: var(--font-display); + font-size: 14px; + font-weight: 700; + letter-spacing: 0.12em; + color: var(--text-primary); + white-space: nowrap; +} + +.app-title .title-accent { + color: var(--mountain); +} + +.header-sep { + width: 1px; + height: 24px; + background: var(--border); + flex-shrink: 0; +} + +.header-right { + display: flex; + align-items: center; + gap: 16px; + margin-left: auto; +} + +.replay-badge { + font-size: 10px; + font-family: var(--font-display); + letter-spacing: 0.1em; + color: #38bdf8; + background: rgba(56, 189, 248, 0.1); + border: 1px solid rgba(56, 189, 248, 0.3); + padding: 3px 8px; + border-radius: 3px; +} + +.back-to-grid-btn { + font-size: 10px; + font-family: var(--font-display); + letter-spacing: 0.08em; + color: #64748b; + background: transparent; + border: 1px solid #1e2a3a; + padding: 3px 10px; + border-radius: 3px; + cursor: pointer; } +.back-to-grid-btn:hover { color: #e2e8f0; border-color: #64748b; } -.App-logo { - height: 40vmin; - pointer-events: none; +.api-status { + font-size: 11px; + font-family: var(--font-display); + letter-spacing: 0.08em; + display: flex; + align-items: center; + gap: 6px; +} + +.api-status-dot { + width: 6px; + height: 6px; + border-radius: 50%; + background: var(--text-dim); +} + +.api-status-dot.ok { + background: var(--progress); + box-shadow: 0 0 6px var(--progress); } -@media (prefers-reduced-motion: no-preference) { - .App-logo { - animation: App-logo-spin infinite 20s linear; - } +.api-status-dot.err { + background: #ef4444; + box-shadow: 0 0 6px #ef4444; +} + +/* ─── MAIN LAYOUT ─── */ +.app-body { + display: grid; + grid-template-columns: 1fr 280px; + flex: 1; + overflow: hidden; +} + +.app-left { + display: flex; + flex-direction: column; + overflow: hidden; + border-right: 1px solid var(--border); } -.App-header { - background-color: #282c34; - min-height: 100vh; +.app-right { display: flex; flex-direction: column; + overflow: hidden; + background: var(--surface); +} + +/* ─── CANVAS ROW ─── */ +.canvas-row { + display: flex; + gap: 0; + padding: 16px; + flex-shrink: 0; + border-bottom: 1px solid var(--border); + overflow-x: auto; +} + +.canvas-wrap { + display: flex; + flex-direction: column; + gap: 8px; + flex: 1; + min-width: 280px; +} + +.canvas-wrap + .canvas-wrap { + margin-left: 16px; +} + +.canvas-label { + font-family: var(--font-display); + font-size: 10px; + font-weight: 500; + letter-spacing: 0.14em; + color: var(--text-dim); + text-transform: uppercase; +} + +.canvas-svg { + display: block; + background: var(--paper-white); +} + +.canvas-3d { + display: block; + background: linear-gradient(180deg, #1a1a2e 0%, #0f101a 100%); + border: 1px solid var(--border); +} + +.canvas-label-row { + display: flex; + align-items: center; + justify-content: space-between; + gap: 10px; +} + +.fold-mode-toggle { + display: inline-flex; + border: 1px solid var(--border); + background: var(--surface); +} + +.fold-mode-btn { + border: none; + background: transparent; + color: var(--text-dim); + font-family: var(--font-display); + font-size: 9px; + letter-spacing: 0.08em; + padding: 3px 7px; + cursor: pointer; +} + +.fold-mode-btn + .fold-mode-btn { + border-left: 1px solid var(--border); +} + +.fold-mode-btn.active { + color: var(--text-primary); + background: #1f2538; +} + +/* ─── STEP FEED ─── */ +.step-feed-section { + flex: 1; + display: flex; + flex-direction: column; + overflow: hidden; +} + +.section-header { + font-family: var(--font-display); + font-size: 10px; + font-weight: 500; + letter-spacing: 0.14em; + color: var(--text-dim); + text-transform: uppercase; + padding: 8px 16px; + border-bottom: 1px solid var(--border); + flex-shrink: 0; +} + +.step-feed { + overflow-y: auto; + flex: 1; + padding: 4px 0; +} + +.step-entry { + display: flex; + flex-direction: column; + gap: 2px; + padding: 8px 16px; + border-bottom: 1px solid var(--border); + cursor: default; + transition: background 0.1s; +} + +.step-entry:hover { + background: var(--surface); +} + +.step-entry.active { + background: var(--surface-2); + border-left: 2px solid var(--valley); + padding-left: 14px; +} + +.step-entry-top { + display: flex; + align-items: center; + gap: 8px; +} + +.step-num { + font-family: var(--font-display); + font-size: 10px; + font-weight: 700; + color: var(--text-dim); + width: 24px; + flex-shrink: 0; +} + +.step-instruction { + font-size: 12px; + color: var(--text-primary); + flex: 1; +} + +.assign-badge { + font-family: var(--font-display); + font-size: 10px; + font-weight: 700; + padding: 1px 5px; + line-height: 1.4; + flex-shrink: 0; +} + +.assign-badge.M { + background: var(--mountain); + color: #0d0d14; +} + +.assign-badge.V { + background: var(--valley); + color: #0d0d14; +} + +.assign-badge.B { + background: var(--border-bright); + color: var(--text-dim); +} + +.step-reward-delta { + font-size: 11px; + color: var(--text-dim); + padding-left: 32px; +} + +.step-reward-delta .delta-positive { + color: var(--progress); +} + +.step-reward-delta .delta-negative { + color: #ef4444; +} + +/* ─── REWARD PANEL ─── */ +.reward-panel { + padding: 12px 16px; + border-bottom: 1px solid var(--border); + flex-shrink: 0; +} + +.reward-row { + display: flex; + align-items: center; + gap: 8px; + margin-bottom: 6px; +} + +.reward-row:last-child { + margin-bottom: 0; +} + +.reward-label { + font-family: var(--font-display); + font-size: 10px; + font-weight: 500; + letter-spacing: 0.06em; + color: var(--text-dim); + width: 72px; + flex-shrink: 0; + text-transform: uppercase; +} + +.reward-track { + flex: 1; + height: 8px; + background: var(--bg); + border: 1px solid var(--border); + overflow: hidden; +} + +.reward-bar { + height: 100%; + transition: width 0.4s ease; +} + +.reward-value { + font-family: var(--font-display); + font-size: 11px; + font-weight: 500; + color: var(--text-primary); + width: 36px; + text-align: right; + flex-shrink: 0; +} + +.reward-value.dim { + color: var(--text-dim); +} + +.reward-divider { + height: 1px; + background: var(--border); + margin: 6px 0; +} + +/* ─── INFO BADGES ─── */ +.info-badges { + padding: 12px 16px; + display: flex; + flex-direction: column; + gap: 8px; +} + +.info-row { + display: flex; + align-items: center; + justify-content: space-between; + gap: 8px; +} + +.info-key { + font-family: var(--font-display); + font-size: 10px; + font-weight: 500; + letter-spacing: 0.06em; + color: var(--text-dim); + text-transform: uppercase; +} + +.info-val { + font-family: var(--font-display); + font-size: 11px; + font-weight: 700; + color: var(--text-primary); +} + +.info-val.bool-true { + color: var(--progress); +} + +.info-val.bool-false { + color: #ef4444; +} + +.info-val.dim { + color: var(--text-dim); +} + +/* ─── TARGET SELECTOR ─── */ +.target-selector { + display: flex; + align-items: center; + gap: 8px; +} + +.target-selector-label { + font-family: var(--font-display); + font-size: 10px; + font-weight: 500; + letter-spacing: 0.10em; + color: var(--text-dim); + text-transform: uppercase; + white-space: nowrap; +} + +.target-select { + background: var(--surface-2); + border: 1px solid var(--border-bright); + color: var(--text-primary); + font-family: var(--font-display); + font-size: 11px; + padding: 4px 8px; + outline: none; + cursor: pointer; + min-width: 180px; +} + +.target-select:focus { + border-color: var(--valley); +} + +optgroup { + background: var(--surface); + color: var(--text-dim); + font-family: var(--font-display); + font-size: 10px; +} + +option { + background: var(--surface-2); + color: var(--text-primary); + font-family: var(--font-display); +} + +/* ─── PLAYER CONTROLS ─── */ +.player-controls { + display: flex; + align-items: center; + gap: 6px; + flex-shrink: 0; +} + +.ctrl-btn { + background: var(--surface-2); + border: 1px solid var(--border-bright); + color: var(--text-primary); + font-family: var(--font-display); + font-size: 11px; + font-weight: 500; + padding: 4px 10px; + cursor: pointer; + white-space: nowrap; + line-height: 1.4; + letter-spacing: 0.04em; + transition: background 0.1s, border-color 0.1s; +} + +.ctrl-btn:hover:not(:disabled) { + background: var(--surface); + border-color: var(--text-dim); +} + +.ctrl-btn:disabled { + opacity: 0.35; + cursor: not-allowed; +} + +.ctrl-btn.play { + border-color: var(--valley); + color: var(--valley); +} + +.ctrl-btn.play:hover:not(:disabled) { + background: rgba(56, 189, 248, 0.1); +} + +.ctrl-step-display { + font-family: var(--font-display); + font-size: 11px; + color: var(--text-dim); + padding: 4px 8px; + border: 1px solid var(--border); + background: var(--bg); + white-space: nowrap; + min-width: 72px; + text-align: center; +} + +/* ─── LOADING / ERROR ─── */ +.app-overlay { + position: fixed; + inset: 0; + display: flex; align-items: center; justify-content: center; - font-size: calc(10px + 2vmin); - color: white; + background: var(--bg); + z-index: 100; +} + +.overlay-message { + font-family: var(--font-display); + font-size: 13px; + letter-spacing: 0.1em; + color: var(--text-dim); + display: flex; + align-items: center; + gap: 12px; } -.App-link { - color: #61dafb; +.pulse-dot { + width: 8px; + height: 8px; + border-radius: 50%; + background: var(--valley); + animation: pulse 1.2s ease-in-out infinite; } -@keyframes App-logo-spin { - from { - transform: rotate(0deg); - } - to { - transform: rotate(360deg); - } +@keyframes pulse { + 0%, 100% { opacity: 0.2; transform: scale(0.8); } + 50% { opacity: 1; transform: scale(1); } +} + +/* ─── MISC ─── */ +.episode-loading { + display: flex; + align-items: center; + justify-content: center; + gap: 8px; + padding: 12px 16px; + font-family: var(--font-display); + font-size: 11px; + color: var(--text-dim); + letter-spacing: 0.08em; } diff --git a/src/App.js b/src/App.js index 37845757234ccb68531c10cf7a2ffc589c47e342..5b5842f6787d369c34cd3e280637d38071b9c2b7 100644 --- a/src/App.js +++ b/src/App.js @@ -1,23 +1,239 @@ -import logo from './logo.svg'; +import { useState, useEffect, useCallback, useRef } from 'react'; import './App.css'; +import CreaseCanvas from './components/CreaseCanvas'; +import RewardPanel from './components/RewardPanel'; +import StepFeed from './components/StepFeed'; +import InfoBadges from './components/InfoBadges'; +import TargetSelector from './components/TargetSelector'; +import PlayerControls from './components/PlayerControls'; +import Fold3DCanvas from './components/Fold3DCanvas'; + +const API_BASE = ''; + +// Read ?ep=
-
- 
- [{"instruction": "Valley fold along horizontal center", "from": [0, 0.5], "to": [1, 0.5], "assignment": "V"}] ',
+ '[{"instruction": "Invalid fold", "from": [0.3, 0.3], "to": [0.7, 0.7], "assignment": "V"}] ',
+ 'this is not valid JSON at all',
+ ]
+ target_names = [dataset_items[0]['target_name']] * 3
+ rewards = reward_fn(test_completions, target_names=target_names)
+ print(f"\nDry run rewards: {rewards}")
+ print("Dry run complete — reward function works.")
+ return
+
+ # Load model via unsloth
+ # VL models (SpatialThinker, Qwen2.5-VL) use FastVisionModel
+ # Text-only models use FastLanguageModel
+ is_vl = _is_vl_model(args.model)
+
+ try:
+ if is_vl:
+ from unsloth import FastVisionModel as ModelLoader
+ print(f"Loading VL model (vision-language): {args.model}")
+ else:
+ from unsloth import FastLanguageModel as ModelLoader
+ print(f"Loading text model: {args.model}")
+ except ImportError:
+ print("ERROR: unsloth not installed. Run: pip install unsloth")
+ print("Or run with --dry_run to test the reward function without a model.")
+ return
+
+ model, tokenizer = ModelLoader.from_pretrained(
+ model_name=args.model,
+ max_seq_length=args.max_seq_length,
+ load_in_4bit=True,
+ )
+
+ model = ModelLoader.get_peft_model(
+ model,
+ r=32,
+ target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+ "gate_proj", "up_proj", "down_proj"],
+ lora_alpha=32,
+ lora_dropout=0,
+ use_gradient_checkpointing="unsloth",
+ )
+
+ # Convert dataset to HuggingFace Dataset format
+ from datasets import Dataset
+
+ # GRPOTrainer expects 'prompt' column and optionally others.
+ # We embed target_name in the dataset so the reward fn can use it.
+ hf_dataset = Dataset.from_list(dataset_items)
+
+ # Build main reward function
+ reward_fn = make_reward_fn(env, args.max_folds)
+
+ from trl import GRPOConfig, GRPOTrainer
+
+ config = GRPOConfig(
+ output_dir=args.output,
+ num_train_epochs=args.epochs,
+ per_device_train_batch_size=args.batch_size,
+ gradient_accumulation_steps=args.grad_accum,
+ learning_rate=args.lr,
+ max_completion_length=512,
+ num_generations=args.n_generations,
+ temperature=1.0,
+ logging_steps=1,
+ report_to="trackio",
+ run_name="origami-grpo",
+ )
+
+ # GRPOTrainer passes all dataset columns as kwargs to reward_funcs.
+ # The 'target_name' column arrives as a list (one per completion in the batch).
+ def wrapped_reward_fn(completions, target_name=None, **kwargs):
+ """Wrapper that extracts target_name from batch columns."""
+ target_names = target_name if isinstance(target_name, list) else [target_name] * len(completions)
+ return reward_fn(completions, target_names=target_names)
+
+ trainer = GRPOTrainer(
+ model=model,
+ config=config,
+ train_dataset=hf_dataset,
+ reward_funcs=[wrapped_reward_fn],
+ tokenizer=tokenizer,
+ )
+
+ print(f"\nStarting GRPO training...")
+ print(f" Model: {args.model}")
+ print(f" Level: {args.level} targets")
+ print(f" Epochs: {args.epochs}")
+ print(f" Generations per prompt: {args.n_generations}")
+ print(f" Output: {args.output}/")
+
+ trainer.train()
+
+ # Save
+ model.save_pretrained(args.output)
+ tokenizer.save_pretrained(args.output)
+ print(f"\nModel saved to {args.output}/")
+
+
+if __name__ == '__main__':
+ main()
diff --git a/train_origami.ipynb b/train_origami.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..75abc7d7a0f567b06b8032c43c8d6e8e1d58180c
--- /dev/null
+++ b/train_origami.ipynb
@@ -0,0 +1,184 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "id": "8smrrb11v84",
+ "source": "# Optigami — Origami RL Training (GRPO)\n\n**Train an LLM to generate valid origami fold sequences using verifiable geometric rewards.**\n\nArchitecture:\n- **Environment**: `env/` — CreaseGraph + Kawasaki/Maekawa/BLB verifiers + target matching\n- **Policy model**: SpatialThinker (Qwen2.5-VL-7B) or vanilla Qwen2.5-VL-7B\n- **Training**: Unsloth GRPO — model generates complete fold sequences, gets terminal reward\n- **Tracking**: Trackio — real-time reward curves on HF Spaces\n\n| Reward Component | Weight | What it measures |\n|---|---|---|\n| `progress` | 0.45 | Geometric crease coverage vs target |\n| `economy` | 0.10 | Penalty for excess creases |\n| `kawasaki` | 0.08 | Kawasaki theorem satisfaction |\n| `maekawa` | 0.07 | Maekawa theorem satisfaction |\n| `blb` | 0.05 | Big-Little-Big lemma |\n| `anchored` | 0.05 | Valid anchor point usage |\n| `completion` | +10.0 | Bonus when target reached |",
+ "metadata": {}
+ },
+ {
+ "cell_type": "markdown",
+ "id": "kn1k9d357j",
+ "source": "## 1. Setup\n\n**GPU**: H100 80GB (Northflank/CoreWeave) or A100/T4 (Colab)\n\nInstall dependencies. Unsloth handles efficient model loading + LoRA.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "d10vqzep5b6",
+ "source": "%%capture\n!pip install unsloth trackio shapely numpy datasets\n!pip install --upgrade trl transformers\n\n# Check GPU\nimport torch\nprint(f\"GPU: {torch.cuda.get_device_name(0)}\")\nprint(f\"VRAM: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "y6wsagz8h",
+ "source": "## 2. Configuration\n\nChoose base model and hyperparameters. Two options:\n- **SpatialThinker** (`OX-PIXL/SpatialThinker-Qwen2.5-VL-7B`): Pre-trained for spatial reasoning via RL\n- **Vanilla Qwen2.5-VL** (`unsloth/Qwen2.5-VL-7B-Instruct`): Standard vision-language model\n\nWe'll compare both to see which learns origami folding faster.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "dh1zapl0w5s",
+ "source": "# ── Config ──────────────────────────────────────────────────────────────────\n# Toggle MODEL_NAME to switch between SpatialThinker and vanilla Qwen2.5-VL\n\nMODEL_NAME = \"OX-PIXL/SpatialThinker-Qwen2.5-VL-7B\"\n# MODEL_NAME = \"unsloth/Qwen2.5-VL-7B-Instruct\" # uncomment for vanilla\n\nMAX_SEQ_LENGTH = 2048\nLORA_R = 32\nLORA_ALPHA = 32\nEPOCHS = 3\nBATCH_SIZE = 2\nGRAD_ACCUM = 4\nLR = 5e-6\nN_GENERATIONS = 8 # completions sampled per prompt (GRPO group size)\nMAX_FOLDS = 8 # max folds per episode\nLEVEL = 1 # target difficulty (1=simple, 2=medium, 3=hard)\nMAX_COMPLETION_LEN = 512\nOUTPUT_DIR = \"origami-grpo\"\n\n# Trackio — set your HF Space ID for live dashboard\nTRACKIO_SPACE_ID = None # e.g. \"your-username/optigami-training\"\n\nprint(f\"Model: {MODEL_NAME}\")\nprint(f\"Config: {EPOCHS} epochs, batch={BATCH_SIZE}, grad_accum={GRAD_ACCUM}, lr={LR}\")\nprint(f\"GRPO: {N_GENERATIONS} generations, max_folds={MAX_FOLDS}, level={LEVEL}\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "o5hhfbp0wb",
+ "source": "## 3. Clone Repo & Test Environment\n\nClone the optigami repo (skip if running locally) and verify the environment works.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "94cemjucczl",
+ "source": "import os\n\n# Clone repo if not already present (Colab/Northflank)\nif not os.path.exists(\"env/environment.py\"):\n !git clone https://huggingface.co/spaces/openenv-community/optigami /content/optigami 2>/dev/null || true\n os.chdir(\"/content/optigami\")\n\n# Verify env/ is accessible\nfrom env.environment import OrigamiEnvironment\nfrom env.rewards import compute_reward\nfrom env.prompts import parse_fold_list\n\nenv = OrigamiEnvironment(mode=\"code_as_policy\", max_steps=MAX_FOLDS)\nprint(f\"Available targets: {env.available_targets()}\")\nprint(f\"Environment mode: {env.mode}\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "id": "2j9mccejyfx",
+ "source": "# ── Dry run: test reward function ───────────────────────────────────────────\n# Verify rewards work before loading the model\n\nimport copy\n\ndef make_reward_fn(env_template, max_folds):\n \"\"\"Reward function: clone env, run completion, return total reward.\"\"\"\n def reward_fn(completions, prompts=None, **kwargs):\n rewards = []\n target_names = kwargs.get(\"target_names\", [None] * len(completions))\n for completion, target_name in zip(completions, target_names):\n try:\n e = env_template.clone()\n e.reset(target_name=target_name)\n _, reward_dict, _, _ = e.step(completion)\n rewards.append(float(reward_dict[\"total\"]))\n except Exception:\n rewards.append(-0.1)\n return rewards\n return reward_fn\n\nreward_fn = make_reward_fn(env, MAX_FOLDS)\n\ntest_completions = [\n '[{\"instruction\": \"Valley fold along horizontal center\", \"from\": [0, 0.5], \"to\": [1, 0.5], \"assignment\": \"V\"}] ',\n '[{\"instruction\": \"Bad fold\", \"from\": [0.3, 0.3], \"to\": [0.7, 0.7], \"assignment\": \"V\"}] ',\n 'not valid JSON',\n]\ntarget_names = [\"half_horizontal\"] * 3\nrewards = reward_fn(test_completions, target_names=target_names)\n\nfor comp, r in zip([\"perfect fold\", \"partial fold\", \"garbage\"], rewards):\n print(f\" {comp}: reward = {r:.3f}\")\nprint(\"\\nReward function OK ✓\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "46gs2p1cy4",
+ "source": "## 4. Load Model + LoRA\n\nLoad the VL model with Unsloth's `FastVisionModel` (4-bit quantized) and apply LoRA adapters.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "82f76od6d2k",
+ "source": "from unsloth import FastVisionModel\n\nmodel, tokenizer = FastVisionModel.from_pretrained(\n model_name=MODEL_NAME,\n max_seq_length=MAX_SEQ_LENGTH,\n load_in_4bit=True,\n)\n\nmodel = FastVisionModel.get_peft_model(\n model,\n r=LORA_R,\n target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\",\n \"gate_proj\", \"up_proj\", \"down_proj\"],\n lora_alpha=LORA_ALPHA,\n lora_dropout=0,\n use_gradient_checkpointing=\"unsloth\",\n)\n\nprint(f\"Model loaded: {MODEL_NAME}\")\nprint(f\"LoRA rank: {LORA_R}, alpha: {LORA_ALPHA}\")\nprint(f\"Trainable params: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "67dyfrj23y",
+ "source": "## 5. Build Dataset\n\nGenerate prompts from all level-appropriate targets. Each prompt embeds the target crease pattern description and asks the model to output `[...] `.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "1msqpzj5fwu",
+ "source": "import random\nfrom datasets import Dataset\n\ndef build_dataset(env, level=1):\n \"\"\"Build training dataset of prompts from env targets.\"\"\"\n all_names = env.available_targets()\n level_names = [\n n for n in all_names\n if env._targets[n].get(\"level\", 1) == level\n ]\n if not level_names:\n level_names = all_names\n\n items = []\n for name in level_names:\n obs = env.reset(target_name=name)\n items.append({\"prompt\": obs[\"prompt\"], \"target_name\": name})\n\n # Repeat each target 10x; ensure at least 50 examples\n repeat = max(10, (50 + len(items) - 1) // len(items))\n items = items * repeat\n random.shuffle(items)\n return items\n\ndataset_items = build_dataset(env, level=LEVEL)\nhf_dataset = Dataset.from_list(dataset_items)\n\nprint(f\"Dataset: {len(dataset_items)} examples\")\nprint(f\"Targets in dataset: {sorted(set(d['target_name'] for d in dataset_items))}\")\nprint(f\"\\nSample prompt (first 300 chars):\\n{dataset_items[0]['prompt'][:300]}...\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "7n3r3nsw8ae",
+ "source": "## 6. Trackio Setup\n\nInitialize Trackio for real-time training visualization. Trackio is a free W&B alternative that deploys a Gradio dashboard to HF Spaces.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "bgru9rsw95b",
+ "source": "import trackio\n\n# Initialize Trackio run\ntrackio_kwargs = {\n \"project_name\": \"optigami\",\n \"run_name\": f\"grpo-{MODEL_NAME.split('/')[-1]}-level{LEVEL}\",\n}\nif TRACKIO_SPACE_ID:\n trackio_kwargs[\"space_id\"] = TRACKIO_SPACE_ID\n\ntrackio.init(**trackio_kwargs)\nprint(f\"Trackio initialized: {trackio_kwargs['run_name']}\")\nif TRACKIO_SPACE_ID:\n print(f\"Dashboard: https://huggingface.co/spaces/{TRACKIO_SPACE_ID}\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "n8aqymlszo",
+ "source": "## 7. GRPO Training\n\nRun GRPO with Trackio logging. The trainer samples `N_GENERATIONS` completions per prompt, computes rewards via the environment, and updates the policy using group-relative advantages.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "ci4imd9ws7v",
+ "source": "from trl import GRPOConfig, GRPOTrainer\n\n# ── Per-component reward functions for detailed logging ─────────────────────\nREWARD_COMPONENTS = [\"kawasaki\", \"maekawa\", \"blb\", \"progress\", \"economy\", \"completion\"]\n\ndef make_component_fn(env_template, component):\n \"\"\"Create a reward function that returns a single component's value.\"\"\"\n def component_fn(completions, target_name=None, **kwargs):\n target_names = target_name if isinstance(target_name, list) else [target_name] * len(completions)\n rewards = []\n for completion, tn in zip(completions, target_names):\n try:\n e = env_template.clone()\n e.reset(target_name=tn)\n _, reward_dict, _, _ = e.step(completion)\n rewards.append(float(reward_dict.get(component, 0.0)))\n except Exception:\n rewards.append(0.0)\n return rewards\n component_fn.__name__ = f\"reward_{component}\"\n return component_fn\n\n# Main reward function (returns total reward)\ndef wrapped_reward_fn(completions, target_name=None, **kwargs):\n \"\"\"Main reward function — extracts target_name from batch columns.\"\"\"\n target_names = target_name if isinstance(target_name, list) else [target_name] * len(completions)\n return reward_fn(completions, target_names=target_names)\n\n# Build list of all reward functions: [total, kawasaki, maekawa, blb, progress, economy, completion]\nall_reward_fns = [wrapped_reward_fn] + [\n make_component_fn(env, c) for c in REWARD_COMPONENTS\n]\n\n# ── GRPO Config ─────────────────────────────────────────────────────────────\nconfig = GRPOConfig(\n output_dir=OUTPUT_DIR,\n num_train_epochs=EPOCHS,\n per_device_train_batch_size=BATCH_SIZE,\n gradient_accumulation_steps=GRAD_ACCUM,\n learning_rate=LR,\n max_completion_length=MAX_COMPLETION_LEN,\n num_generations=N_GENERATIONS,\n temperature=1.0,\n logging_steps=1,\n save_steps=50,\n report_to=\"trackio\",\n run_name=f\"grpo-{MODEL_NAME.split('/')[-1]}-level{LEVEL}\",\n)\n\ntrainer = GRPOTrainer(\n model=model,\n config=config,\n train_dataset=hf_dataset,\n reward_funcs=all_reward_fns,\n tokenizer=tokenizer,\n)\n\nprint(f\"Trainer ready. Starting GRPO training...\")\nprint(f\" Model: {MODEL_NAME}\")\nprint(f\" Dataset: {len(hf_dataset)} examples\")\nprint(f\" Reward functions: total + {REWARD_COMPONENTS}\")\nprint(f\" Logging to: Trackio\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "code",
+ "id": "mwy1i7d509",
+ "source": "# ── Train! ──────────────────────────────────────────────────────────────────\ntrainer.train()\n\n# Save model + tokenizer\nmodel.save_pretrained(OUTPUT_DIR)\ntokenizer.save_pretrained(OUTPUT_DIR)\nprint(f\"\\nTraining complete. Model saved to {OUTPUT_DIR}/\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "54iu59w1ipb",
+ "source": "## 8. Evaluation\n\nRun the trained model on all targets and measure solve rates + average rewards.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "7loiy5mrsi3",
+ "source": "# ── Evaluate trained model on all targets ───────────────────────────────────\nfrom transformers import TextStreamer\n\nFastVisionModel.for_inference(model)\n\neval_results = {}\nn_samples = 5 # generations per target\n\nfor target_name in env.available_targets():\n obs = env.reset(target_name=target_name)\n prompt = obs[\"prompt\"]\n\n # Tokenize prompt\n messages = [{\"role\": \"user\", \"content\": prompt}]\n input_ids = tokenizer.apply_chat_template(messages, return_tensors=\"pt\").to(model.device)\n\n target_rewards = []\n target_solved = 0\n\n for _ in range(n_samples):\n outputs = model.generate(\n input_ids=input_ids,\n max_new_tokens=MAX_COMPLETION_LEN,\n temperature=0.7,\n do_sample=True,\n )\n completion = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)\n\n # Score completion\n e = env.clone()\n e.reset(target_name=target_name)\n try:\n _, reward_dict, _, _ = e.step(completion)\n total = float(reward_dict[\"total\"])\n solved = reward_dict.get(\"completion\", 0) > 0\n except Exception:\n total = -0.1\n solved = False\n\n target_rewards.append(total)\n if solved:\n target_solved += 1\n\n eval_results[target_name] = {\n \"avg_reward\": sum(target_rewards) / len(target_rewards),\n \"max_reward\": max(target_rewards),\n \"solve_rate\": target_solved / n_samples,\n }\n print(f\" {target_name:20s} avg={eval_results[target_name]['avg_reward']:.3f} \"\n f\"max={eval_results[target_name]['max_reward']:.3f} \"\n f\"solved={target_solved}/{n_samples}\")\n\n# Summary\navg_solve = sum(r[\"solve_rate\"] for r in eval_results.values()) / len(eval_results)\navg_reward = sum(r[\"avg_reward\"] for r in eval_results.values()) / len(eval_results)\nprint(f\"\\nOverall: avg_reward={avg_reward:.3f}, solve_rate={avg_solve:.1%}\")\n\n# Log to Trackio\ntrackio.log({\"eval/avg_reward\": avg_reward, \"eval/solve_rate\": avg_solve})\nfor name, res in eval_results.items():\n trackio.log({f\"eval/{name}_reward\": res[\"avg_reward\"], f\"eval/{name}_solved\": res[\"solve_rate\"]})",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "c9dr4ht05r",
+ "source": "## 9. A/B Comparison: SpatialThinker vs Vanilla Qwen2.5-VL\n\nTo compare both models, run this notebook twice:\n1. First run with `MODEL_NAME = \"OX-PIXL/SpatialThinker-Qwen2.5-VL-7B\"`\n2. Second run with `MODEL_NAME = \"unsloth/Qwen2.5-VL-7B-Instruct\"`\n\nBoth runs log to the same Trackio project (`optigami`) with different run names, so you can overlay the reward curves directly in the dashboard.\n\nThe cell below loads saved eval results from both runs for comparison (run after both trainings complete).",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "qwwd4wyuhnq",
+ "source": "# ── Save eval results for comparison ────────────────────────────────────────\nimport json\n\nmodel_tag = MODEL_NAME.split(\"/\")[-1]\neval_path = f\"eval_results_{model_tag}_level{LEVEL}.json\"\n\nwith open(eval_path, \"w\") as f:\n json.dump(eval_results, f, indent=2)\nprint(f\"Eval results saved to {eval_path}\")\n\n# ── Compare (run after both models are trained) ────────────────────────────\nspatial_path = f\"eval_results_SpatialThinker-Qwen2.5-VL-7B_level{LEVEL}.json\"\nvanilla_path = f\"eval_results_Qwen2.5-VL-7B-Instruct_level{LEVEL}.json\"\n\nif os.path.exists(spatial_path) and os.path.exists(vanilla_path):\n with open(spatial_path) as f:\n spatial = json.load(f)\n with open(vanilla_path) as f:\n vanilla = json.load(f)\n\n print(f\"\\n{'Target':<22} {'SpatialThinker':>16} {'Vanilla Qwen':>16} {'Delta':>10}\")\n print(\"-\" * 66)\n for target in sorted(set(list(spatial.keys()) + list(vanilla.keys()))):\n s_r = spatial.get(target, {}).get(\"avg_reward\", 0)\n v_r = vanilla.get(target, {}).get(\"avg_reward\", 0)\n delta = s_r - v_r\n print(f\" {target:<20} {s_r:>14.3f} {v_r:>14.3f} {delta:>+8.3f}\")\n\n s_avg = sum(r[\"avg_reward\"] for r in spatial.values()) / len(spatial)\n v_avg = sum(r[\"avg_reward\"] for r in vanilla.values()) / len(vanilla)\n print(f\"\\n {'OVERALL':<20} {s_avg:>14.3f} {v_avg:>14.3f} {s_avg - v_avg:>+8.3f}\")\n\n s_solve = sum(r[\"solve_rate\"] for r in spatial.values()) / len(spatial)\n v_solve = sum(r[\"solve_rate\"] for r in vanilla.values()) / len(vanilla)\n print(f\" {'Solve Rate':<20} {s_solve:>13.1%} {v_solve:>13.1%} {s_solve - v_solve:>+7.1%}\")\nelse:\n print(f\"Run both models to compare. Looking for:\\n {spatial_path}\\n {vanilla_path}\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ },
+ {
+ "cell_type": "markdown",
+ "id": "812csd43vxk",
+ "source": "## 10. Push to HuggingFace Hub (optional)\n\nUpload the trained LoRA adapter to HF for deployment or further fine-tuning.",
+ "metadata": {}
+ },
+ {
+ "cell_type": "code",
+ "id": "h38kp70n16q",
+ "source": "# ── Push to HF Hub (uncomment and set your repo) ───────────────────────────\n# from huggingface_hub import login\n# login(token=\"hf_...\") # or use HF_TOKEN env var\n#\n# HF_REPO = \"your-username/optigami-grpo-spatialthinker\"\n# model.push_to_hub(HF_REPO)\n# tokenizer.push_to_hub(HF_REPO)\n# print(f\"Model pushed to https://huggingface.co/{HF_REPO}\")\n\ntrackio.finish()\nprint(\"Done! Check your Trackio dashboard for training curves.\")",
+ "metadata": {},
+ "execution_count": null,
+ "outputs": []
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "name": "python",
+ "version": "3.10.0"
+ },
+ "colab": {
+ "provenance": [],
+ "gpuType": "A100"
+ },
+ "accelerator": "GPU"
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
\ No newline at end of file
diff --git a/trainer/mock_env.py b/trainer/mock_env.py
deleted file mode 100644
index e1360d105ad111bd36b6b2f1c99c37a0269ba31f..0000000000000000000000000000000000000000
--- a/trainer/mock_env.py
+++ /dev/null
@@ -1,233 +0,0 @@
-"""
-Mock origami environment for trainer development.
-
-Returns fake PaperState responses so we can iterate on the GRPO loop
-without waiting for the real physics engine. The mock applies geometric
-transforms (vertex rotations around fold lines) but skips energy/strain
-computation — those return plausible dummy values.
-"""
-
-import math
-import numpy as np
-from dataclasses import dataclass, field
-
-
-@dataclass
-class Material:
- name: str = "paper"
- thickness_mm: float = 0.1
- youngs_modulus_gpa: float = 2.0
- max_strain: float = 0.03 # 3%
-
-
-@dataclass
-class PaperState:
- vertices: np.ndarray # (N, 3)
- edges: np.ndarray # (E, 2)
- faces: list[list[int]]
- assignments: list[str] # M/V/B per edge
- fold_angles: np.ndarray # (E,) degrees
-
- rest_lengths: np.ndarray # (E,)
- strain: np.ndarray # (N,)
- energy: float = 0.0
-
- face_orders: list[tuple] = field(default_factory=list)
- num_layers: int = 1
-
- material: Material = field(default_factory=Material)
-
- bounding_box: np.ndarray = field(default_factory=lambda: np.array([1.0, 1.0, 0.0]))
- deployment_ratio: float = 1.0
- is_valid: bool = True
- kawasaki_violation: float = 0.0
- maekawa_violation: float = 0.0
- self_intersections: int = 0
-
- def to_dict(self) -> dict:
- return {
- "width": float(self.bounding_box[0]),
- "height": float(self.bounding_box[1]),
- "material": {
- "name": self.material.name,
- "thickness_mm": self.material.thickness_mm,
- "youngs_modulus_gpa": self.material.youngs_modulus_gpa,
- },
- "vertices": self.vertices.tolist(),
- "edges": self.edges.tolist(),
- "assignments": self.assignments,
- "fold_angles": self.fold_angles.tolist(),
- "num_layers_at_center": self.num_layers,
- "bounding_box": {
- "x": float(self.bounding_box[0]),
- "y": float(self.bounding_box[1]),
- "z": float(self.bounding_box[2]),
- },
- }
-
-
-def create_flat_sheet(width: float = 1.0, height: float = 1.0,
- material: Material | None = None) -> PaperState:
- """Create a flat rectangular sheet with 4 vertices, 5 edges (incl diagonal), 2 faces."""
- verts = np.array([
- [0, 0, 0],
- [width, 0, 0],
- [width, height, 0],
- [0, height, 0],
- ], dtype=float)
-
- edges = np.array([
- [0, 1], [1, 2], [2, 3], [3, 0], # boundary
- [0, 2], # diagonal
- ], dtype=int)
-
- faces = [[0, 1, 2], [0, 2, 3]]
- assignments = ["B", "B", "B", "B", "F"] # boundary + flat diagonal
- fold_angles = np.zeros(len(edges))
- rest_lengths = np.array([np.linalg.norm(verts[e[1]] - verts[e[0]]) for e in edges])
- strain = np.zeros(len(verts))
-
- mat = material or Material()
- return PaperState(
- vertices=verts, edges=edges, faces=faces,
- assignments=assignments, fold_angles=fold_angles,
- rest_lengths=rest_lengths, strain=strain,
- material=mat,
- bounding_box=np.array([width, height, 0.0]),
- )
-
-
-def _rotate_points(points: np.ndarray, axis_point: np.ndarray,
- axis_dir: np.ndarray, angle_rad: float) -> np.ndarray:
- """Rotate points around an arbitrary axis using Rodrigues' formula."""
- k = axis_dir / np.linalg.norm(axis_dir)
- translated = points - axis_point
- cos_a = math.cos(angle_rad)
- sin_a = math.sin(angle_rad)
- rotated = (translated * cos_a +
- np.cross(k, translated) * sin_a +
- k * (np.dot(translated, k).reshape(-1, 1)) * (1 - cos_a))
- return rotated + axis_point
-
-
-def apply_fold_mock(state: PaperState, fold: dict) -> tuple[PaperState, str | None]:
- """
- Apply a single fold operation to the paper state (mock version).
-
- fold = {
- "type": "valley" | "mountain",
- "line": {"start": [x, y], "end": [x, y]},
- "angle": 0-180
- }
-
- Returns (new_state, error_string_or_None).
- """
- fold_type = fold.get("type", "valley")
- line = fold.get("line", {})
- angle_deg = fold.get("angle", 180)
-
- start = np.array(line.get("start", [0, 0]), dtype=float)
- end = np.array(line.get("end", [0, 0]), dtype=float)
-
- if np.allclose(start, end):
- return state, "Fold line has zero length"
-
- if fold_type not in ("valley", "mountain"):
- return state, f"Unknown fold type: {fold_type}"
-
- if not (0 < angle_deg <= 180):
- return state, f"Angle must be in (0, 180], got {angle_deg}"
-
- # Fold direction: valley folds up (+z), mountain folds down (-z)
- sign = 1.0 if fold_type == "valley" else -1.0
- angle_rad = sign * math.radians(angle_deg)
-
- # Determine which vertices are on the "folding" side of the line
- line_dir_2d = end - start
- normal_2d = np.array([-line_dir_2d[1], line_dir_2d[0]]) # perpendicular
-
- new_verts = state.vertices.copy()
- for i, v in enumerate(new_verts):
- point_2d = v[:2] - start
- side = np.dot(point_2d, normal_2d)
- if side > 1e-9: # on the positive side → rotate
- axis_point = np.array([start[0], start[1], 0.0])
- axis_dir = np.array([line_dir_2d[0], line_dir_2d[1], 0.0])
- new_verts[i] = _rotate_points(
- v.reshape(1, -1), axis_point, axis_dir, angle_rad
- ).flatten()
-
- # Update bounding box (clamp near-zero values from floating point)
- bb = np.ptp(new_verts, axis=0) # max - min per axis
- bb = np.where(np.abs(bb) < 1e-10, 0.0, bb)
- # Add minimum thickness per layer (material thickness)
- thickness = state.material.thickness_mm / 1000.0 # convert mm to m
- num_layers = state.num_layers + 1
- bb[2] = max(bb[2], thickness * num_layers)
-
- # Mock strain: small random value per vertex
- new_strain = np.random.uniform(0, 0.01, len(new_verts))
-
- # Mock energy
- new_energy = state.energy + 0.1 * angle_deg / 180.0
-
- # Update assignments — add new edge as M or V
- new_assignments = state.assignments.copy()
-
- # Deployment ratio estimate: each full fold (180°) halves the area in one direction.
- # Partial folds reduce proportionally. This is a mock approximation —
- # the real engine will compute from actual face overlaps.
- fold_factor = angle_deg / 180.0 # 1.0 for full fold, 0.5 for 90°, etc.
- deploy_ratio = state.deployment_ratio * (1.0 - 0.5 * fold_factor)
-
- new_state = PaperState(
- vertices=new_verts,
- edges=state.edges.copy(),
- faces=state.faces.copy(),
- assignments=new_assignments,
- fold_angles=state.fold_angles.copy(),
- rest_lengths=state.rest_lengths.copy(),
- strain=new_strain,
- energy=new_energy,
- material=state.material,
- bounding_box=bb,
- deployment_ratio=deploy_ratio,
- num_layers=state.num_layers + 1,
- is_valid=True,
- kawasaki_violation=0.0,
- maekawa_violation=0.0,
- self_intersections=0,
- )
- return new_state, None
-
-
-def execute_fold_strategy(strategy_fn, paper_state: PaperState,
- max_folds: int = 20) -> tuple[PaperState, list[dict], str | None]:
- """
- Execute a fold_strategy function against the mock environment.
-
- Returns (final_state, applied_folds, error_or_None).
- """
- state_dict = paper_state.to_dict()
- try:
- folds = strategy_fn(state_dict)
- except Exception as e:
- return paper_state, [], f"Strategy function raised: {e}"
-
- if not isinstance(folds, list):
- return paper_state, [], "Strategy must return a list of fold dicts"
-
- applied = []
- current = paper_state
- for i, fold in enumerate(folds):
- if i >= max_folds:
- break
- if not isinstance(fold, dict):
- return current, applied, f"Fold {i} is not a dict"
-
- current, error = apply_fold_mock(current, fold)
- if error:
- return current, applied, f"Fold {i} failed: {error}"
- applied.append(fold)
-
- return current, applied, None
diff --git a/trainer/prompts.py b/trainer/prompts.py
deleted file mode 100644
index cde29f9d14221f613066bdae84371bc3af2136f0..0000000000000000000000000000000000000000
--- a/trainer/prompts.py
+++ /dev/null
@@ -1,169 +0,0 @@
-"""
-Prompt templates for origami fold strategy generation.
-
-The LLM receives a task description and paper state, then generates
-a fold_strategy(paper_state) function that returns fold operations.
-"""
-
-SYSTEM_PROMPT = """\
-You are an origami engineer. You design fold patterns for real-world applications \
-like solar panel packing, deployable shelters, and medical stents.
-
-You will be given a folding task with material constraints. Write a Python function \
-`fold_strategy(paper_state)` that returns a list of fold operations to achieve the goal.
-
-Rules:
-- Only use native Python (no imports except math, itertools, functools)
-- Each fold: {"type": "valley"|"mountain", "line": {"start": [x,y], "end": [x,y]}, "angle": 0-180}
-- Fold lines must intersect the paper boundaries
-- Fewer folds is better (efficiency matters)
-- Respect material strain limits
-- Output ONLY the function in ```python ... ``` backticks\
-"""
-
-
-TASK_TEMPLATES = {
- "half_fold": {
- "name": "half_fold",
- "prompt": """\
-TASK: Fold a {width}m x {height}m {material} sheet in half to minimize one dimension.
-
-MATERIAL: {material} (thickness: {thickness_mm}mm, max strain: {max_strain_pct}%)
-CONSTRAINTS: Maximum {max_folds} fold operations.
-TARGET: Deployment ratio <= 0.5 (folded area is half or less of original)
-
-CURRENT STATE:
- Sheet: {width}m x {height}m, flat (0 folds applied)
- Bounding box: {width}m x {height}m x 0.0m
-
-Write a fold_strategy(paper_state) function that returns a list of fold operations.
-Each fold: {{"type": "valley"|"mountain", "line": {{"start": [x,y], "end": [x,y]}}, "angle": 0-180}}
-
-```python
-def fold_strategy(paper_state):
- # Your code here
- return [...]
-```""",
- "target_ratio": 0.5,
- "max_folds": 3,
- },
-
- "letter_fold": {
- "name": "letter_fold",
- "prompt": """\
-TASK: Fold a {width}m x {height}m {material} sheet into thirds (like a letter).
-
-MATERIAL: {material} (thickness: {thickness_mm}mm, max strain: {max_strain_pct}%)
-CONSTRAINTS: Maximum {max_folds} fold operations.
-TARGET: Deployment ratio <= 0.33
-
-CURRENT STATE:
- Sheet: {width}m x {height}m, flat (0 folds applied)
-
-Write a fold_strategy(paper_state) function that returns a list of fold operations.
-Each fold: {{"type": "valley"|"mountain", "line": {{"start": [x,y], "end": [x,y]}}, "angle": 0-180}}
-
-```python
-def fold_strategy(paper_state):
- # Your code here
- return [...]
-```""",
- "target_ratio": 0.33,
- "max_folds": 5,
- },
-
- "solar_panel": {
- "name": "solar_panel",
- "prompt": """\
-TASK: Fold a {width}m x {height}m Mylar sheet to minimize packed volume for a solar panel.
-The folded panel must be deployable (unfold cleanly to near-original area).
-
-MATERIAL: Mylar (thickness: 0.05mm, Young's modulus: 4 GPa, max strain: 3%)
-CONSTRAINTS:
- - Maximum {max_folds} fold operations
- - Must pack into bounding box <= 15cm x 15cm x 5cm
- - Must deploy to >= 80% of original area
- - No self-intersections
-
-TARGET: Deployment ratio <= 0.05 (95% volume reduction)
-
-CURRENT STATE:
- Sheet: {width}m x {height}m, flat (0 folds applied)
- Bounding box: {width}m x {height}m x 0.0m
-
-HINT: Consider tessellated patterns like Miura-ori — alternating mountain and valley
-folds in a grid create a highly compact, single-DOF deployable structure.
-
-Write a fold_strategy(paper_state) function that returns a list of fold operations.
-Each fold: {{"type": "valley"|"mountain", "line": {{"start": [x,y], "end": [x,y]}}, "angle": 0-180}}
-
-```python
-def fold_strategy(paper_state):
- # Your code here
- return [...]
-```""",
- "target_ratio": 0.05,
- "max_folds": 20,
- },
-
- "stent_fold": {
- "name": "stent_fold",
- "prompt": """\
-TASK: Fold a {width}m x {height}m Nitinol sheet into a compact cylinder for a medical stent.
-
-MATERIAL: Nitinol (thickness: 0.1mm, Young's modulus: 75 GPa, max strain: 8%)
-CONSTRAINTS:
- - Maximum {max_folds} fold operations
- - Compressed diameter: 3mm
- - Deployed diameter: 10mm
- - Must be radially deployable
-
-TARGET: Minimize packed cross-section while maintaining deployability.
-
-Write a fold_strategy(paper_state) function that returns a list of fold operations.
-
-```python
-def fold_strategy(paper_state):
- # Your code here
- return [...]
-```""",
- "target_ratio": 0.1,
- "max_folds": 15,
- },
-}
-
-
-# Default task configs for each level
-TASK_CONFIGS = {
- "half_fold": {
- "width": 1.0, "height": 1.0, "material": "paper",
- "thickness_mm": 0.1, "max_strain_pct": 3, "max_folds": 3,
- },
- "letter_fold": {
- "width": 1.0, "height": 1.0, "material": "paper",
- "thickness_mm": 0.1, "max_strain_pct": 3, "max_folds": 5,
- },
- "solar_panel": {
- "width": 1.0, "height": 1.0, "material": "mylar",
- "thickness_mm": 0.05, "max_strain_pct": 3, "max_folds": 20,
- },
- "stent_fold": {
- "width": 0.1, "height": 0.03, "material": "nitinol",
- "thickness_mm": 0.1, "max_strain_pct": 8, "max_folds": 15,
- },
-}
-
-
-def build_prompt(task_name: str = "half_fold", **overrides) -> str:
- """Build a complete user prompt for a given task."""
- task = TASK_TEMPLATES[task_name]
- config = {**TASK_CONFIGS[task_name], **overrides}
- return task["prompt"].format(**config)
-
-
-def get_task_target_ratio(task_name: str) -> float:
- return TASK_TEMPLATES[task_name]["target_ratio"]
-
-
-def get_task_max_folds(task_name: str) -> int:
- return TASK_TEMPLATES[task_name]["max_folds"]
diff --git a/trainer/rewards.py b/trainer/rewards.py
deleted file mode 100644
index 34d006beaf41550c29f47403e30694f491b957fe..0000000000000000000000000000000000000000
--- a/trainer/rewards.py
+++ /dev/null
@@ -1,388 +0,0 @@
-"""
-Reward functions for origami GRPO training.
-
-Three reward functions following the 2048 pattern:
- 1. code_valid — Does the generated code parse and produce fold instructions?
- 2. physically_valid — Are the folds geometrically/physically valid?
- 3. fold_quality — How good is the folding solution (compactness, efficiency)?
-
-Lexicographic gating (from SpatialThinker): if code doesn't parse,
-all downstream rewards are 0. This prevents reward hacking.
-"""
-
-import ast
-import sys
-import math
-import traceback
-from typing import Callable
-
-# Use real engine if available, fall back to mock
-try:
- from engine.paper import Paper
- from engine.fold_engine import execute_fold_strategy
- from engine.materials import Material, get_material
- from engine.validation import validate_paper
- from engine.metrics import compute_metrics
-
- def _create_sheet(width, height, material):
- return Paper.create_flat_sheet(width, height, material)
-
- USE_REAL_ENGINE = True
-except ImportError:
- from trainer.mock_env import (
- PaperState as Paper, create_flat_sheet, execute_fold_strategy, Material
- )
-
- def _create_sheet(width, height, material):
- return create_flat_sheet(width, height, material)
-
- def validate_paper(p):
- from types import SimpleNamespace
- return SimpleNamespace(
- is_valid=p.is_valid, kawasaki_valid=True, maekawa_valid=True,
- kawasaki_violation=p.kawasaki_violation,
- maekawa_violation=p.maekawa_violation,
- self_intersection_count=p.self_intersections,
- )
-
- def compute_metrics(p, orig):
- return {
- "deployment_ratio": p.deployment_ratio,
- "fold_count": sum(1 for a in p.assignments if a in ("M", "V")),
- "max_strain": float(p.strain.max()) if len(p.strain) > 0 else 0.0,
- }
-
- USE_REAL_ENGINE = False
-
-
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-def extract_function(text: str) -> str | None:
- """Extract a Python function from triple-backtick code blocks."""
- if text.count("```") < 2:
- return None
- first = text.find("```") + 3
- second = text.find("```", first)
- fx = text[first:second].strip()
- fx = fx.removeprefix("python\n").removeprefix("python\r\n")
- # Find the def statement
- def_idx = fx.find("def ")
- if def_idx == -1:
- return None
- fx = fx[def_idx:]
- if fx.startswith("def fold_strategy("):
- return fx
- return None
-
-
-def check_imports_stdlib_only(code: str) -> tuple[bool, str]:
- """Check that code only imports from Python stdlib."""
- try:
- tree = ast.parse(code)
- except SyntaxError as e:
- return False, f"syntax error: {e}"
-
- ALLOWED_MODULES = {
- "math", "itertools", "functools", "collections", "copy",
- "operator", "typing", "random", "heapq", "bisect",
- }
-
- for node in ast.walk(tree):
- if isinstance(node, ast.Import):
- for alias in node.names:
- root = alias.name.split(".")[0]
- if root not in ALLOWED_MODULES:
- return False, f"non-stdlib import: {alias.name}"
- elif isinstance(node, ast.ImportFrom):
- if node.module:
- root = node.module.split(".")[0]
- if root not in ALLOWED_MODULES:
- return False, f"non-stdlib import: {node.module}"
-
- return True, "ok"
-
-
-def create_sandboxed_function(code: str) -> Callable:
- """
- Execute the function code in a restricted namespace.
- Returns the fold_strategy function object.
- """
- allowed_builtins = {
- "range", "len", "int", "float", "str", "list", "dict", "tuple",
- "set", "bool", "abs", "min", "max", "sum", "sorted", "reversed",
- "enumerate", "zip", "map", "filter", "round", "isinstance",
- "True", "False", "None", "print",
- }
- safe_builtins = {k: __builtins__[k] if isinstance(__builtins__, dict)
- else getattr(__builtins__, k)
- for k in allowed_builtins
- if (k in __builtins__ if isinstance(__builtins__, dict)
- else hasattr(__builtins__, k))}
- safe_builtins["__import__"] = __import__ # needed for stdlib imports
-
- namespace = {"__builtins__": safe_builtins}
- exec(code, namespace)
-
- if "fold_strategy" not in namespace:
- raise ValueError("No fold_strategy function defined")
-
- return namespace["fold_strategy"]
-
-
-# ---------------------------------------------------------------------------
-# State for strategy execution
-# ---------------------------------------------------------------------------
-
-# Current task config (set by train.py before training starts)
-if USE_REAL_ENGINE:
- _default_material = get_material("paper")
-else:
- _default_material = Material()
-
-_current_task = {
- "width": 1.0,
- "height": 1.0,
- "material": _default_material,
- "target_ratio": 0.5,
- "max_folds": 3,
-}
-
-PRINT_EVERY = 5
-_print_counter = 0
-
-
-def set_task_config(width=1.0, height=1.0, material=None,
- target_ratio=0.5, max_folds=3):
- global _current_task
- _current_task = {
- "width": width,
- "height": height,
- "material": material or Material(),
- "target_ratio": target_ratio,
- "max_folds": max_folds,
- }
-
-
-# ---------------------------------------------------------------------------
-# Reward 1: code_valid
-# ---------------------------------------------------------------------------
-
-def code_valid(completions, **kwargs) -> list[float]:
- """
- Does the generated code parse as valid Python and produce a callable?
-
- +1.0 — valid function that can be created
- -0.5 — correct structure but exec/sandbox fails
- -2.0 — no function found or syntax error
- -20.0 — non-stdlib imports (heavy penalty)
- """
- scores = []
- for completion in completions:
- response = completion[0]["content"]
- function_code = extract_function(response)
-
- if function_code is None:
- scores.append(-2.0)
- continue
-
- ok, info = check_imports_stdlib_only(function_code)
- if not ok:
- if "syntax error" in info:
- scores.append(-2.0)
- else:
- scores.append(-20.0) # non-stdlib imports
- continue
-
- try:
- create_sandboxed_function(function_code)
- scores.append(1.0)
- except Exception:
- scores.append(-0.5)
-
- return scores
-
-
-# ---------------------------------------------------------------------------
-# Reward 2: physically_valid
-# ---------------------------------------------------------------------------
-
-def physically_valid(completions, **kwargs) -> list[float]:
- """
- Are the folds physically possible?
-
- +1.0 — all folds valid, no violations
- -2.0 — per Kawasaki/Maekawa violation
- -5.0 — any self-intersection
- -1.0 — strain exceeds material limit
- 0.0 — function broken / can't run
- """
- scores = []
- for completion in completions:
- response = completion[0]["content"]
- function_code = extract_function(response)
-
- if function_code is None:
- scores.append(0.0)
- continue
-
- ok, info = check_imports_stdlib_only(function_code)
- if not ok:
- scores.append(0.0)
- continue
-
- try:
- strategy_fn = create_sandboxed_function(function_code)
- except Exception:
- scores.append(0.0)
- continue
-
- try:
- paper = _create_sheet(
- _current_task["width"],
- _current_task["height"],
- _current_task["material"],
- )
- original = paper
- final_state, applied, error = execute_fold_strategy(
- strategy_fn, paper, _current_task["max_folds"]
- )
-
- if error:
- scores.append(0.0)
- continue
-
- if len(applied) == 0:
- scores.append(0.0)
- continue
-
- # Score based on validity using engine validation
- val = validate_paper(final_state)
- metrics = compute_metrics(final_state, original)
-
- score = 1.0
- score -= 2.0 * val.kawasaki_violation
- score -= 2.0 * val.maekawa_violation
- if val.self_intersection_count > 0:
- score -= 5.0
- max_strain = metrics.get("max_strain", 0.0)
- if max_strain > _current_task["material"].max_strain:
- score -= 1.0
-
- scores.append(score)
-
- except TimeoutError:
- scores.append(-1.0)
- except Exception:
- scores.append(0.0)
-
- return scores
-
-
-# ---------------------------------------------------------------------------
-# Reward 3: fold_quality
-# ---------------------------------------------------------------------------
-
-def fold_quality(completions, **kwargs) -> list[float]:
- """
- How good is the folding solution?
-
- +20.0 * compactness — main reward (1 - deployment_ratio)
- +10.0 bonus — if meets target ratio
- -0.5 per fold — efficiency penalty
- -3.0 * overstrain — material stress penalty
- -1.0 — timeout
- -3.0 — exception
- 0.0 — function broken
- """
- global _print_counter
- scores = []
-
- for completion in completions:
- response = completion[0]["content"]
- function_code = extract_function(response)
-
- should_print = (_print_counter % PRINT_EVERY == 0)
- _print_counter += 1
-
- if should_print:
- print(f"\n--- Strategy (sample {_print_counter}) ---")
- print(function_code if function_code else "[no function extracted]")
-
- if function_code is None:
- scores.append(0.0)
- continue
-
- ok, info = check_imports_stdlib_only(function_code)
- if not ok:
- scores.append(0.0)
- continue
-
- try:
- strategy_fn = create_sandboxed_function(function_code)
- except Exception:
- scores.append(0.0)
- continue
-
- try:
- paper = _create_sheet(
- _current_task["width"],
- _current_task["height"],
- _current_task["material"],
- )
- original = paper
- final_state, applied, error = execute_fold_strategy(
- strategy_fn, paper, _current_task["max_folds"]
- )
-
- if error:
- if should_print:
- print(f"Error: {error}")
- scores.append(0.0)
- continue
-
- num_folds = len(applied)
- if num_folds == 0:
- scores.append(0.0)
- continue
-
- # Use engine metrics
- metrics = compute_metrics(final_state, original)
- deploy_ratio = metrics.get("deployment_ratio", 1.0)
- max_strain = metrics.get("max_strain", 0.0)
-
- # Compactness: main reward signal
- compactness = 1.0 - deploy_ratio
- score = 20.0 * compactness
-
- # Bonus for meeting target
- if deploy_ratio <= _current_task["target_ratio"]:
- score += 10.0
-
- # Fold efficiency penalty
- score -= 0.5 * num_folds
-
- # Strain penalty
- mat_limit = _current_task["material"].max_strain
- if max_strain > mat_limit:
- score -= 3.0 * (max_strain / mat_limit)
-
- if should_print:
- print(f"Folds: {num_folds}, Ratio: {deploy_ratio:.3f}, "
- f"Compactness: {compactness:.3f}, Score: {score:.2f}")
- bb = metrics.get("bounding_box", {})
- print(f"BBox: {bb.get('x',0):.3f} x {bb.get('y',0):.3f} x {bb.get('z',0):.3f}")
-
- scores.append(score)
-
- except TimeoutError:
- if should_print:
- print("Timeout!")
- scores.append(-1.0)
- except Exception as e:
- if should_print:
- print(f"Exception: {e}")
- scores.append(-3.0)
-
- return scores
diff --git a/trainer/train.py b/trainer/train.py
deleted file mode 100644
index a781ad7f931d9bedc0227c1144fe42c073965acf..0000000000000000000000000000000000000000
--- a/trainer/train.py
+++ /dev/null
@@ -1,208 +0,0 @@
-"""
-Origami GRPO Training Script
-
-Usage (Colab with T4/A100):
- python trainer/train.py
-
-Or in a notebook:
- %run trainer/train.py
-
-Requires: unsloth, trl>=0.22.2, transformers>=4.56.2, trackio, datasets
-"""
-
-import os
-import sys
-
-# Ensure project root is on path
-PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-if PROJECT_ROOT not in sys.path:
- sys.path.insert(0, PROJECT_ROOT)
-
-from trainer.prompts import build_prompt, SYSTEM_PROMPT, get_task_target_ratio, get_task_max_folds
-from trainer.rewards import code_valid, physically_valid, fold_quality, set_task_config
-
-try:
- from engine.materials import get_material
- Material = type(get_material("paper")) # get the Material class
-except ImportError:
- from trainer.mock_env import Material
- def get_material(name):
- return Material()
-
-# ============================================================================
-# Config
-# ============================================================================
-
-MODEL_NAME = "unsloth/Qwen2.5-7B-Instruct"
-MAX_SEQ_LENGTH = 2048
-LORA_RANK = 4
-
-# Start with the simplest task
-TASK_NAME = "half_fold"
-
-# GRPO hyperparameters (from 2048 reference, adapted for origami)
-LEARNING_RATE = 2e-4
-MAX_STEPS = 600
-NUM_GENERATIONS = 2
-TEMPERATURE = 1.0
-BATCH_SIZE = 1
-GRAD_ACCUM = 1
-DATASET_SIZE = 1000 # replicated prompt dataset
-
-
-def main():
- # ========================================================================
- # 1. Load model with Unsloth
- # ========================================================================
- from unsloth import FastLanguageModel
-
- print(f"Loading model: {MODEL_NAME}")
- model, tokenizer = FastLanguageModel.from_pretrained(
- model_name=MODEL_NAME,
- load_in_4bit=True,
- max_seq_length=MAX_SEQ_LENGTH,
- )
-
- # ========================================================================
- # 2. Apply LoRA adapters
- # ========================================================================
- model = FastLanguageModel.get_peft_model(
- model,
- r=LORA_RANK,
- target_modules=[
- "q_proj", "k_proj", "v_proj", "o_proj",
- "gate_proj", "up_proj", "down_proj",
- ],
- lora_alpha=LORA_RANK * 2,
- use_gradient_checkpointing="unsloth",
- random_state=3407,
- )
-
- # ========================================================================
- # 3. Build prompt and dataset
- # ========================================================================
- user_prompt = build_prompt(TASK_NAME)
- target_ratio = get_task_target_ratio(TASK_NAME)
- max_folds = get_task_max_folds(TASK_NAME)
-
- # Configure reward functions with task parameters
- set_task_config(
- width=1.0,
- height=1.0,
- material=get_material("paper"),
- target_ratio=target_ratio,
- max_folds=max_folds,
- )
-
- # Create replicated prompt dataset (same pattern as 2048)
- from datasets import Dataset
-
- dataset = Dataset.from_list([
- {
- "prompt": [
- {"role": "system", "content": SYSTEM_PROMPT},
- {"role": "user", "content": user_prompt},
- ],
- }
- ] * DATASET_SIZE)
-
- # Calculate prompt token length for max_completion_length
- prompt_tokens = tokenizer.apply_chat_template(
- [
- {"role": "system", "content": SYSTEM_PROMPT},
- {"role": "user", "content": user_prompt},
- ],
- add_generation_prompt=True,
- tokenize=True,
- )
- max_prompt_length = len(prompt_tokens) + 1
- max_completion_length = MAX_SEQ_LENGTH - max_prompt_length
- print(f"Prompt tokens: {max_prompt_length}, Max completion: {max_completion_length}")
-
- # ========================================================================
- # 4. Test inference before training
- # ========================================================================
- print("\n=== Pre-training inference test ===")
- text = tokenizer.apply_chat_template(
- [
- {"role": "system", "content": SYSTEM_PROMPT},
- {"role": "user", "content": user_prompt},
- ],
- tokenize=False,
- add_generation_prompt=True,
- )
-
- from transformers import TextStreamer
- _ = model.generate(
- **tokenizer(text, return_tensors="pt").to("cuda"),
- temperature=TEMPERATURE,
- max_new_tokens=min(512, max_completion_length),
- streamer=TextStreamer(tokenizer, skip_prompt=True),
- )
-
- # ========================================================================
- # 5. Configure GRPO training
- # ========================================================================
- from trl import GRPOConfig, GRPOTrainer
-
- training_args = GRPOConfig(
- temperature=TEMPERATURE,
- learning_rate=LEARNING_RATE,
- weight_decay=0.001,
- warmup_ratio=0.1,
- lr_scheduler_type="linear",
- optim="adamw_8bit",
- logging_steps=1,
- per_device_train_batch_size=BATCH_SIZE,
- gradient_accumulation_steps=GRAD_ACCUM,
- num_generations=NUM_GENERATIONS,
- max_prompt_length=max_prompt_length,
- max_completion_length=max_completion_length,
- max_steps=MAX_STEPS,
- save_steps=100,
- report_to="trackio",
- output_dir="outputs",
- )
-
- # ========================================================================
- # 6. Create trainer and start training
- # ========================================================================
- trainer = GRPOTrainer(
- model=model,
- processing_class=tokenizer,
- reward_funcs=[
- code_valid, # Reward 1: valid Python?
- physically_valid, # Reward 2: physically possible folds?
- fold_quality, # Reward 3: how good is the solution?
- ],
- args=training_args,
- train_dataset=dataset,
- )
-
- print(f"\n=== Starting GRPO training: {TASK_NAME} ===")
- print(f"Steps: {MAX_STEPS}, Generations: {NUM_GENERATIONS}, LR: {LEARNING_RATE}")
- trainer.train()
-
- # ========================================================================
- # 7. Post-training inference
- # ========================================================================
- print("\n=== Post-training inference ===")
- _ = model.generate(
- **tokenizer(text, return_tensors="pt").to("cuda"),
- temperature=TEMPERATURE,
- max_new_tokens=min(1024, max_completion_length),
- streamer=TextStreamer(tokenizer, skip_prompt=True),
- )
-
- # ========================================================================
- # 8. Save model (optional)
- # ========================================================================
- save_path = "outputs/origami-fold-lora"
- print(f"\nSaving LoRA adapter to {save_path}")
- model.save_pretrained(save_path)
- tokenizer.save_pretrained(save_path)
- print("Done!")
-
-
-if __name__ == "__main__":
- main()
diff --git a/training/__init__.py b/training/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/training/demo.py b/training/demo.py
new file mode 100644
index 0000000000000000000000000000000000000000..3e72332772ce3edfa9e8d68412b2c56b942189ce
--- /dev/null
+++ b/training/demo.py
@@ -0,0 +1,251 @@
+"""
+training/demo.py — Run 8 zero-shot rollouts and stream them to the grid viewer.
+
+Usage:
+ cd /path/to/optigami
+ python -m training.demo
+
+Then open: http://localhost:9001/viewer/training.html
+
+Each of the 8 "strategies" is a heuristic that mimics what a pretrained LLM might
+produce for different tasks — varying from near-optimal to poor. This exercises
+the full broadcast → grid viewer pipeline without requiring an LLM API key.
+"""
+from __future__ import annotations
+
+import asyncio
+import time
+import uuid
+from typing import Callable
+
+import uvicorn
+
+from server.app import app, broadcast
+from training.runner import run_batch
+
+
+# ── 8 zero-shot heuristic strategies ──────────────────────────────────────────
+# Each is a callable: paper_state (dict) → fold_dict
+# These represent the range of strategies a pretrained LLM might generate.
+
+def strategy_perfect_half(paper_state: dict) -> dict:
+ """Valley fold exactly at horizontal midline — optimal for half_fold."""
+ return {"type": "valley", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0}
+
+
+def strategy_slight_offset(paper_state: dict) -> dict:
+ """Valley fold slightly off-center — almost optimal."""
+ return {"type": "valley", "line": {"start": [0.0, 0.48], "end": [1.0, 0.48]}, "angle": 180.0}
+
+
+def strategy_thirds(paper_state: dict) -> dict:
+ """Letter fold at one-third — wrong for half_fold, generates interesting geometry."""
+ fold_count = paper_state.get("fold_count", 0)
+ positions = [0.333, 0.667]
+ if fold_count >= len(positions):
+ return {"type": "stop", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 0.0}
+ return {
+ "type": "valley" if fold_count == 0 else "mountain",
+ "line": {"start": [0.0, positions[fold_count]], "end": [1.0, positions[fold_count]]},
+ "angle": 180.0,
+ }
+
+
+def strategy_vertical(paper_state: dict) -> dict:
+ """Vertical fold — gets compactness but in wrong dimension for target_box."""
+ return {"type": "valley", "line": {"start": [0.5, 0.0], "end": [0.5, 1.0]}, "angle": 180.0}
+
+
+def strategy_mountain(paper_state: dict) -> dict:
+ """Mountain fold at midline — same geometry, different assignment."""
+ return {"type": "mountain", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0}
+
+
+def strategy_accordion(paper_state: dict) -> dict:
+ """Accordion 3-fold — overfolds, achieves high compactness but more folds."""
+ fold_count = paper_state.get("fold_count", 0)
+ positions = [0.25, 0.5, 0.75]
+ assignments = ["valley", "mountain", "valley"]
+ if fold_count >= len(positions):
+ return {"type": "stop", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 0.0}
+ return {
+ "type": assignments[fold_count],
+ "line": {"start": [0.0, positions[fold_count]], "end": [1.0, positions[fold_count]]},
+ "angle": 180.0,
+ }
+
+
+def strategy_diagonal(paper_state: dict) -> dict:
+ """Diagonal fold — achieves compactness but irregular bounding box."""
+ return {"type": "valley", "line": {"start": [0.0, 0.0], "end": [1.0, 1.0]}, "angle": 180.0}
+
+
+def strategy_quarter(paper_state: dict) -> dict:
+ """Two perpendicular folds — 4x compactness for quarter_fold task."""
+ fold_count = paper_state.get("fold_count", 0)
+ if fold_count == 0:
+ return {"type": "valley", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0}
+ if fold_count == 1:
+ return {"type": "valley", "line": {"start": [0.5, 0.0], "end": [0.5, 1.0]}, "angle": 180.0}
+ return {"type": "stop", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 0.0}
+
+
+STRATEGIES: list[tuple[str, Callable]] = [
+ ("perfect_half", strategy_perfect_half),
+ ("slight_offset", strategy_slight_offset),
+ ("thirds_fold", strategy_thirds),
+ ("vertical_fold", strategy_vertical),
+ ("mountain_fold", strategy_mountain),
+ ("accordion_3", strategy_accordion),
+ ("diagonal", strategy_diagonal),
+ ("quarter_fold", strategy_quarter),
+]
+
+
+# ── Demo runner ────────────────────────────────────────────────────────────────
+
+async def run_demo(task_name: str = "half_fold", delay_s: float = 0.5) -> None:
+ """Wait for server to be ready, then fire 8 episodes."""
+ # Give uvicorn time to bind and call startup hook (sets broadcast._loop)
+ await asyncio.sleep(1.5)
+
+ batch_id = 1
+ names, fns = zip(*STRATEGIES)
+ ep_ids = [f"ep_{name}" for name in names]
+
+ print(f"\n[demo] Starting batch {batch_id} — task: {task_name}")
+ print(f"[demo] Open http://localhost:9001/viewer/training.html\n")
+
+ # Signal grid to clear and show G=8
+ await broadcast.start_batch(batch_id, len(fns))
+
+ await asyncio.sleep(delay_s)
+
+ # Run all 8 episodes in the thread pool; broadcast_fn fires into this loop
+ results = await asyncio.gather(*[
+ asyncio.to_thread(
+ _run_one,
+ fn,
+ task_name,
+ ep_id,
+ broadcast.publish,
+ )
+ for fn, ep_id in zip(fns, ep_ids)
+ ])
+
+ scores = [r["score"] for r in results]
+ best_idx = max(range(len(scores)), key=lambda i: scores[i])
+
+ await broadcast.finish_batch(batch_id, scores, best_episode_id=ep_ids[best_idx])
+
+ print("\n[demo] Results:")
+ for name, result in zip(names, results):
+ print(f" {name:20s} score={result['score']:+.2f} status={result['status']}")
+ print(f"\n[demo] Best: {names[best_idx]} (score={scores[best_idx]:+.2f})")
+ print("\n[demo] Grid viewer running. Press Ctrl+C to stop.\n")
+
+
+def _run_one(
+ strategy_fn: Callable,
+ task_name: str,
+ ep_id: str,
+ broadcast_fn: Callable,
+) -> dict:
+ """Thin wrapper: adds a small sleep between steps so the viewer can animate."""
+ from server.models import OrigamiAction
+ from server.origami_environment import OrigamiEnvironment
+
+ env = OrigamiEnvironment()
+ obs = env.reset(task_name=task_name)
+
+ broadcast_fn(ep_id, {
+ "type": "episode_update",
+ "episode_id": ep_id,
+ "task_name": task_name,
+ "step": 0,
+ "observation": _obs_dict(obs),
+ })
+
+ max_steps = env._task.get("max_folds", 10) if env._task else 10
+ status = "done"
+
+ for step_idx in range(max_steps):
+ if obs.done:
+ break
+
+ time.sleep(0.3) # pace so the viewer can animate each step
+
+ fold_dict = strategy_fn(obs.paper_state)
+
+ if fold_dict.get("type") == "stop":
+ break
+
+ action = OrigamiAction(
+ fold_type=fold_dict["type"],
+ fold_line=fold_dict["line"],
+ fold_angle=float(fold_dict.get("angle", 180.0)),
+ )
+ obs = env.step(action)
+
+ broadcast_fn(ep_id, {
+ "type": "episode_update",
+ "episode_id": ep_id,
+ "task_name": task_name,
+ "step": step_idx + 1,
+ "observation": _obs_dict(obs),
+ })
+
+ if obs.done:
+ break
+ else:
+ status = "timeout"
+
+ score = obs.reward if obs.reward is not None else env._total_reward or 0.0
+
+ broadcast_fn(ep_id, {
+ "type": "episode_done",
+ "episode_id": ep_id,
+ "status": status,
+ "score": float(score),
+ "final_metrics": obs.metrics,
+ })
+
+ return {
+ "episode_id": ep_id,
+ "score": float(score),
+ "final_metrics": obs.metrics,
+ "status": status,
+ }
+
+
+def _obs_dict(obs) -> dict:
+ try:
+ return obs.model_dump()
+ except AttributeError:
+ return {
+ "paper_state": getattr(obs, "paper_state", {}),
+ "metrics": getattr(obs, "metrics", {}),
+ "fold_history": getattr(obs, "fold_history", []),
+ "done": getattr(obs, "done", False),
+ "reward": getattr(obs, "reward", None),
+ }
+
+
+# ── Entry point ────────────────────────────────────────────────────────────────
+
+async def _main() -> None:
+ config = uvicorn.Config(app, host="0.0.0.0", port=9001, log_level="warning")
+ server = uvicorn.Server(config)
+
+ # Run demo concurrently with the uvicorn server
+ await asyncio.gather(
+ server.serve(),
+ run_demo(task_name="half_fold"),
+ )
+
+
+if __name__ == "__main__":
+ try:
+ asyncio.run(_main())
+ except KeyboardInterrupt:
+ print("\n[demo] Stopped.")
diff --git a/training/demo_llm.py b/training/demo_llm.py
new file mode 100644
index 0000000000000000000000000000000000000000..6c05a78d6af5225de45957451b1bc6245b944299
--- /dev/null
+++ b/training/demo_llm.py
@@ -0,0 +1,232 @@
+"""
+training/demo_llm.py — 8 rollouts using Claude as the zero-shot fold strategist.
+
+Usage:
+ cd /path/to/optigami
+ ANTHROPIC_API_KEY=sk-... python -m training.demo_llm
+
+Each of the 8 episodes calls Claude (claude-haiku-4-5) once per fold step.
+Claude sees the current paper_state metrics and decides the next fold.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import os
+import re
+import time
+from typing import Any
+
+import anthropic
+import uvicorn
+
+from server.app import app, broadcast
+from server.models import OrigamiAction
+from server.origami_environment import OrigamiEnvironment
+from server.tasks import get_task_by_name
+
+
+TASK_NAME = "half_fold"
+NUM_EPISODES = 8
+MODEL = "claude-haiku-4-5-20251001"
+
+
+# ── LLM strategy factory ───────────────────────────────────────────────────────
+
+def make_llm_strategy(client: anthropic.Anthropic, task: dict, episode_num: int):
+ """Return a strategy_fn for one episode. Each episode gets its own call history."""
+ history: list[dict[str, Any]] = []
+
+ def strategy(paper_state: dict) -> dict:
+ fold_count = paper_state.get("fold_count", 0)
+ compactness = paper_state.get("compactness", 0)
+ bb = paper_state.get("bounding_box", [1, 1, 0])
+ target_box = task.get("target_box", [1, 0.5, 0.02])
+ max_folds = task.get("max_folds", 3)
+
+ user_msg = f"""You are folding a {task['width']}x{task['height']} sheet of {task['material']}.
+Task: {task['description']}
+Target box to fit inside: {target_box}
+Max folds allowed: {max_folds}
+
+Current state (fold {fold_count}/{max_folds}):
+ compactness: {compactness:.3f} (1.0 = fully packed, 0.0 = flat)
+ bounding_box: [{bb[0]:.3f}, {bb[1]:.3f}, {bb[2]:.4f}]
+ fits_target_box: {paper_state.get('fits_target_box', False)}
+
+Choose the next fold. Respond with ONLY valid JSON, no other text:
+{{
+ "type": "valley" or "mountain" or "stop",
+ "line": {{"start": [x, y], "end": [x, y]}},
+ "angle": 180
+}}
+
+Coordinates are normalized 0-1. Use "stop" if done."""
+
+ history.append({"role": "user", "content": user_msg})
+
+ response = client.messages.create(
+ model=MODEL,
+ max_tokens=120,
+ messages=history,
+ )
+ reply = response.content[0].text.strip()
+ history.append({"role": "assistant", "content": reply})
+
+ # Extract JSON — handle markdown code blocks
+ match = re.search(r'\{[^{}]+\}', reply, re.DOTALL)
+ if not match:
+ return {"type": "stop", "line": {"start": [0, 0.5], "end": [1, 0.5]}, "angle": 0.0}
+
+ fold_dict = json.loads(match.group())
+ # Normalize: ensure required keys
+ fold_dict.setdefault("type", "valley")
+ fold_dict.setdefault("line", {"start": [0.0, 0.5], "end": [1.0, 0.5]})
+ fold_dict.setdefault("angle", 180.0)
+ return fold_dict
+
+ return strategy
+
+
+# ── Episode runner ─────────────────────────────────────────────────────────────
+
+def run_episode_llm(
+ strategy_fn,
+ task_name: str,
+ ep_id: str,
+ broadcast_fn,
+) -> dict:
+ env = OrigamiEnvironment()
+ obs = env.reset(task_name=task_name)
+ task = env._task or {}
+
+ broadcast_fn(ep_id, {
+ "type": "episode_update",
+ "episode_id": ep_id,
+ "task_name": task_name,
+ "step": 0,
+ "observation": _obs_dict(obs),
+ })
+
+ max_steps = task.get("max_folds", 5)
+ status = "done"
+
+ for step_idx in range(max_steps):
+ if obs.done:
+ break
+
+ # Build a flat paper_state dict for the LLM (add metrics inline)
+ ps = dict(obs.paper_state)
+ ps.update(obs.metrics) # compactness, fits_target_box, etc.
+ ps["fold_count"] = step_idx
+
+ try:
+ fold_dict = strategy_fn(ps)
+ except Exception as exc:
+ broadcast_fn(ep_id, {
+ "type": "episode_done", "episode_id": ep_id,
+ "status": "error", "score": 0.0,
+ "final_metrics": obs.metrics, "error": str(exc),
+ })
+ return {"episode_id": ep_id, "score": 0.0, "status": "error"}
+
+ if fold_dict.get("type") == "stop":
+ break
+
+ time.sleep(0.4) # pace for viewer animation
+
+ action = OrigamiAction(
+ fold_type=fold_dict["type"],
+ fold_line=fold_dict["line"],
+ fold_angle=float(fold_dict.get("angle", 180.0)),
+ )
+ obs = env.step(action)
+
+ broadcast_fn(ep_id, {
+ "type": "episode_update",
+ "episode_id": ep_id,
+ "task_name": task_name,
+ "step": step_idx + 1,
+ "observation": _obs_dict(obs),
+ })
+
+ if obs.done:
+ break
+ else:
+ status = "timeout"
+
+ score = obs.reward if obs.reward is not None else (env._total_reward or 0.0)
+ broadcast_fn(ep_id, {
+ "type": "episode_done",
+ "episode_id": ep_id,
+ "status": status,
+ "score": float(score),
+ "final_metrics": obs.metrics,
+ })
+
+ return {"episode_id": ep_id, "score": float(score), "status": status}
+
+
+def _obs_dict(obs) -> dict:
+ try:
+ return obs.model_dump()
+ except AttributeError:
+ return {
+ "paper_state": getattr(obs, "paper_state", {}),
+ "metrics": getattr(obs, "metrics", {}),
+ "fold_history": getattr(obs, "fold_history", []),
+ "done": getattr(obs, "done", False),
+ "reward": getattr(obs, "reward", None),
+ }
+
+
+# ── Main ──────────────────────────────────────────────────────────────────────
+
+async def run_demo() -> None:
+ api_key = os.environ.get("ANTHROPIC_API_KEY")
+ if not api_key:
+ raise RuntimeError("Set ANTHROPIC_API_KEY environment variable")
+
+ client = anthropic.Anthropic(api_key=api_key)
+ task = get_task_by_name(TASK_NAME)
+
+ await asyncio.sleep(1.5) # wait for server startup
+
+ print(f"\n[llm-demo] Model: {MODEL}")
+ print(f"[llm-demo] Task: {TASK_NAME} — {task['description']}")
+ print(f"[llm-demo] Open: http://localhost:9001/viewer/training.html\n")
+
+ await broadcast.start_batch(1, NUM_EPISODES)
+
+ ep_ids = [f"ep_{i:02d}" for i in range(NUM_EPISODES)]
+ strategies = [make_llm_strategy(client, task, i) for i in range(NUM_EPISODES)]
+
+ # Run all episodes concurrently (each makes its own Claude API calls)
+ results = await asyncio.gather(*[
+ asyncio.to_thread(run_episode_llm, fn, TASK_NAME, ep_id, broadcast.publish)
+ for fn, ep_id in zip(strategies, ep_ids)
+ ])
+
+ scores = [r["score"] for r in results]
+ best_idx = max(range(len(scores)), key=lambda i: scores[i])
+
+ await broadcast.finish_batch(1, scores, best_episode_id=ep_ids[best_idx])
+
+ print("\n[llm-demo] Results:")
+ for i, result in enumerate(results):
+ print(f" ep_{i:02d} score={result['score']:+.2f} status={result['status']}")
+ print(f"\n[llm-demo] Best: ep_{best_idx:02d} (score={scores[best_idx]:+.2f})")
+ print("\n[llm-demo] Press Ctrl+C to stop.\n")
+
+
+async def _main() -> None:
+ config = uvicorn.Config(app, host="0.0.0.0", port=9001, log_level="warning")
+ server = uvicorn.Server(config)
+ await asyncio.gather(server.serve(), run_demo())
+
+
+if __name__ == "__main__":
+ try:
+ asyncio.run(_main())
+ except KeyboardInterrupt:
+ print("\n[llm-demo] Stopped.")
diff --git a/training/runner.py b/training/runner.py
new file mode 100644
index 0000000000000000000000000000000000000000..bf4de2f41ef42d781b5af158dafc08e77db49d34
--- /dev/null
+++ b/training/runner.py
@@ -0,0 +1,191 @@
+"""
+TrainingRunner — parallel episode executor for GRPO training.
+
+Each episode runs in a ThreadPoolExecutor thread.
+After every env.step(), observations are pushed to the broadcast server (fire-and-forget).
+"""
+from __future__ import annotations
+
+import uuid
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Any, Callable, Optional
+
+from server.models import OrigamiAction
+from server.origami_environment import OrigamiEnvironment
+
+
+BroadcastFn = Callable[[str, dict], None]
+
+
+def run_episode(
+ strategy_fn: Callable[[dict], dict],
+ task_name: str,
+ ep_id: Optional[str] = None,
+ broadcast_fn: Optional[BroadcastFn] = None,
+ max_steps: Optional[int] = None,
+) -> dict:
+ """Run a single origami episode with a given strategy function.
+
+ Args:
+ strategy_fn: Callable that receives paper_state dict and returns a fold dict:
+ {"type": "valley"|"mountain"|"pleat"|"crimp"|"stop",
+ "line": {"start": [x, y], "end": [x, y]},
+ "angle": 180.0}
+ task_name: Name of the task (from server/tasks.py)
+ ep_id: Episode identifier for broadcast; auto-generated if None
+ broadcast_fn: Optional callback(ep_id, data) for live streaming
+ max_steps: Override task's max_folds if provided
+
+ Returns:
+ dict with keys: episode_id, score, final_metrics, fold_history, status
+ """
+ ep_id = ep_id or str(uuid.uuid4())[:8]
+ env = OrigamiEnvironment()
+
+ obs = env.reset(task_name=task_name)
+
+ if broadcast_fn:
+ broadcast_fn(ep_id, {
+ "type": "episode_update",
+ "episode_id": ep_id,
+ "task_name": task_name,
+ "step": 0,
+ "observation": _obs_to_dict(obs),
+ })
+
+ step_limit = max_steps or env._task.get("max_folds", 20) if env._task else 20
+ status = "done"
+
+ for step_idx in range(step_limit):
+ if obs.done:
+ break
+
+ # Strategy generates a fold dict
+ try:
+ fold_dict = strategy_fn(obs.paper_state)
+ except Exception as exc:
+ status = "error"
+ if broadcast_fn:
+ broadcast_fn(ep_id, {
+ "type": "episode_done",
+ "episode_id": ep_id,
+ "status": "error",
+ "score": obs.reward or 0.0,
+ "final_metrics": obs.metrics,
+ "error": str(exc),
+ })
+ break
+
+ fold_type = fold_dict.get("type", "valley")
+ fold_line = fold_dict.get("line", {"start": [0, 0.5], "end": [1, 0.5]})
+ fold_angle = float(fold_dict.get("angle", 180.0))
+
+ action = OrigamiAction(
+ fold_type=fold_type,
+ fold_line=fold_line,
+ fold_angle=fold_angle,
+ )
+ obs = env.step(action)
+
+ if broadcast_fn:
+ broadcast_fn(ep_id, {
+ "type": "episode_update",
+ "episode_id": ep_id,
+ "task_name": task_name,
+ "step": step_idx + 1,
+ "observation": _obs_to_dict(obs),
+ })
+
+ if obs.done:
+ break
+ else:
+ status = "timeout"
+
+ score = obs.reward if obs.reward is not None else (env._total_reward or 0.0)
+
+ if broadcast_fn:
+ broadcast_fn(ep_id, {
+ "type": "episode_done",
+ "episode_id": ep_id,
+ "status": status,
+ "score": float(score),
+ "final_metrics": obs.metrics,
+ })
+
+ return {
+ "episode_id": ep_id,
+ "score": float(score),
+ "final_metrics": obs.metrics,
+ "fold_history": obs.fold_history,
+ "status": status,
+ }
+
+
+def run_batch(
+ strategy_fns: list[Callable[[dict], dict]],
+ task_name: str,
+ broadcast_fn: Optional[BroadcastFn] = None,
+ batch_id: Optional[int] = None,
+ max_workers: int = 8,
+) -> list[dict]:
+ """Run G episodes in parallel with a ThreadPoolExecutor.
+
+ Args:
+ strategy_fns: List of G strategy callables (one per completion)
+ task_name: Task to use for all episodes
+ broadcast_fn: Optional broadcast callback, called after each step
+ batch_id: Batch identifier for broadcast
+ max_workers: Max parallel threads (bounded by G)
+
+ Returns:
+ List of episode result dicts, in same order as strategy_fns
+ """
+ n = len(strategy_fns)
+ ep_ids = [f"ep_{(batch_id or 0):04d}_{i:02d}" for i in range(n)]
+ workers = min(max_workers, n)
+
+ results: list[dict] = [{}] * n
+
+ with ThreadPoolExecutor(max_workers=workers) as pool:
+ futures = {
+ pool.submit(
+ run_episode,
+ fn,
+ task_name,
+ ep_ids[i],
+ broadcast_fn,
+ ): i
+ for i, fn in enumerate(strategy_fns)
+ }
+
+ for future in as_completed(futures):
+ idx = futures[future]
+ try:
+ results[idx] = future.result()
+ except Exception as exc:
+ results[idx] = {
+ "episode_id": ep_ids[idx],
+ "score": 0.0,
+ "final_metrics": {},
+ "fold_history": [],
+ "status": "error",
+ "error": str(exc),
+ }
+
+ return results
+
+
+def _obs_to_dict(obs) -> dict:
+ """Convert OrigamiObservation to a JSON-serializable dict."""
+ try:
+ return obs.model_dump()
+ except AttributeError:
+ return {
+ "task": obs.task if hasattr(obs, "task") else {},
+ "paper_state": obs.paper_state if hasattr(obs, "paper_state") else {},
+ "metrics": obs.metrics if hasattr(obs, "metrics") else {},
+ "fold_history": obs.fold_history if hasattr(obs, "fold_history") else [],
+ "done": obs.done if hasattr(obs, "done") else False,
+ "reward": obs.reward if hasattr(obs, "reward") else None,
+ "error": obs.error if hasattr(obs, "error") else None,
+ }
diff --git a/viewer/training.html b/viewer/training.html
new file mode 100644
index 0000000000000000000000000000000000000000..672931688a8de725f6f4e0373f421ecfef68eccc
--- /dev/null
+++ b/viewer/training.html
@@ -0,0 +1,761 @@
+
+
+
+
+
+OPTIGAMI — TRAINING GRID VIEWER
+
+
+
+
+
+
+
+
+
+
+
+
+
- Edit src/App.js and save to reload.
-
+
+
+ OPTIGAMI RL
+
+
+ {isReplayMode ? (
+ <>
+ REPLAY — {REPLAY_EP_ID}
+
+
+ ) : (
+ setSelectedTarget(name)}
+ />
+ )}
+
+
+
+
);
}
diff --git a/src/App.test.js b/src/App.test.js
index 1f03afeece5ac28064fa3c73a29215037465f789..2e69bb5de9e34aea82df890e820f68617509faf2 100644
--- a/src/App.test.js
+++ b/src/App.test.js
@@ -1,8 +1 @@
-import { render, screen } from '@testing-library/react';
-import App from './App';
-
-test('renders learn react link', () => {
- render(
+
+
+ {apiStatus === 'ok' ? 'API OK' : apiStatus === 'err' ? 'API ERR' : 'CONNECTING'}
+
+
+
+
+
+
+
+
+
+
+ TASK — {targetDef ? targetDef.name.replace(/_/g, ' ').toUpperCase() : '—'}
+
+
+
+
+ {currentStep === 0 ? 'INITIAL STATE' : `STEP ${currentStep} / ${totalSteps}`}
+
+
+
+
+
+ 3D FOLD PREVIEW
+
+
+
+ FOLD SEQUENCE
+ {episodeLoading ? (
+
+
+ FETCHING EPISODE...
+
+ ) : (
+
+
+ METRICS
+ EPISODE INFO
+
+
+ );
+}
diff --git a/src/components/PlayerControls.js b/src/components/PlayerControls.js
new file mode 100644
index 0000000000000000000000000000000000000000..74e01e39e9b87bc9b848d7b0783d29a662e37d61
--- /dev/null
+++ b/src/components/PlayerControls.js
@@ -0,0 +1,54 @@
+export default function PlayerControls({
+ playing,
+ onPlay,
+ onPause,
+ onNext,
+ onPrev,
+ onReset,
+ currentStep,
+ totalSteps,
+ disabled,
+}) {
+ const atStart = currentStep === 0;
+ const atEnd = currentStep >= totalSteps;
+
+ return (
+
+ format
+
+
+ anchored
+
+
+ novelty
+
+
+ completion
+
+
+ regression
+
+
+ efficiency
+
+ {targetDef && (
+ <>
+
+ level
+ LVL {targetDef.level}
+
+
+ creases
+ {targetDef.n_creases}
+
+
+ task
+
+ {(targetDef.name || '').replace(/_/g, ' ').toUpperCase()}
+
+
+
+ )}
+
+
+
+
+ {disabled ? '—/—' : `${currentStep} / ${totalSteps}`}
+
+
+
+
+ );
+}
diff --git a/src/components/RewardPanel.js b/src/components/RewardPanel.js
new file mode 100644
index 0000000000000000000000000000000000000000..bf1acfd24fd4863ed7cd025359b1c123e21142ef
--- /dev/null
+++ b/src/components/RewardPanel.js
@@ -0,0 +1,106 @@
+const METRIC_FIELDS = [
+ {
+ key: 'progress',
+ label: 'progress',
+ color: 'var(--progress)',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'delta',
+ label: 'delta',
+ color: '#6ec6ff',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? (v >= 0 ? '+' : '') + v.toFixed(3) : '—'),
+ },
+ {
+ key: 'kawasaki',
+ label: 'kawasaki',
+ color: 'var(--validity)',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'maekawa',
+ label: 'maekawa',
+ color: 'var(--validity)',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'blb',
+ label: 'blb',
+ color: 'var(--validity)',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'assignment_accuracy',
+ label: 'assign acc',
+ color: '#ffb74d',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'economy',
+ label: 'economy',
+ color: '#81c784',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'degree_sanity',
+ label: 'deg sanity',
+ color: 'var(--validity)',
+ normalize: (v) => Math.min(Math.max(v || 0, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+ {
+ key: 'total',
+ label: 'total',
+ color: '#e0e0e0',
+ normalize: (v) => Math.min(Math.max((v || 0) / 1.5, 0), 1),
+ format: (v) => (v != null ? v.toFixed(3) : '—'),
+ },
+];
+
+function RewardRow({ label, color, pct, formattedValue, isDash, inverted }) {
+ const barColor = inverted && pct > 0 ? 'var(--validity)' : color;
+ return (
+
+ {label}
+
+ );
+}
+
+export default function RewardPanel({ metrics }) {
+ return (
+
+
+
+
+ {formattedValue}
+
+
+ {METRIC_FIELDS.map(({ key, label, color, normalize, format, inverted }) => {
+ const raw = metrics ? metrics[key] : undefined;
+ const isDash = raw === null || raw === undefined;
+ const pct = isDash ? 0 : normalize(raw) * 100;
+ return (
+
+ );
+}
diff --git a/src/components/StepFeed.js b/src/components/StepFeed.js
new file mode 100644
index 0000000000000000000000000000000000000000..0171a35f65e14616707dfd9d8cbbc6da49b79709
--- /dev/null
+++ b/src/components/StepFeed.js
@@ -0,0 +1,94 @@
+import { useEffect, useRef } from 'react';
+
+function compactnessDelta(step, prevStep) {
+ if (!step || !step.metrics) return null;
+ const curr = step.metrics.compactness;
+ if (curr == null) return null;
+ if (prevStep && prevStep.metrics && prevStep.metrics.compactness != null) {
+ return curr - prevStep.metrics.compactness;
+ }
+ return curr;
+}
+
+function foldAssignment(fold) {
+ if (!fold) return null;
+ const t = fold.type || '';
+ if (t === 'valley') return 'V';
+ if (t === 'mountain') return 'M';
+ if (t === 'pleat') return 'P';
+ if (t === 'crimp') return 'C';
+ return t.charAt(0).toUpperCase() || null;
+}
+
+function foldLabel(fold) {
+ if (!fold) return '';
+ const type = fold.type || 'fold';
+ const angle = fold.angle != null ? ` ${fold.angle}°` : '';
+ return `${type.toUpperCase()} FOLD${angle}`;
+}
+
+export default function StepFeed({ steps, currentStep }) {
+ const feedRef = useRef(null);
+ const activeRef = useRef(null);
+
+ useEffect(() => {
+ if (activeRef.current) {
+ activeRef.current.scrollIntoView({ block: 'nearest', behavior: 'smooth' });
+ }
+ }, [currentStep]);
+
+ if (!steps || steps.length === 0) {
+ return (
+
+
+ );
+ }
+
+ return (
+
+ NO STEPS LOADED
+
+
+ {steps.map((step, idx) => {
+ const stepNum = idx + 1;
+ const isActive = currentStep === stepNum;
+ const delta = compactnessDelta(step, idx > 0 ? steps[idx - 1] : null);
+ const asgn = foldAssignment(step.fold);
+ const label = foldLabel(step.fold);
+ const m = step.metrics || {};
+
+ return (
+
+ );
+}
diff --git a/src/components/TargetSelector.js b/src/components/TargetSelector.js
new file mode 100644
index 0000000000000000000000000000000000000000..b7822c56bf51d27c2b4773cb378bae8edb89d9fb
--- /dev/null
+++ b/src/components/TargetSelector.js
@@ -0,0 +1,38 @@
+function groupByLevel(targets) {
+ const levels = {};
+ Object.values(targets).forEach(t => {
+ if (!levels[t.level]) levels[t.level] = [];
+ levels[t.level].push(t);
+ });
+ return levels;
+}
+
+export default function TargetSelector({ targets, selected, onChange }) {
+ const levels = groupByLevel(targets);
+ const sortedLevels = Object.keys(levels).sort((a, b) => Number(a) - Number(b));
+
+ return (
+
+
+ );
+ })}
+
+ #{stepNum}
+ {label}
+ {asgn && (
+ {asgn}
+ )}
+
+ {delta !== null && (
+
+ {'\u0394'} compact:{' '}
+ = 0 ? 'delta-positive' : 'delta-negative'}>
+ {delta >= 0 ? '+' : ''}{delta.toFixed(3)}
+
+ {m.max_strain != null && (
+
+ {' '}| strain: {m.max_strain.toFixed(4)}
+ {m.is_valid != null && (
+ | {m.is_valid ? '✓' : '✗'}
+ )}
+
+ )}
+
+ )}
+
+ TARGET
+
+
+ );
+}
diff --git a/src/index.css b/src/index.css
index ec2585e8c0bb8188184ed1e0703c4c8f2a8419b0..bd577de0611e61395f5ac151db2c3d23b69df803 100644
--- a/src/index.css
+++ b/src/index.css
@@ -1,13 +1,34 @@
-body {
+@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@300;400;500;700&family=IBM+Plex+Mono:wght@300;400;500&display=swap');
+
+*, *::before, *::after {
+ box-sizing: border-box;
margin: 0;
- font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', 'Roboto', 'Oxygen',
- 'Ubuntu', 'Cantarell', 'Fira Sans', 'Droid Sans', 'Helvetica Neue',
- sans-serif;
+ padding: 0;
+}
+
+body {
+ background: #0d0d14;
+ color: #f8fafc;
+ font-family: 'IBM Plex Mono', monospace;
+ font-size: 13px;
+ line-height: 1.5;
-webkit-font-smoothing: antialiased;
- -moz-osx-font-smoothing: grayscale;
+ overflow-x: hidden;
+}
+
+::-webkit-scrollbar {
+ width: 4px;
+ height: 4px;
+}
+
+::-webkit-scrollbar-track {
+ background: #0d0d14;
+}
+
+::-webkit-scrollbar-thumb {
+ background: #2a2a3a;
}
-code {
- font-family: source-code-pro, Menlo, Monaco, Consolas, 'Courier New',
- monospace;
+::-webkit-scrollbar-thumb:hover {
+ background: #3a3a5a;
}
diff --git a/src/reportWebVitals.js b/src/reportWebVitals.js
index 5253d3ad9e6be6690549cb255f5952337b02401d..2950497bb9ec2dc32d6e4488060a6535651e2fbd 100644
--- a/src/reportWebVitals.js
+++ b/src/reportWebVitals.js
@@ -1,13 +1 @@
-const reportWebVitals = onPerfEntry => {
- if (onPerfEntry && onPerfEntry instanceof Function) {
- import('web-vitals').then(({ getCLS, getFID, getFCP, getLCP, getTTFB }) => {
- getCLS(onPerfEntry);
- getFID(onPerfEntry);
- getFCP(onPerfEntry);
- getLCP(onPerfEntry);
- getTTFB(onPerfEntry);
- });
- }
-};
-
-export default reportWebVitals;
+export default function reportWebVitals() {}
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/tests/test_graph.py b/tests/test_graph.py
new file mode 100644
index 0000000000000000000000000000000000000000..20b261c270a2d9158b500e4a56f9b630901d800d
--- /dev/null
+++ b/tests/test_graph.py
@@ -0,0 +1,115 @@
+import numpy as np
+import pytest
+from env.graph import CreaseGraph, VERTEX_TOL
+
+
+def test_init_boundary():
+ g = CreaseGraph()
+ assert len(g.vertices) == 4
+ assert len(g.edges) == 4
+ for eid, (v1, v2, assignment) in g.edges.items():
+ assert assignment == 'B'
+ assert g.interior_vertices() == []
+
+
+def test_add_vertex_dedup():
+ g = CreaseGraph()
+ id1 = g.add_vertex(0.5, 0.5)
+ id2 = g.add_vertex(0.5, 0.5)
+ assert id1 == id2
+
+
+def test_add_vertex_dedup_near():
+ g = CreaseGraph()
+ id1 = g.add_vertex(0.5, 0.5)
+ id2 = g.add_vertex(0.5 + VERTEX_TOL * 0.5, 0.5)
+ assert id1 == id2
+
+
+def test_cyclic_order():
+ g = CreaseGraph()
+ center_id = g.add_vertex(0.5, 0.5)
+
+ right_id = g.add_vertex(0.8, 0.5) # 0 degrees
+ top_id = g.add_vertex(0.5, 0.8) # 90 degrees
+ left_id = g.add_vertex(0.2, 0.5) # 180 degrees
+ bottom_id = g.add_vertex(0.5, 0.2) # 270 degrees / -90 degrees
+
+ e_right = g.add_edge(center_id, right_id, 'M')
+ e_top = g.add_edge(center_id, top_id, 'M')
+ e_left = g.add_edge(center_id, left_id, 'M')
+ e_bottom = g.add_edge(center_id, bottom_id, 'M')
+
+ cyclic = g.get_cyclic_edges(center_id)
+ # Sorted by angle ascending: right(0), top(90), left(180), bottom(-90 → 270)
+ # arctan2 for bottom gives -pi/2 which sorts before 0 in ascending order
+ # So actual ascending order: bottom(-pi/2), right(0), top(pi/2), left(pi)
+ assert len(cyclic) == 4
+
+ def edge_angle(eid):
+ ev1, ev2, _ = g.edges[eid]
+ other_id = ev2 if ev1 == center_id else ev1
+ ox, oy = g.vertices[other_id]
+ cx, cy = g.vertices[center_id]
+ return float(np.arctan2(oy - cy, ox - cx))
+
+ angles = [edge_angle(eid) for eid in cyclic]
+ assert angles == sorted(angles), "Edges should be sorted by ascending angle"
+
+ assert e_right in cyclic
+ assert e_top in cyclic
+ assert e_left in cyclic
+ assert e_bottom in cyclic
+
+ # Verify specific order: bottom < right < top < left in angle space
+ pos = {eid: i for i, eid in enumerate(cyclic)}
+ assert pos[e_bottom] < pos[e_right] < pos[e_top] < pos[e_left]
+
+
+def test_interior_vertices_empty():
+ g = CreaseGraph()
+ assert g.interior_vertices() == []
+
+
+def test_interior_vertices_with_crease_intersection():
+ g = CreaseGraph()
+ center_id = g.add_vertex(0.5, 0.5)
+ assert center_id in g.interior_vertices()
+
+
+def test_split_edge():
+ g = CreaseGraph()
+ # Find the bottom boundary edge (0,0)-(1,0) which is edge 0: v0-v1
+ original_edge_id = None
+ for eid, (v1, v2, assignment) in g.edges.items():
+ x1, y1 = g.vertices[v1]
+ x2, y2 = g.vertices[v2]
+ if {(x1, y1), (x2, y2)} == {(0.0, 0.0), (1.0, 0.0)}:
+ original_edge_id = eid
+ original_v1 = v1
+ original_v2 = v2
+ break
+
+ assert original_edge_id is not None
+
+ mid_id = g.add_vertex(0.5, 0.0)
+ eid1, eid2 = g.split_edge(original_edge_id, mid_id)
+
+ assert original_edge_id not in g.edges
+
+ assert eid1 in g.edges
+ assert eid2 in g.edges
+
+ _, _, a1 = g.edges[eid1]
+ _, _, a2 = g.edges[eid2]
+ assert a1 == 'B'
+ assert a2 == 'B'
+
+ def edge_vertex_set(eid):
+ v1, v2, _ = g.edges[eid]
+ return {v1, v2}
+
+ assert mid_id in edge_vertex_set(eid1)
+ assert mid_id in edge_vertex_set(eid2)
+ assert original_v1 in edge_vertex_set(eid1) or original_v1 in edge_vertex_set(eid2)
+ assert original_v2 in edge_vertex_set(eid1) or original_v2 in edge_vertex_set(eid2)
diff --git a/tests/test_openenv_adapter.py b/tests/test_openenv_adapter.py
new file mode 100644
index 0000000000000000000000000000000000000000..7e10e1c02eced478cccc36ea44cbc80f4d301fdb
--- /dev/null
+++ b/tests/test_openenv_adapter.py
@@ -0,0 +1,49 @@
+from openenv_runtime.environment import OpenEnvOrigamiEnvironment
+from openenv_runtime.models import OrigamiAction, OrigamiObservation
+
+
+def test_openenv_reset_returns_observation():
+ env = OpenEnvOrigamiEnvironment(mode="step", max_steps=8)
+ obs = env.reset(target_name="half_horizontal")
+
+ assert isinstance(obs, OrigamiObservation)
+ assert obs.done is False
+ assert obs.target_name == "half_horizontal"
+ assert hasattr(obs, "prompt") and obs.prompt
+
+
+def test_openenv_step_single_fold_completes_simple_target():
+ env = OpenEnvOrigamiEnvironment(mode="step", max_steps=8)
+ env.reset(target_name="half_horizontal")
+
+ action = OrigamiAction(
+ from_point=[0.0, 0.5],
+ to_point=[1.0, 0.5],
+ assignment="V",
+ )
+ obs = env.step(action)
+
+ assert obs.reward is not None
+ assert obs.reward > 1.0
+ assert obs.done is True
+
+
+def test_openenv_action_alias_from_to():
+ """Client can send 'from' and 'to' (JSON-friendly) instead of from_point/to_point."""
+ action = OrigamiAction.model_validate({
+ "from": [0.0, 0.5],
+ "to": [1.0, 0.5],
+ "assignment": "V",
+ })
+ assert action.from_point == [0.0, 0.5]
+ assert action.to_point == [1.0, 0.5]
+
+
+def test_openenv_state_contains_target_and_step():
+ env = OpenEnvOrigamiEnvironment(mode="step", max_steps=8)
+ env.reset(target_name="half_horizontal")
+
+ state = env.state()
+ assert "paper" in state
+ assert state.get("target") == "half_horizontal"
+ assert state.get("step") == 0
diff --git a/tests/test_paper_state.py b/tests/test_paper_state.py
new file mode 100644
index 0000000000000000000000000000000000000000..86507660b7736857eabe81759c9326015dde397a
--- /dev/null
+++ b/tests/test_paper_state.py
@@ -0,0 +1,97 @@
+import pytest
+from env.paper_state import PaperState, UNIT_SQUARE_CORNERS
+from env.graph import VERTEX_TOL
+
+
+def test_single_crease_no_interior_vertices():
+ paper = PaperState()
+ result = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'V')
+ assert result['valid'] is True
+ interior = paper.graph.interior_vertices()
+ assert interior == [], f"Expected no interior vertices, got {interior}"
+
+
+def test_anchor_points_initial():
+ paper = PaperState()
+ anchors = paper.anchor_points()
+ for corner in UNIT_SQUARE_CORNERS:
+ assert any(
+ abs(ax - corner[0]) < VERTEX_TOL and abs(ay - corner[1]) < VERTEX_TOL
+ for ax, ay in anchors
+ ), f"Corner {corner} not found in anchor_points"
+
+
+def test_anchor_points_grow():
+ paper = PaperState()
+ result = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'V')
+ assert result['valid'] is True
+
+ anchors = paper.anchor_points()
+
+ def has_point(px, py):
+ return any(abs(ax - px) < VERTEX_TOL and abs(ay - py) < VERTEX_TOL for ax, ay in anchors)
+
+ assert has_point(0.0, 0.5), "(0, 0.5) should be in anchor_points after crease"
+ assert has_point(1.0, 0.5), "(1, 0.5) should be in anchor_points after crease"
+
+
+def test_invalid_assignment():
+ paper = PaperState()
+ result = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'X')
+ assert result['valid'] is False
+ assert 'invalid_assignment' in result['errors']
+
+
+def test_fold_history():
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'M')
+ assert len(paper.fold_history) == 1
+
+
+def test_unanchored_returns_false_anchored():
+ paper = PaperState()
+ result = paper.add_crease([0.3, 0.3], [0.7, 0.7], 'M')
+ assert result['anchored'] is False
+
+
+def test_crease_edges_returned():
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'M')
+ edges = paper.crease_edges()
+ assert len(edges) >= 1
+ for e in edges:
+ assert e['assignment'] in ('M', 'V')
+ assert 'v1' in e
+ assert 'v2' in e
+
+
+def test_two_intersecting_creases():
+ paper = PaperState()
+ r1 = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'M')
+ r2 = paper.add_crease([0.5, 0.0], [0.5, 1.0], 'V')
+ assert r1['valid'] is True
+ assert r2['valid'] is True
+ interior = paper.graph.interior_vertices()
+ assert len(interior) >= 1
+ coords = [paper.graph.vertices[vid] for vid in interior]
+ assert any(abs(x - 0.5) < VERTEX_TOL and abs(y - 0.5) < VERTEX_TOL for x, y in coords)
+
+
+def test_duplicate_crease_detection():
+ """Adding the same crease twice → first duplicate=False, second duplicate=True."""
+ paper = PaperState()
+ r1 = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'M')
+ assert r1['valid'] is True
+ assert r1['duplicate'] is False
+
+ r2 = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'V') # same endpoints, different assignment
+ assert r2['valid'] is True
+ assert r2['duplicate'] is True
+
+
+def test_error_returns_have_duplicate_false():
+ """Error return paths include duplicate=False."""
+ paper = PaperState()
+ r_invalid = paper.add_crease([0.0, 0.5], [1.0, 0.5], 'X')
+ assert r_invalid['valid'] is False
+ assert r_invalid['duplicate'] is False
diff --git a/tests/test_verifier.py b/tests/test_verifier.py
new file mode 100644
index 0000000000000000000000000000000000000000..11cfca8d85db157628a8d754b771ae5dc52252a9
--- /dev/null
+++ b/tests/test_verifier.py
@@ -0,0 +1,314 @@
+import pytest
+import numpy as np
+from env.graph import CreaseGraph
+from env.paper_state import PaperState
+from env.verifier import (
+ check_kawasaki_at_vertex,
+ check_maekawa_at_vertex,
+ check_blb_at_vertex,
+ geometric_crease_coverage,
+ check_all_vertices,
+ check_degree_sanity,
+)
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def make_cross_graph(center_coords=(0.5, 0.5), assignment='M') -> tuple[CreaseGraph, int]:
+ """
+ Degree-4 vertex at center with 4 spokes pointing N/S/E/W.
+ All spokes have the given assignment.
+ """
+ g = CreaseGraph()
+ cx, cy = center_coords
+ vid = g.add_vertex(cx, cy)
+
+ neighbors = [
+ (0.0, cy), # left (180°)
+ (1.0, cy), # right (0°)
+ (cx, 0.0), # down (-90°)
+ (cx, 1.0), # up (90°)
+ ]
+ for nx, ny in neighbors:
+ nid = g.add_vertex(nx, ny)
+ g.add_edge(vid, nid, assignment)
+
+ return g, vid
+
+
+# ---------------------------------------------------------------------------
+# Kawasaki tests
+# ---------------------------------------------------------------------------
+
+def test_kawasaki_no_interior_vertices():
+ paper = PaperState()
+ paper.add_crease([0, 0.5], [1, 0.5], 'V')
+ assert paper.graph.interior_vertices() == []
+ result = check_all_vertices(paper.graph)
+ assert result['kawasaki'] == 1.0
+ assert result['n_interior'] == 0
+
+
+def test_kawasaki_valid_degree4_vertex():
+ """Equal 90° sectors → alternating sum = 0 → Kawasaki satisfied."""
+ g, vid = make_cross_graph()
+ ok, err = check_kawasaki_at_vertex(vid, g)
+ assert ok == True
+ assert err == pytest.approx(0.0, abs=1e-9)
+
+
+def test_kawasaki_invalid_vertex():
+ """
+ Manually construct a degree-4 vertex whose sectors are 60°,120°,80°,100°.
+ Alternating sum = 60 - 120 + 80 - 100 = -80° ≠ 0 → should fail.
+ """
+ g = CreaseGraph()
+ cx, cy = 0.5, 0.5
+ vid = g.add_vertex(cx, cy)
+
+ # Place neighbours at specific angles so sectors are exactly as desired.
+ # Sectors are measured CCW between consecutive rays.
+ # We choose ray angles (from center) in ascending arctan2 order:
+ # a0 = 0°
+ # a1 = 60° (sector0 = 60°)
+ # a2 = 180° (sector1 = 120°)
+ # a3 = 260° = -100° (sector2 = 80°)
+ # sector3 (wraparound to a0) = 360° - 260° = 100°
+ # alt_sum = 60 - 120 + 80 - 100 = -80° → |alt_sum| ≈ 1.396 rad
+ r = 0.3
+ angles_deg = [0.0, 60.0, 180.0, 260.0]
+ for deg in angles_deg:
+ rad = np.deg2rad(deg)
+ nx = cx + r * np.cos(rad)
+ ny = cy + r * np.sin(rad)
+ nid = g.add_vertex(nx, ny)
+ g.add_edge(vid, nid, 'M')
+
+ ok, err = check_kawasaki_at_vertex(vid, g)
+ assert ok == False
+ expected_err = abs(np.deg2rad(60 - 120 + 80 - 100))
+ assert err == pytest.approx(expected_err, abs=1e-6)
+
+
+# ---------------------------------------------------------------------------
+# Maekawa tests
+# ---------------------------------------------------------------------------
+
+def test_maekawa_excludes_boundary():
+ """
+ Boundary edges at a vertex should NOT count toward M/V tally.
+ A corner vertex has only boundary edges; Maekawa should return True
+ (fewer than 4 fold edges → vacuously satisfied).
+ """
+ g = CreaseGraph()
+ corner_id = 0 # vertex (0,0)
+ assert check_maekawa_at_vertex(corner_id, g) is True
+
+
+def test_maekawa_valid():
+ """3 M + 1 V → |3-1| = 2 → True."""
+ g = CreaseGraph()
+ cx, cy = 0.5, 0.5
+ vid = g.add_vertex(cx, cy)
+
+ r = 0.3
+ angles_deg = [0.0, 90.0, 180.0, 270.0]
+ assignments = ['M', 'M', 'M', 'V']
+ for deg, asgn in zip(angles_deg, assignments):
+ rad = np.deg2rad(deg)
+ nid = g.add_vertex(cx + r * np.cos(rad), cy + r * np.sin(rad))
+ g.add_edge(vid, nid, asgn)
+
+ assert check_maekawa_at_vertex(vid, g) is True
+
+
+def test_maekawa_invalid():
+ """2 M + 2 V → |2-2| = 0 → False."""
+ g = CreaseGraph()
+ cx, cy = 0.5, 0.5
+ vid = g.add_vertex(cx, cy)
+
+ r = 0.3
+ angles_deg = [0.0, 90.0, 180.0, 270.0]
+ assignments = ['M', 'M', 'V', 'V']
+ for deg, asgn in zip(angles_deg, assignments):
+ rad = np.deg2rad(deg)
+ nid = g.add_vertex(cx + r * np.cos(rad), cy + r * np.sin(rad))
+ g.add_edge(vid, nid, asgn)
+
+ assert check_maekawa_at_vertex(vid, g) is False
+
+
+# ---------------------------------------------------------------------------
+# BLB tests
+# ---------------------------------------------------------------------------
+
+def test_blb_no_violations_equal_sectors():
+ """Equal 90° sectors → no strict local minimum → BLB returns []."""
+ g, vid = make_cross_graph()
+ violations = check_blb_at_vertex(vid, g)
+ assert violations == []
+
+
+def test_blb_violation_detected():
+ """
+ Create a vertex with a strict local-minimum sector whose bounding edges
+ share the same MV assignment → BLB violation.
+
+ Use angles 0°, 10°, 180°, 270° so sector[0]=10° is the strict local min
+ relative to sector[3] (90°) and sector[1] (170°). The two bounding edges
+ are at 0° and 10°; assign both 'M' → violation.
+ """
+ g = CreaseGraph()
+ cx, cy = 0.5, 0.5
+ vid = g.add_vertex(cx, cy)
+
+ r = 0.3
+ # angles ascending (arctan2 order): 0°, 10°, 180°, 270° (= -90°)
+ # sorted arctan2: -90°, 0°, 10°, 180°
+ # sectors: 90°, 10°, 170°, 90° (sum=360°)
+ # sector at index 1 (between 0° and 10°) = 10° is strict local min (90 > 10 < 170)
+ angles_deg = [0.0, 10.0, 180.0, 270.0]
+ edge_ids = []
+ for deg in angles_deg:
+ rad = np.deg2rad(deg)
+ nid = g.add_vertex(cx + r * np.cos(rad), cy + r * np.sin(rad))
+ eid = g.add_edge(vid, nid, 'M')
+ edge_ids.append(eid)
+
+ violations = check_blb_at_vertex(vid, g)
+ assert len(violations) > 0
+
+
+def test_blb_no_violation_when_opposite_assignments():
+ """
+ Same geometry as above but with opposite assignments on the two edges
+ bounding the small sector → no BLB violation.
+ """
+ g = CreaseGraph()
+ cx, cy = 0.5, 0.5
+ vid = g.add_vertex(cx, cy)
+
+ r = 0.3
+ angles_deg = [0.0, 10.0, 180.0, 270.0]
+ # sorted arctan2: -90°(270°), 0°, 10°, 180°
+ # small sector is between 0° and 10° (index 1 and 2 in sorted order)
+ # assign them opposite assignments
+ assignments_by_angle = {
+ 0.0: 'M',
+ 10.0: 'V',
+ 180.0: 'M',
+ 270.0: 'V',
+ }
+ for deg in angles_deg:
+ rad = np.deg2rad(deg)
+ nid = g.add_vertex(cx + r * np.cos(rad), cy + r * np.sin(rad))
+ g.add_edge(vid, nid, assignments_by_angle[deg])
+
+ violations = check_blb_at_vertex(vid, g)
+ assert violations == []
+
+
+# ---------------------------------------------------------------------------
+# Coverage tests
+# ---------------------------------------------------------------------------
+
+def test_coverage_exact_match():
+ """Add exact crease matching target → coverage = 1.0, economy = 1.0."""
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'M')
+
+ target = [{'v1': (0.0, 0.5), 'v2': (1.0, 0.5), 'assignment': 'M'}]
+ coverage, economy, assignment_accuracy = geometric_crease_coverage(paper, target)
+ assert coverage == pytest.approx(1.0)
+ assert economy == pytest.approx(1.0)
+ assert assignment_accuracy == pytest.approx(1.0)
+
+
+def test_coverage_no_match():
+ """No creases added → coverage = 0.0."""
+ paper = PaperState()
+ target = [{'v1': (0.0, 0.5), 'v2': (1.0, 0.5), 'assignment': 'M'}]
+ coverage, economy, assignment_accuracy = geometric_crease_coverage(paper, target)
+ assert coverage == pytest.approx(0.0)
+ assert assignment_accuracy == pytest.approx(1.0) # vacuous: no positional matches
+
+
+def test_coverage_excess_penalty():
+ """
+ Target has 1 crease. Add 3 non-intersecting creases, one matching target.
+ coverage = 1.0, economy = 1 - 2/1 → clamped to 0.0 (economy < 1.0).
+ Uses non-intersecting extras to avoid PaperState edge splitting the target crease.
+ """
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'M') # matches target (midpoint 0.5,0.5)
+ paper.add_crease([0.0, 0.3], [0.5, 0.3], 'V') # extra, no intersection
+ paper.add_crease([0.0, 0.7], [0.5, 0.7], 'V') # extra, no intersection
+
+ target = [{'v1': (0.0, 0.5), 'v2': (1.0, 0.5), 'assignment': 'M'}]
+ coverage, economy, assignment_accuracy = geometric_crease_coverage(paper, target)
+ assert coverage == pytest.approx(1.0)
+ assert economy < 1.0
+ assert assignment_accuracy == pytest.approx(1.0)
+
+
+def test_coverage_assignment_mismatch_partial_credit():
+ """
+ Position matches but assignment doesn't → coverage = 0.5 (partial),
+ assignment_accuracy = 0.0 (no correct assignments among matched).
+ """
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'V') # paper has V, target expects M
+
+ target = [{'v1': (0.0, 0.5), 'v2': (1.0, 0.5), 'assignment': 'M'}]
+ coverage, economy, assignment_accuracy = geometric_crease_coverage(paper, target)
+ assert coverage == pytest.approx(0.5)
+ assert assignment_accuracy == pytest.approx(0.0)
+
+
+# ---------------------------------------------------------------------------
+# check_degree_sanity tests
+# ---------------------------------------------------------------------------
+
+def test_check_degree_sanity_no_interior():
+ """No interior vertices → returns 1.0 (vacuous)."""
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'V')
+ assert check_degree_sanity(paper.graph) == pytest.approx(1.0)
+
+
+def test_check_degree_sanity_even_degree():
+ """Degree-4 interior vertex → even degree → returns 1.0."""
+ g, _ = make_cross_graph()
+ assert check_degree_sanity(g) == pytest.approx(1.0)
+
+
+def test_check_degree_sanity_odd_degree():
+ """Degree-3 interior vertex → odd degree → returns 0.0."""
+ g = CreaseGraph()
+ cx, cy = 0.5, 0.5
+ vid = g.add_vertex(cx, cy)
+ r = 0.3
+ for deg in [0.0, 120.0, 240.0]:
+ rad = np.deg2rad(deg)
+ nid = g.add_vertex(cx + r * np.cos(rad), cy + r * np.sin(rad))
+ g.add_edge(vid, nid, 'M')
+ assert check_degree_sanity(g) == pytest.approx(0.0)
+
+
+# ---------------------------------------------------------------------------
+# check_all_vertices vacuous test
+# ---------------------------------------------------------------------------
+
+def test_check_all_vertices_vacuous():
+ """Single horizontal crease → no interior vertices → all scores = 1.0."""
+ paper = PaperState()
+ paper.add_crease([0.0, 0.5], [1.0, 0.5], 'V')
+ result = check_all_vertices(paper.graph)
+ assert result['kawasaki'] == 1.0
+ assert result['maekawa'] == 1.0
+ assert result['blb'] == 1.0
+ assert result['n_interior'] == 0
+ assert result['per_vertex'] == []
diff --git a/train.py b/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..3af33977c2fba88baafb85e61a449fe96890b86d
--- /dev/null
+++ b/train.py
@@ -0,0 +1,261 @@
+"""
+OrigamiRL — GRPO Training Script
+Code-as-policy: model generates complete fold sequence, gets terminal reward.
+
+Base model: SpatialThinker (Qwen2.5-VL-7B fine-tuned for spatial reasoning)
+or any Unsloth-compatible model.
+
+Usage:
+ python train.py
+ python train.py --model unsloth/Qwen2.5-VL-7B-Instruct --epochs 3
+ python train.py --model OX-PIXL/SpatialThinker-Qwen2.5-VL-7B --epochs 3
+ python train.py --dry_run # test rewards without GPU
+ python train.py --no_semantic # use coordinate-based prompts instead of semantic
+"""
+import argparse
+import random
+
+# VL (vision-language) model identifiers — use FastVisionModel for these
+_VL_MODEL_PATTERNS = ['VL', 'vl', 'Vision', 'vision', 'SpatialThinker', 'SpaceThinker']
+
+
+def _is_vl_model(model_name: str) -> bool:
+ return any(p in model_name for p in _VL_MODEL_PATTERNS)
+
+
+def parse_args():
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--model', default='unsloth/Qwen2.5-VL-7B-Instruct',
+ help='Base model. Use unsloth/Qwen2.5-VL-7B-Instruct or '
+ 'OX-PIXL/SpatialThinker-Qwen2.5-VL-7B for spatial reasoning')
+ parser.add_argument('--max_seq_length', type=int, default=2048)
+ parser.add_argument('--epochs', type=int, default=3)
+ parser.add_argument('--batch_size', type=int, default=2)
+ parser.add_argument('--grad_accum', type=int, default=4)
+ parser.add_argument('--lr', type=float, default=5e-6)
+ parser.add_argument('--n_generations', type=int, default=8)
+ parser.add_argument('--max_folds', type=int, default=8)
+ parser.add_argument('--output', default='origami-grpo')
+ parser.add_argument('--level', type=int, default=1, help='Target difficulty level (1-3)')
+ parser.add_argument('--no_semantic', action='store_true',
+ help='Disable semantic prompts; use coordinate-based target (default: semantic)')
+ parser.add_argument('--dry_run', action='store_true', help='Test reward function without training')
+ return parser.parse_args()
+
+
+def build_dataset(env, level: int = 1, max_folds: int = 8, semantic: bool = True) -> list[dict]:
+ """
+ Build a training dataset of prompts from available targets.
+ Each item: {'prompt': str, 'target_name': str}
+ Repeats each target multiple times to give enough training steps.
+ When semantic=True, uses get_semantic_description for task descriptions.
+ """
+ from env.prompts import get_semantic_description, code_as_policy_prompt
+
+ all_names = env.available_targets()
+
+ # Filter by level; fall back to all targets if none match
+ level_names = [
+ name for name in all_names
+ if env._targets[name].get('level', 1) == level
+ ]
+ if not level_names:
+ level_names = all_names
+
+ items = []
+ for name in level_names:
+ target = env._targets[name]
+ if semantic:
+ desc = get_semantic_description(name, target)
+ prompt = code_as_policy_prompt(target, max_folds=max_folds, semantic_description=desc)
+ else:
+ prompt = code_as_policy_prompt(target, max_folds=max_folds, semantic_description=None)
+ items.append({'prompt': prompt, 'target_name': name})
+
+ # Repeat each target 10x; ensure at least 50 examples
+ repeat = max(10, (50 + len(items) - 1) // len(items))
+ items = items * repeat
+
+ random.shuffle(items)
+ return items
+
+
+def make_reward_fn(env_template, max_folds: int):
+ """
+ Returns a reward function compatible with trl GRPOTrainer.
+
+ Signature: reward_fn(completions, prompts=None, **kwargs) -> list[float]
+
+ For each completion:
+ 1. Clone the environment (fresh paper state)
+ 2. Reset to the target embedded in the prompt (use target_name from kwargs if available)
+ 3. Execute the completion as a fold sequence
+ 4. Return the total reward
+ """
+ def reward_fn(completions, prompts=None, **kwargs):
+ rewards = []
+ target_names = kwargs.get('target_names', [None] * len(completions))
+
+ for completion, target_name in zip(completions, target_names):
+ try:
+ env = env_template.clone()
+ env.reset(target_name=target_name)
+ _, reward_dict, _, _ = env.step(completion)
+ rewards.append(float(reward_dict['total']))
+ except Exception:
+ rewards.append(-0.1)
+
+ return rewards
+
+ return reward_fn
+
+
+def make_detailed_reward_fns(env_template, max_folds: int) -> list:
+ """
+ Returns a list of reward functions, one per reward component.
+ Used for detailed W&B logging of each component separately.
+ Components: kawasaki, maekawa, blb, progress, economy, completion
+ """
+ components = ['kawasaki', 'maekawa', 'blb', 'progress', 'economy', 'completion']
+
+ def make_component_fn(component: str):
+ def component_fn(completions, prompts=None, **kwargs):
+ rewards = []
+ target_names = kwargs.get('target_names', [None] * len(completions))
+
+ for completion, target_name in zip(completions, target_names):
+ try:
+ env = env_template.clone()
+ env.reset(target_name=target_name)
+ _, reward_dict, _, _ = env.step(completion)
+ rewards.append(float(reward_dict.get(component, 0.0)))
+ except Exception:
+ rewards.append(0.0)
+
+ return rewards
+
+ component_fn.__name__ = f'reward_{component}'
+ return component_fn
+
+ return [make_component_fn(c) for c in components]
+
+
+def main():
+ args = parse_args()
+
+ # Import here to allow dry_run without GPU
+ from env.environment import OrigamiEnvironment
+
+ env = OrigamiEnvironment(mode='code_as_policy', max_steps=args.max_folds)
+
+ # Build dataset
+ use_semantic = not args.no_semantic
+ dataset_items = build_dataset(env, level=args.level, max_folds=args.max_folds, semantic=use_semantic)
+ print(f"Dataset: {len(dataset_items)} examples from level-{args.level} targets (semantic={'on' if use_semantic else 'off'})")
+ print(f"Targets: {env.available_targets()}")
+
+ # Dry run: test reward function without loading model
+ if args.dry_run:
+ reward_fn = make_reward_fn(env, args.max_folds)
+ test_completions = [
+ 'OPTIGAMI
+
+ IDLE
+
+ BATCH —
+ EPISODES 0
+ AVG REWARD —
+
+ WS
+