Add RL training environment with OpenEnv backend

- origami_server/: Full OpenEnv environment with physics simulator
- engine/: FOLD parser, analytical fold simulator (BFS + rotation
transforms), chamfer-distance shape matching
- environment.py: OrigamiEnvironment (reset/step/state)
- tasks.py: 7 tasks (triangle, half_fold, quarter_fold, letter_fold,
waterbomb, accordion, miura_ori)
- app.py: FastAPI server with /tasks endpoints
- training/: GRPO training scripts for Colab
- reward.py: valid_fold + shape_match reward functions
- train_grpo.py: Unsloth + TRL GRPOTrainer pipeline
- Dual-service Docker: Next.js UI on :7860, OpenEnv API on :8000

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

.dockerignore

@@ -3,3 +3,4 @@ node_modules
 .git
 __pycache__
 *.md
+outputs/

Dockerfile

@@ -1,39 +1,54 @@
 FROM node:20-alpine AS base
 
-# --- Dependencies ---
+# --- Node dependencies ---
 FROM base AS deps
 WORKDIR /app
 COPY package.json package-lock.json ./
 RUN npm ci
 
-# --- Build ---
+# --- Next.js build ---
 FROM base AS builder
 WORKDIR /app
 COPY --from=deps /app/node_modules ./node_modules
-COPY . .
-
+COPY app/ ./app/
+COPY components/ ./components/
+COPY hooks/ ./hooks/
+COPY lib/ ./lib/
+COPY public/ ./public/
+COPY package.json package-lock.json tsconfig.json next.config.ts postcss.config.mjs eslint.config.mjs .eslintrc.json next-env.d.ts metadata.json ./
 ENV NEXT_TELEMETRY_DISABLED=1
 RUN npm run build
 
-# --- Runner ---
-FROM base AS runner
+# --- Final runner: Node + Python ---
+FROM node:20-alpine AS runner
 WORKDIR /app
 
+# Install Python + pip for the OpenEnv API
+RUN apk add --no-cache python3 py3-pip py3-numpy py3-scipy
+
 ENV NODE_ENV=production
 ENV NEXT_TELEMETRY_DISABLED=1
 ENV PORT=7860
 ENV HOSTNAME=0.0.0.0
 
-RUN addgroup --system --gid 1001 nodejs
-RUN adduser --system --uid 1001 nextjs
-
+# Copy Next.js standalone build
 COPY --from=builder /app/public ./public
-COPY --from=builder --chown=nextjs:nodejs /app/.next/standalone ./
-COPY --from=builder --chown=nextjs:nodejs /app/.next/static ./.next/static
+COPY --from=builder /app/.next/standalone ./
+COPY --from=builder /app/.next/static ./.next/static
+
+# Copy Python RL environment
+COPY origami_server/ ./origami_server/
+COPY training/ ./training/
+COPY requirements.txt ./
+
+# Install Python deps
+RUN pip3 install --no-cache-dir --break-system-packages -r requirements.txt
 
-USER nextjs
+# Start script: run both services
+COPY start.sh ./
+RUN chmod +x start.sh
 
 EXPOSE 7860
+EXPOSE 8000
 
-# Start the Next.js standalone server
-CMD ["node", "server.js"]
+CMD ["./start.sh"]

README.md

@@ -7,5 +7,5 @@ sdk: docker
 app_port: 7860
 pinned: false
 license: mit
-short_description: Interactive 3D origami simulator with AI pattern generation
+short_description: Interactive 3D origami simulator + RL training environment
 ---

origami_server/app.py

@@ -0,0 +1,62 @@
+"""FastAPI entry point — OpenEnv create_app() + custom endpoints."""
+
+import os
+from pathlib import Path
+
+from fastapi.responses import HTMLResponse
+from fastapi.staticfiles import StaticFiles
+
+from openenv.core.env_server.http_server import create_app
+
+from .environment import OrigamiEnvironment
+from .models import OrigamiAction, OrigamiObservation
+
+app = create_app(
+    OrigamiEnvironment,
+    OrigamiAction,
+    OrigamiObservation,
+    env_name="origami_env",
+)
+
+from .tasks import TASKS
+
+
+@app.get("/tasks")
+def get_tasks():
+    return {
+        name: {
+            "name": task["name"],
+            "description": task["description"],
+            "difficulty": task["difficulty"],
+            "paper": task["paper"],
+            "target_fold": task["target_fold"],
+        }
+        for name, task in TASKS.items()
+    }
+
+
+@app.get("/tasks/{task_name}")
+def get_task_detail(task_name: str):
+    if task_name not in TASKS:
+        from fastapi import HTTPException
+
+        raise HTTPException(status_code=404, detail=f"Task '{task_name}' not found")
+    task = TASKS[task_name]
+    return {
+        "name": task["name"],
+        "description": task["description"],
+        "difficulty": task["difficulty"],
+        "paper": task["paper"],
+        "target_fold": task["target_fold"],
+    }
+
+
+def main():
+    import uvicorn
+
+    port = int(os.environ.get("PORT", 8000))
+    uvicorn.run(app, host="0.0.0.0", port=port)
+
+
+if __name__ == "__main__":
+    main()

origami_server/engine/__init__.py

@@ -0,0 +1,3 @@
+from .simulate import simulate
+from .fold_parser import parse_fold, validate_fold
+from .shape_match import compute_shape_match

origami_server/engine/fold_parser.py

@@ -0,0 +1,176 @@
+"""FOLD JSON parsing and validation.
+
+Validates LLM-generated FOLD crease patterns before simulation.
+FOLD spec: https://github.com/edemaine/fold
+"""
+
+from typing import Any
+
+import numpy as np
+
+
+def validate_fold(fold_data: dict[str, Any]) -> tuple[bool, str]:
+    """Validate a FOLD JSON object. Returns (is_valid, error_message)."""
+
+    # Required fields
+    for key in ("vertices_coords", "edges_vertices", "edges_assignment"):
+        if key not in fold_data:
+            return False, f"Missing required field: {key}"
+
+    verts = fold_data["vertices_coords"]
+    edges = fold_data["edges_vertices"]
+    assignments = fold_data["edges_assignment"]
+
+    # Must have at least 3 vertices (a triangle)
+    if len(verts) < 3:
+        return False, f"Need at least 3 vertices, got {len(verts)}"
+
+    # Must have at least 3 edges
+    if len(edges) < 3:
+        return False, f"Need at least 3 edges, got {len(edges)}"
+
+    # Edges and assignments must match length
+    if len(edges) != len(assignments):
+        return False, (
+            f"edges_vertices ({len(edges)}) and "
+            f"edges_assignment ({len(assignments)}) must match length"
+        )
+
+    # Fold angles must match if present
+    if "edges_foldAngle" in fold_data:
+        angles = fold_data["edges_foldAngle"]
+        if len(angles) != len(edges):
+            return False, (
+                f"edges_foldAngle ({len(angles)}) must match "
+                f"edges_vertices ({len(edges)})"
+            )
+
+    # Validate vertex coordinates (2D or 3D)
+    num_verts = len(verts)
+    for i, v in enumerate(verts):
+        if not isinstance(v, (list, tuple)) or len(v) < 2:
+            return False, f"Vertex {i} must be [x, y] or [x, y, z], got {v}"
+
+    # Validate edge indices
+    for i, e in enumerate(edges):
+        if not isinstance(e, (list, tuple)) or len(e) != 2:
+            return False, f"Edge {i} must be [v1, v2], got {e}"
+        v1, v2 = e
+        if v1 < 0 or v1 >= num_verts or v2 < 0 or v2 >= num_verts:
+            return False, f"Edge {i} references invalid vertex: {e}"
+        if v1 == v2:
+            return False, f"Edge {i} is degenerate (same vertex): {e}"
+
+    # Validate assignments
+    valid_assignments = {"M", "V", "B", "F", "U", "C"}
+    for i, a in enumerate(assignments):
+        if a not in valid_assignments:
+            return False, f"Edge {i} has invalid assignment '{a}'"
+
+    # Must have at least one fold crease (M or V)
+    has_fold = any(a in ("M", "V") for a in assignments)
+    if not has_fold:
+        return False, "No fold creases (M or V) found"
+
+    # Must have boundary edges
+    has_boundary = any(a == "B" for a in assignments)
+    if not has_boundary:
+        return False, "No boundary edges (B) found"
+
+    return True, ""
+
+
+def parse_fold(fold_data: dict[str, Any]) -> dict[str, np.ndarray]:
+    """Parse validated FOLD JSON into numpy arrays for simulation.
+
+    Returns dict with:
+        vertices: (N, 3) float64 — vertex positions (z=0 for 2D input)
+        edges: (E, 2) int — edge vertex indices
+        assignments: list[str] — edge type per edge
+        fold_angles: (E,) float64 — target fold angle per edge (radians)
+        faces: (F, 3) int — triangulated face vertex indices
+    """
+    verts = fold_data["vertices_coords"]
+
+    # Ensure 3D (add z=0 if 2D)
+    vertices = np.zeros((len(verts), 3), dtype=np.float64)
+    for i, v in enumerate(verts):
+        vertices[i, 0] = v[0]
+        vertices[i, 1] = v[1]
+        if len(v) > 2:
+            vertices[i, 2] = v[2]
+
+    edges = np.array(fold_data["edges_vertices"], dtype=np.int32)
+    assignments = list(fold_data["edges_assignment"])
+
+    # Fold angles: default based on assignment if not provided
+    if "edges_foldAngle" in fold_data:
+        fold_angles = np.array(fold_data["edges_foldAngle"], dtype=np.float64)
+    else:
+        fold_angles = np.zeros(len(edges), dtype=np.float64)
+        for i, a in enumerate(assignments):
+            if a == "V":
+                fold_angles[i] = 180.0
+            elif a == "M":
+                fold_angles[i] = -180.0
+
+    # Convert degrees to radians for simulation
+    fold_angles_rad = np.radians(fold_angles)
+
+    # Triangulate faces
+    if "faces_vertices" in fold_data:
+        raw_faces = fold_data["faces_vertices"]
+        faces = _triangulate_faces(raw_faces)
+    else:
+        faces = _compute_faces(vertices, edges)
+
+    return {
+        "vertices": vertices,
+        "edges": edges,
+        "assignments": assignments,
+        "fold_angles": fold_angles_rad,
+        "faces": faces,
+    }
+
+
+def _triangulate_faces(raw_faces: list[list[int]]) -> np.ndarray:
+    """Fan-triangulate polygon faces into triangles."""
+    triangles = []
+    for face in raw_faces:
+        if len(face) < 3:
+            continue
+        for i in range(1, len(face) - 1):
+            triangles.append([face[0], face[i], face[i + 1]])
+    if not triangles:
+        return np.zeros((0, 3), dtype=np.int32)
+    return np.array(triangles, dtype=np.int32)
+
+
+def _compute_faces(vertices: np.ndarray, edges: np.ndarray) -> np.ndarray:
+    """Compute triangulated faces from vertices and edges using adjacency."""
+    from collections import defaultdict
+
+    adj = defaultdict(set)
+    for v1, v2 in edges:
+        adj[v1].add(v2)
+        adj[v2].add(v1)
+
+    triangles = set()
+    for v1, v2 in edges:
+        common = adj[v1] & adj[v2]
+        for v3 in common:
+            tri = tuple(sorted([v1, v2, v3]))
+            triangles.add(tri)
+
+    if not triangles:
+        # Fallback: create faces using Delaunay on 2D projection
+        from scipy.spatial import Delaunay
+
+        pts_2d = vertices[:, :2]
+        try:
+            tri = Delaunay(pts_2d)
+            return tri.simplices.astype(np.int32)
+        except Exception:
+            return np.zeros((0, 3), dtype=np.int32)
+
+    return np.array(list(triangles), dtype=np.int32)

origami_server/engine/shape_match.py

@@ -0,0 +1,95 @@
+"""Shape matching for reward computation.
+
+Computes similarity between the LLM's folded shape and the target shape.
+Like AlphaFold's RMSD but for origami vertex positions.
+"""
+
+import numpy as np
+from scipy.spatial.distance import cdist
+
+
+def compute_shape_match(
+    predicted: np.ndarray,
+    target: np.ndarray,
+) -> float:
+    """Compute shape similarity between predicted and target positions.
+
+    Uses chamfer distance normalized by bounding box diagonal.
+    Aligns shapes by centering before comparison.
+
+    Args:
+        predicted: (N, 3) predicted vertex positions.
+        target: (M, 3) target vertex positions.
+
+    Returns:
+        Similarity score in [0, 1]. 1.0 = perfect match.
+    """
+    if len(predicted) == 0 or len(target) == 0:
+        return 0.0
+
+    # Center both point clouds
+    pred_centered = predicted - predicted.mean(axis=0)
+    target_centered = target - target.mean(axis=0)
+
+    # Try multiple rotations and pick best match
+    best_score = 0.0
+    for rotation in _get_alignment_rotations():
+        rotated = pred_centered @ rotation.T
+        score = _chamfer_similarity(rotated, target_centered)
+        best_score = max(best_score, score)
+
+    return best_score
+
+
+def _chamfer_similarity(a: np.ndarray, b: np.ndarray) -> float:
+    """Chamfer distance converted to similarity score."""
+    d = cdist(a, b)
+
+    # Forward: for each point in a, min distance to b
+    forward = d.min(axis=1).mean()
+    # Backward: for each point in b, min distance to a
+    backward = d.min(axis=0).mean()
+    chamfer = (forward + backward) / 2.0
+
+    # Normalize by bounding box diagonal of target
+    all_pts = np.vstack([a, b])
+    bbox_diag = np.linalg.norm(all_pts.max(axis=0) - all_pts.min(axis=0))
+    if bbox_diag < 1e-12:
+        return 1.0 if chamfer < 1e-12 else 0.0
+
+    similarity = max(0.0, 1.0 - chamfer / bbox_diag)
+    return similarity
+
+
+def _get_alignment_rotations() -> list[np.ndarray]:
+    """Generate rotation matrices for alignment search.
+
+    Identity + 90 deg rotations around each axis + mirrors (15 total).
+    """
+    I = np.eye(3)
+    rotations = [I]
+
+    # 90 deg rotations around Z axis
+    for k in range(1, 4):
+        angle = k * np.pi / 2
+        c, s = np.cos(angle), np.sin(angle)
+        rotations.append(np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]]))
+
+    # 90 deg rotations around X axis
+    for k in range(1, 4):
+        angle = k * np.pi / 2
+        c, s = np.cos(angle), np.sin(angle)
+        rotations.append(np.array([[1, 0, 0], [0, c, -s], [0, s, c]]))
+
+    # 90 deg rotations around Y axis
+    for k in range(1, 4):
+        angle = k * np.pi / 2
+        c, s = np.cos(angle), np.sin(angle)
+        rotations.append(np.array([[c, 0, s], [0, 1, 0], [-s, 0, c]]))
+
+    # Mirrors
+    rotations.append(np.diag([-1.0, 1.0, 1.0]))
+    rotations.append(np.diag([1.0, -1.0, 1.0]))
+    rotations.append(np.diag([1.0, 1.0, -1.0]))
+
+    return rotations

origami_server/engine/simulate.py

@@ -0,0 +1,150 @@
+"""Origami fold simulator — analytical rotation with cumulative transforms.
+
+BFS from face 0 through the face adjacency graph. Each face accumulates
+a rotation transform (R, t) such that: folded_pos = R @ flat_pos + t.
+When crossing a fold edge, the fold rotation is composed with the parent
+face's transform. Non-fold edges inherit the parent's transform directly.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+from .fold_parser import parse_fold
+
+
+@dataclass
+class SimResult:
+    """Result of a fold simulation."""
+
+    positions: np.ndarray  # (N, 3) final vertex positions
+    converged: bool
+    steps_taken: int
+    max_strain: float
+    total_energy: float
+
+
+def simulate(
+    fold_data: dict,
+    crease_percent: float = 1.0,
+    max_steps: int = 500,
+    params: dict | None = None,
+) -> SimResult:
+    """Simulate a FOLD crease pattern and return final 3D positions."""
+    parsed = parse_fold(fold_data)
+    flat_pos = parsed["vertices"].copy()
+    edges = parsed["edges"]
+    assignments = parsed["assignments"]
+    fold_angles = parsed["fold_angles"]
+    faces = parsed["faces"]
+    positions = flat_pos.copy()
+
+    if len(faces) == 0:
+        return SimResult(
+            positions=positions, converged=True,
+            steps_taken=0, max_strain=0.0, total_energy=0.0,
+        )
+
+    face_adj = _build_face_adjacency(faces)
+
+    crease_map: dict[tuple[int, int], float] = {}
+    for i, (v1, v2) in enumerate(edges):
+        key = (min(int(v1), int(v2)), max(int(v1), int(v2)))
+        if assignments[i] in ("M", "V"):
+            crease_map[key] = fold_angles[i] * crease_percent
+
+    n_faces = len(faces)
+    face_R = [None] * n_faces
+    face_t = [None] * n_faces
+
+    face_R[0] = np.eye(3)
+    face_t[0] = np.zeros(3)
+
+    visited = [False] * n_faces
+    visited[0] = True
+
+    placed: set[int] = set()
+    for vi in faces[0]:
+        placed.add(int(vi))
+
+    queue = [0]
+    while queue:
+        fi = queue.pop(0)
+        face = faces[fi]
+
+        for j in range(len(face)):
+            v1, v2 = int(face[j]), int(face[(j + 1) % len(face)])
+            edge_key = (min(v1, v2), max(v1, v2))
+
+            for fj in face_adj.get(edge_key, []):
+                if visited[fj]:
+                    continue
+                visited[fj] = True
+                queue.append(fj)
+
+                angle = crease_map.get(edge_key, 0.0)
+
+                if abs(angle) > 1e-10:
+                    p1 = positions[v1].copy()
+                    axis = positions[v2] - p1
+                    axis_len = np.linalg.norm(axis)
+                    if axis_len > 1e-12:
+                        axis_unit = axis / axis_len
+                        fold_rot = Rotation.from_rotvec(
+                            angle * axis_unit,
+                        ).as_matrix()
+                    else:
+                        fold_rot = np.eye(3)
+
+                    face_R[fj] = fold_rot @ face_R[fi]
+                    face_t[fj] = fold_rot @ (face_t[fi] - p1) + p1
+                else:
+                    face_R[fj] = face_R[fi].copy()
+                    face_t[fj] = face_t[fi].copy()
+
+                for vi in faces[fj]:
+                    vi_int = int(vi)
+                    if vi_int not in placed:
+                        positions[vi_int] = face_R[fj] @ flat_pos[vi_int] + face_t[fj]
+                        placed.add(vi_int)
+
+    max_strain = _compute_strain(positions, parsed)
+
+    return SimResult(
+        positions=positions,
+        converged=True,
+        steps_taken=1,
+        max_strain=max_strain,
+        total_energy=0.0,
+    )
+
+
+def _build_face_adjacency(
+    faces: np.ndarray,
+) -> dict[tuple[int, int], list[int]]:
+    """Map each edge (sorted vertex pair) to list of face indices."""
+    adj: dict[tuple[int, int], list[int]] = {}
+    for fi, face in enumerate(faces):
+        n = len(face)
+        for j in range(n):
+            v1, v2 = int(face[j]), int(face[(j + 1) % n])
+            key = (min(v1, v2), max(v1, v2))
+            if key not in adj:
+                adj[key] = []
+            adj[key].append(fi)
+    return adj
+
+
+def _compute_strain(positions: np.ndarray, parsed: dict) -> float:
+    """Compute max axial strain across all edges."""
+    edges = parsed["edges"]
+    vertices_flat = parsed["vertices"]
+    max_strain = 0.0
+    for v1, v2 in edges:
+        rest = np.linalg.norm(vertices_flat[v2] - vertices_flat[v1])
+        curr = np.linalg.norm(positions[v2] - positions[v1])
+        if rest > 1e-12:
+            strain = abs(curr - rest) / rest
+            max_strain = max(max_strain, strain)
+    return max_strain

origami_server/environment.py

@@ -0,0 +1,158 @@
+"""Origami RL Environment — OpenEnv Environment subclass.
+
+Single-shot episodes: LLM submits a FOLD crease pattern, physics simulates it,
+reward = shape similarity to target. Like AlphaFold for origami.
+"""
+
+import uuid
+from typing import Any, Optional
+
+import numpy as np
+from openenv.core import Environment
+
+from .engine.fold_parser import validate_fold
+from .engine.shape_match import compute_shape_match
+from .engine.simulate import SimResult, simulate
+from .models import OrigamiAction, OrigamiObservation, OrigamiState
+from .tasks import get_task
+
+
+class OrigamiEnvironment(
+    Environment[OrigamiAction, OrigamiObservation, OrigamiState]
+):
+    """Origami folding environment.
+
+    Episode flow:
+        1. reset(task_name="triangle") -> returns task description + target info
+        2. step(OrigamiAction(fold_data={...})) -> simulates, scores, returns done=True
+
+    Single action per episode. The action IS the complete crease pattern.
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS = True
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self._state = OrigamiState()
+        self._task: dict = {}
+        self._target_positions: np.ndarray = np.zeros((0, 3))
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> OrigamiObservation:
+        """Start a new episode with a target shape task."""
+        self._state = OrigamiState(
+            episode_id=episode_id or str(uuid.uuid4()),
+            step_count=0,
+        )
+
+        task_name = kwargs.get("task_name", "triangle")
+        self._task = get_task(task_name)
+        self._state.task_name = self._task["name"]
+
+        target_fold = self._task["target_fold"]
+        try:
+            target_result = simulate(target_fold, crease_percent=1.0)
+            self._target_positions = target_result.positions
+        except Exception:
+            self._target_positions = np.zeros((0, 3))
+
+        return OrigamiObservation(
+            done=False,
+            reward=None,
+            task=self._task_info(),
+            fold_data={},
+            final_positions=[],
+            target_positions=self._target_positions.tolist(),
+            shape_similarity=0.0,
+            max_strain=0.0,
+            is_stable=True,
+            error=None,
+        )
+
+    def step(
+        self,
+        action: OrigamiAction,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> OrigamiObservation:
+        """Evaluate the LLM's crease pattern.
+
+        1. Validate FOLD data
+        2. Run physics simulation (creasePercent=1.0)
+        3. Compare final shape to target
+        4. Return observation with reward = similarity * 20
+        """
+        self._state.step_count += 1
+        fold_data = action.fold_data
+
+        is_valid, error_msg = validate_fold(fold_data)
+        if not is_valid:
+            self._state.is_stable = False
+            return OrigamiObservation(
+                done=True,
+                reward=-2.0,
+                task=self._task_info(),
+                fold_data=fold_data,
+                final_positions=[],
+                target_positions=self._target_positions.tolist(),
+                shape_similarity=0.0,
+                max_strain=0.0,
+                is_stable=False,
+                error=f"Invalid FOLD data: {error_msg}",
+            )
+
+        try:
+            result: SimResult = simulate(fold_data, crease_percent=1.0)
+        except Exception as e:
+            self._state.is_stable = False
+            return OrigamiObservation(
+                done=True,
+                reward=-2.0,
+                task=self._task_info(),
+                fold_data=fold_data,
+                final_positions=[],
+                target_positions=self._target_positions.tolist(),
+                shape_similarity=0.0,
+                max_strain=0.0,
+                is_stable=False,
+                error=f"Simulation error: {str(e)}",
+            )
+
+        similarity = compute_shape_match(
+            result.positions, self._target_positions
+        )
+        reward = similarity * 20.0
+
+        self._state.shape_similarity = similarity
+        self._state.is_stable = result.converged
+
+        return OrigamiObservation(
+            done=True,
+            reward=reward,
+            task=self._task_info(),
+            fold_data=fold_data,
+            final_positions=result.positions.tolist(),
+            target_positions=self._target_positions.tolist(),
+            shape_similarity=similarity,
+            max_strain=result.max_strain,
+            is_stable=result.converged,
+            error=None,
+        )
+
+    @property
+    def state(self) -> OrigamiState:
+        return self._state
+
+    def _task_info(self) -> dict:
+        if not self._task:
+            return {}
+        return {
+            "name": self._task.get("name", ""),
+            "description": self._task.get("description", ""),
+            "difficulty": self._task.get("difficulty", 0),
+            "paper": self._task.get("paper", {}),
+        }

origami_server/models.py

@@ -0,0 +1,48 @@
+"""OpenEnv types for the Origami RL environment.
+
+OrigamiAction: LLM submits a FOLD crease pattern.
+OrigamiObservation: Result of simulating that pattern against a target.
+OrigamiState: Internal episode state.
+"""
+
+from typing import Any, Optional
+
+from openenv.core import Action, Observation, State
+from pydantic import Field
+
+
+class OrigamiAction(Action):
+    """LLM submits a FOLD crease pattern as its action.
+
+    The fold_data dict must contain:
+      - vertices_coords: [[x, y], ...] — 2D vertex positions on flat paper
+      - edges_vertices: [[v1, v2], ...] — edge connectivity
+      - edges_assignment: ["B"|"M"|"V", ...] — boundary/mountain/valley
+      - edges_foldAngle: [angle, ...] — target fold angles in degrees
+        (optional — defaults from assignment: M=-180, V=+180, B=0)
+    """
+
+    fold_data: dict[str, Any] = Field(
+        ..., description="FOLD-format crease pattern JSON"
+    )
+
+
+class OrigamiObservation(Observation):
+    """Result of simulating the LLM's crease pattern."""
+
+    task: dict[str, Any] = Field(default_factory=dict)
+    fold_data: dict[str, Any] = Field(default_factory=dict)
+    final_positions: list[list[float]] = Field(default_factory=list)
+    target_positions: list[list[float]] = Field(default_factory=list)
+    shape_similarity: float = 0.0
+    max_strain: float = 0.0
+    is_stable: bool = True
+    error: Optional[str] = None
+
+
+class OrigamiState(State):
+    """Internal state for an origami episode."""
+
+    task_name: str = ""
+    shape_similarity: float = 0.0
+    is_stable: bool = True

origami_server/tasks.py

@@ -0,0 +1,218 @@
+"""Task definitions for origami RL training.
+
+Each task defines a target shape as a reference FOLD crease pattern.
+The LLM must discover a crease pattern that folds into the same shape.
+"""
+
+TASKS: dict[str, dict] = {
+    "triangle": {
+        "name": "triangle",
+        "description": "Fold the paper in half diagonally to make a triangle",
+        "difficulty": 1,
+        "paper": {"width": 1.0, "height": 1.0},
+        "target_fold": {
+            "vertices_coords": [[0, 0], [1, 0], [1, 1], [0, 1]],
+            "edges_vertices": [[0, 1], [1, 2], [2, 3], [3, 0], [0, 2]],
+            "edges_assignment": ["B", "B", "B", "B", "V"],
+            "edges_foldAngle": [0, 0, 0, 0, 180],
+            "faces_vertices": [[0, 1, 2], [0, 2, 3]],
+        },
+    },
+    "half_fold": {
+        "name": "half_fold",
+        "description": "Fold the paper in half horizontally",
+        "difficulty": 1,
+        "paper": {"width": 1.0, "height": 1.0},
+        "target_fold": {
+            "vertices_coords": [
+                [0, 0], [1, 0], [1, 1], [0, 1], [0, 0.5], [1, 0.5],
+            ],
+            "edges_vertices": [
+                [0, 1], [1, 5], [5, 2], [2, 3], [3, 4], [4, 0],
+                [4, 5],
+            ],
+            "edges_assignment": ["B", "B", "B", "B", "B", "B", "V"],
+            "edges_foldAngle": [0, 0, 0, 0, 0, 0, 180],
+            "faces_vertices": [[0, 1, 5, 4], [4, 5, 2, 3]],
+        },
+    },
+    "quarter_fold": {
+        "name": "quarter_fold",
+        "description": "Fold the paper into quarters (two perpendicular folds)",
+        "difficulty": 2,
+        "paper": {"width": 1.0, "height": 1.0},
+        "target_fold": {
+            "vertices_coords": [
+                [0, 0], [0.5, 0], [1, 0],
+                [0, 0.5], [0.5, 0.5], [1, 0.5],
+                [0, 1], [0.5, 1], [1, 1],
+            ],
+            "edges_vertices": [
+                [0, 1], [1, 2], [2, 5], [5, 8], [8, 7], [7, 6], [6, 3], [3, 0],
+                [1, 4], [4, 7],
+                [3, 4], [4, 5],
+            ],
+            "edges_assignment": [
+                "B", "B", "B", "B", "B", "B", "B", "B",
+                "V", "V", "V", "V",
+            ],
+            "edges_foldAngle": [
+                0, 0, 0, 0, 0, 0, 0, 0,
+                180, 180, 180, 180,
+            ],
+            "faces_vertices": [
+                [0, 1, 4, 3],
+                [1, 2, 5, 4],
+                [3, 4, 7, 6],
+                [4, 5, 8, 7],
+            ],
+        },
+    },
+    "letter_fold": {
+        "name": "letter_fold",
+        "description": "Tri-fold the paper like a letter (two parallel folds)",
+        "difficulty": 2,
+        "paper": {"width": 1.0, "height": 1.0},
+        "target_fold": {
+            "vertices_coords": [
+                [0, 0], [1, 0],
+                [0, 1/3], [1, 1/3],
+                [0, 2/3], [1, 2/3],
+                [0, 1], [1, 1],
+            ],
+            "edges_vertices": [
+                [0, 1], [1, 3], [3, 5], [5, 7], [7, 6], [6, 4], [4, 2], [2, 0],
+                [2, 3],
+                [4, 5],
+            ],
+            "edges_assignment": [
+                "B", "B", "B", "B", "B", "B", "B", "B",
+                "V", "M",
+            ],
+            "edges_foldAngle": [
+                0, 0, 0, 0, 0, 0, 0, 0,
+                180, -180,
+            ],
+            "faces_vertices": [
+                [0, 1, 3, 2],
+                [2, 3, 5, 4],
+                [4, 5, 7, 6],
+            ],
+        },
+    },
+    # --- Optigami patterns (from lib/patterns.ts) ---
+    "waterbomb": {
+        "name": "waterbomb",
+        "description": "Create a waterbomb base with valley folds on diagonals and mountain folds on midlines of a square sheet",
+        "difficulty": 3,
+        "paper": {"width": 2.0, "height": 2.0},
+        "target_fold": {
+            "vertices_coords": [
+                [-1, 1], [0, 1], [1, 1],
+                [-1, 0], [0, 0], [1, 0],
+                [-1, -1], [0, -1], [1, -1],
+            ],
+            "edges_vertices": [
+                [0, 1], [1, 2], [2, 5], [5, 8], [8, 7], [7, 6], [6, 3], [3, 0],
+                [0, 4], [2, 4], [6, 4], [8, 4],
+                [1, 4], [3, 4], [5, 4], [7, 4],
+            ],
+            "edges_assignment": [
+                "B", "B", "B", "B", "B", "B", "B", "B",
+                "V", "V", "V", "V",
+                "M", "M", "M", "M",
+            ],
+            "edges_foldAngle": [
+                0, 0, 0, 0, 0, 0, 0, 0,
+                180, 180, 180, 180,
+                -180, -180, -180, -180,
+            ],
+            "faces_vertices": [
+                [0, 1, 4], [1, 2, 4],
+                [2, 5, 4], [5, 8, 4],
+                [8, 7, 4], [7, 6, 4],
+                [6, 3, 4], [3, 0, 4],
+            ],
+        },
+    },
+    "accordion": {
+        "name": "accordion",
+        "description": "Make an accordion (zig-zag) fold with alternating mountain and valley creases like a paper fan",
+        "difficulty": 2,
+        "paper": {"width": 2.0, "height": 2.0},
+        "target_fold": {
+            "vertices_coords": [
+                [-1, 1], [-0.5, 1], [0, 1], [0.5, 1], [1, 1],
+                [-1, -1], [-0.5, -1], [0, -1], [0.5, -1], [1, -1],
+            ],
+            "edges_vertices": [
+                [0, 1], [1, 2], [2, 3], [3, 4],
+                [5, 6], [6, 7], [7, 8], [8, 9],
+                [0, 5], [4, 9],
+                [1, 6], [2, 7], [3, 8],
+            ],
+            "edges_assignment": [
+                "B", "B", "B", "B",
+                "B", "B", "B", "B",
+                "B", "B",
+                "V", "M", "V",
+            ],
+            "edges_foldAngle": [
+                0, 0, 0, 0,
+                0, 0, 0, 0,
+                0, 0,
+                180, -180, 180,
+            ],
+            "faces_vertices": [
+                [0, 1, 6, 5], [1, 2, 7, 6],
+                [2, 3, 8, 7], [3, 4, 9, 8],
+            ],
+        },
+    },
+    "miura_ori": {
+        "name": "miura_ori",
+        "description": "Create a Miura-ori tessellation fold pattern on a 2x2 grid with offset zigzag vertices",
+        "difficulty": 3,
+        "paper": {"width": 2.0, "height": 2.0},
+        "target_fold": {
+            "vertices_coords": [
+                [-1, 1], [0, 1.2], [1, 1],
+                [-1, 0], [0, 0.2], [1, 0],
+                [-1, -1], [0, -0.8], [1, -1],
+            ],
+            "edges_vertices": [
+                [0, 1], [1, 2], [2, 5], [5, 8], [8, 7], [7, 6], [6, 3], [3, 0],
+                [1, 4], [4, 7],
+                [3, 4], [4, 5],
+            ],
+            "edges_assignment": [
+                "B", "B", "B", "B", "B", "B", "B", "B",
+                "M", "M",
+                "V", "M",
+            ],
+            "edges_foldAngle": [
+                0, 0, 0, 0, 0, 0, 0, 0,
+                -180, -180,
+                180, -180,
+            ],
+            "faces_vertices": [
+                [0, 1, 4, 3], [1, 2, 5, 4],
+                [3, 4, 7, 6], [4, 5, 8, 7],
+            ],
+        },
+    },
+}
+
+
+def get_task(name: str | None = None) -> dict:
+    """Get a task by name. Defaults to 'triangle'."""
+    if name is None:
+        name = "triangle"
+    if name not in TASKS:
+        raise ValueError(f"Unknown task '{name}'. Available: {list(TASKS.keys())}")
+    return TASKS[name]
+
+
+def list_tasks() -> list[str]:
+    """List all available task names."""
+    return list(TASKS.keys())

requirements.txt

@@ -0,0 +1,6 @@
+fastapi>=0.100.0
+numpy>=1.24.0
+openenv-core[core]>=0.2.1
+pydantic>=2.0.0
+scipy>=1.10
+uvicorn>=0.23.0

start.sh

@@ -0,0 +1,6 @@
+#!/bin/sh
+# Start the OpenEnv API in the background on port 8000
+python3 -m uvicorn origami_server.app:app --host 0.0.0.0 --port 8000 &
+
+# Start Next.js frontend on port 7860 (foreground)
+node server.js

training/reward.py

@@ -0,0 +1,119 @@
+"""GRPO reward functions for origami RL training.
+
+Two reward functions (matching the 2048 pattern):
+1. valid_fold: Does the LLM output parse as valid FOLD JSON?
+2. shape_match: Simulate and compare to target shape.
+"""
+
+import json
+import re
+from typing import Any
+
+import numpy as np
+
+from origami_server.engine.fold_parser import validate_fold
+from origami_server.engine.shape_match import compute_shape_match
+from origami_server.engine.simulate import simulate
+from origami_server.tasks import get_task
+
+
+def extract_fold_json(response: str) -> dict | None:
+    """Extract FOLD JSON from LLM response text.
+
+    Looks for JSON between ```json ... ``` or raw JSON object.
+    """
+    # Try fenced code block first
+    match = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", response, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(1))
+        except json.JSONDecodeError:
+            pass
+
+    # Try raw JSON object
+    match = re.search(r"\{[^{}]*\"vertices_coords\"[^{}]*\}", response, re.DOTALL)
+    if match:
+        try:
+            return json.loads(match.group(0))
+        except json.JSONDecodeError:
+            pass
+
+    # Try parsing the whole response
+    try:
+        data = json.loads(response.strip())
+        if isinstance(data, dict) and "vertices_coords" in data:
+            return data
+    except (json.JSONDecodeError, ValueError):
+        pass
+
+    return None
+
+
+def valid_fold(completions: list, **kwargs: Any) -> list[float]:
+    """Reward 1: Does the LLM output parse as valid FOLD JSON?
+
+    +1.0  valid FOLD JSON with correct structure
+    -0.5  parseable JSON but invalid FOLD structure
+    -2.0  not parseable as JSON at all
+    """
+    scores = []
+    for completion in completions:
+        response = completion[0]["content"]
+        fold_data = extract_fold_json(response)
+
+        if fold_data is None:
+            scores.append(-2.0)
+            continue
+
+        is_valid, error = validate_fold(fold_data)
+        if is_valid:
+            scores.append(1.0)
+        else:
+            scores.append(-0.5)
+
+    return scores
+
+
+def shape_match(
+    completions: list,
+    task_name: str = "triangle",
+    **kwargs: Any,
+) -> list[float]:
+    """Reward 2: Simulate the fold and compare to target shape.
+
+    Score = similarity * 20.0 (range: 0 to 20)
+    -1.0  if simulation fails/diverges
+    -2.0  if FOLD data is invalid
+    """
+    task = get_task(task_name)
+    target_fold = task["target_fold"]
+
+    # Pre-compute target positions
+    try:
+        target_result = simulate(target_fold, crease_percent=1.0)
+        target_positions = target_result.positions
+    except Exception:
+        return [0.0] * len(completions)
+
+    scores = []
+    for completion in completions:
+        response = completion[0]["content"]
+        fold_data = extract_fold_json(response)
+
+        if fold_data is None:
+            scores.append(-2.0)
+            continue
+
+        is_valid, error = validate_fold(fold_data)
+        if not is_valid:
+            scores.append(-1.0)
+            continue
+
+        try:
+            result = simulate(fold_data, crease_percent=1.0)
+            similarity = compute_shape_match(result.positions, target_positions)
+            scores.append(similarity * 20.0)
+        except Exception:
+            scores.append(-1.0)
+
+    return scores

training/train_grpo.py

@@ -0,0 +1,128 @@
+"""GRPO training script for origami RL.
+
+Follows the OpenEnv + Unsloth pattern:
+- LLM generates FOLD JSON crease patterns
+- Two reward functions: valid_fold + shape_match
+- GRPOTrainer from TRL handles the RL loop
+
+Usage (Colab):
+    python -m training.train_grpo --task triangle --max_steps 600
+"""
+
+import argparse
+
+PROMPT_TEMPLATE = """You are an origami designer. Generate a FOLD-format crease pattern
+that, when folded, produces the target shape described below.
+
+Target: {description}
+Paper size: {width} x {height}
+
+Output a JSON object with these exact fields:
+- vertices_coords: [[x, y], ...] — 2D positions on the flat paper (0 to {width} for x, 0 to {height} for y)
+- edges_vertices: [[v1, v2], ...] — pairs of vertex indices forming edges
+- edges_assignment: ["B"|"M"|"V", ...] — B=boundary, M=mountain fold, V=valley fold
+- edges_foldAngle: [angle, ...] — fold angles in degrees (M: negative like -180, V: positive like 180, B: 0)
+
+Rules:
+- Boundary edges (B) must outline the paper rectangle
+- At least one fold crease (M or V) must exist
+- Mountain fold angles are negative (-180 to 0)
+- Valley fold angles are positive (0 to 180)
+- All vertex indices in edges must be valid (0 to N-1)
+
+Output ONLY the JSON object wrapped in ```json ... ``` markers."""
+
+
+def build_prompt(task: dict) -> str:
+    return PROMPT_TEMPLATE.format(
+        description=task["description"],
+        width=task["paper"]["width"],
+        height=task["paper"]["height"],
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(description="GRPO training for origami RL")
+    parser.add_argument("--task", default="triangle", help="Task name")
+    parser.add_argument("--max_steps", type=int, default=600)
+    parser.add_argument("--num_generations", type=int, default=4)
+    parser.add_argument("--model", default="unsloth/Qwen3-8B-unsloth-bnb-4bit")
+    parser.add_argument("--lr", type=float, default=2e-4)
+    args = parser.parse_args()
+
+    # --- These imports are heavy, only load when actually training ---
+    from datasets import Dataset
+    from trl import GRPOConfig, GRPOTrainer
+    from unsloth import FastLanguageModel
+
+    from origami_server.tasks import get_task
+    from training.reward import shape_match, valid_fold
+
+    task = get_task(args.task)
+    prompt_text = build_prompt(task)
+
+    # Build dataset (1000 copies of same prompt, like 2048)
+    dataset = Dataset.from_list(
+        [
+            {
+                "prompt": [{"role": "user", "content": prompt_text}],
+                "answer": 0,
+            }
+        ]
+        * 1000
+    )
+
+    # Load model with LoRA
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        model_name=args.model,
+        load_in_4bit=True,
+        max_seq_length=2048,
+    )
+
+    model = FastLanguageModel.get_peft_model(
+        model,
+        r=8,
+        target_modules=[
+            "q_proj", "k_proj", "v_proj", "o_proj",
+            "gate_proj", "up_proj", "down_proj",
+        ],
+        lora_alpha=16,
+        use_gradient_checkpointing="unsloth",
+    )
+
+    # Wrap shape_match to inject task_name
+    def shape_match_reward(completions, **kwargs):
+        return shape_match(completions, task_name=args.task, **kwargs)
+
+    # GRPO config
+    training_args = GRPOConfig(
+        temperature=1.0,
+        learning_rate=args.lr,
+        weight_decay=0.001,
+        warmup_ratio=0.1,
+        lr_scheduler_type="linear",
+        optim="adamw_8bit",
+        logging_steps=1,
+        per_device_train_batch_size=1,
+        gradient_accumulation_steps=1,
+        num_generations=args.num_generations,
+        max_prompt_length=1024,
+        max_completion_length=1024,
+        max_steps=args.max_steps,
+        save_steps=100,
+        output_dir="outputs",
+    )
+
+    trainer = GRPOTrainer(
+        model=model,
+        processing_class=tokenizer,
+        reward_funcs=[valid_fold, shape_match_reward],
+        args=training_args,
+        train_dataset=dataset,
+    )
+
+    trainer.train()
+
+
+if __name__ == "__main__":
+    main()