Custom advantage computation

actor.advantage selects the whole advantage computation or individual stages. With the default configuration, existing GAE, reward normalization, and PRM settings retain their behavior. The configuration contains import paths and JSON-serializable keyword arguments; each actor worker resolves its functions at initialization.

actor:
  advantage:
    compute:
      implementation: default
      kwargs: {}
    outcome_reward_shaping: null
    estimator: null
    process_adv_shaping: null

The default pipeline applies length penalties, outcome shaping, reward scaling and normalization, temporal estimation, advantage normalization, then process advantage shaping. Loss aggregation and optimizer minibatch scheduling are unchanged.

Stage	Built-in implementations	Default
outcome_reward_shaping	identity	Unmodified outcome reward
estimator	token_gae, turn_gae	Follows actor.gae_timestep_unit
process_adv_shaping	additive, gvpo, process_weighted	Existing PRM selection, otherwise additive

Every stage also accepts an importable dotted Python function path and kwargs. auto selects the stage default. Process auto does not accept kwargs. Explicit process selection wins over call-local PRM selection; explicit GVPO kwargs override legacy GVPO values, while unspecified values inherit them. Existing PRM settings remain supported.

actor:
  token_rewards_as_adv: true
  advantage:
    estimator:
      implementation: turn_gae
      kwargs: {}
    process_adv_shaping:
      implementation: gvpo
      kwargs:
        negative_scale: 0.2
        zero_penalty: 0.4

Turn GAE requires rollout turn_ids. Its token KL penalties affect actor advantages but not critic returns. Token GAE includes KL in the recurrence. A custom GAE lambda receives lengths in the selected estimator’s timestep unit, including when the estimator overrides gae_timestep_unit. Existing GAELambdaContext functions remain supported; GAELambdaInputs is an alias for the same type.

Stage functions

Outcome functions accept OutcomeRewardInputs from areal.trainer.ppo.advantage. Its data contains rollout tensors after the length penalty, config is the actor configuration, and meta contains trajectory and logical rollout grouping. Return OutcomeResult(rewards=...) with one floating-point reward per tensor row:

from areal.trainer.ppo.advantage import OutcomeResult, OutcomeRewardInputs


def scale_outcome(inputs: OutcomeRewardInputs, *, factor: float) -> OutcomeResult:
    return OutcomeResult(rewards=inputs.data["rewards"] * factor)

Use implementation: my_package.my_module.scale_outcome and kwargs: {factor: 0.5}. The module must be available on every actor worker. Functions must preserve tensor row order, grouping, and device. Do not modify the input mapping or its tensors in place.

When reward normalization uses explicit logical rollout references, custom shaping must also return OutcomeResult.normalization_rewards: one shaped reference per logical rollout, before bias, scaling, and clipping. Read the raw references and row mapping from meta.rollout_groups. Without explicit references, equal rows within a logical rollout allow reference inference; differing rows require explicit shaped references. meta.traj_group_sizes counts tensor rows per prompt, while meta.logical_group_sizes counts logical rollouts per prompt. Shaped references require complete metadata produced by concat_batch; they do not change membership or row counts. Distributed plugins should use framework metadata rather than constructing partial metadata themselves.

Temporal functions accept TemporalInputs from areal.trainer.ppo.advantage_algorithms and return TemporalResult(advantages, returns, total_rewards). Inputs use next-token prediction alignment. Outcome, KL, and folded process rewards are separate channels. The explicit bootstrap_values come from each trajectory’s actual final valid token. Callable lambda functions can be resolved with compute_gae_lambda in areal.trainer.ppo.gae, passing the estimator’s timestep unit. total_rewards is the token-aligned monitoring channel.

Process functions accept ProcessAdvantageInputs and return a tensor of advantages. They receive normalized advantages, aligned direct token rewards (possibly None), and the loss mask. They do not change critic returns. Built-in GVPO and process-weighted shaping require token_rewards_as_adv=True; process-weighted shaping also requires mask_no_eos_with_zero=False.

Replacing the full computation

Set compute.implementation to a dotted function path accepting AdvantageInputs and returning AdvantageResult. A full override cannot be combined with stage settings. The default compute does not accept kwargs; configure its stages instead.

AdvantageResult supplies advantages, returns, loss mask, aligned log probabilities, outcome rewards, total rewards, and KL rewards. Outputs must be finite floating-point tensors on the expected device and with the expected shape; masks must be binary. The identity outcome stage also preserves legacy integer task rewards. The optional extras mapping may contain only input_token_loss_mask, is_truncated, mopd_behavior_logprobs, prox_logp, and ref_logp. It cannot replace input IDs, values, original rewards, or other rollout fields. A full override owns alignment and must supply the derived metadata needed by its training mode, including immutable behavior log probabilities when combining with MOPD.