refactor(openenv): simplify runtime environment and models, extend server API

openenv_runtime/environment.py

@@ -1,183 +1,53 @@
-from __future__ import annotations
+"""
+OpenEnv adapter for Optigami.
 
-from typing import Any, Optional
+Thin wrapper around env.environment.OrigamiEnvironment that adapts it to the
+OpenEnv protocol (Action/Observation types).
+"""
+from env.environment import OrigamiEnvironment as _Env
 
-from openenv.core.env_server.interfaces import Environment
+from .models import OrigamiAction, OrigamiObservation
 
-from env.environment import OrigamiEnvironment
 
-from .models import OrigamiAction, OrigamiObservation, OrigamiState
+class OpenEnvOrigamiEnvironment:
+    """
+    OpenEnv-compatible wrapper for env.environment.OrigamiEnvironment.
 
+    Converts between env's dict-based API and OpenEnv's Action/Observation types.
+    """
 
-class OpenEnvOrigamiEnvironment(Environment[OrigamiAction, OrigamiObservation, OrigamiState]):
-    """OpenEnv adapter over the existing OrigamiEnvironment implementation."""
+    def __init__(self, mode: str = "step", max_steps: int = 8, targets_dir=None):
+        self._env = _Env(mode=mode, max_steps=max_steps, targets_dir=targets_dir)
 
-    SUPPORTS_CONCURRENT_SESSIONS = True
-
-    def __init__(
-        self,
-        default_mode: str = "step",
-        max_steps: int = 8,
-        targets_dir: Optional[str] = None,
-    ):
-        super().__init__()
-        self.default_mode = default_mode
-        self.max_steps = max_steps
-        self.targets_dir = targets_dir
-        self._env: Optional[OrigamiEnvironment] = None
-        self._episode_id: Optional[str] = None
-
-    def _new_env(self, mode: Optional[str] = None) -> OrigamiEnvironment:
-        return OrigamiEnvironment(
-            mode=mode or self.default_mode,
-            max_steps=self.max_steps,
-            targets_dir=self.targets_dir,
-        )
-
-    def reset(
-        self,
-        seed: Optional[int] = None,
-        episode_id: Optional[str] = None,
-        **kwargs: Any,
-    ) -> OrigamiObservation:
-        del seed  # deterministic seed plumbing can be added later
-
-        mode = kwargs.get("mode", self.default_mode)
-        target_name = kwargs.get("target_name")
-
-        self._env = self._new_env(mode=mode)
-        self._episode_id = episode_id
+    def reset(self, target_name=None, **kwargs):
         obs_dict = self._env.reset(target_name=target_name)
+        return self._obs_dict_to_model(obs_dict, reward=None, done=False)
 
+    def step(self, action: OrigamiAction, **kwargs):
+        action_dict = {
+            "from": action.from_point,
+            "to": action.to_point,
+            "assignment": action.assignment,
+        }
+        obs_dict, reward, done, info = self._env.step(action_dict)
+        reward_val = reward.get("total", 0.0) if isinstance(reward, dict) else reward
+        return self._obs_dict_to_model(obs_dict, reward=reward_val, done=done)
+
+    def _obs_dict_to_model(self, obs_dict: dict, reward=None, done=False) -> OrigamiObservation:
         return OrigamiObservation(
-            done=False,
-            reward=None,
-            metadata={"available_targets": self._env.available_targets()},
             prompt=obs_dict.get("prompt", ""),
-            target_name=obs_dict.get("target_name"),
+            target_name=obs_dict.get("target_name", ""),
             step=obs_dict.get("step", 0),
-            paper_state=self._paper_state_snapshot(),
-            info=self._env._info(),
-            reward_components={},
-        )
-
-    def step(
-        self,
-        action: OrigamiAction,
-        timeout_s: Optional[float] = None,
-        **kwargs: Any,
-    ) -> OrigamiObservation:
-        del timeout_s, kwargs
-
-        if self._env is None:
-            self.reset(target_name=action.target_name)
-
-        assert self._env is not None
-
-        if action.target_name and action.target_name != self._env.target_name:
-            self.reset(target_name=action.target_name, mode=self._env.mode)
-
-        try:
-            if action.mode == "sequence":
-                if not action.completion:
-                    return self._error_observation("sequence mode requires completion")
-
-                seq_env = self._new_env(mode="code_as_policy")
-                seq_env.reset(target_name=self._env.target_name)
-                obs_dict, reward_dict, done, info = seq_env.step(action.completion)
-                self._env = seq_env
-            else:
-                if action.fold is not None:
-                    fold_payload = {
-                        "from": list(action.fold.from_point),
-                        "to": list(action.fold.to_point),
-                        "assignment": action.fold.assignment,
-                        "instruction": action.fold.instruction,
-                    }
-                    env_action: Any = fold_payload
-                elif action.completion:
-                    env_action = action.completion
-                else:
-                    return self._error_observation("single mode requires fold or completion")
-
-                obs_dict, reward_dict, done, info = self._env.step(env_action)
-
-            total = reward_dict.get("total") if isinstance(reward_dict, dict) else None
-            return OrigamiObservation(
-                done=bool(done),
-                reward=float(total) if isinstance(total, (int, float)) else None,
-                metadata={"target_name": self._env.target_name},
-                prompt=obs_dict.get("prompt", ""),
-                target_name=obs_dict.get("target_name", self._env.target_name),
-                step=obs_dict.get("step", self._env.step_count),
-                paper_state=self._paper_state_snapshot(),
-                info=info or {},
-                reward_components=reward_dict or {},
-            )
-        except Exception as exc:  # pragma: no cover - defensive path
-            return self._error_observation(str(exc))
-
-    @property
-    def state(self) -> OrigamiState:
-        if self._env is None:
-            tmp_env = self._new_env(mode=self.default_mode)
-            return OrigamiState(
-                episode_id=self._episode_id,
-                step_count=0,
-                mode=tmp_env.mode,
-                target_name=None,
-                paper={},
-                last_reward={},
-                available_targets=tmp_env.available_targets(),
-            )
-
-        env_state = self._env.state()
-        return OrigamiState(
-            episode_id=self._episode_id,
-            step_count=env_state.get("step", self._env.step_count),
-            mode=env_state.get("mode", self._env.mode),
-            target_name=env_state.get("target", self._env.target_name),
-            paper=env_state.get("paper", {}),
-            last_reward=self._env.last_reward or {},
-            available_targets=self._env.available_targets(),
+            paper_fold_json=obs_dict.get("paper_fold_json", {}),
+            reward=reward,
+            done=done,
         )
 
-    def close(self) -> None:
-        if self._env is not None:
-            self._env.close()
-            self._env = None
+    def state(self):
+        return self._env.state()
 
-    def _paper_state_snapshot(self) -> dict[str, Any]:
-        if self._env is None or self._env.paper is None:
-            return {"vertices": {}, "edges": [], "anchor_points": []}
+    def close(self):
+        self._env.close()
 
-        graph = self._env.paper.graph
-        return {
-            "vertices": {str(k): [float(v[0]), float(v[1])] for k, v in graph.vertices.items()},
-            "edges": [
-                {
-                    "id": int(eid),
-                    "v1": [float(graph.vertices[v1][0]), float(graph.vertices[v1][1])],
-                    "v2": [float(graph.vertices[v2][0]), float(graph.vertices[v2][1])],
-                    "assignment": assignment,
-                }
-                for eid, (v1, v2, assignment) in graph.edges.items()
-            ],
-            "anchor_points": [
-                [float(x), float(y)] for (x, y) in self._env.paper.anchor_points()
-            ],
-        }
 
-    def _error_observation(self, message: str) -> OrigamiObservation:
-        return OrigamiObservation(
-            done=False,
-            reward=-0.1,
-            metadata={"error": True},
-            prompt="",
-            target_name=self._env.target_name if self._env else None,
-            step=self._env.step_count if self._env else 0,
-            paper_state=self._paper_state_snapshot(),
-            info=self._env._info() if self._env else {},
-            reward_components={"format": 0.0, "total": -0.1, "error": message},
-            error=message,
-        )
+__all__ = ["OpenEnvOrigamiEnvironment"]

openenv_runtime/models.py

@@ -1,63 +1,53 @@
-from __future__ import annotations
+"""
+OpenEnv Pydantic models for the env/ stack.
 
-from typing import Any, Literal, Optional
+Matches the env/environment data shape: observations with prompt, target_name,
+step, paper_fold_json; actions as fold dicts with from/to/assignment.
+"""
+from typing import Optional
 
-from pydantic import BaseModel, Field, field_validator
+from pydantic import ConfigDict, Field
 
 from openenv.core.env_server.types import Action, Observation, State
 
 
-class OrigamiFold(BaseModel):
-    """Single fold action payload for step-level execution."""
-
-    from_point: list[float] = Field(..., description="Fold line start [x, y]")
-    to_point: list[float] = Field(..., description="Fold line end [x, y]")
-    assignment: Literal["M", "V"] = Field(..., description="Mountain or valley")
-    instruction: str = Field(default="", description="Optional natural language instruction")
-
-    @field_validator("from_point", "to_point")
-    @classmethod
-    def _validate_point(cls, point: list[float]) -> list[float]:
-        if len(point) != 2:
-            raise ValueError("Point must contain exactly 2 coordinates")
-        return [float(point[0]), float(point[1])]
+class OrigamiAction(Action):
+    """One fold operation — from_point, to_point, assignment."""
 
+    model_config = ConfigDict(populate_by_name=True)
 
-class OrigamiAction(Action):
-    """
-    OpenEnv action for Optigami.
-
-    Modes:
-    - single: execute one fold (pass `fold` or JSON `completion` for a single-fold object)
-    - sequence: execute a full <folds>[...]</folds> completion in one step
-    """
-
-    mode: Literal["single", "sequence"] = Field(default="single")
-    fold: Optional[OrigamiFold] = Field(default=None)
-    completion: Optional[str] = Field(default=None)
-    target_name: Optional[str] = Field(
-        default=None,
-        description="Optional target override; reset to this target before stepping",
+    from_point: list[float] = Field(
+        alias="from",
+        description="[x, y] start point of the crease",
+    )
+    to_point: list[float] = Field(
+        alias="to",
+        description="[x, y] end point of the crease",
+    )
+    assignment: str = Field(
+        description="'M' (mountain) or 'V' (valley)",
     )
 
 
 class OrigamiObservation(Observation):
-    """OpenEnv observation payload returned by Optigami."""
-
-    prompt: str = Field(default="")
-    target_name: Optional[str] = Field(default=None)
-    step: int = Field(default=0)
-    paper_state: dict[str, Any] = Field(default_factory=dict)
-    info: dict[str, Any] = Field(default_factory=dict)
-    reward_components: dict[str, float | int | str] = Field(default_factory=dict)
-    error: Optional[str] = Field(default=None)
+    """Observation from env.environment — prompt, target, step, paper state."""
+
+    prompt: str = Field(default="", description="LLM prompt for the current step")
+    target_name: str = Field(default="", description="Name of the target (.fold stem)")
+    step: int = Field(default=0, ge=0, description="Current step index")
+    paper_fold_json: dict = Field(
+        default_factory=dict,
+        description="Graph edges (crease pattern state)",
+    )
 
 
 class OrigamiState(State):
-    """OpenEnv state payload for Optigami."""
+    """Server-side episode state."""
+
+    paper: dict = Field(default_factory=dict, description="Paper state")
+    target: Optional[str] = Field(default=None, description="Target name")
+    step: int = Field(default=0, ge=0, description="Step count")
+    mode: str = Field(default="step", description="'step' or 'code_as_policy'")
+
 
-    mode: str = Field(default="step")
-    target_name: Optional[str] = Field(default=None)
-    paper: dict[str, Any] = Field(default_factory=dict)
-    last_reward: dict[str, Any] = Field(default_factory=dict)
-    available_targets: list[str] = Field(default_factory=list)
+__all__ = ["OrigamiAction", "OrigamiObservation", "OrigamiState"]

openenv_server/app.py

@@ -1,12 +1,25 @@
 from __future__ import annotations
 
+import asyncio
 import json
 from pathlib import Path
 
 import numpy as np
+from fastapi import HTTPException, WebSocket
 from fastapi.responses import HTMLResponse, JSONResponse
 from fastapi.staticfiles import StaticFiles
 
+from openenv.core.env_server.http_server import create_app
+
+from env.environment import OrigamiEnvironment
+from openenv_runtime.environment import OpenEnvOrigamiEnvironment
+from openenv_runtime.models import OrigamiAction, OrigamiObservation
+from server.training_broadcast import TrainingBroadcastServer
+
+
+# ---------------------------------------------------------------------------
+# Numpy-safe JSON response
+# ---------------------------------------------------------------------------
 
 def _np_default(obj):
     if isinstance(obj, np.bool_):
@@ -23,56 +36,150 @@ def _np_default(obj):
 class NumpyJSONResponse(JSONResponse):
     def render(self, content) -> bytes:
         return json.dumps(content, default=_np_default).encode("utf-8")
-from openenv.core.env_server.http_server import create_app
 
-from openenv_runtime.environment import OpenEnvOrigamiEnvironment
-from openenv_runtime.models import OrigamiAction, OrigamiObservation
+
+# ---------------------------------------------------------------------------
+# Episode registry for replay
+# ---------------------------------------------------------------------------
+
+_episode_registry: dict[str, dict] = {}
 
 
+# ---------------------------------------------------------------------------
+# OpenEnv app + training broadcast server
+# ---------------------------------------------------------------------------
+
 app = create_app(
-    env=lambda: OpenEnvOrigamiEnvironment(),
+    env=lambda: OpenEnvOrigamiEnvironment(mode="step"),
     action_cls=OrigamiAction,
     observation_cls=OrigamiObservation,
     env_name="optigami",
 )
 
+broadcast = TrainingBroadcastServer()
+
+
+def _ensure_broadcast_loop():
+    """Set broadcast loop on first use (replaces deprecated on_event('startup'))."""
+    if broadcast._loop is None or broadcast._loop.is_closed():
+        try:
+            broadcast._loop = asyncio.get_running_loop()
+        except RuntimeError:
+            pass
+
+
+@app.middleware("http")
+async def _set_broadcast_loop(request, call_next):
+    """Ensure broadcast has event loop before handling requests."""
+    _ensure_broadcast_loop()
+    return await call_next(request)
+
 
 # ---------------------------------------------------------------------------
-# Demo fold sequences — new format: type, line {start, end}, angle
-# ---------------------------------------------------------------------------
-
-DEMO_SEQUENCES: dict[str, list[dict]] = {
-    "half_fold": [
-        {"type": "valley", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0},
-    ],
-    "quarter_fold": [
-        {"type": "valley", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0},
-        {"type": "valley", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0},
-    ],
-    "letter_fold": [
-        {"type": "valley", "line": {"start": [0.0, 0.333], "end": [1.0, 0.333]}, "angle": 180.0},
-        {"type": "mountain", "line": {"start": [0.0, 0.667], "end": [1.0, 0.667]}, "angle": 180.0},
-    ],
-    "map_fold": [
-        {"type": "valley", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0},
-        {"type": "mountain", "line": {"start": [0.5, 0.0], "end": [0.5, 1.0]}, "angle": 180.0},
-    ],
-    "solar_panel": [
-        {"type": "valley", "line": {"start": [0.0, 0.25], "end": [1.0, 0.25]}, "angle": 180.0},
-        {"type": "mountain", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 180.0},
-        {"type": "valley", "line": {"start": [0.0, 0.75], "end": [1.0, 0.75]}, "angle": 180.0},
-    ],
-    "shelter_wall": [
-        {"type": "valley", "line": {"start": [0.0, 0.333], "end": [1.0, 0.333]}, "angle": 180.0},
-        {"type": "valley", "line": {"start": [0.0, 0.667], "end": [1.0, 0.667]}, "angle": 180.0},
-    ],
-    "stent": [
-        {"type": "valley", "line": {"start": [0.0, 0.25], "end": [1.0, 0.25]}, "angle": 90.0},
-        {"type": "mountain", "line": {"start": [0.0, 0.5], "end": [1.0, 0.5]}, "angle": 90.0},
-        {"type": "valley", "line": {"start": [0.0, 0.75], "end": [1.0, 0.75]}, "angle": 90.0},
-        {"type": "stop", "line": {"start": [0.0, 0.0], "end": [1.0, 1.0]}, "angle": 0.0},
-    ],
-}
+# Health endpoint
+# ---------------------------------------------------------------------------
+
+@app.get("/health", include_in_schema=True)
+async def health():
+    return {"status": "ok"}
+
+
+# ---------------------------------------------------------------------------
+# Episode replay endpoint
+# ---------------------------------------------------------------------------
+
+@app.get("/episode/replay/{ep_id}", include_in_schema=True, response_class=NumpyJSONResponse)
+async def replay_episode(ep_id: str):
+    if ep_id not in _episode_registry:
+        raise HTTPException(status_code=404, detail="Episode not found")
+    return NumpyJSONResponse(_episode_registry[ep_id])
+
+
+# ---------------------------------------------------------------------------
+# Training grid viewer WebSocket
+# ---------------------------------------------------------------------------
+
+@app.websocket("/ws/training")
+async def training_ws(websocket: WebSocket):
+    """Read-only spectator WebSocket for the training grid viewer."""
+    _ensure_broadcast_loop()
+    await broadcast.connect_spectator(websocket)
+
+
+# ---------------------------------------------------------------------------
+# Helper: extract crease folds from .fold target
+# ---------------------------------------------------------------------------
+
+def _target_to_folds(target: dict) -> list[dict]:
+    """Extract crease folds from a target .fold dict (edges with M or V)."""
+    verts = target.get("vertices_coords", [])
+    edges_v = target.get("edges_vertices", [])
+    edges_a = target.get("edges_assignment", [])
+    folds = []
+    for (v1, v2), ass in zip(edges_v, edges_a):
+        if ass in ("M", "V") and v1 < len(verts) and v2 < len(verts):
+            p1 = verts[v1]
+            p2 = verts[v2]
+            folds.append({"from": p1, "to": p2, "assignment": ass})
+    return folds
+
+
+def _graph_state_to_fold(paper_dict: dict) -> dict:
+    """Convert internal graph state dict to FOLD-format arrays for the frontend.
+
+    Input format (from env.state()['paper']):
+        vertices: {id: (x, y), ...}
+        edges: {id: (v1_id, v2_id, assignment), ...}  (only M/V)
+
+    Output format (FOLD):
+        vertices_coords: [[x, y, 0], ...]
+        edges_vertices: [[i, j], ...]
+        edges_assignment: ['M'|'V'|'B', ...]
+        faces_vertices: [[i, j, k], ...]  (Delaunay triangulation for 3D)
+    """
+    raw_verts = paper_dict.get("vertices", {})
+    raw_edges = paper_dict.get("edges", {})
+
+    if not raw_verts:
+        return {}
+
+    sorted_ids = sorted(raw_verts.keys(), key=lambda k: int(k) if isinstance(k, (int, str)) else k)
+    id_to_idx = {vid: idx for idx, vid in enumerate(sorted_ids)}
+
+    vertices_coords = []
+    for vid in sorted_ids:
+        xy = raw_verts[vid]
+        vertices_coords.append([float(xy[0]), float(xy[1]), 0.0])
+
+    edges_vertices = []
+    edges_assignment = []
+    for eid in sorted(raw_edges.keys(), key=lambda k: int(k) if isinstance(k, (int, str)) else k):
+        v1_id, v2_id, asgn = raw_edges[eid]
+        if v1_id in id_to_idx and v2_id in id_to_idx:
+            edges_vertices.append([id_to_idx[v1_id], id_to_idx[v2_id]])
+            edges_assignment.append(asgn)
+
+    faces_vertices = _triangulate_vertices(vertices_coords)
+
+    return {
+        "vertices_coords": vertices_coords,
+        "edges_vertices": edges_vertices,
+        "edges_assignment": edges_assignment,
+        "faces_vertices": faces_vertices,
+    }
+
+
+def _triangulate_vertices(vertices_coords: list) -> list:
+    """Delaunay triangulate the 2D vertex set for 3D mesh rendering."""
+    if len(vertices_coords) < 3:
+        return []
+    try:
+        from scipy.spatial import Delaunay
+        pts = np.array([[v[0], v[1]] for v in vertices_coords])
+        tri = Delaunay(pts)
+        return tri.simplices.tolist()
+    except Exception:
+        return [[0, 1, 2], [0, 2, 3]] if len(vertices_coords) >= 4 else []
 
 
 # ---------------------------------------------------------------------------
@@ -81,68 +188,62 @@ DEMO_SEQUENCES: dict[str, list[dict]] = {
 
 @app.get("/targets", include_in_schema=True, response_class=NumpyJSONResponse)
 def get_targets():
-    """Return available task names and metadata for the frontend."""
-    from server.tasks import get_task_by_name, available_task_names
-
+    """Return available target names and metadata from env/targets/*.fold."""
+    env = OrigamiEnvironment()
+    names = env.available_targets()
     result: dict[str, dict] = {}
-    for name in available_task_names():
-        t = get_task_by_name(name)
+    for name in names:
+        target = env._targets.get(name, {})
         result[name] = {
             "name": name,
-            "level": t.get("difficulty", 1),
-            "description": t.get("description", ""),
-            "n_creases": t.get("max_folds", 3),
-            "difficulty": t.get("difficulty", 1),
-            "material": t.get("material", "paper"),
+            "level": target.get("level", 1),
+            "description": target.get("description", ""),
+            "n_creases": len([a for a in target.get("edges_assignment", []) if a in ("M", "V")]),
+            "difficulty": target.get("level", 1),
+            "material": "paper",
         }
     return NumpyJSONResponse(result)
 
 
 @app.get("/episode/demo", include_in_schema=True, response_class=NumpyJSONResponse)
-def demo_episode(target: str = "half_fold"):
-    """Return a pre-solved demo episode for the given task."""
-    from server.origami_environment import OrigamiEnvironment
-    from server.models import OrigamiAction as NewOrigamiAction
-    from server.tasks import get_task_by_name
+def demo_episode(target: str = "half_horizontal"):
+    """Return a pre-solved demo episode for the given .fold target."""
+    env = OrigamiEnvironment(mode="step")
+    targets = env.available_targets()
+    if target not in targets:
+        target = targets[0] if targets else "half_horizontal"
 
-    # Fall back to half_fold if target not found
-    folds = DEMO_SEQUENCES.get(target, DEMO_SEQUENCES["half_fold"])
-
-    env = OrigamiEnvironment()
-    obs = env.reset(task_name=target)
+    t = env._targets.get(target, {})
+    folds = _target_to_folds(t)
 
+    obs_dict = env.reset(target_name=target)
     steps: list[dict] = []
 
     for i, fold_dict in enumerate(folds):
-        if fold_dict.get("type") == "stop":
-            break
-
-        action = NewOrigamiAction(
-            fold_type=fold_dict["type"],
-            fold_line=fold_dict["line"],
-            fold_angle=float(fold_dict.get("angle", 180.0)),
-        )
-
-        obs = env.step(action)
+        obs_dict, reward, done, info = env.step(fold_dict)
+        graph = env.paper.graph
+        all_edges = {eid: (v1, v2, a) for eid, (v1, v2, a) in graph.edges.items()}
+        fold_state = _graph_state_to_fold({
+            "vertices": dict(graph.vertices),
+            "edges": all_edges,
+        })
 
         steps.append({
             "step": i + 1,
             "fold": fold_dict,
-            "paper_state": obs.paper_state,
-            "metrics": obs.metrics,
-            "done": obs.done,
+            "paper_state": fold_state,
+            "metrics": reward if isinstance(reward, dict) else {"total": reward},
+            "done": done,
         })
-
-        if obs.done:
+        if done:
             break
 
-    task_def = get_task_by_name(target) if target else {}
-
     return NumpyJSONResponse({
         "task_name": target,
-        "task": task_def,
+        "task": {"name": target, "level": t.get("level", 1), "description": t.get("description", "")},
+        "target_crease": t,
         "steps": steps,
-        "final_metrics": obs.metrics if steps else {},
+        "final_metrics": steps[-1]["metrics"] if steps else {},
     })