Multi-Agent REINFORCE

REINFORCE is a class of policy-gradient methods that optimize the policy directly using sampled returns. These methods are widely used for LLM fine-tuning because of their simplicity and efficiency. Examples include GRPO, Dr. GRPO, RLOO, ReMax, TreeRPO, and REINFORCE++. REINFORCE can be extended to multi-agent settings, where multiple LLM agents respond at each turn and their joint responses receive a shared reward.

All MAREINFORCE trainers share one training loop. At each rollout node, every agent samples $$ K $$ responses, the trainer scores the index-aligned joint responses with one shared reward, and each agent weights its own log-probabilities by the resulting shared advantages before a buffered optimizer step.

MAREINFORCE training flow Every joint sample earns one shared return; its advantage weights each agent's own policy gradient.

MAREINFORCE#

The policy-gradient surrogate for Multi-Agent REINFORCE (MAREINFORCE), before advantage normalization or reference regularization, can be expressed as:

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\sum_{t=0}^{H-1} R_t \cdot \log \pi_{\theta_i}(a_{i,t}\mid h_{i,t})\Bigg],

where

\( R_t \)

is the return at turn

\( t \)

and

\( H \)

is the horizon (i.e., the number of dialog turns). The expectation is taken over initial observations from the dataset

\mathcal{D}

and the joint action history of all episodes following policy

\boldsymbol{\pi}_{\boldsymbol{\theta}}

REINFORCE methods do not use a critic model for value estimation. Their policy-gradient estimates can have high variance due to stochastic sampling and long-term credit assignment. Common approaches to reducing this variance include using an action-independent baseline or averaging over more samples, such as $$ K $$ samples at each joint history $\mathbf{h}_t$ .

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K} \sum_{t=0}^{H-1} \left(R^{k}_t - b(\mathbf{h}_t)\right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg],

where the baseline

b(\mathbf{h}_t)

is action-independent.

MAGRPO#

Multi-Agent Group-Relative Policy Optimization (MAGRPO) is an instantiation of MAREINFORCE inspired by GRPO, where the group-average baseline is the mean return of $$ K $$ samples. In the default aligned joint mode, its surrogate before advantage normalization is:

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H-1} \left(R^{k}_t - \frac{1}{K}\sum_{l=1}^{K}R^{l}_t\right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg].

The sampled group mean includes the sample being updated, so it is not the action-independent baseline assumed in the general expression above. In multi-turn training, each branch’s return is its immediate reward plus the mean return of its child node, with no discount applied. In cross mode, each agent’s completion instead receives the sum of returns over all joint combinations containing that completion.

By default, advantage_normalization=True then centers and scales the advantages within each node’s generation group using the population standard deviation, with a minimum denominator of 1e-6. If reference KL regularization is enabled, the sampled reference penalty is subtracted from returns before the baseline and normalization are applied.

The current MAGRPOTrainer loss slice omits the final completion token, so its implemented sequence score differs from the full sequence log-probability in the surrogate above. The canonical setting is joint_mode="cross"; the accepted "crossed" alias generates Cartesian combinations but does not apply the same per-completion return aggregation.

MAGRPOConfig Parameters

num_agents: Configured number of agents, which should match the trainer setup
num_turns: Number of turns per episode
num_train_epochs: Number of training epochs
agent_learning_rate: Learning rate
logging_steps: Logging interval in steps
num_generations: Number of generations to sample per prompt for each agent
max_new_tokens: Maximum number of new tokens to generate
temperature: Temperature for sampling
top_p: Top-p for sampling
top_k: Top-k for sampling
discount: Retained configuration field; the current MAGRPO implementation computes undiscounted returns and does not use this value
joint_mode: Joint action composition (aligned for index-aligned, cross for Cartesian product)
early_termination_threshold: A training branch stops expanding when its node’s mean immediate reward exceeds this threshold (default -0.2; None disables early termination)
rollout_buffer_size: Number of node samples to buffer per agent before an update
train_batch_size: Number of node samples per optimizer step; defaults to rollout_buffer_size
advantage_normalization: Whether to center and scale advantages within each node’s generation group (default True)
eval_interval: Evaluation runs at batch indices divisible by N, including the first batch of each epoch; 0 disables periodic evaluation
eval_num_samples: Number of samples to evaluate per evaluation run
eval_batch_size: Evaluation DataLoader batch size
external_prompt_passthrough: Whether to use external prompts directly in multi-turn training
advantage_mode: Baseline mode (mean, max, rloo, raw)
parallel_training: Agent execution mode (none or mp)
agent_devices: Agent device placement; required explicitly for parallel_training="mp"
reference_kl_enabled: Whether to enable a frozen-reference KL penalty (default False)
reference_kl_coef: Weight of the reference penalty (default 0.1)
reference_devices: Optional device placement for the reference models

MAGRPOTrainer Parameters

agent_model or agents: Model identifier string for homogeneous agents, or a list of model identifiers or preloaded models; preloaded models are supplied through agents
num_agents: Number of agents created from agent_model (default 2), configured independently of args. When only agents is supplied, its length determines the agent count
tokenizer: Shared tokenizer or one per agent; optional for model identifiers, required for preloaded models
reward_func: Callable that returns a list of floats (required)
reward_processor: Optional processor to apply to rewards (e.g., scaling)
formatters: Single callable or list of per-agent callables for prompt formatting
args: Instance of MAGRPOConfig (optional)
train_dataset: Dataset required when calling train()
eval_dataset: Evaluation dataset (optional)
model_config: Model configuration dictionary (optional)
wandb_config: Configuration for Weights & Biases logging (optional)
external_transition: Function providing transitions between turns; required when num_turns > 1 (the default is 2)
eval_logger: Optional task-metrics callback, called as eval_logger(agent_completions_turns=..., batch_items=..., prompts=...); returns one metrics dictionary per evaluated item and is used with eval_aggregator
eval_aggregator: Optional callback, called as eval_aggregator(metrics, num_turns=...); the returned turn_k/... metrics are logged under eval/. Both evaluation callbacks must be supplied for custom metrics

MAGRPO collects samples into a rollout buffer and recomputes their log-probabilities during updates, without importance sampling or ratio clipping. Samples may therefore predate earlier updates from the same buffer or rollout tree. Since the trainer has no critic model, value_clip_range does not apply.

The trainer uses a fixed training DataLoader batch size of 1 and requires at least num_generations=2, including for the raw-return variant. By default, train_batch_size=rollout_buffer_size: each agent averages the losses of its buffered node samples for one optimizer step. Smaller mini-batches produce multiple optimizer steps; a partial buffer is flushed at the end of each epoch.

Other Variants#

CoMLRL also provides other MAREINFORCE variants with different baselines:

MARLOO: Multi-Agent REINFORCE Leave-One-Out. The baseline is the mean return of the other generations in the same node’s sample group, not the other agents.

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H-1} \left(R^{k}_t - \frac{1}{K-1}\sum_{l=1, l\neq k}^{K}R^{l}_t\right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg].

MAReMax: Multi-Agent REINFORCE with Group Max (MAReMaxConfig and MAReMaxTrainer). The baseline is the maximum sampled group return at the node; this implementation does not generate a separate greedy baseline response.

J(\theta_i) = \mathbb{E}_{\mathbf{o}_0 \sim \mathcal{D}, \mathbf{h}_t \sim \boldsymbol{\pi}_{\boldsymbol{\theta}}} \Bigg[\frac{1}{K}\sum_{k=1}^{K}\sum_{t=0}^{H-1} \left(R^{k}_t - \mathrm{max}_l\, R^l_t \right) \cdot \log \pi_{\theta_i}(a^{k}_{i,t}\mid h_{i,t})\Bigg].

MAREINFORCETrainer, MARLOOTrainer, and MAReMaxTrainer derive from comlrl.trainers.reinforce.MAGRPOTrainer. Their corresponding configuration classes inherit MAGRPOConfig, with default advantage_mode values of raw, rloo, and max, respectively. They share the same trainer interface and configuration fields.

The four trainers differ only in the baseline subtracted from the shared return. For one node with $$ K=4 $$ joint responses, the animation below compares their advantages before normalization: MAREINFORCE keeps the raw return, as in REINFORCE; MAGRPO subtracts the group mean, as in GRPO; MARLOO subtracts the mean of the other $$ K-1 $$ returns, as in RLOO, which scales the MAGRPO advantages by $$ K/(K-1) $$ ; and MAReMax subtracts the group maximum, so no response receives a positive advantage.

Baselines across MAREINFORCE variants Each row subtracts a different baseline from the same shared returns; shaded bars mark the advantages before normalization.

The variant formulas describe advantages before normalization. All variants inherit advantage_normalization=True, which re-centers even raw returns and max-baseline advantages. Except for the standard-deviation floor and numerical precision, the four modes then yield the same normalized advantages. Users can retain the baseline differences shown above, including the raw-return MAREINFORCE estimator, by setting advantage_normalization=False.