OpenEnv documentation

The Task API

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.4.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

The Task API

This guide explains the Task API: the optional discovery layer that lets a dataset-backed environment publish the set of tasks it can run, so a trainer or evaluator can enumerate them and choose which one each episode uses.

It exists to answer a common question: OpenEnv’s step loop gives you one episode at a time, so how does a training script find out that an environment holds 7,595 test problems, and how does it ask for problem number 12?

The Short Answer

The Task API has two halves:

  • Server side: your environment optionally implements five methods — list_splits(), list_tasks(), num_tasks(), get_task(), get_task_range(). These are described by the TaskProvider protocol.
  • HTTP side: when those methods exist, the environment server automatically exposes them as routes under /{env_name}/…. You do not register anything.

Task discovery is metadata only. It never starts an episode. Selecting a task happens the usual way, through reset():

env.list_splits()            # ["train", "test"]
env.num_tasks("test")        # 7595
env.get_task("test", 12)     # {"id": "test-12", "index": 12, "split": "test"}
env.reset(split="test", index=12)   # <- this is what starts the episode

The API is shaped to be compatible with ORS/OpenReward task and split conventions, which is why environments imported with openenv import get it for free.

When You Need It

Implement the Task API when your environment is dataset-backed — when an episode means “run one row of a dataset” and a trainer needs to iterate, shard, or shuffle over those rows.

You do not need it for environments where every episode is generated or self-contained: Echo, Snake, a game, a REPL. Those environments simply omit the methods, and the routes report 501 Not Implemented.

The Protocol

TaskProvider is a typing.Protocol, not a base class. You do not inherit from it — you just define the methods on your Environment subclass and the server discovers them by name.

from typing import Any, Optional

class TaskProvider(Protocol):
    def list_splits(self) -> list[Any]: ...
    def list_tasks(self, split: str) -> list[Any]: ...
    def num_tasks(self, split: str) -> int: ...
    def get_task(self, split: str, index: int) -> Any: ...
    def get_task_range(
        self,
        split: str,
        start: Optional[int] = None,
        stop: Optional[int] = None,
    ) -> list[Any]: ...

Each method may be sync or async; the server awaits the result when it is awaitable, so you can back task discovery with async dataset clients.

A “task spec” is deliberately untyped (Any). Return whatever your environment needs — a dict, a Pydantic model, a dataclass. The server converts it to JSON, handling Pydantic models, dataclasses, and plain objects. In practice a positional stub such as {"id": "test-12", "index": 12, "split": "test"} is enough, because reset() is what actually loads the row.

Two rules that are easy to miss

  1. Task methods must be side-effect-free. They are discovery, not control. They must not mutate episode state, consume a stream, or advance a cursor.
  2. They must work on a freshly constructed environment. Each task route builds a short-lived environment instance, calls the one method, and closes it again. Nothing you set up in reset() is available, so read configuration in __init__ (or lazily inside the method) rather than relying on episode state.

HTTP Routes

The routes are registered by HTTPEnvServer and appear in the server’s OpenAPI schema under the Task API tag. All of them are namespaced by environment name except /list_environments.

MethodRouteBodyResponse
GET/list_environments["latex_ocr_env"]
GET/{env_name}/splits[{"name": "train", "type": "train"}, …]
POST/{env_name}/tasks{"split": "test"}{"tasks": [...], "env_name": "latex_ocr_env"}
POST/{env_name}/num_tasks{"split": "test"}{"num_tasks": 7595}
POST/{env_name}/task{"split": "test", "index": 12}{"task": {...}}
POST/{env_name}/task_range{"split": "test", "start": 0, "stop": 32}{"tasks": [...]}

The request bodies are ListTasksRequest, NumTasksRequest, GetTaskRequest, and GetTaskRangeRequest. start and stop follow Python slice semantics: start is inclusive, stop is exclusive, and both may be omitted.

These are HTTP-only. There are no WebSocket message types for task discovery — the /ws session protocol stays focused on reset, step, state, and close.

env_name

{env_name} is the name passed to create_app(..., env_name="latex_ocr_env"). If you do not pass one, it defaults to the environment factory’s class name (for example LatexOCREnvironment). Matching is case-insensitive, and any other name returns 404. Pass env_name explicitly for any environment that implements the Task API — the generated URL is part of your public surface, and a class rename should not break a trainer.

from openenv.core.env_server import create_app

app = create_app(
    LatexOCREnvironment,
    LatexOCRAction,
    LatexOCRObservation,
    env_name="latex_ocr_env",
)

Split normalization

list_splits() may return plain strings, dicts, or Pydantic models. The server normalizes each entry to {"name": ..., "type": ...}:

  • a dict or model is passed through as-is (after JSON conversion)
  • a string becomes {"name": s, "type": s} when s is train, validation, or test
  • any other string becomes {"name": s, "type": "validation"}

So returning ["train", "holdout"] yields [{"name": "train", "type": "train"}, {"name": "holdout", "type": "validation"}]. Return dicts yourself if you need to control the type for a custom split name.

Error semantics

ConditionStatus
{env_name} does not match the server’s environment404 Not Found
Environment does not define the method, or it raises NotImplementedError501 Not Implemented
Method raises IndexError (index out of range)400 Bad Request

Raise IndexError from get_task() for an out-of-range index so callers get a 400 rather than a 500. Reserve NotImplementedError for capabilities the environment genuinely does not support — for example, random access on a streaming split.

Implementing It

A dataset-backed environment reads its dataset in __init__, answers discovery from metadata, and loads the actual row in reset():

from typing import Any, Optional
from uuid import uuid4

from openenv.core.env_server import Environment
from openenv.core.env_server.types import State


class LatexOCREnvironment(Environment):
    SPLITS = ["train", "test"]

    def __init__(self, dataset_name: str = "unsloth/LaTeX_OCR") -> None:
        super().__init__()
        self.dataset_name = dataset_name

    # --- Task API: discovery only, no side effects ---

    def list_splits(self) -> list[str]:
        return self.SPLITS

    def num_tasks(self, split: str) -> int:
        return len(self._load_split(split))

    def list_tasks(self, split: str) -> list[dict[str, Any]]:
        return [
            {"id": f"{split}-{i}", "index": i}
            for i in range(self.num_tasks(split))
        ]

    def get_task(self, split: str, index: int) -> dict[str, Any]:
        n = self.num_tasks(split)
        if index < 0 or index >= n:
            raise IndexError(f"index {index} out of range for split {split} (n={n})")
        return {"id": f"{split}-{index}", "index": index, "split": split}

    def get_task_range(
        self, split: str, start: Optional[int] = None, stop: Optional[int] = None
    ) -> list[dict[str, Any]]:
        n = self.num_tasks(split)
        start = 0 if start is None else start
        stop = n if stop is None else min(stop, n)
        return [
            {"id": f"{split}-{i}", "index": i, "split": split}
            for i in range(start, stop)
        ]

    # --- Episode: this is where a task is actually loaded ---

    def reset(
        self,
        split: str = "test",
        index: Optional[int] = None,
        seed: Optional[int] = None,
        episode_id: Optional[str] = None,
        **kwargs: Any,
    ) -> LatexOCRObservation:
        if split not in self.list_splits():
            raise ValueError(f"unknown split {split!r}; expected {self.list_splits()}")
        self._state = State(episode_id=episode_id or str(uuid4()), step_count=0)
        row = self._load_row(split, index, seed)
        return LatexOCRObservation(image_base64=row["image"], done=False)

Selecting a task in reset()

ResetRequest allows extra fields, and both the HTTP /reset route and the WebSocket reset message forward unknown keys to your reset() — after filtering them against its signature. So naming the parameters split and index in reset() is all it takes for this to work:

result = env.reset(split="test", index=12)

The corresponding wire calls are POST /reset with {"split": "test", "index": 12}, or a WebSocket {"type": "reset", "data": {"split": "test", "index": 12}}. The server filters incoming keys against your reset() signature, so a key it does not declare is dropped silently unless the signature also has **kwargs. A misspelled parameter therefore shows up as “the environment ignored my selection” rather than an error — worth checking first when a task selection appears not to take effect.

When index is omitted, pick a row from seed so episodes stay reproducible.

Streaming and very large splits

list_tasks() and get_task_range() can be asked for more rows than you want to materialize. Two defenses are worth building in:

  • return positional stubs rather than real rows, and cap how many you generate (list_tasks() returning a bounded preview is fine — num_tasks() still reports the honest total)
  • if a split is streamed sequentially and cannot be randomly accessed, raise ValueError or NotImplementedError from reset(index=...) and say which mode does support indexing

Report the true count from num_tasks() even when list_tasks() is truncated. A trainer that shards work by num_tasks() needs the real denominator.

Consuming the Task API

The core clients (EnvClient, SyncEnvClient) do not ship Task API methods, because task specs are environment-specific. Environment clients add thin HTTP helpers alongside the inherited Gym-style reset()/step():

import requests
from urllib.parse import urljoin

ENV_NAME = "latex_ocr_env"


class LatexOCREnv(EnvClient[LatexOCRAction, LatexOCRObservation, State]):
    def list_splits(self) -> list[str]:
        resp = requests.get(urljoin(self._http_base(), f"{ENV_NAME}/splits"), timeout=30)
        resp.raise_for_status()
        return [s["name"] for s in resp.json()]

    def num_tasks(self, split: str) -> int:
        resp = requests.post(
            urljoin(self._http_base(), f"{ENV_NAME}/num_tasks"),
            json={"split": split},
            timeout=60,
        )
        resp.raise_for_status()
        return int(resp.json()["num_tasks"])

Note that the Task API is served over HTTP while episodes typically run over the WebSocket /ws session. A client that was constructed from a ws:// URL has to derive the HTTP base URL before calling these routes.

From a training or evaluation loop, discovery and episode control then compose naturally:

with LatexOCREnv.from_docker_image("latex-ocr-env:latest") as env:
    total = env.num_tasks("test")
    for index in range(total):
        result = env.reset(split="test", index=index)
        prediction = model(result.observation.image_base64)
        result = env.step(LatexOCRAction(latex=prediction))
        print(index, result.reward)

Or with curl, against a locally running server:

curl http://localhost:8000/list_environments
curl http://localhost:8000/latex_ocr_env/splits
curl -X POST http://localhost:8000/latex_ocr_env/num_tasks \
  -H 'Content-Type: application/json' -d '{"split": "test"}'
curl -X POST http://localhost:8000/latex_ocr_env/task_range \
  -H 'Content-Type: application/json' -d '{"split": "test", "start": 0, "stop": 4}'

Imported Environments

openenv import generates wrappers that already implement the Task API by delegating to the source environment’s own task and split methods, and whose reset() accepts split / index / task_spec. If you are bringing in an ORS/OpenReward or Prime Intellect Verifiers environment, you get task discovery without writing any of the above — see the CLI reference for openenv import.

Checklist

Before shipping a dataset-backed environment:

  1. All five methods defined, or a deliberate NotImplementedError for the ones you cannot support.
  2. Every method works on a fresh instance, with no reliance on reset() having run.
  3. No side effects — no cursor advanced, no stream consumed, no episode state mutated.
  4. env_name passed explicitly to create_app().
  5. get_task() raises IndexError for out-of-range indices.
  6. num_tasks() reports the honest total, even if list_tasks() truncates.
  7. reset() accepts split and index with the same names the trainer will use, and falls back to seed-driven selection when index is omitted.

Related Reading

Update on GitHub