File size: 4,993 Bytes
99d2ff3 e2284e6 99d2ff3 1ec7127 99d2ff3 e2284e6 99d2ff3 1ec7127 99d2ff3 e2284e6 f776c3a 99d2ff3 e2284e6 99d2ff3 e2284e6 99d2ff3 1ec7127 99d2ff3 9de2ed4 99d2ff3 9de2ed4 99d2ff3 6423c48 | 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 | from fastapi import FastAPI, HTTPException
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel
from typing import Dict, Any
import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
from models import Action, Observation, State
from server.llm_env import LLMEnv
app = FastAPI(title="LLM Control OpenEnv")
# Global in-memory store for environments and session state
envs: Dict[str, LLMEnv] = {}
completed_episodes: Dict[str, Dict[str, Any]] = {}
# Default global environment initialized with a reset state
default_env = LLMEnv()
default_env.reset()
class ResetRequest(BaseModel):
task: str = "easy"
class StepRequest(BaseModel):
action: Action
episode_id: str | None = None
class GraderRequest(BaseModel):
episode_id: str
# Default global environment to satisfy simple paths
default_env = LLMEnv()
@app.get("/", response_class=HTMLResponse)
async def serve_gui():
path = os.path.join(os.path.dirname(__file__), "index.html")
try:
with open(path, "r") as f:
return f.read()
except FileNotFoundError:
return "GUI index.html not found. Check the root directory."
@app.post("/reset")
async def reset(req: ResetRequest = ResetRequest()):
if req.task not in ["easy", "medium", "hard"]:
raise HTTPException(status_code=400, detail="Invalid task")
env = LLMEnv(task=req.task)
obs = env.reset()
state = env.state
envs[state.episode_id] = env
# Also set default env to the latest reset for easy single-agent testing
global default_env
default_env = env
return {
"observation": obs.model_dump(),
"state": state.model_dump()
}
@app.post("/step")
async def step(req: StepRequest):
# Retrieve env
env = default_env
if req.episode_id and req.episode_id in envs:
env = envs[req.episode_id]
obs, reward, done, info = env.step(req.action)
if done:
# Save cumulative reward for grading
completed_episodes[env.state.episode_id] = {
"reward": env.state.cumulative_reward,
"bounds": env._reward_bounds()
}
return {
"observation": obs.model_dump(),
"reward": reward,
"done": done,
"info": info
}
@app.get("/state", response_model=State)
async def get_state(episode_id: str | None = None):
env = default_env
if episode_id and episode_id in envs:
env = envs[episode_id]
return env.state
@app.post("/baseline")
async def run_baseline():
import subprocess
try:
# baseline.py should be in the directory above server
baseline_path = os.path.join(os.path.dirname(__file__), "baseline.py")
result = subprocess.run([sys.executable, baseline_path], capture_output=True, text=True, check=True)
# Parse the output to return the dict
# We expect JSON or eval-able output from baseline, or simply look at the final prints
# But this implies we should structure baseline.py to just run the tasks
# Or we can just run the baseline logic directly here if we want API.
# For safety, let's just execute it and return the raw output or parse a standard format.
# We'll just run our logic from baseline script here directly if the subprocess is too complex,
# but the prompt says POST /baseline runs baseline.py, so we will return stdout.
# Actually, let's try to extract JSON from the stdout.
import json
out = result.stdout.strip().splitlines()[-1]
# assume last line is valid JSON dict
scores = json.loads(out)
return scores
except subprocess.CalledProcessError as e:
raise HTTPException(status_code=500, detail=f"Baseline failed: {e.stderr}")
except Exception as e:
raise HTTPException(status_code=500, detail=f"Baseline error: {str(e)}")
@app.post("/grader")
async def grader(req: GraderRequest):
if req.episode_id not in completed_episodes:
# Check active envs
if req.episode_id in envs:
env = envs[req.episode_id]
r = env.state.cumulative_reward
b_min, b_max = env._reward_bounds()
score = (norm * 0.998) + 0.001
return {"score": score}
raise HTTPException(status_code=404, detail="Episode not found or not finished")
data = completed_episodes[req.episode_id]
r = data["reward"]
b_min, b_max = data["bounds"]
norm = (r - b_min) / (b_max - b_min)
# Smoothly scale to strictly (0, 1) to avoid manual edge limits
score = (norm * 0.998) + 0.001
return {"score": score}
@app.get("/tasks")
async def get_tasks():
return {
"tasks": ["easy", "medium", "hard"],
"action_schema": Action.model_json_schema()
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=7860)
|