LLM Proximal Policy Optimization (LLM PPO)¶

PPO (Proximal Policy Optimization) is a policy-gradient method that keeps updates inside a clipped trust region. LLMPPO adapts this idea to causal language models and is designed for both single-turn and multi-turn fine-tuning.

In AgileRL, the implementation is turn-aware:

Turn-level credit assignment: each generated turn is treated as one RL action, with discounting across turns.
Actor-critic optimization: policy and value adapters are updated jointly, with clipped policy/value losses plus entropy regularization.
Single-turn and multi-turn parity: single-turn prompting is treated as the special case where all action tokens belong to turn 0.

This algorithm can therefore be used in multi-turn agentic finetuning or single-turn reasoning tasks.

Variance Reduction¶

LLM PPO reduces the variance of its policy gradient with a learned value baseline: a value head is trained alongside the policy to predict expected return, and the advantage is computed as return minus the value estimate (GAE-style when discounting across turns). Compared with the group-relative normalization used by GRPO and the Return Batch Normalization (ReBN) used by LLM REINFORCE, this is the most expressive variance reducer (state-conditioned, no group requirement) but also the most expensive: an extra adapter must be trained, the baseline is biased while the value head is catching up, and value-fit pathologies are an extra failure mode to debug.

Example¶

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from agilerl.algorithms import LLMPPO

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-0.5B",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B")

agent = LLMPPO(
    actor_network=model,
    pad_token_id=tokenizer.eos_token_id,
    pad_token=tokenizer.eos_token,
    device="cuda" if torch.cuda.is_available() else "cpu",
    batch_size=8,
    update_epochs=1,
    clip_coef=0.2,
    max_output_tokens=128,
    max_model_len=1024,
)

Training¶

The typical training entry point is train_llm_rollout in agilerl.training.llm. Single-turn reasoning is the max_turns=1 case of the same function.

from agilerl.training.llm import train_llm_rollout
from agilerl.llm_envs import RolloutHarness

def reward_fn(completion: str, answer: str, question: str) -> float:
    del question
    return float(answer.lower() in completion.lower())

# 1) Single-turn / reasoning datasets (a prompt dataset is just an env, max_turns=1)
class QADataset:
    """Single-turn dataset env: serve a question on reset, score it on step."""

    def __init__(self, questions, answers, reward_fn, prompt_builder,
                 test_questions=None, test_answers=None):
        self.questions, self.answers = questions, answers
        self.test_questions, self.test_answers = test_questions, test_answers
        self.reward_fn, self.prompt_builder = reward_fn, prompt_builder
        self._cursor, self._split = 0, ""

    @property
    def dataset_size(self) -> int:
        return len(self.questions)

    def reset(self, seed=None, *, row_index=None, evaluation=None):
        if evaluation and self.test_questions is not None:
            qs, ans, split = self.test_questions, self.test_answers, "eval"
        else:
            qs, ans, split = self.questions, self.answers, "train"
        if row_index is None:
            if split != self._split:
                self._cursor, self._split = 0, split
            row_index, self._cursor = self._cursor, self._cursor + 1
        self._q, self._a = qs[row_index % len(qs)], ans[row_index % len(ans)]
        return self.prompt_builder(self._q), {}

    def step(self, action):
        return "", float(self.reward_fn(action, self._a, self._q)), True, False, {}

env_factory = lambda: RolloutHarness.local(
    QADataset(
        questions=["2+2?", "Capital of France?"],
        answers=["4", "Paris"],
        reward_fn=reward_fn,
        prompt_builder=lambda question: f"Q: {question}\nA:",
        test_questions=["3+3?"],
        test_answers=["6"],
    ),
    tokenizer,
    max_turns=1,
    pad_id=tokenizer.eos_token_id,
    apply_chat_template=True,
    max_model_len=1024,
)

trained_pop = train_llm_rollout(
    pop=[agent],
    max_turns=1,
    env_factory=env_factory,
    max_steps=2000,
    evaluation_interval=50,
)

# 2) Multi-turn text environments (each rollout drives its own in-process env)
class ToyRolloutEnv:
    def reset(self, seed=None):
        del seed
        return "Start: What is 2+2?", {}

    def step(self, action: str):
        reward = 1.0 if "4" in action else 0.0
        return "Done.", reward, True, False, {"correct": bool(reward)}

env_factory = lambda: RolloutHarness.local(
    ToyRolloutEnv(),
    tokenizer,
    max_turns=4,
    pad_id=tokenizer.eos_token_id,
    max_model_len=1024,
    max_output_tokens=128,
)

trained_pop = train_llm_rollout(
    pop=[agent],
    max_turns=4,
    env_factory=env_factory,
    max_steps=2000,
    evaluation_interval=50,
)

Saving and Loading Agents¶

To save an agent, use the save_llm_checkpoint function:

from agilerl.utils.utils import save_llm_checkpoint

save_llm_checkpoint(agent, "path/to/checkpoint")

As with other AgileRL LLM algorithms, loading is done with Hugging Face from_pretrained APIs for the base model and adapter.

Parameters¶

class agilerl.algorithms.ppo_llm.PPO(*args: Any, **kwargs: Any)¶

Turn-level PPO for LLM finetuning with actor/reference adapters.

Each generation sequence (turn) is treated as a single RL action. GAE discounts between turns, not between tokens within a turn. Single-turn is the special case where all action tokens share turn 0.

Parameters:

pad_token_id (int) – Token id used for sequence padding.
pad_token (str) – Padding token string.
model_name (str | None, optional) – HF model name or local path used when building internally.
actor_network (PreTrainedModel | PeftModel | None, optional) – Pre-built actor model. If omitted, model_name is used.
model_config (dict[str, Any] | None, optional) – Extra kwargs passed when constructing a model from model_name.
hp_config (HyperparameterConfig | None, optional) – Hyperparameter mutation configuration.
index (int, optional) – Population index used by evolutionary workflows.
batch_size (int, optional) – Batch size used for PPO updates.
beta (float, optional) – KL penalty coefficient against the reference policy.
vf_coef (float, optional) – Value loss coefficient.
clip_coef (float, optional) – PPO clipping coefficient.
gamma (float, optional) – Discount factor across turns.
gae_lambda (float, optional) – GAE lambda used for turn-level advantage estimation.
critic_warmup_steps (int, optional) – First learn steps that update only the critic, on Monte Carlo returns (GAE lambda = 1), leaving the policy fixed. critic_warmup_steps_done counts them and is checkpointed.
lr_actor (float, optional) – Actor learning rate.
lr_critic (float | None, optional) – Critic/value-head learning rate. If None, lr_actor is used.
max_grad_norm (float, optional) – Gradient clipping norm.
share_grad_clip (bool, optional) – Clip actor and critic grads by their combined norm.
update_epochs (int, optional) – Number of PPO epochs per update.
temperature (float, optional) – Sampling temperature for generation.
repetition_penalty (float, optional) – Repetition penalty used during generation.
top_p (float, optional) – Nucleus sampling threshold.
top_k (int, optional) – Top-k sampling threshold.
min_p (float, optional) – Minimum probability cutoff for sampling.
use_separate_reference_adapter (bool, optional) – Whether to keep a separate reference adapter.
calc_position_embeddings (bool, optional) – Whether to compute position embeddings.
micro_batch_size_per_gpu (int | None, optional) – Optional target micro-batch size per GPU.
mini_batch_size (int | None, optional) – Per-rank trajectories covered by one optimizer step. None uses batch_size / world_size. gradient_accumulation_steps is derived as mini_batch_size / micro_batch_size_per_gpu.
max_output_tokens (int | None, optional) – Maximum newly generated tokens per completion.
min_output_tokens (int | None, optional) – Minimum newly generated tokens per completion.
max_model_len (int, optional) – Maximum model context length.
max_row_tokens (int | None, optional) – Longest training row (prompt plus completion) the memory estimate sizes for. None uses max_model_len.
hf_generate_chunk_size (int | None, optional) – Number of prompts per HuggingFace generation chunk. Ignored when colocated.
lora_config (LoraConfig | None, optional) – LoRA configuration.
cosine_lr_schedule_config (CosineLRScheduleConfig | None, optional) – Warmup-cosine schedule stepped once per learn call; the actor’s starts after critic_warmup_steps.
fsdp_config (FSDPConfig | None, optional) – FSDP2 sharding settings for distributed runs, defaults to None
device (str, optional) – Device for accelerated computing, ‘cpu’ or ‘cuda’, defaults to ‘cpu’
wrap (bool, optional) – Whether to wrap models for distributed execution.
clone (bool, optional) – Whether this instance is being created as a clone.
offload_trainer_during_rollout (bool, optional) – For colocated vLLM, offload the trainer’s own base to CPU during rollout (and bring it back for the training step) so the rollout engine and the trainer never both hold a base on the GPU. Defaults to True; inert without colocated vLLM, and disabled under FSDP2 sharding.
vllm_config (VLLMConfig | None, optional) – vLLM runtime configuration.
seed (int, optional) – Random seed.
turn_level_clip (bool, optional) – Legacy gate for per-turn ratio clipping, honored only when importance_sampling_level="auto". Superseded by importance_sampling_level.
importance_sampling_level (Literal["auto", "token", "turn", "trajectory"], optional) – IS / ratio-pooling level for the policy surrogate, orthogonal to advantage_granularity. "token" clips per token; "turn" pools the ratio per turn; "trajectory" pools over the whole completion; the paired advantage is pooled to the same bucket. "auto" (default) uses the GAE granularity when turn_level_clip is set, else token. Turn/trajectory pooling couples a unit’s tokens and cannot be token-chunked in the fused kernel, so set use_liger_loss=False there (the standard path is always memory-bounded).
advantage_granularity (Literal["turn", "token", "auto"], optional) – PPO action granularity. "turn" enforces turn-level updates, "token" enforces token-level updates, and "auto" uses token-level only when all samples are single-turn.
turn_ratio_pooling (Literal["sum", "mean"], optional) – Reduction used to pool per-token log-ratios into a per-turn ratio when the importance-sampling level is "turn" (the default "auto" resolves to turn for multi-turn batches); ignored at token/trajectory level. "sum" (default) yields the product ratio per turn — the standard, paper-aligned per-turn importance weight. "mean" yields a length-normalized geometric-mean ratio (GSPO-style); reach for it on long or highly variable-length turns, where the product ratio is far outside the clip band on every turn and saturates the clipped surrogate — length-normalizing keeps the per-turn ratio in range so the surrogate stays informative.
turn_value_reduction (str, optional) – Aggregation used to map token critic values to turn values. "mean" reproduces existing behavior, "final_value" uses the final action token value in each turn.
whiten_advantages (bool, optional) – Whether to whiten computed advantages before PPO optimization.
gradient_checkpointing (bool, optional) – Enable gradient checkpointing.
torch_compiler (str | None, optional) – Optional torch compile mode.
cast_logprobs_to_fp32 (bool, optional) – When True (default), run the per-token log-prob reduction (gather / logsumexp) in fp32 before casting back to the input dtype, for numerically stable log-probs. False runs it in the input dtype, saving a little memory at the cost of a per-token bf16 quantisation error that can bias importance-sampling ratios.
chunk_rows (int | None, optional) – Primary chunk-size setting for fused logit tiles. Applies to both standard and Liger paths.
use_liger_loss (bool, optional) – Use the Liger fused policy loss, defaults to True (requires liger-kernel). Recommended for PPO: via AgileRL’s LigerFusedLinearPolicyLossFunction (not the upstream Liger GRPO kernel), it is roughly memory-neutral with a mild speedup that grows with sequence length (~1.1x at long sequences) at token-level IS. Separate from the Liger model patches (fused RMSNorm/RoPE/SwiGLU), which apply whenever liger-kernel is installed.
fuse_actor_critic_pass (bool | None, optional) – True runs the actor and critic rows of a micro-batch in one forward and backward; False runs them as two passes, so only one pass’s activations are live. None (default) fuses when the memory estimate of the fused pass fits this GPU’s total memory, and always fuses off CUDA.
quantization_config (BitsAndBytesConfig | None, optional) – Optional transformers.BitsAndBytesConfig for loading the base model in 4-/8-bit (QLoRA). lm_head is kept unquantized so the fused-linear-logprob path stays numerically exact.
activation_offload (bool, optional) – When True, run the training forward inside torch.autograd.graph.save_on_cpu so tensors saved for backward live in pinned host RAM instead of GPU memory. Trades PCIe bandwidth for GPU memory (the win grows with sequence length); a no-op during rollout / reference forwards.
moe_lora_recompute (bool | None, optional) – Recompute routed-expert LoRA activations in backward on frozen packed base weights. None (default) recomputes only outside activation-checkpointed blocks.
vllm_importance_sampling_correction (bool, optional) – When True (default) and colocated, correct the rollout/trainer log-prob mismatch by weighting each training token by clamp(exp(trainer - sampling), max=vllm_importance_sampling_cap). Active only for training rollouts; inert on the HuggingFace path and at eval.
vllm_importance_sampling_cap (float, optional) – Upper clamp on the vLLM importance-sampling ratio (default 2.0), bounding the correction weight to limit variance from outlier tokens. Must be > 0.
vllm_max_logprob_gap (float, optional) – Log a warning when a learn step’s mean |trainer - vLLM| per-token log-prob gap exceeds this, in nats (default 0.1). bf16 engines usually sit near 0.01-0.03.
vllm_max_clip_fraction (float, optional) – Log a warning when the fraction of action tokens whose trainer/vLLM ratio reaches vllm_importance_sampling_cap exceeds this (default 0.02). bf16 engines usually stay under 0.01.
use_sequence_packing (bool, optional) – Opt in to padding-free sequence packing for the gradient forward. Only honoured under a FlashAttention-2 / FlexAttention backend, otherwise inert; the fused actor+critic pass packs only under FlexAttention, and the no-grad reference/old-value pass stays padded.
lora_tower_target_modules (list[str] | None, optional) – Tower and connector modules of a VL model to adapt as well. lora_config.target_modules always land on the language model only. Defaults to None.
rollout_strip_multimodal_towers (bool | list[VllmModality], optional) – strip_multimodal_towers of a vLLM engine outside this process that serves the adapter; tower LoRA skips the towers it leaves unbuilt. A colocated engine uses vllm_config instead. Defaults to False.
kl_clamp (float | None, optional) – Per-token bound on the K3 KL penalty. Tokens where the policy sits below the reference with K3 above the bound (10 is about 2.6 nats below) get no KL gradient and add nothing to the loss. The kl metric stays unbounded; kl_clamp_frac reports the share of action tokens past the bound. None disables the bound, defaults to 10.0. Must be > 0.
logit_bias (dict[int, float] | None, optional) – Token id to logit offset vLLM adds when it samples; the trainer’s log-probs add the same. None adds none.

add_scores(scores: Sequence[float | list[float]]) → None¶

Add scores to the metrics.

Parameters:: scores (Sequence[float | list[float]]) – List of scores (or per-agent score rows) to add.

apply_lr() → None¶: Set lr / lr_critic as the optimizer’s peaks at the schedule’s current step.

clean_up() → None¶: Clean up the algorithm.

clone(index: int | None = None, wrap: bool = True) → Self¶

Create a clone of the algorithm.

QLoRA clones rebuild the base via from_pretrained and transfer only adapter (+ value head) weights. FSDP2 clones copy a rank-0 CPU full state dict onto a fresh CPU actor, then shard. The dense full model is never placed on GPU.

Parameters:

index (int | None, optional) – The index of the clone, defaults to None
wrap (bool, optional) – Unused. Clones always call wrap_models(). Kept so tournament / multi-frequency can pass wrap=False.

Returns:

A clone of the algorithm

Return type:

EvolvableAlgorithm

configure_batch_size_per_process(batch_size: int, micro_batch_size_per_gpu: int | None, mini_batch_size: int | None, group_size: int = 1) → None¶

Derive per-process batch sizes and gradient accumulation steps.

batch_size is the global collect size (prompt groups for GRPO-family). Each data-parallel replica holds (batch_size / dp_size) * group_size samples, where dp_size is the process-group world folded by tensor-parallel degree (world_size / tp). Unset mini_batch_size is that per-rank count. Unset micro_batch_size_per_gpu uses the mini-batch. gradient_accumulation_steps is mini_batch_size / micro_batch_size_per_gpu.

Parameters:

batch_size (int) – Global collect size across ranks.
micro_batch_size_per_gpu (int | None) – Samples per forward/backward pass.
mini_batch_size (int | None) – Per-rank samples per optimizer step.
group_size (int) – Completions per prompt (GRPO-family).

static copy_attributes(agent: IndividualT, clone: IndividualT, exclude: Iterable[str] = ()) → IndividualT¶

Copy the non-evolvable attributes of the algorithm to a clone.

Parameters:

agent (EvolvableAlgorithm) – The algorithm to copy attributes from.
clone (EvolvableAlgorithm) – The clone of the algorithm.
exclude (Iterable[str]) – Attribute names to leave on clone / agent.

Returns:

The clone of the algorithm.

Return type:

EvolvableAlgorithm

property current_lr: float¶: Actor learning rate the next learn call trains with.

property current_lr_critic: float¶: Critic learning rate the next learn call trains with; lr is the peak without lr_critic.

eval_policy_network_ids() → set[int]¶

Return the id of every evaluation network in the agent’s policy group.

Returns:: Identities of the policy’s evaluation networks.
Return type:: set[int]

evolvable_attributes(networks_only: bool = False) → dict[str, Any]¶

Return the attributes related to the evolvable networks in the algorithm. Includes attributes that are either EvolvableModule or ModuleDict objects, as well as the optimizers associated with the networks.

Parameters:: networks_only (bool, optional) – If True, only include evolvable networks, defaults to False
Returns:: A dictionary of network attributes.
Return type:: dict[str, Any]

finalize_training_step(num_steps: int) → None¶

Close the agent’s training block, storing any captured GraMa scores.

Parameters:: num_steps (int) – Number of steps taken during the training step.
Returns:: None.
Return type:: None

property fitness: list[float | ndarray[tuple[int, ...], dtype[_ScalarType_co]]]¶: Fitness history (scalars, or per-sub-agent rows for multi-agent).

get_action(obs: list[RolloutPrompt] | RolloutPrompt, training: bool = True, **kwargs: Any) → ActionResult¶

Generate completion tokens for each prompt in the batch.

Parameters:

obs (LLMObsType) – A single prompt dict or a list of HF-style prompt dicts.
training (bool) – If False, use near-deterministic decoding where applicable.
kwargs (Any) – Additional keyword arguments accepted for base-class compatibility.

Returns:

An ActionResult of per-prompt completion token IDs and masks. When the vLLM sampling-mismatch correction is enabled (training rollouts on the vLLM path), sampling_logps carries the captured per-row sampling logprobs; otherwise it is None.

Return type:

ActionResult

Return the dimension of the action space as it pertains to the underlying networks (i.e. the output size of the networks).

Parameters:: action_space (spaces.Space or list[spaces.Space].) – The action space of the environment.
Returns:: The dimension of the action space.
Return type:: int | dict[str, int] | tuple[int | dict[str, int], …]

get_eval_modules(cloning: bool = True) → tuple[dict[str, EvolvableModule], dict[str, EvolvableModule]]¶

Get the offsprings of all of the evaluation modules in the individual.

Parameters:: cloning (bool, optional) – Whether to clone each evaluation module before returning it, defaults to True.
Returns:: Tuple of offspring policy and the rest of the evaluation modules
Return type:: tuple[dict[str, EvolvableModule], dict[str, EvolvableModule]]

get_lr_names() → list[str | tuple[str, str]]¶: Return the learning-rate attribute name(s) of each optimizer.

get_policy() → EvolvableModuleProtocol¶: Return the policy network of the algorithm.

Return the dimension of the state space as it pertains to the underlying networks (i.e. the input size of the networks).

Parameters:: observation_space (spaces.Space or list[spaces.Space].) – The observation space of the environment.
Returns:: The dimension of the state space.
Return type:: tuple[int, …] | dict[str, tuple[int, …]]

property hp_config: HyperparameterConfig¶: Return the hyperparameter configuration for Evo-HPO mutations.

property index: int¶: Return the index of the algorithm.

init_training_step(capture_grama: bool = False) → None¶

Open the agent’s training block: metrics tracking, and GraMa capture.

Hooks are registered afresh each cycle, so they follow the agent through architecture mutations, checkpoint reloads and accelerator re-wrapping. Opening a block implicitly closes one that an earlier call left open.

Parameters:: capture_grama (bool) – Whether to register GraMa capture hooks for this training step. Defaults to False since the LLM finetuners never run ReGraMa.
Returns:: None.
Return type:: None

static inspect_attributes(agent: EvolvableAlgorithmProtocol | AgentWrapperProtocol[Any], input_args_only: bool = False, exclude: Iterable[str] = ()) → dict[str, Any]¶

Inspect and retrieve the attributes of the current object, excluding attributes related to the underlying evolvable networks (i.e. EvolvableModule, torch.optim.Optimizer) and with an option to include only the attributes that are input arguments to the constructor.

Parameters:

input_args_only (bool) – If True, only include attributes that are input arguments to the constructor. Defaults to False.
exclude (Iterable[str], optional) – Extra attribute names to drop from the result, on top of the standard exclusions below. For a caller-specific reason to leave an attribute out of its own view.

Returns:

A dictionary of attribute names and their values.

Return type:

dict[str, Any]

Update actor and critic adapters; a critic warmup step updates only the critic.

Parameters:

experiences (LLMRolloutExperiences) – (token_ids, action_masks, rewards). For single-turn, rewards is a flat tensor of scalars; for multi-turn, shape [batch, max_turns] per-turn rewards.
turn_ids (torch.Tensor | None) – Optional [batch, seq_len - 1] turn index per generated token, including tokens outside the action mask; -1 elsewhere. If None, all action tokens are turn 0.
sampling_logps (list[torch.Tensor | None] | None) – Optional per-row flat vLLM sampling logprobs (one 1-D tensor per trajectory, generated tokens only; concatenated across turns for multi-turn) for the vLLM sampling-mismatch correction. vLLM reports them under the penalised distribution it sampled from, and the policy carries no penalties, so the correction is the exact policy-to-behaviour ratio. Parallel to the stacked token_ids rows. None disables the correction for this update.
episode_segments (list[EpisodeSegments | None] | None) – Optional segment layout per trajectory, parallel to the stacked token_ids rows (None for an unsegmented trajectory). Each segment runs its forwards as its own row; values, GAE and returns span the whole episode. The update keeps the optimizer steps of the unsegmented batch: each step accumulates a window of segment rows, padded so every rank runs the same micro-batches.
pixel_values (torch.Tensor | None) – Optional vision rows of every trajectory, stacked in token_ids row order. The reference, actor and critic forwards each see the vision rows of the rows they run.
pixel_image_counts (Sequence[int] | None) – Optional vision rows per trajectory, needed when trajectories hold different numbers of images.
image_token_id (int | None) – Token id the VL forward scatters one image feature row into. Required with pixel_values and episode_segments, to cut the filler rows down to one image.
sampling_supports (list[torch.Tensor | None] | None) – Optional per-row (n, K) token ids vLLM could sample at each generated token, -1 padded, parallel to sampling_logps. The policy, old-policy and reference log-probs of those tokens are normalized over their support, the distribution vLLM sampled from.

Returns:

Mean training metrics across PPO minibatch updates, actor and critic gradient norms averaged over optimizer steps, the pre-update value fit, critic_warmup, row padding stats and learn_phase_<phase>_s wall seconds.

Return type:

dict[str, float]

classmethod load(path: str, device: str | device = 'cpu', accelerator: Accelerator | None = None) → Self¶

Load an algorithm from a checkpoint.

Parameters:

path (string) – Location to load checkpoint from.
device (str, optional) – Device to load the algorithm on, defaults to ‘cpu’
accelerator (Accelerator | None, optional) – Accelerator object for distributed computing, defaults to None

Returns:

An instance of the algorithm

Return type:

SingleAgentAlgorithm

load_checkpoint(path: str, load_optimizer: bool = False, overwrite_reference_adapter: bool | None = None, overwrite_critic_adapter: bool = False, restore_config: bool = True, restore_hyperparameters: bool = True) → None¶

Load adapter weights and algorithm state from a checkpoint directory.

Adapter roles restored on load:

actor — the trained policy. Always loaded.

reference — loaded from the checkpoint’s reference/ adapter when it has one; otherwise the checkpoint’s actor is copied onto reference so SFT -> DPO -> GRPO chains work out of the box.

critic — loaded from the checkpoint’s critic/ adapter when it has one, otherwise left at its fresh LoRA init. Set overwrite_critic_adapter to seed it from the actor.

The checkpoint’s LoRA config must match the live algorithm’s config; a mismatch raises ValueError (re-create the agent with the checkpoint’s LoRA config to load it).

The same flow applies to plain, DDP and FSDP2 runs:

lora_only=T -> PEFT adapter dirs are loaded into the live
adapters.

lora_only=F -> the full actor state_dict is restored from
attributes.pt.

When load_optimizer=True the optimizer state is restored from attributes.pt; if the checkpoint contains no optimizer state (saved with save_optimizer=False), a UserWarning is emitted and a freshly-initialised optimizer is used. The LR schedule always resumes at the checkpoint’s learn step, with this instance’s lr / lr_critic (after any hyperparameter restore) as its peaks.

Parameters:

path (str) – Directory containing a checkpoint written by save_checkpoint().
load_optimizer (bool) – If True also load the optimizer state so training can resume.
overwrite_reference_adapter (bool | None) – Copy the checkpoint’s actor onto reference even when it has a reference/ adapter. None copies only when the checkpoint has no reference adapter.
overwrite_critic_adapter (bool) – Seed critic from the checkpoint’s actor.
restore_config (bool) – If False, keep this instance’s algorithm settings and restore only training state: step count, scores, fitness, reference_update_tracker, rng and, with restore_hyperparameters, the registry’s mutable hyperparameters.
restore_hyperparameters (bool) – If False, keep this instance’s values for the registry’s mutable hyperparameters (e.g. lr). Only a run that mutates them needs the checkpoint’s values.

load_weights(path: str, overwrite_reference_adapter: bool | None = None, overwrite_critic_adapter: bool = False) → None¶

Load only the LoRA adapters (and value head) from a checkpoint directory.

Parameters:

path (str) – Directory containing a checkpoint written by save_checkpoint().
overwrite_reference_adapter (bool | None) – See load_checkpoint().
overwrite_critic_adapter (bool) – See load_checkpoint().

property lr_schedule_step: int¶: Learn steps the LR schedule has taken; 0 without a schedule.

property mut: str | None¶: Return the mutation object of the algorithm.

mutation_hook() → None¶: Execute the hooks registered with the algorithm.

classmethod population(size: int, device: str | device = 'cpu', resume_from_checkpoint: str | None = None, **kwargs: Any) → list[Self]¶

Create a population of LLM algorithms.

Builds agent 0 fully (loading the model from disk), then clones for agents 1..N. Under FSDP2 / QLoRA this uses adapter-only clone(); otherwise the actor is copied via clone_llm().

Parameters:

size (int) – The size of the population.
device (DeviceType) – Torch device. Defaults to "cpu".
resume_from_checkpoint (str | None) – Path to checkpoint to resume from.
kwargs (Any) – Additional keyword arguments to pass to the algorithm constructor.

Returns:

A list of LLM algorithms.

Return type:

list[LLMAlgorithm]

Preprocess observations (dummy) for forward pass through neural network.

Parameters:: observation (torch.Tensor[float] or dict[str, torch.Tensor[float]]) – Observations of environment
Returns:: Preprocessed observations
Return type:: torch.Tensor[float] or dict[str, torch.Tensor[float]] or tuple[torch.Tensor[float], …]

recompile() → None¶

Recompile evolvable modules with torch.compile.

Iterates over evolvable_attributes and compiles each one. Skipped for distributed runs, matching _initialize_actors().

register_mutation_hook(hook: LambdaType | MethodType) → None¶

Parameters:: hook (MutationHook) – The hook to be executed after mutation.

register_network_group(group: NetworkGroup) → None¶

Set the evaluation network for the algorithm.

Parameters:: name (str) – The name of the evaluation network.

reinit_optimizers(optimizer: OptimizerConfig | None = None) → None¶

Reinitialize the optimizers of an algorithm. If no optimizer is passed, all optimizers are reinitialized.

Parameters:: optimizer (OptimizerConfig | None, optional) – The optimizer to reinitialize, defaults to None, in which case all optimizers are reinitialized.

save_checkpoint(path: str, lora_only: bool = True, save_optimizer: bool = True) → None¶

Save adapter weights and algorithm state to a directory.

AgileRL never persists base-model weights when lora_only=True for LLM algorithms: a checkpoint is a directory containing

<adapter>/adapter_model.safetensors + adapter_config.json — one subdirectory per adapter in selected_adapters (always actor, plus reference / critic when those adapters are configured). Written only when lora_only=True.

attributes.pt — algorithm hyperparameters, plus (optionally) the actor state dict and/or optimizer state dict. Always present.

The same format is written for plain, DDP and FSDP2 runs:

lora_only=T, save_optimizer=T -> PEFT adapter dirs on disk +
optimizer state in attributes.pt

lora_only=T, save_optimizer=F -> PEFT adapter dirs only lora_only=F, save_optimizer=T -> full actor state_dict +

optimizer state in attributes.pt

lora_only=F, save_optimizer=F -> full actor state_dict in attributes.pt

FSDP2-sharded parameters and optimizer state are gathered to full tensors before saving, so checkpoints are rank-count independent.

Parameters:

path (str) – Directory to write the checkpoint into.
lora_only (bool) – If True (default) only adapter weights are written to disk via save_pretrained; the base model is shared across checkpoints and not serialised. If False, the full actor state dict is persisted into attributes.pt.
save_optimizer (bool) – If True (default) also persist the optimizer and LR scheduler state in attributes.pt so training can resume.

property scores: list[float | list[float]]¶: Per-episode scores (per-group score rows for multi-agent metrics).

select_adapter(adapter_name: str) → Generator[None, None, None]¶

Temporarily switch adapter; restores the actor adapter on exit.

Parameters:: adapter_name (str) – Name of the adapter to activate (“actor”, “critic”, “reference”).

set_reference_policy(reference_update_tracker: int) → None¶

Update the reference policy when the tracker advances past the stored value.

Base weights are immutable in AgileRL’s LoRA-only training: with use_separate_reference_adapter=True the actor adapter is copied onto the reference adapter; without one the implicit reference (the base model with adapters disabled) cannot move, so the update request is acknowledged with a one-time warning and the KL anchor stays the initial policy.

Parameters:: reference_update_tracker (int) – The reference policy update tracker

set_training_mode(training: bool) → None¶

Set the training mode of the algorithm.

Parameters:: training (bool) – If True, set the algorithm to training mode.

snapshot_checkpoint(lora_only: bool = True, save_optimizer: bool = True) → LLMCheckpointSnapshot¶

Copy the checkpoint of save_checkpoint() into host memory.

All ranks must call this together (FSDP2 gathers are collective). The snapshot holds the weights and optimizer state of this step: later training does not change it, so LLMCheckpointSnapshot.write() may run on another thread while training continues.

Parameters:

lora_only (bool) – See save_checkpoint().
save_optimizer (bool) – See save_checkpoint().

Returns:

Host copy of the checkpoint; empty off the main process.

Return type:

LLMCheckpointSnapshot

property steps: int¶: Cumulative global step count.

test(env: RolloutHarness, loop: int = 1, *args: Any, **kwargs: Any) → ndarray¶

Return fitness (test) score of the llm on the test sub-set.

Parameters:

env (RolloutHarness) – Tokenized rollout episode environment (single- or multi-turn).
loop (int) – Number of outer test iterations (episodes).

Returns:

Zero-dimensional array holding the mean per-step reward, which is also recorded in the agent’s fitness history.

Return type:

np.ndarray

Move experiences to the device.

Parameters:: experiences (tuple[torch.Tensor[float], ...]) – Experiences to move to device
Returns:: Experiences on the device
Return type:: tuple[torch.Tensor[float], …]

unrolled_eval_networks() → list[tuple[str | None, Module]]¶

Return the agent’s evaluation networks as (network_id, network) pairs.

Returns:: One (network_id, network) pair per measured network.
Return type:: list[tuple[str | None, torch.nn.Module]]

unwrap_models() → None¶: Unwraps the models in the algorithm from the accelerator.

update_existing_adapter(checkpoint_dir: str, adapter_name: str) → None¶

Overwrite weights of an existing adapter in-place without creating new parameters.

Parameters:

checkpoint_dir (str) – Checkpoint directory
adapter_name (str.) – Adapter name

Returns:

None

Return type:

None

static update_lr(optimizer: OptimizerWrapper, lr: float | tuple[float, float | None], scheduler_config: CosineLRScheduleConfig | None = None, schedule_step: int = 0) → LambdaLR | None¶

Set the peak learning rate of each param group and rebuild the schedule.

Parameters:

optimizer (OptimizerWrapper) – LLM optimizer.
lr (float | tuple[float, float | None]) – Learning rate value, or actor/critic pair; a None critic uses the actor rate.
scheduler_config (CosineLRScheduleConfig | None) – Scheduler configuration; None holds lr constant.
schedule_step (int) – Learn steps the schedule has already taken.

Returns:

Scheduler at schedule_step when scheduler_config is set.

Return type:

LambdaLR | None

use_adapter(adapter_name: str) → None¶

Switch the active PEFT adapter, handling all side-effects.

For “reference”: switches adapter and freezes reference params (never trained). For all others: switches adapter and restores requires_grad=True on all training adapter LoRA params so distributed gradient hooks keep firing.

Parameters:: adapter_name (str) – Name of the adapter to activate (“actor”, “critic”, “reference”).

wrap_models() → None¶

Prepare the actor for training.

Places the actor (dense .to(device) or FSDP2 shard) then builds the optimizer and LR scheduler on those parameters. FSDP2 wraps each transformer block with activation checkpointing before fully_shard when gradient_checkpointing is on. Data-parallel prepare_actor uses HuggingFace gradient_checkpointing_enable.