Spaces:

anugrah55
/

cernenv-trainer

Sleeping

File size: 6,118 Bytes

"""FastAPI app exposing ``CERNCollisionEnvironment`` over the OpenEnv HTTP API.



We delegate the standard OpenEnv routes (``/reset``, ``/step``, ``/state``,

``/schema``, ``/health``, ``/mcp``) to ``create_fastapi_app`` and add a

human-friendly landing page at ``/`` so the Hugging Face Space preview

shows the project description instead of a 404.

"""

from __future__ import annotations

import os
from typing import Optional

from fastapi.responses import HTMLResponse
from openenv.core.env_server import create_fastapi_app

from models import CollisionObservation, ExperimentAction
from server.environment import CERNCollisionEnvironment


_LANDING_PAGE = """\

<!doctype html>

<html lang=en>

<head>

  <meta charset=utf-8>

  <title>CERNenv — LHC Discovery RL Environment</title>

  <meta name=viewport content="width=device-width,initial-scale=1">

  <style>

    body { font-family: ui-sans-serif, system-ui, -apple-system, sans-serif;

           margin: 2rem auto; max-width: 780px; color:#111; padding: 0 1rem; line-height:1.5 }

    h1 { margin-bottom: .25rem }

    h2 { margin-top: 2rem; border-bottom: 1px solid #eee; padding-bottom: .25rem }

    code { background:#f4f4f4; padding:.05rem .35rem; border-radius:4px; font-size:.95em }

    pre  { background:#0e1116; color:#e6edf3; padding:1rem; border-radius:6px; overflow-x:auto }

    pre code { background:transparent; color:inherit; padding:0 }

    .pill { display:inline-block; padding:.1rem .55rem; border-radius:999px;

            background:#e8f0ff; color:#1d4ed8; font-size:.85em; margin-right:.25rem }

    .muted { color:#666 }

    a { color:#1d4ed8 }

    table { border-collapse:collapse; margin: .5rem 0 }

    th,td { text-align:left; padding:.25rem .9rem .25rem 0; vertical-align: top }

    th { color:#444; font-weight:600 }

  </style>

</head>

<body>

  <h1>⚛️ CERNenv</h1>

  <p class=muted>An LHC (Large Hadron Collider) particle-discovery RL environment for autonomous physicist agents — built for the Meta OpenEnv Hackathon.</p>



  <p>

    <span class=pill>OpenEnv</span>

    <span class=pill>POMDP</span>

    <span class=pill>16 action types</span>

    <span class=pill>3 difficulty levels</span>

    <span class=pill>HTTP + WebSocket</span>

  </p>



  <h2>What this is</h2>

  <p>

    A Large Language Model (LLM) agent plays a high-energy physicist running an

    analysis at the LHC. Each step it picks one structured action — configure

    the beam, allocate luminosity, set a trigger, collect collisions, fit a

    resonance, estimate significance, submit a discovery claim, and so on —

    and receives a noisy detector-style observation. The latent particle

    (mass, decay channel, branching ratios, width) is hidden ground truth.

    Reward decomposes into per-step shaping + a dominant terminal calibration

    against the truth particle.

  </p>



  <h2>API</h2>

  <table>

    <tr><th><code>GET  /health</code></th><td>liveness probe</td></tr>

    <tr><th><code>GET  /schema</code></th><td>JSON schemas for actions, observations, state</td></tr>

    <tr><th><code>POST /reset</code></th><td>start a new episode (e.g. <code>{"seed": 7, "scenario": "easy_diphoton_160"}</code>)</td></tr>

    <tr><th><code>POST /step</code></th><td>execute one action (<code>{"action": {"action_type": ..., "parameters": {...}, "justification": "..."}}</code>)</td></tr>

    <tr><th><code>GET  /state</code></th><td>current public state snapshot</td></tr>

    <tr><th><code>GET  /docs</code></th><td>interactive Swagger UI</td></tr>

    <tr><th><code>GET  /metadata</code></th><td>environment metadata</td></tr>

  </table>



  <h2>Quickstart</h2>

  <pre><code># reset

curl -X POST $URL/reset \\

  -H 'Content-Type: application/json' \\

  -d '{"seed": 7, "scenario": "easy_diphoton_160"}'



# step

curl -X POST $URL/step \\

  -H 'Content-Type: application/json' \\

  -d '{"action": {"action_type": "configure_beam",

                 "parameters": {"sqrt_s_tev": 13.0},

                 "justification": "set 13 TeV"}}'</code></pre>



  <h2>Companion Spaces</h2>

  <ul>

    <li>📓 Trainer (Unsloth + LoRA + GRPO on A100): <a href="https://huggingface.co/spaces/anugrah55/cernenv-trainer">anugrah55/cernenv-trainer</a></li>

    <li>🎯 Trained adapters (LoRA): <a href="https://huggingface.co/anugrah55/cernenv-grpo-qwen2.5-3b">anugrah55/cernenv-grpo-qwen2.5-3b</a></li>

  </ul>



  <p class=muted style="margin-top:2rem">CERNenv · OpenEnv-compatible · BSD-3-Clause</p>

</body>

</html>

"""


def make_env_factory(

    max_steps: int,

    default_difficulty: Optional[str],

):
    def factory() -> CERNCollisionEnvironment:
        return CERNCollisionEnvironment(
            max_steps=max_steps,
            default_difficulty=default_difficulty,
        )

    return factory


def build_app(

    *,

    max_steps: int = 40,

    default_difficulty: Optional[str] = None,

):
    """Construct the FastAPI app with a per-session environment factory.



    The OpenEnv-provided routes (`/reset`, `/step`, `/state`, `/schema`,

    `/health`, `/mcp`) come from ``create_fastapi_app``. We then mount a

    friendly landing page at ``/`` so the Space preview is informative.

    """
    factory = make_env_factory(max_steps=max_steps, default_difficulty=default_difficulty)
    fa_app = create_fastapi_app(factory, ExperimentAction, CollisionObservation)

    @fa_app.get("/", response_class=HTMLResponse, include_in_schema=False)
    def landing() -> HTMLResponse:  # pragma: no cover - trivial
        return HTMLResponse(_LANDING_PAGE)

    return fa_app


app = build_app(
    max_steps=int(os.getenv("CERNENV_MAX_STEPS", "40")),
    default_difficulty=os.getenv("CERNENV_DEFAULT_DIFFICULTY") or None,
)


def main() -> None:  # pragma: no cover - CLI entrypoint
    import uvicorn

    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", "8000"))
    uvicorn.run("server.app:app", host=host, port=port, log_level="info")


if __name__ == "__main__":  # pragma: no cover
    main()