Training¶

Trainers¶

The Trainer classes provide a high-level, manifest-driven interface for running AgileRL evolutionary training. See the Trainers guide for usage examples and manifest reference.

Abstract base trainer for AgileRL evolutionary training.

Parameters:

algorithm (AlgoSpec | str) – An algorithm spec or a string algorithm name.
environment (EnvSpecType | str) – A gymnasium.Env instance, a PettingZoo ParallelEnv instance, or an env-name string.
training (TrainingSpec) – Training loop parameters (max steps, population size, etc.).
mutation (MutationSpec | None) – Mutation probabilities and RL-HP ranges.
selection_strategy (SelectionStrategySpec | None) – Selection strategy driving evolutionary HPO — a TournamentSelectionSpec or a MultiFrequencySelectionSpec (MF-PBT).
replay_buffer (ReplayBufferType | None) – Replay buffer configuration. Off-policy algorithms auto-create a default buffer when this is None.
resume_from_checkpoint (str | None) – Checkpoint to continue an interrupted run from, restoring optimizer state and the hyperparameters it belongs to. Mutually exclusive with load_weights_from.
load_weights_from (str | None) – Checkpoint to warm-start a new run from, taking only the weights. Mutually exclusive with resume_from_checkpoint.
device (str | torch.device) – Torch device (e.g. "cpu", "cuda").
accelerator (Accelerator | None) – Accelerator instance.
hp_config (HyperparameterConfig | None) – Hyperparameter config for HPO. Falls back to the one the mutation spec describes.
actor_network (EvolvableModule | None) – Pre-built actor to hand every agent’s constructor instead of building one from the spec.
critic_network (EvolvableModule | None) – Pre-built critic, for the algorithms that take one.
actor_networks (ModuleDict | None) – Pre-built per-agent actors, for multi-agent algorithms.
critic_networks (ModuleDict | list[EvolvableModule] | None) – Pre-built critics, for the algorithms that take several: per-agent for multi-agent algorithms, a pair for TD3.
kwargs – Accepts the deprecated tournament alias for selection_strategy.

classmethod from_manifest(manifest: str | Path | dict[str, Any] | TrainingManifest, **kwargs: Any) → Self¶

Instantiate a Trainer from a JSON-style manifest or a TrainingManifest instance.

The manifest supplies the algorithm, environment, and training configuration; any trainer-specific construction arguments are passed through as keyword arguments (e.g. device, accelerator, and resume_from_checkpoint for LocalTrainer, or client and api_key for ArenaTrainer).

Parameters:

manifest (str | Path | dict[str, Any] | TrainingManifest) – Path to a YAML/JSON file, or a raw dict, or a TrainingManifest instance.
kwargs – Trainer-specific construction arguments forwarded to the subclass constructor.

Returns:

A fully configured Trainer instance.

Return type:

SelfTrainerT

abstract train() → TrainingLoopReturn | dict[str, Any]¶

Run the training loop.

LocalTrainer runs training locally and returns a tuple of (population, fitnesses) where population is the final evolved population and fitnesses contains each agent’s fitness from the final evaluation round.
ArenaTrainer submits a job to Arena and returns the API response as a dict.

Returns:: The training result, whose type depends on the trainer.
Return type:: TrainingLoopReturn | dict[str, Any]

class agilerl.training.trainer.LocalTrainer(algorithm: AlgoSpec | str, environment: EnvSpecType | str, training: TrainingSpec | None = None, mutation: MutationSpec | None = None, selection_strategy: SelectionStrategySpec | None = None, replay_buffer: ReplayBufferType | None = None, *, resume_from_checkpoint: str | None = None, load_weights_from: str | None = None, hpo: bool = False, device: str | torch.device = 'cpu', accelerator: Accelerator | None = None, hp_config: HyperparameterConfig | None = None, actor_network: EvolvableModule | None = None, critic_network: EvolvableModule | None = None, actor_networks: ModuleDict | None = None, critic_networks: ModuleDict | list[EvolvableModule] | None = None, **kwargs: Any)¶

Local trainer that streamlines the AgileRL evolutionary training process.

Automatically builds the components necessary for RL training with evolutionary HPO from a series of Pydantic models that validate the specified training configuration, and dispatches to the algorithm-specific training loop through LocalTrainer.train(). Handles all of the RL training paradigms available in AgileRL.

Parameters:

algorithm (AlgorithmSpec | str) – An :class:`AlgorithmSpec instance or a string algorithm name.
environment (gym.Env | ParallelEnv) – An RL environment following Gymnasium or PettingZoo API.
training (TrainingSpec | None) – Training parameters. Defaults to TrainingSpec() (1M steps, single agent, no HPO).
mutation (MutationSpec | Mutations | None) – Mutation probabilities and RL hyperparameter ranges. When hp_config is omitted, ranges come from mutation.rl_hp_selection.
selection_strategy (SelectionStrategySpec | None) – Selection strategy driving evolutionary HPO: a TournamentSelectionSpec or a MultiFrequencySelectionSpec (MF-PBT).
replay_buffer (ReplayBufferSpec | ReplayBuffer | None) – Replay buffer configuration.
hpo (bool) – Whether to enable evolutionary HPO using default mutation probabilities, tournament selection, and RL hyperparameters to mutate. Defaults to False.
resume_from_checkpoint (str | None) – Checkpoint to continue an interrupted run from, restoring optimizer state and the hyperparameters it belongs to. Mutually exclusive with load_weights_from.
load_weights_from (str | None) – Checkpoint to warm-start a new run from, taking only the weights. Mutually exclusive with resume_from_checkpoint.
device (str) – Torch device string (e.g. "cpu", "cuda").
accelerator (Accelerator | None) – Accelerator instance.
kwargs – Accepts the deprecated tournament alias for selection_strategy.

train(verbose: bool = True, save_elite: bool = False, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, checkpoint_steps: int | None = None, checkpoint_path: str | None = None, overwrite_checkpoints: bool = False, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None) → TrainingLoopReturn¶

Run a local training job given the passed configuration.

Parameters:

verbose (bool) – If True, print verbose output. Defaults to True.
save_elite (bool) – If True, save the elite agent. Defaults to False.
elite_path (str | None) – The path to save the elite agent. Defaults to None.
wb (bool) – If True, enable Weights & Biases logging. Defaults to False.
tensorboard (bool) – If True, enable TensorBoard logging. Defaults to False.
tensorboard_log_dir (str | None) – The path to save the TensorBoard logs. Defaults to None, which will use the default TensorBoard log directory tensorboard_logs.
checkpoint_steps (int | None) – The number of steps between checkpoints. Defaults to None.
checkpoint_path (str | None) – The path to save the checkpoints. Defaults to None.
overwrite_checkpoints (bool) – If True, overwrite the checkpoint. Defaults to False.
wandb_api_key (str | None) – The Weights & Biases API key. Defaults to None.
wandb_kwargs (dict[str, Any] | None) – The Weights & Biases keyword arguments. Defaults to None.

Returns:

A tuple of (population, fitnesses) where population is the final evolved population and fitnesses contains each agent’s fitness from the final evaluation round.

Return type:

TrainingLoopReturn

Submits AgileRL training jobs to the Arena RLOps platform.

Parameters:

algorithm (AlgoSpec | str) – An :class:`AlgorithmSpec instance or a string algorithm name.
environment (ArenaEnvSpec | str) – An :class:`ArenaEnvSpec instance or a string env name.
training (TrainingSpec) – Training loop parameters.
client (ArenaClient | None) – An authenticated ArenaClient. One is created automatically using the provided API key. Defaults to None.
api_key (str | None) – The Arena API key. Defaults to None.
mutation (MutationSpec | None) – Mutation probabilities and RL-HP ranges. Defaults to None.
selection_strategy (TournamentSelectionSpec | None) – Tournament selection configuration. Arena runs tournament selection only, so MF-PBT is rejected. Defaults to None.
replay_buffer (ReplayBufferType | None) – Replay buffer configuration. Defaults to None.
kwargs – Accepts the deprecated tournament alias for selection_strategy.

Build the manifest and submit the training job to Arena.

Parameters:

resource_id (str | int | None) – Arena cluster type or resource id for the job.
num_nodes (int | None) – The number of nodes to use for training.
project (str | None) – The project to submit the experiment to.
experiment_name (str | None) – The name of the experiment to submit.
reward_file (str | Path | bytes | None) – Python reward module for reasoning dataset jobs.
completion (str | None) – Optional model completion for reward validation.

Returns:

Arena API response.

Return type:

dict[str, Any]

Training Functions¶

If you are using a Gym-style environment, our on- and off-policy training functions return a population of trained agents and logged training metrics.

agilerl.training.train_off_policy.train_off_policy(env: Env | VectorEnv, env_name: str, algo: str, pop: list[DQN | RainbowDQN | DDPG | TD3], memory: ReplayBuffer | PrioritizedReplayBuffer | MultiStepReplayBuffer, init_hp: dict[str, Any] | None = None, mut_p: dict[str, Any] | None = None, max_steps: int = 1000000, evo_steps: int = 10000, eval_steps: int | None = None, eval_loop: int = 1, learning_delay: int = 0, eps_start: float = 1.0, eps_end: float = 0.01, eps_decay: float = 0.999, target: float | None = None, n_step_memory: MultiStepReplayBuffer | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, checkpoint: int | None = None, checkpoint_path: str | None = None, overwrite_checkpoints: bool = False, save_elite: bool = False, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, verbose: bool = True, accelerator: Accelerator | None = None, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None) → tuple[list[DQN | RainbowDQN | DDPG | TD3], list[float]]¶

Run the general online off-policy RL training; returns trained population of agents and their fitnesses.

Parameters:

env (Gym-style environment) – The environment to train in. Can be vectorized.
env_name (str) – Environment name
algo (str) – RL algorithm name
pop (list[SingleAgentAlgorithm]) – Population of agents
memory (object) – Experience Replay Buffer
init_hp (dict, optional) – Dictionary containing initial hyperparameters, defaults to None
mut_p (dict, optional) – Dictionary containing mutation parameters, defaults to None
max_steps (int, optional) – Maximum number of steps in environment, defaults to 1000000
evo_steps (int, optional) – Evolution frequency (steps), defaults to 10000
eval_steps (int, optional) – Number of evaluation steps per episode. If None, will evaluate until environment terminates or truncates. Defaults to None
eval_loop (int, optional) – Number of evaluation episodes, defaults to 1
learning_delay (int, optional) – Steps in environment before starting learning, defaults to 0
eps_start (float, optional) – Maximum exploration - initial epsilon value, defaults to 1.0
eps_end (float, optional) – Minimum exploration - final epsilon value, defaults to 0.1
eps_decay (float, optional) – Epsilon decay per episode, defaults to 0.995
target (float, optional) – Target score for early stopping, defaults to None
n_step_memory (object, optional) – Multi-step Experience Replay Buffer to be used alongside Prioritized ERB, defaults to None
selection_strategy (object, optional) – selection strategy driving population evolution. A TournamentSelection or MultiFrequencySelection (MF-PBT) object, defaults to None
tournament (object, optional) – Deprecated alias for selection_strategy (a TournamentSelection object), defaults to None
mutation (object, optional) – Mutation object, defaults to None
checkpoint (int, optional) – Checkpoint frequency (steps), defaults to None
checkpoint_path (str, optional) – Location to save checkpoint, defaults to None
overwrite_checkpoints (bool, optional) – Overwrite previous checkpoints during training, defaults to False
save_elite (bool, optional) – Boolean flag indicating whether to save elite member at the end of training, defaults to False
elite_path (str, optional) – Location to save elite agent, defaults to None
wb (bool, optional) – Weights & Biases tracking, defaults to False
tensorboard (bool, optional) – TensorBoard tracking, defaults to False
tensorboard_log_dir (str, optional) – Directory for TensorBoard logs, defaults to None
verbose (bool, optional) – Display training stats, defaults to True
accelerator (accelerate.Accelerator(), optional) – Accelerator for distributed computing, defaults to None
wandb_api_key (str, optional) – API key for Weights & Biases, defaults to None
wandb_kwargs (dict, optional) – Additional kwargs to pass to wandb.init()

Returns:

Trained population of agents and their fitnesses

Return type:

tuple[list[SingleAgentAlgorithm], list[float]]

agilerl.training.train_on_policy.train_on_policy(env: Env | VectorEnv, env_name: str, algo: str, pop: list[PPO], init_hp: dict[str, Any] | None = None, mut_p: dict[str, Any] | None = None, max_steps: int = 1000000, evo_steps: int = 10000, eval_steps: int | None = None, eval_loop: int = 1, target: float | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, checkpoint: int | None = None, checkpoint_path: str | None = None, overwrite_checkpoints: bool = False, save_elite: bool = False, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, verbose: bool = True, accelerator: Accelerator | None = None, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None, collect_rollouts_fn: Callable[[...], tuple[list[float], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], ndarray[tuple[int, ...], dtype[_ScalarType_co]], dict[str, Any]]] | None = None) → tuple[list[PPO], list[float]]¶

Run the general on-policy RL training; returns trained population of agents and their fitnesses.

Parameters:

env (Gym-style environment) – The environment to train in. Can be vectorized.
env_name (str) – Environment name
algo (str) – RL algorithm name
pop (list[SingleAgentAlgorithm]) – Population of agents
init_hp (dict, optional) – Dictionary containing initial hyperparameters, defaults to None
mut_p (dict, optional) – Dictionary containing mutation parameters, defaults to None
max_steps (int, optional) – Maximum number of steps in environment, defaults to 1000000
evo_steps (int, optional) – Evolution frequency (steps), defaults to 10000
eval_steps (int, optional) – Number of evaluation steps per episode. If None, will evaluate until environment terminates or truncates. Defaults to None
eval_loop (int, optional) – Number of evaluation episodes, defaults to 1
target (float, optional) – Target score for early stopping, defaults to None
selection_strategy (object, optional) – selection strategy driving population evolution. A TournamentSelection or MultiFrequencySelection (MF-PBT) object, defaults to None
tournament (object, optional) – Deprecated alias for selection_strategy (a TournamentSelection object), defaults to None
mutation (object, optional) – Mutation object, defaults to None
checkpoint (int, optional) – Checkpoint frequency (steps), defaults to None
checkpoint_path (str, optional) – Location to save checkpoint, defaults to None
overwrite_checkpoints (bool, optional) – Overwrite previous checkpoints during training, defaults to False
save_elite (bool, optional) – Boolean flag indicating whether to save elite member at the end of training, defaults to False
elite_path (str, optional) – Location to save elite agent, defaults to None
wb (bool, optional) – Weights & Biases tracking, defaults to False
tensorboard (bool, optional) – TensorBoard tracking, defaults to False
tensorboard_log_dir (str, optional) – Directory for TensorBoard logs, defaults to None
verbose (bool, optional) – Display training stats, defaults to True
accelerator (accelerate.Accelerator(), optional) – Accelerator for distributed computing, defaults to None
wandb_api_key (str, optional) – API key for Weights & Biases, defaults to None
wandb_kwargs – Additional kwargs to pass to wandb.init()
collect_rollouts_fn (Callable or None, optional) – Optional function used to collect rollouts. If None and agents use a rollout buffer, a default function will be selected based on whether the agent is recurrent.

Returns:

Trained population of agents and their fitnesses

Return type:

list[SingleAgentAlgorithm], list[float]

If you are training on static, offline data, you can use our offline RL training function.

agilerl.training.train_offline.train_offline(env: VectorEnv, env_name: str, algo: str, pop: list[CQN], memory: ReplayBuffer, init_hp: dict[str, Any] | None = None, mut_p: dict[str, Any] | None = None, max_steps: int = 1000000, evo_steps: int = 10000, eval_steps: int | None = None, eval_loop: int = 1, target: float | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, checkpoint: int | None = None, checkpoint_path: str | None = None, overwrite_checkpoints: bool = False, save_elite: bool = False, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, verbose: bool = True, accelerator: Accelerator | None = None, dataset: File | Mapping[Any, Any] | None = None, minari_dataset_id: str | None = None, remote: bool = False, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None) → tuple[list[CQN], list[float]]¶

Run the general offline RL training; returns trained population of agents and their fitnesses.

Parameters:

env (gym.vector.VectorEnv) – The vectorized environment used to evaluate the population
env_name (str) – Environment name
algo (str) – RL algorithm name
pop (list[CQN]) – Population of agents
memory (ReplayBuffer) – Experience Replay Buffer
init_hp (dict, optional) – Dictionary containing initial hyperparameters, defaults to None
mut_p (dict, optional) – Dictionary containing mutation parameters, defaults to None
max_steps (int, optional) – Maximum number of steps in environment, defaults to 1000000
evo_steps (int, optional) – Evolution frequency (steps), defaults to 10000
eval_steps (int, optional) – Number of evaluation steps per episode. If None, will evaluate until environment terminates or truncates. Defaults to None
eval_loop (int, optional) – Number of evaluation episodes, defaults to 1
target (float, optional) – Target score for early stopping, defaults to None
selection_strategy (object, optional) – selection strategy driving population evolution. A TournamentSelection or MultiFrequencySelection (MF-PBT) object, defaults to None
tournament (object, optional) – Deprecated alias for selection_strategy (a TournamentSelection object), defaults to None
mutation (object, optional) – Mutation object, defaults to None
checkpoint (int, optional) – Checkpoint frequency (steps), defaults to None
checkpoint_path (str, optional) – Location to save checkpoint, defaults to None
overwrite_checkpoints (bool, optional) – Overwrite previous checkpoints during training, defaults to False
save_elite (bool, optional) – Boolean flag indicating whether to save elite member at the end of training, defaults to False
elite_path (str, optional) – Location to save elite agent, defaults to None
wb (bool, optional) – Weights & Biases tracking, defaults to False
tensorboard (bool, optional) – TensorBoard tracking, defaults to False
tensorboard_log_dir (str, optional) – Directory for TensorBoard logs, defaults to None
verbose (bool, optional) – Display training stats, defaults to True
accelerator (accelerate.Accelerator(), optional) – Accelerator for distributed computing, defaults to None
dataset (OfflineDataset | None, optional) – Open HDF5 handle or in-memory array mapping, defaults to None
minari_dataset_id (str, optional) – Minari dataset ID for loading data, defaults to None
remote (bool, optional) – Load Minari dataset from remote, defaults to False
wandb_api_key (str, optional) – API key for Weights & Biases, defaults to None
wandb_kwargs (dict, optional) – Additional kwargs to pass to wandb.init()

Returns:

Trained population of agents and their fitnesses

Return type:

tuple[list[CQN], list[float]]

The multi-agent off-policy and on-policy training functions handle PettingZoo-style environments and multi-agent algorithms.

agilerl.training.train_multi_agent_off_policy.train_multi_agent_off_policy(env: ParallelEnv | AsyncPettingZooVecEnv, env_name: str, algo: str, pop: list[MADDPG | MATD3], memory: ReplayBuffer, sum_scores: bool = True, init_hp: dict[str, Any] | None = None, mut_p: dict[str, Any] | None = None, max_steps: int = 50000, evo_steps: int = 25, eval_steps: int | None = None, eval_loop: int = 1, learning_delay: int = 0, target: float | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, checkpoint: int | None = None, checkpoint_path: str | None = None, overwrite_checkpoints: bool = False, save_elite: bool = False, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, verbose: bool = True, accelerator: Accelerator | None = None, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None) → tuple[list[MADDPG | MATD3], list[float] | list[dict[str, float]]]¶

Run the general off-policy multi-agent RL training; returns trained population of agents and their fitnesses.

Parameters:

env (Gym-style environment) – The environment to train in. Can be vectorized.
env_name (str) – Environment name
algo (str) – RL algorithm name
pop (list[MADDPG | MATD3]) – Population of agents
memory (ReplayBuffer) – Experience Replay Buffer
sum_scores (bool, optional) – Boolean flag indicating whether to sum sub-agents scores, typically True for co-operative environments, defaults to True
init_hp (dict) – Dictionary containing initial hyperparameters.
mut_p (dict, optional) – Dictionary containing mutation parameters, defaults to None
max_steps (int, optional) – Maximum number of steps in environment, defaults to 50000
evo_steps (int, optional) – Evolution frequency (steps), defaults to 25
eval_steps (int, optional) – Number of evaluation steps per episode. If None, will evaluate until environment terminates or truncates. Defaults to None
eval_loop (int, optional) – Number of evaluation episodes, defaults to 1
learning_delay (int, optional) – Steps in environment before starting learning, defaults to 0
target (float, optional) – Target score for early stopping, defaults to None
selection_strategy (object, optional) – selection strategy driving population evolution. A TournamentSelection or MultiFrequencySelection (MF-PBT) object, defaults to None
tournament (object, optional) – Deprecated alias for selection_strategy (a TournamentSelection object), defaults to None
mutation (object, optional) – Mutation object, defaults to None
checkpoint (int, optional) – Checkpoint frequency (steps), defaults to None
checkpoint_path (str, optional) – Location to save checkpoint, defaults to None
overwrite_checkpoints (bool, optional) – Overwrite previous checkpoints during training, defaults to False
save_elite (bool, optional) – Boolean flag indicating whether to save elite member at the end of training, defaults to False
elite_path (str, optional) – Location to save elite agent, defaults to None
wb (bool, optional) – Weights & Biases tracking, defaults to False
tensorboard (bool, optional) – TensorBoard tracking, defaults to False
tensorboard_log_dir (str, optional) – Directory for TensorBoard logs, defaults to None
verbose (bool, optional) – Display training stats, defaults to True
accelerator (accelerate.Accelerator(), optional) – Accelerator for distributed computing, defaults to None
wandb_api_key (str, optional) – API key for Weights & Biases, defaults to None
wandb_kwargs (dict, optional) – Additional kwargs to pass to wandb.init()

Returns:

Trained population of agents and their fitnesses. Fitnesses are per-agent dictionaries when sum_scores is False.

Return type:

tuple[list[MADDPG | MATD3], list[float] | list[dict[str, float]]]

agilerl.training.train_multi_agent_on_policy.train_multi_agent_on_policy(env: ParallelEnv | PettingZooVecEnv, env_name: str, algo: str, pop: list[IPPO], sum_scores: bool = True, init_hp: dict[str, Any] | None = None, mut_p: dict[str, Any] | None = None, max_steps: int = 50000, evo_steps: int = 25, eval_steps: int | None = None, eval_loop: int = 1, target: float | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, checkpoint: int | None = None, checkpoint_path: str | None = None, overwrite_checkpoints: bool = False, save_elite: bool = False, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, verbose: bool = True, accelerator: Accelerator | None = None, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None) → tuple[list[IPPO], list[float] | list[dict[str, float]]]¶

Run the general on-policy multi-agent RL training; returns trained population of agents and their fitnesses.

Parameters:

env (Gym-style environment) – The environment to train in. Can be vectorized.
env_name (str) – Environment name
algo (str) – RL algorithm name
pop (list[object]) – Population of agents
sum_scores (bool, optional) – Boolean flag indicating whether to sum sub-agents scores, typically True for co-operative environments, defaults to True
init_hp (dict) – Dictionary containing initial hyperparameters.
mut_p (dict, optional) – Dictionary containing mutation parameters, defaults to None
max_steps (int, optional) – Maximum number of steps in environment across the entire population, defaults to 50000
evo_steps (int, optional) – Evolution frequency (steps), defaults to 25
eval_steps (int, optional) – Number of evaluation steps per episode. If None, will evaluate until environment terminates or truncates. Defaults to None
eval_loop (int, optional) – Number of evaluation episodes, defaults to 1
target (float, optional) – Target score for early stopping, defaults to None
selection_strategy (object, optional) – selection strategy driving population evolution. A TournamentSelection or MultiFrequencySelection (MF-PBT) object, defaults to None
tournament (object, optional) – Deprecated alias for selection_strategy (a TournamentSelection object), defaults to None
mutation (object, optional) – Mutation object, defaults to None
checkpoint (int, optional) – Checkpoint frequency (steps), defaults to None
checkpoint_path (str, optional) – Location to save checkpoint, defaults to None
overwrite_checkpoints (bool, optional) – Overwrite previous checkpoints during training, defaults to False
save_elite (bool, optional) – Boolean flag indicating whether to save elite member at the end of training, defaults to False
elite_path (str, optional) – Location to save elite agent, defaults to None
wb (bool, optional) – Weights & Biases tracking, defaults to False
tensorboard (bool, optional) – TensorBoard tracking, defaults to False
tensorboard_log_dir (str, optional) – Directory for TensorBoard logs, defaults to None
verbose (bool, optional) – Display training stats, defaults to True
accelerator (accelerate.Accelerator(), optional) – Accelerator for distributed computing, defaults to None
wandb_api_key (str, optional) – API key for Weights & Biases, defaults to None
wandb_kwargs (dict, optional) – Additional kwargs to pass to wandb.init()

Finally, if you are training a LLM, you can use our LLM training functions. train_llm_rollout runs online RL over rollout (generate-and-score) environments and should be used with GRPO, PPO or REINFORCE; it drives multi-turn rollouts, and single-turn reasoning is the max_turns=1 case. train_llm_dataset runs offline, teacher-forced training over a DatasetEnv dataloader; the algorithm of the population selects the regime, with DPO for pairwise preference data and SFT for supervised fine-tuning on static data.

agilerl.training.llm.train_llm_rollout(pop: list[SupportedRollout], max_turns: int, env_factory: Callable[[], RolloutHarness], init_hp: dict[str, Any] | None = None, max_steps: int = 32768, save_elite: bool | None = None, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, csv: bool = False, csv_log_dir: str | None = None, evo_steps: int | None = None, checkpoint_steps: int | None = None, checkpoint_path: str | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None, evaluation_interval: int = 50, eval_loop: int = 1, max_reward: float | None = None, verbose: bool = True, max_wall_seconds: float | None = None, io_timeout_s: float | None = 600.0, completion_logging: CompletionLoggingConfig | None = None) → tuple[list[SupportedRollout], Any]¶

Train a population of LLM agents over rollout (generate-and-score) environments.

Collects token-level episodes (reset returns (obs, info), repeated get_action / step (full completion tensor), then get_episode_data), then runs turn-level updates. For a RolloutHarness with max_model_len set, a trajectory whose cumulative prompt would overflow the context is stopped with truncated=True.

Parameters:

pop (list[SupportedRollout]) – Population of LLMPPO, LLMREINFORCE or GRPO agents to finetune.
max_turns (int) – Maximum interaction turns per episode.
env_factory (Callable[[], RolloutHarness]) – Zero-arg factory that returns a fresh env for each trajectory rollout. Required to ensure trajectory state isolation.
init_hp (dict, optional) – Initial hyperparameters.
max_steps (int) – Progress-bar budget in sample steps, defaults to 32768.
save_elite (bool, optional) – Whether to save the elite checkpoint, defaults to None.
elite_path (str, optional) – Directory for checkpoints, defaults to None.
wb (bool, optional) – Whether to log to Weights and Biases, defaults to False.
tensorboard (bool, optional) – Whether to log to TensorBoard, defaults to False.
tensorboard_log_dir (str, optional) – Directory for TensorBoard event files, defaults to None.
csv (bool, optional) – Whether to log aggregate metrics to CSV, defaults to False.
csv_log_dir (str, optional) – Path for the CSV file, defaults to None.
evo_steps (int, optional) – Steps between evolution (requires a selection strategy and mutation).
checkpoint_steps (int, optional) – Save checkpoint every N outer iterations when no evolution.
checkpoint_path (str, optional) – Directory for periodic checkpoints; falls back to elite_path.
selection_strategy (SelectionStrategyProtocol | None, optional) – Selection strategy driving evolution, defaults to None.
tournament (TournamentSelection, optional) – Deprecated alias for selection_strategy, defaults to None.
mutation (Mutations, optional) – Mutation operator for evolution, defaults to None.
wandb_api_key (str, optional) – W&B API key, defaults to None.
wandb_kwargs (dict, optional) – Additional kwargs forwarded to wandb.init().
evaluation_interval (int, optional) – How often to call agent.test on a fresh environment from env_factory.
eval_loop (int, optional) – Episodes averaged per evaluation for the HPO fitness score; defaults to 1 (matching the other trainers). Raise it for less noisy tournament selection at higher eval cost.
max_reward (float, optional) – If set, adds accuracy metric vs this threshold.
verbose (bool) – Progress bar and periodic train summaries, defaults to True.
max_wall_seconds (float | None) – Stop after this wall-clock duration (seconds); None disables.
io_timeout_s (float | None) – Backstop deadline for one concurrent round of env round-trips; a hung env or stalled transport raises TimeoutError rather than blocking the batch forever. Defaults to 600 s; None disables it. Forwarded to RolloutCollector.
completion_logging (CompletionLoggingConfig | None) – Sample decoded prompts, completions and rewards each iteration. Written every interval iterations to the console (verbose), a W&B completions table (wb) and jsonl_path. One write per interval step includes every population member. The latest samples are logged at error level if training raises. None disables it.

Returns:

The finetuned population and its last recorded fitnesses.

Return type:

tuple[list[SupportedRollout], Any]

agilerl.training.llm.train_llm_dataset(pop: list[SupportedDataset], env: DatasetEnv | None = None, env_fn: Callable[[], DatasetEnv] | None = None, init_hp: dict[str, Any] | None = None, save_elite: bool | None = None, elite_path: str | None = None, wb: bool = False, tensorboard: bool = False, tensorboard_log_dir: str | None = None, csv: bool = False, csv_log_dir: str | None = None, evo_steps: int | None = None, checkpoint_steps: int | None = None, checkpoint_path: str | None = None, selection_strategy: SelectionStrategyProtocol | None = None, tournament: TournamentSelection | None = None, mutation: Mutations | None = None, wandb_api_key: str | None = None, wandb_kwargs: dict[str, Any] | None = None, evaluation_interval: int = 10, verbose: bool = True, max_steps: int | None = None, num_epochs: int | None = None) → tuple[list[SupportedDataset], Any]¶

Train a population of DPO or SFT agents over a DatasetEnv dataloader.

Each training step draws a labelled batch from the dataset environment. The algorithm of pop[0] selects the regime: DPO minimises a pairwise preference loss over chosen/rejected pairs, while SFT minimises the response cross-entropy. Both share evolution, checkpointing, and metrics logging.

Parameters:

pop (list[SupportedDataset]) – Population of DPO or SFT agents to finetune.
env (DatasetEnv | None) – Shared dataset environment that yields labelled batches.
env_fn (Callable[[], DatasetEnv] | None) – Optional factory that creates one dataset environment per agent.
init_hp (dict[str, Any] | None) – Initial hyperparameters for logging and defaults.
save_elite (bool | None) – Whether to save the elite checkpoint during evolution.
elite_path (str | None) – Path used for checkpoint saving.
wb (bool) – Whether to log metrics to Weights and Biases.
tensorboard (bool) – Whether to log to TensorBoard.
tensorboard_log_dir (str | None) – Directory for TensorBoard event files.
csv (bool) – Whether to log aggregate metrics to CSV.
csv_log_dir (str | None) – Path for the CSV file.
evo_steps (int | None) – Number of outer iterations between evolution steps.
checkpoint_steps (int | None) – Number of iterations between checkpoint saves when evolution is disabled.
checkpoint_path (str | None) – Directory for periodic checkpoints; falls back to elite_path.
selection_strategy (SelectionStrategyProtocol | None) – Selection strategy driving evolution.
tournament (TournamentSelection | None) – Deprecated alias for selection_strategy.
mutation (Mutations | None) – Mutation operator used during evolution.
wandb_api_key (str | None) – Optional W&B API key.
wandb_kwargs (dict[str, Any] | None) – Additional kwargs forwarded to wandb.init().
evaluation_interval (int) – Frequency (iterations) for evaluation.
verbose (bool) – Whether to print periodic training summaries.
max_steps (int | None) – Maximum step budget; defaults to dataset-driven length.
num_epochs (int | None) – Number of epochs to run; takes precedence over max_steps.

Returns:

The finetuned population and its last recorded fitnesses.

Return type:

tuple[list[SupportedDataset], Any]