RL algorithms

The RL algorithms rLLM supports — advantage estimators and auxiliary losses — and how to build your own.

Modules: rllm.trainer.algorithms.advantage (advantage estimators) · rllm.trainer.algorithms.aux_loss (auxiliary losses) · rllm.trainer.algorithms.config (AlgorithmConfig)

In rLLM an RL “algorithm” is assembled from two composable, backend-agnostic pieces, both configured under rllm.algorithm:

1. An advantage estimator — turns per-trajectory rewards into advantages (the GRPO family).
2. Zero or more auxiliary losses — extra loss terms on tokens the policy gradient ignores (for example ECHO’s loss on environment-observation tokens).

The backend then applies a policy loss function (ppo, importance_sampling, …) using those advantages. The same registries feed all three training backends (verl, tinker, fireworks), so flipping algorithms is a one-line config change that works everywhere.

How advantages are computed

Advantages are computed on TrajectoryGroups, partitioned by group_role. The estimator hook is scalar reward per trajectory in, scalar advantage per trajectory out — the unified trainer broadcasts each trajectory’s advantage across its response tokens. This single contract is what lets the same estimator code serve every backend.

Select the estimator with adv_estimator:

rllm:
  algorithm:
    adv_estimator: grpo
    norm_adv_by_std_in_grpo: true   # GRPO only; false = center by mean only

See Advantage estimator for the full data contract, per-role estimator maps, and what the scalar hook intentionally cannot express (GAE and other per-token / critic-based estimators — for those, see Pre-computing advantage).

Built-in advantage estimators

Estimator	adv_estimator	What it computes
GRPO (default)	grpo	Center rewards within each group by the group mean; optionally divide by the group std (norm_adv_by_std_in_grpo, default true). No value network.
REINFORCE	reinforce	No baseline — the reward is the advantage.
REINFORCE++ baseline	reinforce_plus_plus_baseline	Per-group mean baseline, then whiten by the role-batch std.
PRPO	prpo	Center and normalize across the whole batch (flatten all groups), not per-group.
RLOO	rloo	Leave-one-out baseline per group: each trajectory is scored against the mean of the others.
ECHO	echo	GRPO’s advantage unchanged, plus an environment-observation auxiliary loss (see below).

All six live in one registry and run on every backend. (The verl backend can also be pointed at its native estimators, but the rLLM-native set above is the portable path.)

ECHO and auxiliary losses

ECHO keeps GRPO’s advantage exactly and adds a cross-entropy auxiliary loss on the environment-observation tokens (tool / terminal output) that the policy conditions on but the policy gradient never trains. For tool-use and terminal agents — where rollouts are dominated by observation tokens and many rollouts fail — this turns every rollout, including the failures, into dense supervision at no extra rollout cost.

adv_estimator=echo defaults the env-loss weight λ to the paper’s 0.05. Tune it explicitly (productive range 0.01–0.05; 0.0 reproduces plain GRPO):

# any backend — verl / tinker / fireworks
... rllm.algorithm.adv_estimator=echo
... rllm.algorithm.adv_estimator=echo rllm.algorithm.env_loss_coef=0.03

The auxiliary-loss framework

env_loss_coef is shorthand for the general aux_losses mechanism. An auxiliary loss is a named loss applied to a token subset (a “mask”). These two configs are equivalent:

rllm:
  algorithm:
    env_loss_coef: 0.05                       # shorthand
    # — or, the general form —
    aux_losses:
      - { type: env_prediction, coef: 0.05 }  # ECHO's loss, by name

env_prediction is the one built-in client (length-normalized cross-entropy on observation tokens). You can list several aux losses; each is added to the policy loss with its own coef.

The same config runs on every backend, with different mechanics and cost:

Backend	How the aux term is applied	Metric to watch
verl	Folded into GRPO’s single forward/backward pass — free and exact.	actor/aux_env_prediction_loss
tinker / fireworks	A second, gradient-accumulated cross_entropy pass over the same rollouts — no extra rollouts, one extra backward.	train/aux_*

Because loss normalization differs across managed services, λ (the coef) may need per-backend retuning. Watch the metric above to confirm the auxiliary loss is actually falling.

Building custom RL algorithms

Both extension points are registries — register a function or a class and reference it by name from config. No backend code to touch.

A custom advantage estimator

Decorate a function with @register_rllm_adv_estimator("name"). It is scalar-per-trajectory in, scalar-per-trajectory out; **kwargs carries traj_groups when you need per-trajectory metadata.

import numpy as np
from rllm.trainer.algorithms.advantage import register_rllm_adv_estimator

@register_rllm_adv_estimator("batch_mean_baseline")
def batch_mean_baseline(rewards, algorithm_config, **kwargs):
    # rewards: list of 1-D arrays (one per TrajectoryGroup)
    batch_mean = np.mean(np.concatenate(rewards)) if rewards else 0.0
    adv = [group - batch_mean for group in rewards]
    return adv, adv          # (advantages_by_group, returns_by_group)

rllm:
  algorithm:
    adv_estimator: batch_mean_baseline

The full contract, role-level estimator maps, and a complete worked example (porting OPO from verl) are in Advantage estimator.

A custom auxiliary loss

Subclass AuxiliaryLoss, choose a token mask, and register it. The default weighting is uniform (every masked token gets coef); override weight(ctx) to return a per-token tensor instead.

from rllm.trainer.algorithms import register_aux_loss, AuxiliaryLoss, MASK_ACTION

@register_aux_loss("entropy_bonus")
class EntropyBonus(AuxiliaryLoss):
    mask = MASK_ACTION          # or MASK_OBSERVATION
    # def weight(self, ctx):    # optional — return a per-token tensor; None = uniform `coef`
    #     ...

rllm:
  algorithm:
    aux_losses:
      - { type: entropy_bonus, coef: 0.01 }

The same class runs on all three backends (verl folds it into the policy pass; tinker/fireworks add the extra cross-entropy pass). Dynamic per-token weights and auxiliary losses that need a separate model call (e.g. a teacher branch) are an in-progress extension — see the design note design/auxiliary-losses.md.

Need a per-token advantage signal (GAE, reward-to-go, a custom detector) that the scalar estimator hook can’t express? Compute it in your workflow and feed it through use_precomputed_advantage to bypass the estimator entirely.

Configuration quick reference

All under rllm.algorithm (full reference):

Field	Default	Purpose
adv_estimator	grpo	Global advantage estimator (registry name).
norm_adv_by_std_in_grpo	true	GRPO: divide centered rewards by group std (false = center only).
env_loss_coef	0.05 if adv_estimator=echo, else 0.0	Shorthand for aux_losses: [{type: env_prediction, coef: …}].
aux_losses	[]	List of {type, coef} auxiliary-loss specs.
loss_fn	backend default	Policy loss function (ppo, importance_sampling, …).

Per-role estimator overrides are set via traj_group_adv_estimator_map on the AgentTrainer constructor — see Advantage estimator.

Next steps

Advantage estimator

The estimator hook contract, role-level maps, and the OPO worked example

Configuration

Every AlgorithmConfig field

Capability matrix

Which algorithms and losses each backend supports

Backends

verl vs tinker vs fireworks