TASK_API.md

# OpenEnv Task API — How It Works

A thin, optional layer that lets a **dataset-backed** environment expose its tasks
(splits + indexed rows) so a trainer can *enumerate and select* which task each
episode runs. Added in #726 (core ≥ 0.4.x).

---

## The contract

Implement any of these on your `Environment` (duck-typed via the `TaskProvider`
protocol — none are required; unimplemented ones return HTTP `501`):

| Method | Returns | HTTP endpoint |
|---|---|---|
| `list_splits()` | `["train", "test"]` | `GET  /{env}/splits` |
| `num_tasks(split)` | `int` | `POST /{env}/num_tasks` |
| `list_tasks(split)` | `[{id, index}, …]` | `POST /{env}/tasks` |
| `get_task(split, index)` | `{id, index, …}` | `POST /{env}/task` |
| `get_task_range(split, start, stop)` | `[…]` (slice) | `POST /{env}/task_range` |

Task *selection* is separate: it happens through `reset(split=…, index=…)`.
The Task API only **describes** the dataset; `reset` **binds** an episode to a row.

```
Task API  ── describes ──▶  "what tasks exist"   (list/count/peek)
reset()   ── binds ──────▶  "run THIS task"       (episode starts)
step()    ── grades ─────▶  reward vs hidden target
```

---

## Architecture

```mermaid
flowchart LR
    T[Trainer / Client] -->|HTTP GET/POST| API[Task API routes]
    T -->|WebSocket reset/step| EP[Episode routes]
    subgraph Server [OpenEnv FastAPI server]
        API --> F["_env_factory()"]
        EP  --> S[Session env instance]
    end
    F --> D[(Dataset)]
    S --> D
    D -. loaded once .- CACHE[[module-level cache]]
```

Two doors into the same environment: **HTTP** for task discovery, **WebSocket**
for the episode loop.

---

## Request flow (the important gotcha)

Every Task API call spins up a **fresh, throwaway** environment instance:

```mermaid
sequenceDiagram
    participant C as Client
    participant R as Route handler
    participant E as Env instance
    C->>R: POST /env/num_tasks {split}
    R->>E: env = _env_factory()   %% NEW instance
    R->>E: env.num_tasks(split)
    E-->>R: 7632
    R->>E: env.close()            %% destroyed
    R-->>C: {"num_tasks": 7632}
```

> ⚠️ **Consequence:** if your `__init__` loads the dataset, it reloads on
> *every* discovery call. **Load the dataset in a module-level / process cache**,
> not per-instance. (This env uses `@lru_cache` on `_load_split`.)

The episode loop is different — a WebSocket **session** holds one env instance
across `reset`/`step`:

```mermaid
sequenceDiagram
    participant C as Client
    participant S as Session env
    C->>S: reset(split="test", index=0)
    S-->>C: obs{ image, target hidden }
    C->>S: step(action{ latex })
    S-->>C: obs{ reward, done, target revealed }
```

---

## End-to-end RL workflow

```mermaid
flowchart TD
    A[list_splits] --> B[num_tasks split]
    B --> C{sample index i}
    C --> D[reset split, i]
    D --> E[observation: image, no target]
    E --> F[policy / VLM → action]
    F --> G[step action]
    G --> H[reward = rubric pred, hidden target]
    H --> I{more tasks?}
    I -->|yes| C
    I -->|no| J[aggregate reward]
```

---

## Design properties (why it's shaped this way)

- **Optional** — envs without a dataset simply don't implement it (graceful `501`).
- **Stateless discovery** — task metadata is derivable from `(split, index)`, so
  discovery needs no live session and scales horizontally.
- **Ground truth stays server-side** — `reset` never ships the target; only
  `step`'s result reveals it. The agent can't cheat.
- **Framework-neutral** — the same shape backs the ORS / Verifiers importers.

---

## Improvements worth making

| Area | Gap today | Suggested improvement |
|---|---|---|
| **Discovery cost** | fresh instance + `.close()` per call | cache a lightweight metadata-only provider; skip full env init for discovery |
| **Payload size** | `list_tasks` returns *all* rows in one response | prefer `get_task_range` pagination; cap/soft-limit `list_tasks` |
| **Task identity** | `id` is positional (`test-0`) | add a content hash / stable ID so shuffles & re-splits stay reproducible |
| **Filtering** | only split + index | add tag/difficulty/length filters (e.g. `list_tasks(split, where=…)`) |
| **Determinism** | random `reset` seeded ad hoc | standardize a `seed → index` mapping for reproducible curricula |
| **Schema** | tasks are free-form dicts | publish a typed task schema per env (validation + tooling) |
| **Discovery of prompt** | image/prompt only via `reset` | let `get_task` optionally include a lightweight preview |

---

## Where it can fail / doesn't fit

- **Huge / streaming datasets** — `num_tasks` needs a length; pure streaming
  datasets have no random index. Needs a materialized index or row count.
- **Per-instance dataset load** — the #1 footgun: 380 MB re-downloaded on every
  discovery call. Must cache at module scope.
- **Dynamic / generated tasks** — procedurally generated or infinite task spaces
  don't map to `(split, index)`; `num_tasks` is ill-defined.
- **Multi-step / stateful episodes** — the API indexes *starting states*; if a
  "task" is a whole interactive trajectory (env resets mid-episode), indexing is fuzzy.
- **Large task specs over HTTP** — `list_tasks` on a 68k-row split returns a giant
  JSON blob; clients should page with `get_task_range`.
- **Index drift** — positional IDs break if the dataset is re-uploaded/shuffled;
  runs aren't reproducible across dataset versions.
- **Non-tabular data** — video/3D/interactive-web tasks may not reduce to a row
  index cleanly.
- **Auth/gated datasets** — the server must carry credentials; discovery fails
  silently as `501`/`500` if the dataset can't load.
- **Concurrency** — if the env isn't `SUPPORTS_CONCURRENT_SESSIONS`, parallel
  rollouts serialize; discovery instances still multiply.

---

## TL;DR

> The Task API is a **read-only catalog** (`splits → count → row`) over a
> dataset, exposed via HTTP; `reset(split, index)` is what actually **runs** a
> task. It's optional, stateless, and keeps ground truth hidden — but assumes a
> **finite, indexable, cacheable** dataset. Cache the data at process scope,
> paginate large splits, and it fits supervised-style RL cleanly; it strains on
> streaming, generated, or deeply stateful task spaces.