Group Relative Policy Optimization (GRPO)¶

GRPO (Group Relative Policy Optimization) is an elegant simplification of PPO (Proximal Policy Optimization) that makes reinforcement learning more computationally efficient, especially for large language models.

The two key innovations are:

Eliminating the critic network: Instead of training a separate value function to estimate expected rewards (which requires additional compute and memory), GRPO normalizes rewards across a batch of samples. It calculates advantage by subtracting the mean reward from each sample’s reward and dividing by the standard deviation.
Group-based evaluation: GRPO generates multiple outputs using the same policy, evaluates them as a group, and then updates the model. This approach reduces variance in the training signal by smoothing out the randomness inherent in probabilistic environments.

These changes are particularly valuable for LLM training because they reduce computational overhead by removing the need for a separate critic model, provide more stable gradient updates in environments with sparse or noisy rewards, and they simplify implementation while maintaining or improving performance.

In AgileRL, GRPO can be used for single-turn reasoning tasks or multi-turn agentic finetuning. In the multi-turn case, rollouts are still treated as a bandit problem, with environment generated tokens masked and reward signal calculated from cumulative episode reward.

The objective is selected via the loss_type argument, which accepts "grpo" (the default token-level PPO-style clipped surrogate), "gspo" (sequence-level importance ratio, see GSPO) and "cispo" (clamped importance-weighted log-prob objective, see CISPO). The CISPO and GSPO classes are thin subclasses that pin loss_type to the matching variant.

Variance Reduction¶

GRPO replaces PPO’s learned value head with group-relative normalization: for each prompt, group_size rollouts are drawn and their returns are z-scored within the group to form the advantage. The upside is that there is no critic to train, fit or tune, which is attractive for LLM scale; the downside is that the baseline degenerates when the group’s returns collapse (e.g. all rollouts succeed or all fail), and the quality of the variance reduction is tied to the group size. Compare with the learned value baseline used by LLM PPO and Return Batch Normalization (ReBN) used by LLM REINFORCE.

Example¶

For more details on how to set up GRPO and use it for training, check out the tutorial.

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

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

agent = GRPO(
  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,
  group_size=8,
)

Saving and Loading Agents¶

To save an agent, use the save_llm_checkpoint function:

from agilerl.utils.utils import save_llm_checkpoint

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

To load a trained model, you must use the HuggingFace .from_pretrained method, AgileRL is compatible with HuggingFace and Peft models:

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch

base_model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B",
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-3B")
model = PeftModel.from_pretrained(base_model, "path/to/model/directory")

Parameters¶

class agilerl.algorithms.grpo.GRPO(*args: Any, **kwargs: Any)¶

Group Relative Policy Optimization (GRPO).

Paper: https://arxiv.org/pdf/2402.03300

Parameters:

pad_token_id (int) – Pad token id
pad_token (str) – Pad token
model_name (str, optional) – Model name
actor_network (PreTrainedModel | PeftModel | None) – HuggingFace LLM
model_config (dict[str, Any], optional) – Model configuration, to be used when creating the model from a name or path
hp_config (HyperparameterConfig, optional) – RL hyperparameter mutation configuration, defaults to None, whereby algorithm mutations are disabled.
index (int, optional) – Index to keep track of object instance during tournament selection and mutation, defaults to 0
batch_size (int, optional) – Mini-batch size for learning, defaults to 16
beta (float, optional) – Beta coefficient, controls the strength of the KL divergence penalty, defaults to 0.001
lr (float, optional) – Learning rate for optimizer, defaults to 5e-7
clip_coef (float | tuple[float, float], optional) – Surrogate clipping coefficient as either a symmetric scalar (mapped to [1-clip_coef, 1+clip_coef]) or an explicit ratio tuple (clip_coef_min, clip_coef_max).
max_grad_norm (float, optional) – Maximum norm for gradient clipping, defaults to 0.1
update_epochs (int, optional) – Number of policy update epochs, defaults to 1
group_size (int, optional) – Group size, defaults to 8
temperature (float, optional) – Temperature, controls randomness of text generation
repetition_penalty (float, optional) – Repetition penalty used during generation, defaults to 1.0
top_p (float, optional) – Top-p nucleus sampling threshold, defaults to 0.95
top_k (int, optional) – Top-k sampling threshold, defaults to 50
min_p (float, optional) – Minimum probability cutoff for sampling, defaults to 0.0
calc_position_embeddings (bool, optional) – Flag indicating whether to calculate position embeddings, defaults to True
micro_batch_size_per_gpu (int, optional) – Trajectories per backward pass on one rank (the memory setting). Optimizer-step cadence comes from mini_batch_size. If None, the full per-rank batch is used in a single forward pass, defaults to None
mini_batch_size (int, optional) – Per-rank trajectories covered by one optimizer step. None uses (batch_size / world_size) * group_size. gradient_accumulation_steps is derived as mini_batch_size / micro_batch_size_per_gpu.
max_output_tokens (int, optional) – Max number of answer tokens, defaults to None
min_output_tokens (int, optional) – Minimum output tokens, defaults to 0
max_model_len (int, optional) – Maximum context window length, defaults to 1024
hf_generate_chunk_size (int | None, optional) – Number of prompts per HuggingFace generation chunk. Ignored when colocated.
lora_config (LoraConfig, optional) – Config for LoRA, defaults to None
cosine_lr_schedule_config (CosineLRScheduleConfig, optional) – Config for cosine lr scheduling, defaults to None
offload_trainer_during_rollout (bool) – 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.
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) – Wrap models for distributed training upon creation, defaults to True
clone (bool, optional) – Flag to indicate if the instantiation is a cloning, defaults to False
vllm_config (VLLMConfig, optional) – Config for VLLM generation, defaults to None
seed (int, optional) – Seed for the random number generator, defaults to 42
gradient_checkpointing (bool, optional) – Flag to indicate if gradient checkpointing should be used, defaults to True
torch_compiler (str | None, optional) – Torch compile mode (e.g. 'default'), defaults to None
use_liger_loss (bool, optional) – Use the Liger fused loss, defaults to False (requires liger-kernel; warns and falls back otherwise). Not recommended for GRPO/CISPO/GSPO: the upstream Liger GRPO kernel shows no speedup over AgileRL’s already memory-bounded standard path and uses slightly more memory. PPO/REINFORCE route use_liger_loss through a different AgileRL liger-based kernel where it does help (see their docs). The Liger model patches (fused RMSNorm/RoPE/SwiGLU) apply whenever liger-kernel is installed and are independent of this flag.
use_kl_advantage_shaping (bool, optional) – Apply KL-based shaping directly to token advantages before PPO clipping, defaults to False.
adv_norm (str, optional) – Advantage normalization mode. "mean_std" divides by standard deviation, "mean_only" only centers, defaults to "mean_std".
loss_type (Literal["grpo", "gspo", "cispo"], optional) – PPO-style loss variant to optimize. One of "grpo", "gspo", or "cispo", defaults to "grpo". This selects the objective: "grpo"/"gspo" use the min-clip surrogate, "cispo" the clamped-weight x log-prob objective. "gspo" is sugar for "grpo" at trajectory level (it forces importance_sampling_level="trajectory").
importance_sampling_level (Literal["token", "turn", "trajectory"] | None, optional) –
Granularity at which the importance ratio is pooled before clipping/weighting, defaults to None (resolves to "token"; loss_type="gspo" forces "trajectory" and warns if a different level was requested explicitly). This is independent of advantage_granularity (the advantage axis).
- "token" — per-token ratio (standard GRPO / CISPO).
- "turn" — pool the per-token log-ratio over each turn (length- normalized geometric mean) and clip/weight per turn. Requires turn_ids in learn().
- "trajectory" — pool over the whole completion (GSPO).
Turn/trajectory pooling couples a unit’s tokens, so it has no fused Liger kernel and runs on the standard (always memory-bounded) path; only token level can use the Liger path when use_liger_loss=True.
advantage_granularity (Literal["auto", "trajectory", "turn"], optional) –
Unit at which the group-relative advantage is computed, independent of importance_sampling_level. Defaults to "auto".
- "trajectory" — one group-relative scalar per completion (standard GRPO), broadcast to all tokens.
- "turn" — group-relative per turn (each turn’s reward normalized within its group), broadcast to that turn’s tokens. Requires turn_ids and per-turn rewards (batch, max_turns) in learn().
- "auto" — "turn" when the batch has per-turn rewards and any sample has more than one turn; otherwise "trajectory". There is no token-level GRPO advantage (group-relative needs a per-unit reward).
Any advantage x IS combination is valid.
use_separate_reference_adapter (bool, optional) – Keep a dedicated reference LoRA adapter whose weights are frozen snapshots of the actor used for the KL-divergence baseline. When False the reference log-probs are obtained by disabling the actor adapter at inference time. Defaults to True.
whiten_advantages (bool, optional) – If True, whiten token-level advantages over valid action positions, defaults to False.
adv_clip_range (float | None, optional) – Optional symmetric clamp range applied to advantages before loss computation, defaults to None.
filter_zero_adv (bool, optional) – If True, filter samples whose absolute advantage is below adv_filter_eps. Single-process runs drop the samples from the update; multi-process runs zero their advantages instead, keeping the collective schedule identical on every rank, defaults to False.
adv_filter_eps (float, optional) – Threshold used with filter_zero_adv; samples with |advantage| <= eps are filtered out, defaults to 0.0.
turn_advantage_trajectory_fallback (bool, optional) – With per-turn advantages, give a (sample, turn) cell whose group has fewer than two members that played the turn the sample’s trajectory advantage instead of zero, defaults to True. Turns nobody played stay at zero either way.
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.
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.
lora_target_scope (str | None, optional) – Optional PEFT LoRA path scope for multimodal models (e.g. "language_model"). Passed to adapt_lora_config_for_model().
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.
use_sequence_packing (bool, optional) – Opt in to padding-free sequence packing for the gradient forward (sequences pack into one varlen / blockmask pass). Only honoured under a FlashAttention-2 / FlexAttention backend, otherwise inert; the no-grad reference/old-logprob pass stays padded.
loss_norm (Literal["micro_batch", "accumulation_window"], optional) – Token population the policy loss is normalized over. "micro_batch" (default) normalizes each micro-batch on its own. "accumulation_window" normalizes by the action tokens of this rank’s samples entering the optimizer step, so a token weighs the same wherever it is in the rank’s gradient-accumulation window; without it a short trajectory’s tokens outweigh a long one’s in proportion to the length ratio. Under data parallelism the gradient all-reduce then averages the per-rank means. Applies to the standard and the fused Liger path alike.

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.

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 rank holds (batch_size / world_size) * group_size samples. Unset mini_batch_size uses micro_batch_size_per_gpu when the class default is "micro_batch" (RL rollout algorithms) and that is set, else the per-rank collect. 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

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, repeat_prompts: bool = True, *args: Any, **kwargs: Any) → ActionResult¶

Return generated completions for each prompt (GRPO groups when training).

Parameters:

obs (LLMObsType) – List of HF-style prompt dicts (this implementation mutates them).
training (bool) – If True, generate with training sampling settings.
repeat_prompts (bool) – If True and training=True, duplicate each prompt self.group_size times (legacy GRPO grouped mode). If False, treat the batch as already expanded trajectories.

Returns:

An ActionResult of completion token IDs, per-sequence action masks, and (when captured) per-completion vLLM sampling logprobs for the mismatch correction.

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 agent network parameters to learn from experiences.

Parameters:

experiences (LLMRolloutExperiences) – (token_ids, action_masks, rewards) stacked batch. For importance_sampling_level="turn" with per-turn rewards, rewards is (batch, max_turns); otherwise it is one scalar per trajectory (per-turn rewards are summed to the episode return).
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 sampling-mismatch correction. Parallel to the stacked token_ids rows. None disables the correction for this update.
turn_ids (torch.Tensor | None) – (batch, seq_len-1) turn index per action token (-1 for non-action tokens), aligned with the action mask. Required when the resolved advantage granularity is "turn" (per-turn group-relative advantages need per-turn rewards). Also consumed by turn-level importance-ratio pooling when importance_sampling_level="turn". Ignored when neither applies.

Returns:

Dict with averaged loss, kl (NaN on the fused path at beta == 0.0), clipfrac and completion_length (plus per-learn advantage stats, the update-loop entropy / kl_ref / kl_old / is_* diagnostics, averaged grad_norm_pre / grad_norm_post, and the vllm_is_* sampling-mismatch metrics when the correction is active).

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) → 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 (and LR scheduler) 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.

Parameters:

path (str) – Directory containing a checkpoint written by save_checkpoint().
load_optimizer (bool) – If True also load the optimizer and LR scheduler 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.

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 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], …]

process_liger_metrics(aux: list[Tensor]) → tuple[Tensor, Tensor]¶

Split the fused kernel’s auxiliary outputs into KL and clip fraction.

The kernel returns [kl, clipfrac] with a KL coefficient and [clipfrac] at beta == 0.0, so the clip-fraction index follows beta. KL is NaN (unmeasured, not zero) without a coefficient.

Parameters:: aux (list[torch.Tensor]) – Auxiliary outputs from LigerFusedLinearGRPOFunction.
Returns:: (kl_or_nan, clipfrac) scalars.
Return type:: tuple[torch.Tensor, torch.Tensor]

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.

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: torch.optim.Optimizer, lr: float | tuple[float, float], scheduler_config: CosineLRScheduleConfig | None = None) → SequentialLR | None¶

Update the learning rate of the optimizer.

Parameters:

optimizer (Optimizer) – Optimizer
lr (float | tuple[float, float]) – Learning rate value, or actor/critic pair.
scheduler_config (CosineLRScheduleConfig | None) – Scheduler configuration

Returns:

A fresh warmup scheduler when scheduler_config is set.

Return type:

SequentialLR | 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.