LLM environments

Every LLM-training environment is reached through the OpenEnv API (the [llm] extra): whatever backs an env — a prompt dataset, plain Python functions, an imported gem / AxonRL env, or a sandboxed VM — a RolloutHarness drives it through a backend, in one of two transports: in-process (a InProcessEnvClient — no server, no sockets) or remote (an OpenEnv server reached over a URL by an RemoteEnvClient). resolve_env() hosts a local env as a URL from the same entrypoint / env_config a manifest would use, and returns (url, server). A server wraps each env in an OpenEnvWrapper, and a standard text interface — TextAction (message) and TextObservation (prompt) — carries the model’s text both ways, so there are no per-env codecs. See Environments (OpenEnv) for a guide.

  • RolloutHarness — the token-level rollout env: it owns tokenisation, the multi-turn loop and the provenance mask, and drives the env through its backend (reset / step). Build it with local() (a local env, in-process), a url (a hosted OpenEnv server reached over a WebSocket session), or from_spec() (resolve a URL or a module:Class entrypoint). Single-turn reasoning is simply max_turns=1.

  • RolloutCollector — runs independent groups of RolloutHarness rollouts over a batch; a shared TaskAssigner hands each group one common task — a dataset row (drawn in per-epoch shuffled order) or a reset seed.

  • DatasetEnv — supervised fine-tuning and preference-optimization (no generation).

  • InProcessEnvClient (in-process) and RemoteEnvClient (one WebSocket session against a hosted server) are the backends a RolloutHarness drives. resolve_env() hosts an entrypoint as a URL (max_concurrent_envs gives each session a fresh env); OpenEnvServer is the server it starts. spec_to_factory() resolves a non-URL spec to an env factory for the in-process path.

apply_chat_template() lives in agilerl.utils.llm_utils and is re-exported here for convenience.

Declaring what an env needs. A manifest’s entrypoint names the env to build — a registry id, or your own module:Class. Optional factory is the module:attr callable that receives that id, as in factory: gem:make with entrypoint: game:Sudoku-v0-easy. env_config is leftover kwargs only. Add env_packages ({"uv": [...]} or {"pip": [...]}) to declare what has to be installed for that import to work, and the first env build installs it here if it is missing. If those requirements cannot resolve alongside the trainer’s own, host the env with resolve_env() and point env_url at the URL it returns.

To drive a real external OpenEnv server (e.g. an env on the HuggingFace Hub), point RolloutHarness at its URL; each rollout opens its own /ws session via RemoteEnvClient and the server backs it with a fresh env instance, so the server’s max_concurrent_envs must cover batch_size * group_size sessions plus one for the lazily built eval env. If that server exposes an MCP tool, pass mcp_tool= to RolloutHarness / RemoteEnvClient to make the tool available during rollouts. In a run manifest, request_timeout_s bounds each message at 300 seconds by default; 0 disables the bound.

One environment per rollout. Each concurrent rollout gets its own env instance, created once and reused for the whole run. For batched training, RolloutCollector calls the env_factory batch_size * group_size times — once per slot — so the count is determined by the batch, at the training layer. A local env factory (lambda: RolloutHarness.local(make_env(), tok)) runs each in-process. To host the same entrypoint as a URL instead:

from agilerl.llm_envs.openenv_server import resolve_env

url, server = resolve_env(
    "game:Sudoku-v0-easy",
    factory="gem:make",
    max_concurrent_envs=batch_size * group_size + 1,
)
env_factory = lambda: RolloutHarness(url, tok, max_turns=50)

max_concurrent_envs must cover batch_size * group_size sessions plus one for the lazily built eval env.

class agilerl.llm_envs.RolloutHarness(env_client: str | Callable[[], str] | EnvClientProtocol, tokenizer: PreTrainedTokenizerBase, max_turns: int = 1, *, timeout_s: float | None = None, mcp_tool: str | None = None, action_field: str = 'message', observation_field: str | None = None, observation_processor: Callable[[Any], str] | None = None, instruction: str = '', strict_chat_template_boundary: bool = True, apply_chat_template: bool = True, chat_template_kwargs: dict[str, Any] | None = None, max_model_len: int | None = None, system_prompt: str | None = None, vision_processor: Callable[..., Mapping[str, Any]] | None = None, tasks: Sequence[Mapping[str, Any]] | None = None, eval_tasks: Sequence[Mapping[str, Any]] | None = None, segment_prompt_tokens: int | None = None, segment_max_images: int | None = None, restart_keep_turns: int = 0, action_error_field: str = '')

Token-level rollout env: tokenisation + turn loop over a text env client.

Assembles the multi-turn transcript and the provenance mask (only policy-generated tokens train, via get_episode_data()); terminates on max_model_len overflow.

class agilerl.llm_envs.RolloutCollector(env_factory: Callable[[...], RolloutHarness], batch_size: int, group_size: int, env_config: dict[str, Any] | None = None, io_timeout_s: float | None = 600.0, base_seed: int | None = None, slot_acquire_timeout_s: float | None = 300.0, rank: int | None = None, world_size: int | None = None, adaptive_task_sampling: bool = False)

Batched in-process collector over one rank’s share of a global batch.

batch_size is the world-wide prompt-group count (same meaning as GRPO BATCH_SIZE). This rank allocates (batch_size / world_size) * group_size RolloutHarness slots. world_size defaults to the process group (1 when not distributed).

Driven lock-step (reset/step/get_trajectories, the colocated path, single-caller) or per-episode (reset_episode/step_episode/ finalize_episode, the async path — a slot is held from reset to finalize). The per-episode methods are synchronous and thread-safe; asyncio callers offload them via asyncio.to_thread. A step_episode whose episode is finalized while its env round-trip is in flight raises rather than applying the stale result to the slot’s next episode.

class agilerl.llm_envs.TaskAssigner(dataset_size: int, *, seed: int | None = None, rank: int = 0, world_size: int = 1, adaptive: bool = False)

Assign each episode a (seed, row_index) task; a GRPO group shares one.

Dataset rows are reshuffled each epoch and split into one equal shard per data-parallel rank; a procedural env has no rows and is seeded instead. With adaptive, each row is instead drawn from the shard with probability proportional to its TaskRowStats weight, fed by record_outcome().

Parameters:

  • dataset_size – Rows in the env’s dataset; 0 for a procedural env.

  • seed – Seed for the per-epoch shuffle or weighted draw (None -> a fixed default).

  • rank – This process’s shard index in [0, world_size).

  • world_size – Number of data-parallel shards (1 = no sharding).

  • adaptive – Draw rows by recent informative-group rate instead of the epoch cycle.

class agilerl.llm_envs.DatasetEnv(train_dataset: Dataset, test_dataset: Dataset, tokenizer: PreTrainedTokenizerBase, *, objective: Literal['preference', 'sft'], response_column: str = 'response', chat_template_kwargs: dict[str, Any] | None = None, data_batch_size_per_gpu: int = 8, max_context_length: int | None = None, min_completion_length: int | None = None, seed: int = 42, rank: int = 0, world_size: int = 1)

Dataset-backed LLM env for SFT and DPO (no generation).

Variables:

  • dataset_size – {"train": N, "test": M} row counts after filtering.

  • num_epochs – Full passes completed over the train split.

class agilerl.llm_envs.openenv_server.OpenEnvWrapper(inner: TextEnvProtocol, *, env_name: str | None = None, owns_inner: bool = False)

Adapt a plain-text env to OpenEnv’s typed Environment ABC.

Translates between the env’s string reset/step and OpenEnv’s typed Action/Observation/State, surfacing dataset_size/tools on the state and forwarding env info into observation metadata.

Parameters:

  • inner – The local env to host.

  • env_name – Name in the OpenEnv metadata; defaults to inner’s class name.

  • owns_inner – If True, close closes inner (the per-session path).

class agilerl.llm_envs.openenv_server.TextAction(*, metadata: ~typing.Dict[str, ~typing.Any] = <factory>, message: str = '')

OpenEnv action carrying the policy’s generated text.

class agilerl.llm_envs.openenv_server.TextObservation(*, done: bool = False, reward: bool | int | float | None = None, metadata: ~typing.Dict[str, ~typing.Any] = <factory>, prompt: str = '', truncated: bool = False, question: ~typing.Any = None, answer: ~typing.Any = None)

OpenEnv observation carrying the next prompt text and optional labels.

truncated is declared because the wire has only done (no terminated/truncated split), so it can travel and reconstruct the 5-tuple.

On reset, prompt is shown to the policy and the labels stay unset. On the terminal step, question / answer are set so rubrics can score.

class agilerl.llm_envs.openenv_server.OpenEnvServer(env: TextEnvProtocol | None = None, *, make_env: Callable[[], TextEnvProtocol] | None = None, host: str = '127.0.0.1', port: int = 0, env_name: str | None = None, max_concurrent_envs: int | None = None, advertise_host: str | None = None)

Serve OpenEnv’s create_app on uvicorn in a background daemon thread.

Pass env to share one env (one session at a time), or make_env + max_concurrent_envs for a fresh env per /ws session. max_concurrent_envs other than None or 1 requires make_env.

Parameters:

  • env – A single shared local env. Exactly one of env/make_env.

  • make_env – Zero-arg factory building a fresh env per session.

  • host – Interface to bind (default loopback).

  • port – TCP port; 0 lets the OS pick one.

  • env_name – Name in the env’s OpenEnv metadata; defaults to its class name.

  • max_concurrent_envs – Max live sessions; set to the group size with make_env.

  • advertise_host – Host clients should dial, when it differs from the bound interface — a server binding 0.0.0.0 advertises its routable address (a pod IP, say), since 0.0.0.0 is not reachable from anywhere.

class agilerl.llm_envs.openenv.RemoteEnvClient(base_url: str | Callable[[], str], *, timeout_s: float | None = None, connect_timeout_s: float = 30.0, mcp_tool: str | None = None, action_field: str = 'message', tasks: Sequence[Mapping[str, Any]] | None = None, eval_tasks: Sequence[Mapping[str, Any]] | None = None)

Drives an environment hosted behind a URL, over one WebSocket session.

The session is the environment instance: the server builds a fresh env when the connection opens and destroys it when the connection closes, which is how concurrent episodes stay isolated. There is no resume, so a connection lost mid-episode fails that episode rather than reconnecting — a fresh connection would silently be a fresh environment, while the transcript carried on as if nothing had changed. Reconnection therefore happens only at the next reset, where a new environment is what the caller wanted anyway.

Calls block the calling thread while the session’s I/O runs on SESSION_LOOP, shared by every session in the process.

Parameters:

  • base_url – Root URL of the env server, or a zero-arg callable returning it — called again on every reconnect, so a host that moved is found again.

  • timeout_s – Per-message timeout; None (default) is unbounded.

  • connect_timeout_s – Timeout for establishing the session.

  • mcp_tool – If set, send text as call_tool(mcp_tool, {arg: text}) for MCP servers.

  • action_field – Action field (or MCP argument name) carrying the text.

  • tasks – Per-row reset kwargs; row_index selects the entry merged into reset.

  • eval_tasks – Held-out per-row reset kwargs used in eval_mode() when set.

class agilerl.llm_envs.openenv.InProcessEnvClient(env: TextEnvProtocol | Environment, *, action_field: str = 'message')

Drives an environment object living in this process — no network involved.

A plain-text env is wrapped in OpenEnvWrapper so it behaves exactly as it would when hosted behind a URL; an OpenEnv Environment (e.g. QADatasetEnv) is used as-is, which keeps its rubric metadata intact. Built by RolloutHarness.local() / RolloutHarness.from_spec().

Parameters:

  • env – The env to drive — plain-text or an OpenEnv Environment.

  • action_field – The action field the model’s text goes into, on the action class the env declares (message on our TextAction, but e.g. code for an env whose ACTION_CLS names it that).

agilerl.llm_envs.openenv_server.resolve_env(spec: str, env_config: dict[str, Any] | None = None, *, factory: str | None = None, host: str = '127.0.0.1', port: int = 0, max_concurrent_envs: int | None = None, advertise_host: str | None = None) → tuple[str, OpenEnvServer | None]

Resolve an env spec to a (url, server) a RolloutHarness can hit.

A URL is already hosted -> (url, None); anything else is built and hosted on a local OpenEnvServer. max_concurrent_envs serves a fresh env per /ws session (one server backs a whole group); unset is one at a time.

agilerl.llm_envs.env_sources.spec_to_factory(spec: str) → Callable[..., TextEnvProtocol]

Resolve a module:attr constructor path to the callable that builds the env.

Registry ids (game:GuessTheNumber-v0-easy, CartPole-v1) are not constructors; pass those as the entrypoint to a separate factory.