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

* Add Pi coding-agent harness (pi_env) reusing opencode_env sandbox + proxy

PiSessionFactory/PiSession/PiConfig mirror the OpenCode primitive: install the
pi CLI (Node bootstrap + npm), write models.json pointing at an OpenAI-compatible
endpoint, run pi --print --mode json headless, and (transparent_proxy mode)
capture per-token logprobs via the shared interception proxy for GRPO.

The sandbox backend (HFSandboxBackend) and interception proxy are imported from
opencode_env for now; the plan is to consolidate both into openenv.core in a
follow-up. Includes builder + factory-lifecycle unit tests.

* Add deployable HTTP env layer to pi_env (parity with opencode_env)

Mirrors opencode_env's second layer so pi_env is both an in-process harness
primitive and a deployable OpenEnv environment: PiEnv MCP client + FastAPI
server (PiEnvironment with the run_rollout tool, task catalog, Gradio UI) +
Dockerfile + openenv.yaml. HF sandbox (image, HF_TOKEN) replaces E2B
(template, E2B_API_KEY); in-sandbox paths are /root; the sandbox backend and
interception proxy are imported from opencode_env. Adds uv.lock and the
pi-env server-image entry to docker-build.yml.

* Add pi_env to the docs (env stub + catalog card + toctree nav)

sync_env_docs.py --fix generates the docs/source/environments/pi.md stub
from the README; the HTML catalog card in environments.md and the
_toctree.yml nav entry are added manually (the script does not manage
those). check-env-docs now passes.

* Format test_pi_runtime.py (usort import order + ruff line wrap)

* Address review: fix docs links, node version gate, black-box model id, pipefail

- README: correct the HF sandbox docs URL, drop redundant/circular links
- .dockerignore: HF_TOKEN, not E2B_API_KEY
- pi_runtime: gate Node >=22.19 (not just major 22); use the full model id and
  select by --provider so a slashed id isn't split; set -o pipefail so the
  piped tee doesn't mask Pi's exit code

* Keep server/client separation in pi_env server; correct setup timing docs

Import the primitive from its defining modules instead of the top-level
pi_env package (which re-exports the client), per the server/client-separation
invariant. Reword the setup docs: it runs at rollout start, not strictly
before the agent.

* Address review: cover cold-bootstrap in timeouts, doc peer-dep, drop dead field

- harness/server: size the sandbox lifetime and MCP-tool timeout to the full
  Node+Pi cold install, not just the agent
- README: note the opencode_env install prerequisite for the primitive path;
  correct the server Dockerfile comment
- config: drop unused request_timeout_ms
- client: document the endpoint shorthand

* Add pre-baked HF sandbox image for pi_env; fix mode docstring

Mirror opencode_env's pre-baked image: bake Node + Pi + proxy deps so rollouts
skip the cold install (the proxy source is uploaded at runtime since it lives
in opencode_env). Also drop the stale '(default)' on transparent_proxy in the
harness docstring (the factory default is black_box).

* Address review: skip verify/reward on failed setup; pipefail proxy-deps; guard UI stderr

- server: only run verify + compute reward when the run is clean (a failed
  setup/agent leaves a half-prepared sandbox); dedupe the endpoint error paren
- harness: set -o pipefail on the proxy-deps install so a failed pip isn't
  masked by tail
- gradio: don't index an empty stderr split (IndexError when exit!=0, no stderr)

* Give setup commands a longer timeout than verify (pip install / downloads)

* Use get_token() for the HF credential check; gate proxy deps vs source separately

- token check via huggingface_hub.get_token() (HF_TOKEN, legacy var, or a
  cached login) instead of guessing env var names
- _start_proxy: skip the pip install when the deps import, independent of the
  proxy source upload, so the pre-baked image (deps baked, source not) skips it

* Drop unused os import in pi_environment (token check now uses get_token)

* Default PiEnv client message timeout to cover a full rollout (server caps at 1500s)

* Retry PiSessionFactory.create() with backoff

Mirror of the OpenCodeSessionFactory retry: session creation (sandbox + Pi
install + proxy + agent) is the flakiest step, and a single transient failure
dropped the rollout as unscorable. create() now retries create_attempts times
with exponential backoff; the per-attempt logic moved to _create_once, which
already tears its own sandbox down on failure, so a retry never leaks.

* Run setup before the agent and skip verify on a non-zero agent exit

Addresses review: run_rollout ran setup after create() had already started the
agent, so setup raced the agent over the workspace. create() now takes
start_agent (default True, so the loop-owning training path is unchanged), the
server passes start_agent=False, runs setup, then calls session.start_agent(),
giving the correct setup then agent then verify order. A non-zero agent exit is
now treated like a timeout, so a crashed agent no longer runs verify or earns a
reward.

AI-assisted.

* Add pi_env demo and training tutorial

- examples/pi_env_simple.py: inference demo hitting the deployed Space
- docs tutorial pi-agent-grpo.md: black-box AsyncGRPO training recipe, links to TRL

Author: sergiopaniego

.github/workflows/docker-build.yml

@@ -112,6 +112,12 @@ 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: pi-sandbox
+            dockerfile: envs/pi_env/hf_image/Dockerfile
+            context: envs/pi_env
           - name: openapp-env
             dockerfile: envs/openapp_env/server/Dockerfile
             context: envs/openapp_env

docs/source/_toctree.yml

@@ -54,6 +54,8 @@
     title: RL Training with an Agentic Harness
   - local: tutorials/sft-warmup
     title: SFT Training with Environments
+  - local: tutorials/pi-agent-grpo
+    title: Training a Real Coding Agent (Pi)
   title: Tutorials
 - isExpanded: false
   sections:
@@ -131,6 +133,8 @@
     title: OpenCode
   - local: environments/pelican_svg
     title: Pelican SVG
+  - local: environments/pi
+    title: Pi
   - local: environments/sophistry_bench_sprint
     title: Sophistry Bench Sprint
   title: Environments

docs/source/environments.md

@@ -266,6 +266,13 @@ The OpenEnv community has built a catalog of ready-to-run environments that cove
         <a href="https://huggingface.co/spaces/sergiopaniego/pelican-svg-env" class="!no-underline border dark:border-gray-700 px-3 py-1 rounded text-sm hover:shadow">🤗 HF</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,215 @@
+<!-- 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/package_reference/sandbox)
+against any OpenAI-compatible LLM endpoint, optionally capturing per-token
+logprobs for GRPO training.
+
+It mirrors `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. The primitive reuses the sandbox backend +
+proxy from `opencode_env`, so install it alongside `pi_env`:
+
+```bash
+pip install "openenv-opencode-env @ git+https://github.com/huggingface/OpenEnv.git#subdirectory=envs/opencode_env"
+```
+
+```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. For faster rollouts use the
+pre-baked image (Node + Pi + proxy deps already installed), built by CI from
+`hf_image/Dockerfile`:
+
+```python
+sandbox_backend=HFSandboxBackend(image="ghcr.io/huggingface/openenv-pi-sandbox:latest")
+```
+
+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`; `pi_env` ships no `sandbox/` of its own.

docs/source/tutorials/index.md

@@ -27,4 +27,5 @@ Already familiar with the basics? These tutorials cover specific workflows in de
 | [Evaluating agents with Inspect AI](evaluation-inspect.md) | Wrap an OpenEnv environment in an Inspect AI `Task`, run it via `InspectAIHarness`, and get a structured `EvalResult`. | No | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/OpenEnv/blob/main/examples/evaluation_inspect.ipynb) |
 | [BrowserGym Harness Rollouts](browsergym-harness.md) | Drive BrowserGym through the OpenEnv harness runtime when a trainer needs token sampling, logprobs, and reward assignment inside the training loop. | Yes | — |
 | [Collecting rollouts for supervised training](sft-warmup.md) | Run a teacher model to collect reward-labeled rollouts, filter them, and fine-tune a student with TRL's `SFTTrainer` as a warm-start for GRPO. | Yes | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/OpenEnv/blob/main/examples/sft_warmup.ipynb) |
+| [Training a Real Coding Agent (Pi)](pi-agent-grpo.md) | Train the actual Pi agent (black-box, loop-owning) with TRL's AsyncGRPOTrainer: a transparent proxy captures each turn's token ids and logprobs while the agent runs its own tool loop. | Yes | — |
 

docs/source/tutorials/pi-agent-grpo.md

@@ -0,0 +1,54 @@
+# Coding Agent Training with TRL (Pi)
+
+This tutorial covers the black-box training path: training the actual
+[`pi`](https://github.com/badlogic/pi-mono) coding agent, with its own planner,
+tools, context management, and stop condition, using TRL's experimental
+`AsyncGRPOTrainer`. The agent owns its loop, and OpenEnv captures what it did.
+
+> [!NOTE]
+> Three GRPO patterns, three tutorials. For a standard `reset()` / `step()`
+> flow where TRL drives the episode, see the
+> [Wordle GRPO tutorial](wordle-grpo.md). For harness rollouts where the
+> trainer still generates each turn (white-box), see the
+> [BrowserGym harness tutorial](browsergym-harness.md). Use this page when you
+> want to train a production agent as-is, without reimplementing its loop.
+
+## How It Works
+
+The full recipe lives in TRL. The moving pieces:
+
+1. Each rollout runs the agent inside an OpenEnv session created by
+   `PiSessionFactory` from
+   [`pi_env`](https://github.com/huggingface/OpenEnv/tree/main/envs/pi_env),
+   in `transparent_proxy` mode. A small proxy inside the sandbox forwards the
+   agent's `/v1/chat/completions` calls to your vLLM server and records each
+   turn's token ids and logprobs to a trace.
+2. When the agent stops, TRL's `HarnessRolloutWorker` reads the trace, rebuilds
+   the per-turn training rows from the recorded ids, and scores the final
+   workspace with the session's `verify()` method (a held-out verifier the
+   agent never sees).
+3. `AsyncGRPOTrainer` trains on those rows, propagating the rollout reward to
+   every trained token through the group-relative advantage. NCCL weight sync
+   keeps the vLLM server on the current policy, so the agent always samples
+   from the model being trained.
+
+Each rollout gets its own isolated session: one sandbox, one proxy port, one
+agent process. Three small functions adapt the recipe to your task:
+`rollout_reward_fn` (outcome to scalar reward), `train_turn_fn` (which turns
+receive gradient), and `agent_turn_fn` (which trace entries are real agent
+turns rather than auxiliary calls like title generation).
+
+## Full Recipe
+
+The reference script trains on competitive-coding problems from
+`agentica-org/DeepCoder-Preview-Dataset`: the agent writes `solution.py`, and
+the verifier runs it against held-out tests, returning the fraction passed. It
+is self-contained, runs the agent in a local subprocess sandbox (no container
+setup needed), needs two GPUs (one serving the policy with vLLM, one
+training).
+
+- [`examples/scripts/openenv/pi.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/openenv/pi.py)
+  in TRL: the complete, runnable script.
+- [`envs/pi_env`](https://github.com/huggingface/OpenEnv/tree/main/envs/pi_env):
+  the OpenEnv side, including the session factory, sandbox backends, and the
+  transparent interception proxy.

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 HF_TOKEN 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