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

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:

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:

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

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.