base env

BaseEnv runs the standard model loop: feed the prompt, let the model respond and call tools, then score the transcript. you implement the dataset and a scoring hook; everything else is optional.

the interface

create_dataset is required, plus at least one of the two scoring hooks. the full surface:

from contextlib import asynccontextmanager
from pathlib import Path

from benchmax.envs import BaseEnv, BaseRollout, DatasetSplit, JsonlDataset, Tool
from benchmax.rewards import extract_completion_text

class MyEnv(BaseEnv):
    max_turns = 1  # bounds the chat loop; max_tool_calls separately bounds tool usage

    async def create_dataset(
        self, split: DatasetSplit, base_dir: Path, *, max_examples: int | None = None
    ) -> JsonlDataset:
        """required. map a split ("train" or "eval") to a finite dataset. every
        example payload contains `prompt_messages`; other fields are opaque data
        your reward can read later. `max_examples` lets validation and capped
        runs load a small slice."""
        return JsonlDataset(
            base_dir / f"{split}.jsonl", row_to_example=..., max_examples=max_examples
        )

    # implement at least one of the two scoring hooks; when both are present
    # the trainer merges their reward dicts.

    async def compute_reward(self, rollout: BaseRollout) -> dict[str, float]:
        """score one completed rollout."""
        answer = extract_completion_text(rollout.messages)
        return {"correct": float(answer == rollout.example_args["answer"])}

    async def compute_group_rewards(self, rollouts):
        """score a completed sibling group jointly (win-rate, ranking,
        diversity)."""
        ...

    # every hook below is optional; the defaults give a tool-less loop.

    async def list_tools(self) -> list[Tool]:
        """declare OpenAI-compatible tools the model may call. defaults to []."""
        return [...]

    async def run_tool(self, rollout_id: str, tool_name: str, **tool_args):
        """execute one tool call; the return value is the tool result the
        model sees."""
        ...

    @asynccontextmanager
    async def rollout_context(self, rollout_id, example):
        """acquire per-rollout resources before the model sees the prompt and
        release them afterwards (workspace files, db seed, a sandbox)."""
        yield

    async def aclose(self):
        """close environment-owned shared resources when the environment is
        done."""

rollout.messages is the full transcript, rollout.example_args is the per-example data from your dataset payload, and rollout.split tells you whether the attempt came from training or evaluation traffic.

there is no separate system-prompt hook: put system messages in each example’s prompt_messages when you build the dataset.

how a rollout works

1. create_dataset produces examples whose prompt_messages seed the conversation (a single prompt can be several messages, including system and chat history)
2. (optional) rollout_context sets up per-rollout state (e.g. copies files to a workspace, seeds a database)
3. the model receives the prompt.
4. the model generates a response, optionally calling tools via list_tools / run_tool, for up to max_turns turns
5. compute_reward scores the final transcript (and compute_group_rewards can add group-relative components once every sibling finishes)
6. rollout_context exits, releasing anything acquired in step 2
7. the model updates to produce higher-scoring responses

this loop repeats across your dataset. the environment stays fixed while the model improves.

reward components are inferred from the dictionaries returned by compute_reward and compute_group_rewards. components may be conditional; the trainer fills a missing component with zero wherever another rollout in the same training batch emitted that component. operational failures return no components and record why in termination_reason, without cancelling successful siblings.

use Python’s standard logging module from any env method. the trainer attributes each log record to its rollout automatically, so logs appear alongside the completion and reward breakdown in the rollout inspector. see logging.

go deeper