"""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()) @app.middleware("http") 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 @app.get("/api/tasks", include_in_schema=False) def tasks(): return [_summary(t, plans_by_task) for t in pack.tasks] @app.get("/api/runs", include_in_schema=False) 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 @app.get("/api/tasks/{task_id}", include_in_schema=False) def task(task_id: str): return pack.public(_task(task_id)) @app.get("/api/tasks/{task_id}/reference", include_in_schema=False) def reference(task_id: str): return _task(task_id).reference @app.get("/api/tasks/{task_id}/timeline", include_in_schema=False) 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") @lru_cache(maxsize=256) 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} @app.get("/api/tasks/{task_id}/plans", include_in_schema=False) 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 @app.get("/api/tasks/{task_id}/replay", include_in_schema=False) 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) @app.get("/api/episodes/{episode_id}/replay", include_in_schema=False) 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}"} @app.get("/api/episodes", include_in_schema=False) def episodes(): return [{k: v for k, v in ep.summary().items() if k != "steps"} | {"steps": len(ep.steps)} for ep in STORE.recent()] @app.get("/api/episodes/{episode_id}/timeline", include_in_schema=False) 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") @app.get("/", include_in_schema=False) 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) @app.get("/healthz", include_in_schema=False) 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()