Add pi_env: Pi coding-agent environment (HF sandbox + transparent-proxy logprobs)

.github/workflows/docker-build.yml

@@ -112,6 +112,9 @@ jobs:
           - name: opencode-sandbox
             dockerfile: envs/opencode_env/sandbox/hf_image/Dockerfile
             context: envs/opencode_env
+          - name: pi-env
+            dockerfile: envs/pi_env/server/Dockerfile
+            context: envs/pi_env
           - name: openapp-env
             dockerfile: envs/openapp_env/server/Dockerfile
             context: envs/openapp_env

docs/source/_toctree.yml

@@ -129,6 +129,8 @@
     title: Agent World Model
   - local: environments/opencode
     title: OpenCode
+  - local: environments/pi
+    title: Pi
   - local: environments/sophistry_bench_sprint
     title: Sophistry Bench Sprint
   title: Environments

docs/source/environments.md

@@ -258,6 +258,13 @@ The OpenEnv community has built a catalog of ready-to-run environments that cove
         <a href="environments/opencode" class="!no-underline border dark:border-gray-700 px-3 py-1 rounded text-sm hover:shadow">📄 Docs</a>
       </div>
     </div>
+    <div class="border dark:border-gray-700 p-5 rounded-lg shadow">
+      <div class="font-bold mb-2">Pi</div>
+      <p class="text-sm"><code>pi_env</code> runs the Pi coding agent inside an isolated Hugging Face sandbox against any OpenAI-compatible LLM endpoint, optionally capturing per-token logprobs.</p>
+      <div class="flex gap-2 mt-3">
+        <a href="environments/pi" class="!no-underline border dark:border-gray-700 px-3 py-1 rounded text-sm hover:shadow">📄 Docs</a>
+      </div>
+    </div>
     <div class="border dark:border-gray-700 p-5 rounded-lg shadow">
       <div class="font-bold mb-2">Sophistry Bench Sprint</div>
       <p class="text-sm"><code>sophistry_bench_sprint_env</code> is a single-turn advocacy reward-hacking environment on QuALITY passages: the policy defends an assigned answer and the reward proxy peaks at 8 <code>&lt;claim&gt;</code> tags, with four weight-0 canaries that detect format hacking.</p>

docs/source/environments/pi.md

@@ -0,0 +1,210 @@
+<!-- openenv-source: pi_env -->
+# Pi Environment for OpenEnv
+
+`pi_env` runs the [Pi](https://github.com/badlogic/pi-mono) coding agent
+inside an isolated [Hugging Face sandbox](https://huggingface.co/docs/huggingface_hub/guides/sandboxes)
+against any OpenAI-compatible LLM endpoint, optionally capturing per-token
+logprobs for GRPO training.
+
+It mirrors [`opencode_env`](../opencode_env): same two-layer design (an
+in-process harness primitive + a deployable HTTP env), same transparent-proxy
+logprob capture, same uniform `(instruction, setup, verify)` Task shape. The
+agent is Pi instead of OpenCode, and the default sandbox backend is Hugging
+Face instead of E2B.
+
+The env is **task-agnostic** — every rollout is configured at call-time
+with a uniform Task shape:
+
+  - **`instruction`** — prompt for the agent
+  - **`setup`** — list of bash commands run *before* the agent (pip install,
+    git clone, file downloads — anything you need staged in the sandbox)
+  - **`verify`** — list of bash commands run *after* the agent (asserts,
+    pytest invocations, score-file writes)
+
+Reward = `passed_verify / total_verify` unless any `verify` command writes
+a float to `/root/logs/verifier/reward.txt` (override).
+
+## In-process primitive (no HTTP)
+
+For trainers that drive a sandbox directly without an HTTP boundary — this is
+what loop-owning GRPO training uses:
+
+```python
+import os
+from pi_env import PiConfig, PiSessionFactory, PiTask, HFSandboxBackend
+
+factory = PiSessionFactory(
+    config=PiConfig(
+        base_url="https://api.openai.com/v1",
+        api_key=os.environ["OPENAI_API_KEY"],
+        model="gpt-4o-mini",
+        sandbox_home="/root",                   # HF sandbox execs as root
+    ),
+    sandbox_backend=HFSandboxBackend(image="python:3.12"),
+    mode="transparent_proxy",                   # captures per-token logprobs
+)
+session = factory.create(task=PiTask(instruction="..."))
+session.wait_for_completion()
+turns = session.fetch_proxy_trace()             # per-turn (tokens, logprobs)
+session.close()
+```
+
+Pi is pointed at the endpoint via a `models.json` provider block written under
+`PI_CODING_AGENT_DIR` (`api: openai-completions`), then launched headless with
+`pi --print --no-session --mode json`. In `transparent_proxy` mode the
+in-sandbox proxy fronts `base_url`, injects `logprobs=true`, and writes each
+turn's `(messages, completion_token_ids, per_token_logps)` to
+`proxy_trace.jsonl`.
+
+### Sandbox backend
+
+`HFSandboxBackend` (from `opencode_env.sandbox`, shared with `opencode_env`)
+runs the agent in a Hugging Face sandbox. `image="python:3.12"` cold-installs
+Node 22 (bootstrapped) + the Pi CLI (`npm install -g @mariozechner/pi-coding-agent`)
++ the proxy's Python deps on every rollout. Any backend satisfying the
+`SandboxBackend` / `SandboxHandle` / `BgJob` protocols in
+`opencode_env.sandbox.base` can be plugged in the same way.
+
+> The sandbox backend and interception proxy live in `opencode_env` for now;
+> the plan is to consolidate both into `openenv.core` so `pi_env` and
+> `opencode_env` share them without a cross-package import.
+
+## Deployed env (HTTP)
+
+The deployed Space exposes:
+
+- **Web UI** at `/web` — pick endpoint, write task, hit Run, watch live phase
+  log + reward + logprobs.
+- **MCP tool API** at `/mcp` — programmatic `run_rollout` calls.
+- **OpenAPI docs** at `/docs`, **health** at `/health`.
+
+```python
+import os
+from pi_env import PiEnv
+
+with PiEnv(base_url="https://<user>-pi-env.hf.space") as env:
+    env.reset()
+    result = env.run_rollout(
+        endpoint="openai",                          # vllm | openai | hf_router
+        api_key=os.environ["OPENAI_API_KEY"],       # or set as a Space secret
+        instruction=(
+            "Create binary_search.py exposing def binary_search(arr, target) -> int "
+            "that returns the index of target in arr, or -1 if absent."
+        ),
+        setup=[],
+        verify=[
+            "test -f /root/workdir/binary_search.py",
+            "python -c \"import sys; sys.path.insert(0, '/root/workdir'); "
+            "import binary_search; "
+            "assert binary_search.binary_search([1,2,3], 2) == 1; print('OK')\"",
+        ],
+        task_id="binary_search_v1",
+    )
+    print("reward:", result.reward)
+    print("turns:", len(result.proxy_turns))
+```
+
+## The MCP Tool: `run_rollout`
+
+Single tool, two ways to specify the LLM endpoint:
+
+**Option A — endpoint shorthand (recommended)**: pass `endpoint="vllm"` (or
+`"openai"` / `"hf_router"`). The server resolves `base_url`, `api_key`, and
+`model` from env vars + catalog defaults. Any explicit field overrides.
+
+**Option B — fully explicit**: pass `base_url` + `api_key` + `model` directly.
+
+| Arg | Type | Default | Notes |
+|---|---|---|---|
+| `endpoint` | `str` | `""` | One of `"vllm"` / `"openai"` / `"hf_router"`. |
+| `base_url` / `api_key` / `model` | `str` | `""` | Override / supply explicitly. |
+| `instruction` | `str` | required | Prompt passed to `pi`. |
+| `setup` | `list[str]` | `[]` | Bash commands run **before** the agent. |
+| `verify` | `list[str]` | `[]` | Bash commands run **after** the agent. |
+| `task_id` | `str` | `""` | Echoed back in result. |
+| `mode` | `str` | `"transparent_proxy"` | Or `"black_box"` (no logprobs). |
+| `disable_thinking` | `bool \| None` | `None` (catalog default) | Inject `chat_template_kwargs.enable_thinking=false`. |
+| `max_tokens_cap` | `int` | `4096` | Per-turn `max_tokens` clamp. |
+| `top_logprobs` | `int` | `5` | HF Router cap is 5; OpenAI 0–20; vLLM unbounded. |
+| `agent_timeout_s` | `float` | `600.0` | Hard wall budget for one `pi` run. |
+| `image` | `str` | `""` | HF sandbox image; blank → `python:3.12` (cold-installs Node + Pi). |
+
+Returns `RolloutResult` JSON with: `reward`, `setup_results[]`,
+`verify_results[]`, `proxy_turns[]`, `files{}`, `agent_log_tail`,
+`proxy_log_tail`, `wall_s`, `agent_exit_code`, `sandbox_id`, `error`.
+
+## Two Operating Modes
+
+| Mode | What it does | Best for |
+|---|---|---|
+| **`transparent_proxy`** (default) | In-sandbox proxy at `localhost:7000` forwards Pi's LLM calls to `base_url`, injects `logprobs=true`, captures per-turn `(messages, completion_tokens, logprobs)` to `proxy_trace.jsonl`. | GRPO / RL training, observability, top-k distillation. |
+| **`black_box`** | No proxy. Pi talks straight to `base_url`. | Smoke tests, eval, SFT data collection. |
+
+## Building the Docker Image
+
+```bash
+cd envs/pi_env
+
+openenv validate           # check pyproject.toml + openenv.yaml + server/app.py + uv.lock
+openenv build -t pi-env    # builds the image (uses server/Dockerfile)
+
+# run locally with an HF token (Sandbox + Jobs access)
+docker run -p 8000:8000 -e HF_TOKEN=hf_... pi-env
+```
+
+Or build directly:
+
+```bash
+docker build -t pi-env -f envs/pi_env/server/Dockerfile envs/pi_env
+```
+
+## Environment Variables
+
+| Variable | Required | Purpose |
+|---|---|---|
+| `HF_TOKEN` | **yes** for any rollout | Hugging Face sandbox credentials. |
+| `MAX_CONCURRENT_ENVS` | no | Env-instance pool size. Default `4`. |
+| `ENABLE_WEB_INTERFACE` | no | Set `false` to disable the `/web` Gradio mount. Default `true`. |
+| `VLLM_URL` / `VLLM_API_KEY` / `VLLM_MODEL` | for `endpoint="vllm"` | OAI-compatible base URL (key defaults to `intercepted`). |
+| `OPENAI_API_KEY` / `OPENAI_BASE_URL` / `OPENAI_MODEL` | for `endpoint="openai"` | Standard OpenAI. |
+| `HF_ROUTER_API_KEY` / `HF_ROUTER_BASE_URL` / `HF_ROUTER_MODEL` | for `endpoint="hf_router"` | HF Router. |
+
+Pick `provider:` suffixes that actually return logprobs:
+**Together / Nscale / Scaleway / SambaNova / Cerebras**. Avoid Novita /
+Hyperbolic / Featherless (silent drop) and Groq (HTTP 400).
+
+## Project Structure
+
+```
+pi_env/
+├── README.md                       # this file
+├── openenv.yaml                    # OpenEnv space spec
+├── pyproject.toml                  # deps + ``server`` entrypoint
+├── __init__.py                     # re-exports primitive + client + models
+│
+├── client.py                       # PiEnv(MCPToolClient)
+├── models.py                       # RolloutResult / RolloutTurn / PiState
+│
+├── config.py                       # PiConfig (primitive)
+├── harness.py                      # PiSession / PiSessionFactory (CLI-only)
+├── pi_runtime.py                   # models.json builder + install/run cmds
+├── task.py                         # PiTask
+│
+└── server/
+    ├── __init__.py
+    ├── app.py                      # FastAPI factory; mounts Gradio at /web
+    ├── pi_environment.py           # MCPEnvironment with single ``run_rollout`` tool
+    ├── gradio_ui.py                # the /web Gradio Blocks UI
+    ├── catalog.py                  # endpoint shorthand resolver
+    └── Dockerfile                  # multi-stage uv build (used by ``openenv build``)
+```
+
+The sandbox backend + interception proxy are imported from
+[`opencode_env.sandbox`](../opencode_env/sandbox); `pi_env` ships no `sandbox/`
+of its own.
+
+## References
+
+- [OpenEnv docs](https://huggingface.co/docs/openenv)
+- [Pi coding agent](https://github.com/badlogic/pi-mono)
+- [Hugging Face sandboxes](https://huggingface.co/docs/huggingface_hub/guides/sandboxes)

envs/pi_env/.dockerignore

@@ -0,0 +1,29 @@
+# Don't bloat the build context. Copy only what the runtime needs.
+.venv
+__pycache__
+**/__pycache__
+*.pyc
+*.pyo
+.pytest_cache
+.mypy_cache
+.ruff_cache
+
+# NEVER ship secrets into the image. The container reads E2B_API_KEY etc.
+# from runtime env vars (HF Space secrets, ``docker run -e``).
+.env
+.env.*
+
+# Test artifacts and local debug output.
+tests/_artifacts/
+*.log
+.coverage
+htmlcov/
+
+# Editor / IDE noise.
+.idea
+.vscode
+.DS_Store
+
+# Git metadata is unused at runtime.
+.git
+.gitignore

envs/pi_env/.gitignore

@@ -0,0 +1,12 @@
+# Tests are dev-only — never ship into the deployable env.
+tests/
+
+# Local debug artifacts (covered above, kept for clarity).
+*.log
+
+# Local secrets / venv (also covered at repo root, kept here for safety).
+.env
+.env.*
+.venv
+__pycache__/
+*.pyc