Environment Specifications

Pydantic models describing the environment for each training scenario.

class agilerl.models.env.GymEnvSpec(*, name: str, num_envs: Annotated[int, Ge(ge=1)] = 16, entrypoint: str | None = None, path: str | None = None, config: dict[str, Any] | None = None, wrappers: list[tuple[Any, dict[str, Any]] | str | Callable[[...], Any]] | None = None, sync: bool = False)

Gym environment specification.

Parameters:

  • entrypoint (str or None) – Entrypoint for the environment, if custom. Defaults to None.

  • path (str or None) – Path to the environment, if custom. Defaults to None.

  • config (dict[str, Any] or None) – Environment configuration, if custom. Defaults to None.

  • wrappers (list[tuple[Any, dict[str, Any]] | str] or None) – Environment wrappers, if custom. Defaults to None.

  • sync (bool) – Use synchronous vectorization instead of async.

static construct_custom_env_fn(entrypoint: str, path: str | None = None, config: dict[str, Any] | None = None, wrappers: list[tuple[Any, dict[str, Any]] | str | Callable[[...], Any]] | None = None) → Callable[[], Any]

Construct a custom environment given the configuration.

Parameters:

  • entrypoint (str or None) – Entrypoint for the environment, if custom. Defaults to None.

  • path (str or None) – Path to the environment, if custom. Defaults to None.

  • config (dict[str, Any] or None) – Environment configuration, if custom. Defaults to None.

  • wrappers (list[tuple[Any, dict[str, Any]] | str] or None) – Environment wrappers, if custom. Defaults to None.

Returns:

Custom environment factory function.

Return type:

Callable[[], Any]

make_env(extra_wrappers: list[type] | None = None) → AsyncVectorEnv | SyncVectorEnv

Instantiate the vectorized environment given the configuration.

Parameters:

extra_wrappers (list[type] or None, optional) – Optional list of wrapper classes to apply to each individual environment before vectorization.

Returns:

Vectorized environment

Return type:

GymEnvType

make_single_env() → Env

Create a single (non-vectorized) environment instance.

Useful for probing the observation/action space without the overhead of spinning up a full vectorized environment.

Returns:

A single gymnasium environment.

Return type:

gymnasium.Env

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class agilerl.models.env.PzEnvSpec(*, name: str, num_envs: Annotated[int, Ge(ge=1)] = 16, entrypoint: str | None = None, path: str | None = None, config: dict[str, Any] | None = None, wrappers: list[tuple[Any, dict[str, Any]] | str | Callable[[...], Any]] | None = None)

PettingZoo environment specification.

Parameters:

  • entrypoint (str or None) – Entrypoint for the environment, if custom. Defaults to None.

  • path (str or None) – Path to the environment, if custom. Defaults to None.

  • config (dict[str, Any] or None) – Environment configuration, if custom. Defaults to None.

  • wrappers (list[WrapperSpec] or None) – Environment wrappers, if custom. Defaults to None.

static construct_custom_env_fn(entrypoint: str, path: str | None = None, config: dict[str, Any] | None = None, wrappers: list[tuple[Any, dict[str, Any]] | str | Callable[[...], Any]] | None = None) → Callable[[], ParallelEnv]

Construct a custom PettingZoo environment factory.

For PettingZoo, we always require an explicit constructor/entrypoint.

Parameters:

  • entrypoint (str or None) – Entrypoint for the environment, if custom. Defaults to None.

  • path (str or None) – Path to the environment, if custom. Defaults to None.

  • config (dict[str, Any] or None) – Environment configuration, if custom. Defaults to None.

  • wrappers (list[WrapperSpec] or None) – Environment wrappers, if custom. Defaults to None.

Returns:

Custom PettingZoo environment factory.

Return type:

Callable[[], ParallelEnv]

make_env(extra_wrappers: list[type] | None = None) → AsyncPettingZooVecEnv

Instantiate vectorized PettingZoo environments from a constructor.

Parameters:

extra_wrappers (list[type] or None, optional) – Optional list of wrapper classes to apply to each individual environment before vectorization.

Returns:

Vectorized PettingZoo environments.

Return type:

AsyncPettingZooVecEnv

make_single_env() → ParallelEnv

Create a single (non-vectorized) PettingZoo environment instance.

Useful for probing the observation/action spaces without the overhead of spinning up a full vectorized environment.

Returns:

A single PettingZoo parallel environment.

Return type:

ParallelEnv

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class agilerl.models.env.OfflineEnvSpec(*, name: str, num_envs: Annotated[int, Ge(ge=1)] = 16, entrypoint: str | None = None, path: str | None = None, config: dict[str, Any] | None = None, wrappers: list[tuple[Any, dict[str, Any]] | str | Callable[[...], Any]] | None = None, sync: bool = False, minari_dataset_id: str | None = None, dataset_path: str | None = None, remote: bool = False, dataset: Any = None)

Environment specification for offline RL training.

Wraps a standard Gymnasium evaluation environment together with the dataset source used to fill the replay buffer before training begins.

Exactly one of minari_dataset_id or dataset_path must be provided. When minari_dataset_id is set, the dataset is loaded via the Minari library. When dataset_path is set, the dataset is loaded from a local HDF5 file.

Parameters:

  • minari_dataset_id (str | None) – Identifier for a Minari dataset (e.g. "cartpole-v0").

  • dataset_path (str | None) – Path to a local HDF5 dataset file.

  • remote (bool) – If True, download the Minari dataset from the remote repository when it is not available locally.

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class agilerl.models.env.LLMEnvSpec(*, env_type: LLMEnvType, dataset: str | None = None, columns: dict[str, str] | None = None, prompt_template: dict[str, Any] | None = None, max_reward: float | None = None, train_test_split: Annotated[float, Ge(ge=0.0), Le(le=1.0)] = 0.9, reward_file_path: str | None = None, reward_fn_name: str | None = None, response_column: str = 'response', env_name: str | None = None, entrypoint: str | None = None, env_config: dict[str, Any] | None = None, max_turns: Annotated[int | None, Ge(ge=1)] = None, data_batch_size_per_gpu: Annotated[int, Ge(ge=1)] = 8, return_raw_completions: bool = False, max_context_length: int | None = None, seed: int | None = None)

Environment specification for LLM reasoning and preference training.

Declaratively captures the dataset, reward function, and prompt template needed to construct a ReasoningGym or PreferenceGym. Fields are aligned with what Arena expects for LLM training jobs.

Parameters:

  • env_type (LLMEnvType) – The type of LLM environment ("reasoning" or "preference").

  • columns (dict[str, str] | None) – Optional mapping from source dataset column names to the names expected by the gym (e.g. {"question": "input", "answer": "output"} for reasoning).

  • prompt_template (dict[str, Any] | None) – Chat-template configuration passed as conversation_template to ReasoningGym.

  • max_reward (float | None) – Maximum achievable reward, forwarded to the LLM training loop for accuracy logging.

  • train_test_split (float) – Fraction of the dataset used for training.

  • reward_file_path (str | None) – Path to a Python file containing the reward function. Required for reasoning environments.

  • dataset (str) – Path to a Parquet dataset file or a HuggingFace dataset. Required for reasoning/preference/sft environments.

  • env_name (str | None) – GEM environment id (e.g. "game:Sudoku-v0-easy"). Mutually exclusive with entrypoint.

  • entrypoint (str | None) – Dotted path to a callable that returns a MultiTurnEnv. Mutually exclusive with env_name.

  • env_config (dict[str, Any] | None) – Keyword arguments forwarded to the entrypoint callable. Only used when entrypoint is set.

  • max_turns (int | None) – Maximum interaction turns per episode. If None for multiturn environments, the value is probed from the environment.

make_env(tokenizer: Any, accelerator: Accelerator | None = None) → ReasoningGym | PreferenceGym | SFTGym

Make the environment for the LLM agent.

For multiturn environments, use make_multiturn_env_factory() instead — the training loop needs a factory, not a single env.

Parameters:

  • tokenizer (Any) – The tokenizer.

  • accelerator (Accelerator | None) – The accelerator.

Returns:

The reasoning or preference gym environment.

Return type:

ReasoningGym | PreferenceGym | SFTGym

make_multiturn_env_factory(tokenizer: Any, *, max_model_len: int | None = None, max_output_tokens: int | None = None) → Callable[[], Any]

Build a factory that creates wrapped multi-turn env instances.

Each call to the returned factory creates a fresh TokenObservationWrapper. The underlying environment is either a GEM environment (env_name) or a custom class resolved from entrypoint.

If max_turns is None, it is probed from a temporary environment instance and stored back on the spec.

Parameters:

  • tokenizer (Any) – The tokenizer (shared across all instances).

  • max_model_len (int | None) – Maximum model context length for sliding-window prompt truncation inside the wrapper.

  • max_output_tokens (int | None) – Maximum newly generated tokens per turn.

Returns:

A zero-argument callable that creates a wrapped env.

Return type:

Callable[[], TokenObservationWrapper]

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

property name: str

Human-readable name: dataset path, GEM env name, or entrypoint.

class agilerl.models.env.BanditEnvSpec(*, name: str = 'BanditEnv', features: DataFrame | str | Path | None = None, targets: DataFrame | str | Path | None = None, entrypoint: str | None = None, path: str | None = None, config: dict[str, Any] | None = None)

Environment specification for contextual bandit training.

Supports two modes:

Dataset mode: provide features and targets (as DataFrames or file paths) to construct a BanditEnv from a labelled dataset.

Custom entrypoint mode: provide an entrypoint (e.g. "my_module:MyBanditEnv") to instantiate an arbitrary bandit environment. The resolved callable is invoked with **config..

Exactly one of (features + targets) or entrypoint must be provided.

Parameters:

  • name (str) – Human-readable name for the environment / dataset.

  • features (pandas.DataFrame | str | Path | None) – Dataset features. A pd.DataFrame or a path to a file.

  • targets (pandas.DataFrame | str | Path | None) – Dataset targets. A pd.DataFrame or a path to a file.

  • entrypoint (str | None) – Dotted path to a callable that returns a bandit environment (e.g. "my_module:MyBanditEnv").

  • path (str | None) – Optional filesystem path added to sys.path before resolving the entrypoint.

  • config (dict[str, Any] | None) – Keyword arguments forwarded to the entrypoint callable.

make_env() → BanditEnvProtocol

Construct a bandit environment.

In dataset mode, returns a BanditEnv built from features and targets. In entrypoint mode, resolves the callable and invokes it with **config.

Returns:

A bandit environment satisfying BanditEnvProtocol.

Return type:

BanditEnvProtocol

model_config: ClassVar[ConfigDict] = {'arbitrary_types_allowed': True}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].