Spaces:
Sleeping
Sleeping
| """The server: OpenEnv's app plus the map replay, a train-graph viewer and their read-only JSON on the same port. | |
| OpenEnv provides `/reset`, `/step`, `/state`, `/schema`, `/ws`, `/mcp`, `/health`, the Task API under `/rail_ops/...` | |
| and its web UI at `/web`, where an "Explore" tab (`gradio_ui.py`) frames the pages below next to OpenEnv's Playground. | |
| We add: | |
| /twin/ the map replay: trains moving on the line's OpenStreetMap track, any plan | |
| /viewer/ a train graph (time against distance) of any plan, live episodes too | |
| /explorer/ the eval explorer: the board and every model's plan on every eval task | |
| /api/tasks task summaries, with the lazy baselines' rewards | |
| /api/runs the committed model runs: each episode's reward and outcome, and a board | |
| /api/tasks/{id} a public task (what the agent sees, structured) | |
| /api/tasks/{id}/reference optimum, floor, baselines and the real outcome - for people; no tool shows them | |
| /api/tasks/{id}/timeline?plan=... paths for the graph: plan = optimal | naive | greedy | greedy_fast | |
| /api/episodes live episodes on this server (newest first) | |
| /api/episodes/{episode_id}/timeline the latest plan an episode checked or submitted, as paths | |
| /api/tasks/{id}/plans the plans the replay can show: optimum, baselines, models' submitted plans | |
| (from results/rollouts/*/index.json, or RAIL_RUNS_DIR), live episodes | |
| /api/tasks/{id}/replay?plan=... everything the replay animates (rail_core.replay) | |
| /api/episodes/{episode_id}/replay the same for the latest plan of a live episode | |
| uvicorn rail_openenv.server:app --port 8000 | |
| `/` is OpenEnv's: it redirects to the web UI (`/web/`), or to `/explorer/` when ENABLE_WEB_INTERFACE is false. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import random | |
| import statistics | |
| from functools import lru_cache | |
| from pathlib import Path | |
| from fastapi import HTTPException | |
| from fastapi.responses import RedirectResponse | |
| from fastapi.staticfiles import StaticFiles | |
| from openenv.core.env_server.http_server import create_app | |
| from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation | |
| from rail_core import TaskPack, evaluate, hhmm, load_pack, parse_plan | |
| from rail_core.baselines import POLICIES | |
| from rail_core.check import network | |
| from rail_core.replay import replay | |
| try: | |
| from .environment import STORE, RailOpsEnvironment, RailState | |
| from .gradio_ui import build_ui | |
| except ImportError: | |
| from environment import STORE, RailOpsEnvironment, RailState | |
| from gradio_ui import build_ui | |
| ENV_NAME = "rail_ops" | |
| WEB = Path(__file__).resolve().parent / "web" | |
| _UP = Path(__file__).resolve().parents | |
| RUNS_DIR = Path(os.environ.get("RAIL_RUNS_DIR") or (_UP[4] if len(_UP) > 4 else _UP[-1]) / "results" / "rollouts") | |
| BASELINES = {"greedy": "greedy dispatcher", "greedy_fast": "greedy dispatcher (turns early)", | |
| "naive": "change nothing"} | |
| TIERS = ("real", "blockage", "loop", "compound") | |
| PROVIDERS = {"together": "Together", "baseten": "Baseten", "novita": "Novita", "deepinfra": "DeepInfra", | |
| "fireworks": "Fireworks"} | |
| def model_name(spec: str) -> str: | |
| """A person's name for a run's model spec: "hf:Qwen/Qwen3.8-2.4T-A95B:together@default" -> "Qwen3.8-2.4T-A95B | |
| (Together)", "claude-sonnet-5-5@high" -> "Claude Sonnet 5.5 (effort high)".""" | |
| model, _, effort = spec.partition("@") | |
| if model.startswith("hf:"): | |
| repo, _, provider = model[3:].partition(":") | |
| name = repo.split("/")[-1] | |
| return f"{name} ({PROVIDERS.get(provider, provider)})" if provider else name | |
| if model.startswith("claude-"): | |
| parts = model.split("-")[1:] | |
| family = " ".join(p.capitalize() for p in parts if not p.isdigit()) | |
| version = ".".join(p for p in parts if p.isdigit()) | |
| name = f"Claude {family} {version}".strip() | |
| return f"{name} (effort {effort})" if effort and effort != "default" else name | |
| return model | |
| def load_runs(runs_dir: Path = RUNS_DIR) -> list[dict]: | |
| """The committed rollout runs (results/rollouts/*/index.json): [{run, episodes}].""" | |
| out = [] | |
| for idx in sorted(runs_dir.glob("*/index.json")) if runs_dir.is_dir() else []: | |
| try: | |
| body = json.loads(idx.read_text()) | |
| except (OSError, ValueError): | |
| continue | |
| out.append({"run": body["run"], "episodes": body.get("episodes", [])}) | |
| return out | |
| def run_plans(runs_dir: Path = RUNS_DIR) -> dict[str, list[dict]]: | |
| """Submitted plans in the committed rollout runs, by task: [{id, label, model, run, reward, plan}].""" | |
| out: dict[str, list[dict]] = {} | |
| for body in load_runs(runs_dir): | |
| for e in body["episodes"]: | |
| if e.get("plan") is None: | |
| continue | |
| out.setdefault(e["task_id"], []).append( | |
| {"id": f"run:{body['run']}:{e['model']}", "label": model_name(e["model"]), | |
| "model": e["model"], "run": body["run"], "reward": e.get("reward"), "plan": e["plan"]}) | |
| for entries in out.values(): # two runs of one model on a task: name the run too | |
| for e in entries: | |
| if sum(x["model"] == e["model"] for x in entries) > 1: | |
| e["label"] = f"{e['label']} ({e['run']})" | |
| return out | |
| def ci(xs, n=2000, seed=0): | |
| """95 % bootstrap interval of the mean, the same as eval/report.py's (so the numbers match the README).""" | |
| xs = sorted(xs) | |
| rng = random.Random(seed) | |
| means = sorted(statistics.mean(rng.choices(xs, k=len(xs))) for _ in range(n)) | |
| return means[int(0.025 * n)], means[int(0.975 * n)] | |
| def runs_summary(runs: list[dict]) -> dict: | |
| """What the eval explorer shows: one row per (run, model) with the board's columns, and every episode's outcome.""" | |
| rows, episodes = [], [] | |
| for body in runs: | |
| by: dict[str, list[dict]] = {} | |
| for e in body["episodes"]: | |
| by.setdefault(e["model"], []).append(e) | |
| for model, es in by.items(): | |
| r = [e["reward"] for e in es] | |
| lo, hi = ci(r) | |
| real = [e for e in es if e["difficulty"] == "real" and e.get("real_cost") is not None] | |
| rows.append({ | |
| "run": body["run"], "model": model, "name": model_name(model), "split": es[0]["split"], | |
| "n": len(es), "mean": statistics.mean(r), "ci": [lo, hi], | |
| "tiers": {t: statistics.mean(x["reward"] for x in es if x["difficulty"] == t) | |
| for t in TIERS if any(x["difficulty"] == t for x in es)}, | |
| "feasible": sum(bool(e["feasible"]) for e in es), | |
| "optimal": sum(1 for e in es if e["feasible"] and e["cost"] == e["optimal_cost"]), | |
| "no_submit": sum(not e["submitted"] for e in es), | |
| "beat_real": sum(1 for e in real if e["feasible"] and e["cost"] < e["real_cost"]), "real": len(real), | |
| "median_turns": statistics.median(e["agent_turns"] for e in es)}) | |
| for e in es: | |
| episodes.append({k: e.get(k) for k in ("model", "task_id", "reward", "feasible", "submitted", "cost", | |
| "optimal_cost", "end_reason", "agent_turns", "checks", | |
| "violations")} | |
| | {"run": body["run"], "plan_id": f"run:{body['run']}:{model}" | |
| if e.get("plan") is not None else None}) | |
| rows.sort(key=lambda x: (x["split"], -x["mean"])) | |
| return {"models": rows, "episodes": episodes} | |
| def _summary(t, plans: dict | None = None) -> dict: | |
| ref = t.reference | |
| return {"task_id": t.task_id, "split": t.split, "day": t.day, "start": hhmm(t.start), "end": hhmm(t.end), | |
| "difficulty": t.difficulty, "runs": len(t.runs), "closures": [c.kind for c in t.closures], | |
| "causes": sorted({c.cause for c in t.closures if c.cause}), | |
| "model_plans": len((plans or {}).get(t.task_id, [])), "optimal_cost": ref.get("optimal_cost"), | |
| "rewards": {"greedy": ref.get("greedy_reward"), "greedy_fast": ref.get("greedy_fast_reward"), | |
| "naive": ref.get("naive_reward"), "real": (ref.get("real") or {}).get("reward_if_legal")}} | |
| def timeline(task, plan) -> dict: | |
| """Every served run as a path of (km, minute) points, for a time-distance graph.""" | |
| net = network(task) | |
| res = evaluate(task, plan) | |
| paths = [] | |
| for x in res.runs: | |
| if x.cancelled: | |
| paths.append({"train": x.train, "set": x.unit, "cancelled": True, "points": []}) | |
| continue | |
| r = task.run(x.train) | |
| pts = [] | |
| for i in sorted(set(x.arr) | set(x.dep)): | |
| km = net.by_code[r.route[i]].km | |
| if i in x.arr: | |
| pts.append([km, x.arr[i]]) | |
| if i in x.dep and x.dep[i] != x.arr.get(i): | |
| pts.append([km, x.dep[i]]) | |
| paths.append({"train": x.train, "set": x.unit, "kind": r.kind, "points": pts, "delay": x.delay, | |
| "unserved": x.unserved, "problems": x.problems}) | |
| return {"task_id": task.task_id, "start": task.start, "end": task.end, | |
| "points": [{"code": p.code, "km": p.km, "kind": p.kind, "tracks": p.tracks} for p in net.points], | |
| "sections": [{"a": s.a, "b": s.b, "single": s.single} for s in net.sections], | |
| "closures": [{"kind": c.kind, "where": c.where, "start": c.start, "end": c.end, "cause": c.cause} | |
| for c in task.closures], | |
| "paths": paths, "feasible": res.feasible, "cost": res.cost, "violations": res.violations[:50], | |
| "optimal_cost": task.reference.get("optimal_cost")} | |
| PAGES = ("/twin", "/explorer", "/viewer") | |
| def mount_api(app, pack: TaskPack) -> None: | |
| plans_by_task = run_plans() | |
| runs = runs_summary(load_runs()) | |
| async def revalidate_pages(request, call_next): | |
| """The pages and their scripts change with every deploy. Without a Cache-Control header browsers guess how | |
| long to reuse a stored copy (often hours), so a new version did not show; `no-cache` makes them ask every time, | |
| and an unchanged file costs a 304.""" | |
| response = await call_next(request) | |
| if request.url.path.startswith(PAGES): | |
| response.headers.setdefault("Cache-Control", "no-cache") | |
| return response | |
| def tasks(): | |
| return [_summary(t, plans_by_task) for t in pack.tasks] | |
| def model_runs(): | |
| return runs | |
| def _task(task_id): | |
| try: | |
| return pack.get(task_id) | |
| except KeyError: | |
| raise HTTPException(404, f"no task {task_id}") from None | |
| def task(task_id: str): | |
| return pack.public(_task(task_id)) | |
| def reference(task_id: str): | |
| return _task(task_id).reference | |
| def task_timeline(task_id: str, plan: str = "optimal"): | |
| t = _task(task_id) | |
| if plan == "optimal": | |
| p, _ = parse_plan(t, t.reference["optimal_plan"]) | |
| elif plan in POLICIES: | |
| p = POLICIES[plan](t) | |
| else: | |
| raise HTTPException(400, "plan must be optimal, naive, greedy or greedy_fast") | |
| return timeline(t, p) | |
| def _plan(t, plan_id: str): | |
| """(plan, parse problems, label) for a plan id.""" | |
| if plan_id == "optimal": | |
| p, probs = parse_plan(t, t.reference["optimal_plan"]) | |
| return p, probs, "CP-SAT optimum" | |
| if plan_id in POLICIES: | |
| return POLICIES[plan_id](t), [], BASELINES.get(plan_id, plan_id) | |
| for e in plans_by_task.get(t.task_id, []): | |
| if e["id"] == plan_id: | |
| p, probs = parse_plan(t, e["plan"]) | |
| return p, probs, e["label"] | |
| raise HTTPException(400, f"unknown plan {plan_id!r}: see /api/tasks/{t.task_id}/plans") | |
| def _replay(task_id: str, plan_id: str) -> dict: | |
| t = _task(task_id) | |
| p, probs, label = _plan(t, plan_id) | |
| return replay(t, p, label, probs) | {"plan_id": plan_id} | |
| def task_plans(task_id: str): | |
| t = _task(task_id) | |
| ref = t.reference | |
| out = [{"id": "optimal", "label": "CP-SAT optimum", "reward": 1.0, "cost": ref.get("optimal_cost")}] | |
| out += [{"id": k, "label": v, "reward": ref.get(f"{k}_reward"), | |
| "cost": ref.get(f"{k}_cost") if ref.get(f"{k}_feasible") else None} for k, v in BASELINES.items()] | |
| out += [{k: e[k] for k in ("id", "label", "reward")} for e in plans_by_task.get(task_id, [])] | |
| out += [{"id": f"episode:{ep.episode_id}", "label": f"live episode {ep.episode_id}", "reward": None} | |
| for ep in STORE.recent() if ep.task.task_id == task_id] | |
| return out | |
| def task_replay(task_id: str, plan: str = "optimal"): | |
| if plan.startswith("episode:"): | |
| return episode_replay(plan.split(":", 1)[1]) | |
| return _replay(task_id, plan) | |
| def episode_replay(episode_id: str): | |
| ep = STORE.get(episode_id) | |
| if ep is None: | |
| raise HTTPException(404, f"no episode {episode_id}") | |
| last = next((s for s in reversed(ep.steps) if s.get("plan") is not None), None) | |
| p, probs = parse_plan(ep.task, last["plan"]) if last else ({}, []) | |
| return replay(ep.task, p, f"live episode {episode_id}", probs) | {"plan_id": f"episode:{episode_id}"} | |
| def episodes(): | |
| return [{k: v for k, v in ep.summary().items() if k != "steps"} | {"steps": len(ep.steps)} | |
| for ep in STORE.recent()] | |
| def episode_timeline(episode_id: str): | |
| ep = STORE.get(episode_id) | |
| if ep is None: | |
| raise HTTPException(404, f"no episode {episode_id}") | |
| last = next((s for s in reversed(ep.steps) if s.get("plan") is not None), None) | |
| p, _ = parse_plan(ep.task, last["plan"]) if last else ({}, []) | |
| return timeline(ep.task, p) | {"episode": {k: v for k, v in ep.summary().items() if k != "steps"}} | |
| if (WEB / "twin").is_dir(): | |
| app.mount("/twin", StaticFiles(directory=WEB / "twin", html=True), name="twin") | |
| if (WEB / "explorer").is_dir(): | |
| app.mount("/explorer", StaticFiles(directory=WEB / "explorer", html=True), name="explorer") | |
| if WEB.is_dir(): | |
| app.mount("/viewer", StaticFiles(directory=WEB, html=True), name="viewer") | |
| def root(): | |
| return RedirectResponse("/explorer/" if (WEB / "explorer" / "index.html").is_file() else "/twin/") | |
| def create_server(pack: TaskPack | None = None): | |
| pack = pack or load_pack() | |
| def factory(): | |
| return RailOpsEnvironment(pack) | |
| os.environ.setdefault("ENABLE_WEB_INTERFACE", "true") | |
| app = create_app(factory, CallToolAction, CallToolObservation, env_name=ENV_NAME, | |
| max_concurrent_envs=int(os.environ.get("MAX_CONCURRENT_ENVS", "64")), | |
| gradio_builder=build_ui, custom_tab_name="Explore", custom_tab_primary=True, | |
| show_default_tab=True, title_override="RailSimEnv v1 Environment", state_cls=RailState) | |
| mount_api(app, pack) | |
| def healthz(): | |
| return {"status": "ok", "tasks": {s: pack.count(s) for s in pack.splits()}} | |
| return app | |
| _app = None | |
| def __getattr__(name: str): | |
| global _app | |
| if name == "app": | |
| if _app is None: | |
| _app = create_server() | |
| return _app | |
| raise AttributeError(name) | |
| def main(host: str = "0.0.0.0", port: int | None = None) -> None: | |
| import uvicorn | |
| port = port or int(os.environ.get("PORT", "8000")) | |
| uvicorn.run(create_server(), host=host, port=port, ws_ping_interval=60, ws_ping_timeout=300) | |
| if __name__ == "__main__": | |
| main() | |