File size: 6,839 Bytes
19abe39
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
"""
FastAPI server for the origami RL environment.
Serves episode data to the React frontend.

Usage: uvicorn server:app --reload --port 8000
"""

try:
    from fastapi import FastAPI
    from fastapi.middleware.cors import CORSMiddleware
    from pydantic import BaseModel
except ImportError:
    print("Run: pip install fastapi uvicorn pydantic")
    raise

from typing import Optional


app = FastAPI(title="OrigamiRL API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # localhost:3000 for React dev
    allow_methods=["*"],
    allow_headers=["*"],
)


class FoldAction(BaseModel):
    from_point: list[float]  # [x, y]
    to_point: list[float]    # [x, y]
    assignment: str          # 'M' or 'V'
    instruction: str = ""


class EpisodeStep(BaseModel):
    step: int
    fold: Optional[FoldAction]
    paper_state: dict        # FOLD JSON of current crease graph
    anchor_points: list[list[float]]
    reward: dict
    done: bool
    info: dict
    prompt: str              # LLM prompt at this step


class EpisodeResult(BaseModel):
    target_name: str
    target: dict             # FOLD JSON of target
    steps: list[EpisodeStep]
    final_reward: dict


@app.get("/")
def health_check():
    """Health check — returns status and available target names."""
    from env.environment import OrigamiEnvironment
    env = OrigamiEnvironment()
    return {"status": "ok", "targets": env.available_targets()}


@app.get("/targets")
def get_targets():
    """Return list of available target names and their metadata."""
    from env.environment import OrigamiEnvironment
    env = OrigamiEnvironment()
    targets = {}
    for name in env.available_targets():
        t = env._targets[name]
        targets[name] = {
            "name": name,
            "level": t.get("level", 1),
            "description": t.get("description", ""),
            "n_creases": sum(1 for a in t["edges_assignment"] if a in ("M", "V")),
        }
    return targets


@app.get("/episode/run")
def run_episode(target: str = "half_horizontal", completion: str = ""):
    """
    Run a code-as-policy episode with a provided completion string.

    If completion is empty, returns the prompt so the caller knows what to send.
    Returns full episode result with all steps.
    """
    from env.environment import OrigamiEnvironment
    from env.prompts import parse_fold_list, code_as_policy_prompt
    from env.rewards import compute_reward, target_crease_edges

    env = OrigamiEnvironment(mode="step")
    obs = env.reset(target_name=target)

    if not completion:
        return {"prompt": obs["prompt"], "steps": [], "target": env.target}

    try:
        folds = parse_fold_list(completion)
    except ValueError as e:
        return {"error": str(e), "steps": []}

    steps = []
    for i, fold in enumerate(folds):
        result = env.paper.add_crease(fold["from"], fold["to"], fold["assignment"])
        reward = compute_reward(env.paper, result, env.target)

        paper_state = {
            "vertices": {str(k): list(v) for k, v in env.paper.graph.vertices.items()},
            "edges": [
                {
                    "id": k,
                    "v1": list(env.paper.graph.vertices[v[0]]),
                    "v2": list(env.paper.graph.vertices[v[1]]),
                    "assignment": v[2],
                }
                for k, v in env.paper.graph.edges.items()
            ],
            "anchor_points": [list(p) for p in env.paper.anchor_points()],
        }

        # Build per-step prompt reflecting current state
        from env.prompts import step_level_prompt
        step_prompt = step_level_prompt(
            target=env.target,
            paper_state=env.paper,
            step=i + 1,
            max_steps=env.max_steps,
            last_reward=reward,
        )

        steps.append({
            "step": i + 1,
            "fold": {
                "from_point": fold["from"],
                "to_point": fold["to"],
                "assignment": fold["assignment"],
                "instruction": fold.get("instruction", ""),
            },
            "paper_state": paper_state,
            "anchor_points": [list(p) for p in env.paper.anchor_points()],
            "reward": reward,
            "done": reward.get("completion", 0) > 0,
            "info": env._info(),
            "prompt": step_prompt,
        })

        if reward.get("completion", 0) > 0:
            break

    return {
        "target_name": target,
        "target": env.target,
        "steps": steps,
        "final_reward": steps[-1]["reward"] if steps else {},
    }


@app.get("/episode/demo")
def demo_episode(target: str = "half_horizontal"):
    """Return a pre-solved demo episode for each target."""
    DEMO_COMPLETIONS = {
        "half_horizontal": '<folds>[{"instruction": "Valley fold along horizontal center line", "from": [0, 0.5], "to": [1, 0.5], "assignment": "V"}]</folds>',
        "half_vertical": '<folds>[{"instruction": "Mountain fold along vertical center line", "from": [0.5, 0], "to": [0.5, 1], "assignment": "M"}]</folds>',
        "diagonal_main": '<folds>[{"instruction": "Valley fold along main diagonal", "from": [0, 0], "to": [1, 1], "assignment": "V"}]</folds>',
        "diagonal_anti": '<folds>[{"instruction": "Mountain fold along anti-diagonal", "from": [1, 0], "to": [0, 1], "assignment": "M"}]</folds>',
        "thirds_h": '<folds>[{"instruction": "Valley fold at one-third height", "from": [0, 0.333], "to": [1, 0.333], "assignment": "V"}, {"instruction": "Valley fold at two-thirds height", "from": [0, 0.667], "to": [1, 0.667], "assignment": "V"}]</folds>',
        "thirds_v": '<folds>[{"instruction": "Mountain fold at one-third width", "from": [0.333, 0], "to": [0.333, 1], "assignment": "M"}, {"instruction": "Mountain fold at two-thirds width", "from": [0.667, 0], "to": [0.667, 1], "assignment": "M"}]</folds>',
        "accordion_3h": '<folds>[{"instruction": "Valley fold at quarter height", "from": [0, 0.25], "to": [1, 0.25], "assignment": "V"}, {"instruction": "Mountain fold at half height", "from": [0, 0.5], "to": [1, 0.5], "assignment": "M"}, {"instruction": "Valley fold at three-quarter height", "from": [0, 0.75], "to": [1, 0.75], "assignment": "V"}]</folds>',
        "accordion_4h": '<folds>[{"instruction": "Valley fold at 0.2", "from": [0, 0.2], "to": [1, 0.2], "assignment": "V"}, {"instruction": "Mountain fold at 0.4", "from": [0, 0.4], "to": [1, 0.4], "assignment": "M"}, {"instruction": "Valley fold at 0.6", "from": [0, 0.6], "to": [1, 0.6], "assignment": "V"}, {"instruction": "Mountain fold at 0.8", "from": [0, 0.8], "to": [1, 0.8], "assignment": "M"}]</folds>',
    }
    completion = DEMO_COMPLETIONS.get(target, DEMO_COMPLETIONS["half_horizontal"])
    return run_episode(target=target, completion=completion)