Multi-LLM PbRL

Multi-agent PbRL trains collaborating LLM agents from preferences between their joint responses. CoMLRL builds the preference data from task items and learns from it with one of two algorithms: MARLHF fits a joint reward model and then trains the policies with RL, while MADPO optimizes the policies directly from the preferences. MARLHFTrainer, MADPOTrainer, and their iterative variants MARLHFIterTrainer and MADPOIterTrainer are available from comlrl.trainers.preference.

All four trainers currently require num_turns=1 and aligned joint responses (joint_mode="aligned"). An iteration of the iterative trainers is a round of data collection and training, not a dialog turn.

Preference Data#

Each agent \( i\in\mathcal{I}=\{1,\ldots,n\} \) samples \( a_{i,t}\sim\pi_{\theta_i}(\cdot\mid h_{i,t}) \) from its local history, and the responses form the joint action \( \mathbf{a}_t=\{a_{1,t},\ldots,a_{n,t}\} \) with joint history \( \mathbf{h}_t \). With a single turn (\( H=1 \)), \( h_{i,0} \) is agent \( i \)’s prompt and \( \mathbf{a}_0 \) is the joint response.

Users provide task items, not preference data. The trainers sample joint responses, score them with an oracle function, and turn the scores into preference pairs.

Oracle Function#

The oracle function scores each joint response with one shared value. By default, it is the task reward reward_func, after the optional reward_processor; it can follow rules or call an LLM reward model, as described in Training Initialization. Both algorithms learn only from which joint response wins, so the score values matter only for choosing pairs. In MARLHFIterTrainer, preference_scoring_reward="reward_model" makes the most recently fitted reward model the oracle for new comparisons, with the task reward as a fallback until one exists.

Preference Criterion#

For each task item, every agent samples preference_num_candidates responses to its prompt, and responses with the same candidate index form a joint response. Of any two joint responses with different scores, the higher-scored one is preferred: it is the winner \( w \), and the other is the loser \( l \). Ties are discarded. pair_selection and preference_pairs_per_sample decide which pairs each task keeps. In iterative training, each comparison is instead between the current agents’ joint response and the comparator’s joint response at the same candidate index, and either side can win.

Data Format#

A pair \( (\mathbf{h}_H^{w},\mathbf{h}_H^{l}) \) holds the two joint histories of one task. Each pair is a record like the following, with one entry per agent in each list:

{
  "prompts": ["<agent 1 prompt>", "<agent 2 prompt>"],
  "winner_completions": ["<agent 1 response>", "<agent 2 response>"],
  "loser_completions": ["<agent 1 response>", "<agent 2 response>"],
  "winner_reward": 0.92,
  "loser_reward": 0.41,
  "candidate_reward_mean": 0.63,
  "raw_rewards": [0.92, 0.41],
  "target_raw_reward": 0.92,
  "comparator_raw_reward": 0.41
}

winner_reward and loser_reward are the oracle scores of the two joint responses, and candidate_reward_mean is the mean score of the candidates compared for the task. The last three fields come from iterative training: raw_rewards holds the winner’s and loser’s scores before reward_processor, and target_raw_reward and comparator_raw_reward hold those of the current agents’ sample and the comparator’s sample. The iterative trainers write each round’s pairs to iteration_0001.json, iteration_0002.json, and so on under preference_replay_dir, as {"iteration": 1, "num_pairs": ..., "pairs": [...]}.

Training Pipeline#

Training a cooperative LLM-based multi-agent system with PbRL alternates between collecting preference data from the current agents and learning from it with MARLHF or MADPO.

• Base trainers (MARLHFTrainer, MADPOTrainer) collect the pairs once, then train on them. If no task yields a non-tied pair, training returns without updating the policies.
• Iterative trainers (MARLHFIterTrainer, MADPOIterTrainer) repeat the cycle for num_iterations rounds. Round \( \nu \) compares the current agents with the comparator, stores the new pairs as \( \mathcal{D}^{(\nu)} \), lets the replay buffer select the training pairs \( \mathcal{B}^{(\nu)} \) from this and earlier rounds, and trains on them. MADPO runs its updates for num_train_epochs per round. MARLHF fits a newly initialized reward model on them and then runs policy learning.

Evaluation always samples from the current policies and scores them with the task reward (reward_func and reward_processor), never with a learned reward model.

Comparator#

In iterative training, the comparator supplies the other side of each comparison. The current agents and the comparator respond to the same task, and every same-index comparison with a clear winner under the oracle function becomes a pair, whichever side wins.

comparator_policy chooses the comparator’s source:

• comparator_policy="current" (default): Active actor weights, sampled with a separate random-number stream.

• comparator_policy="current_copy": Current actor weights, copied to comparator_devices for each collection phase when those differ from the actor devices.

• comparator_policy="model": Explicit models from comparator_model_name or comparator_agents.

• comparator_policy="history": A saved actor checkpoint from comparator_history_k iterations back; the initial policy until one exists.

• comparator_policy="api": A remote model endpoint.

The history source saves policy checkpoints automatically. The two sides may sample different numbers of candidates; only indices present on both sides are compared, and pair_selection="comparator_reward" does not assume the comparator wins. By default, the comparator generates one response per role; with comparator_generation_mode="centralized", one comparator model writes every role’s response in a single joint response, which is parsed into role responses before pairs are built.

Model sources can be model repositories or local checkpoints, as described in Training Entities:

comparator_policy=model
comparator_model_name=/path/to/comparator-checkpoint
comparator_generation_mode=centralized
comparator_centralized_agent_index=0
comparator_devices=cuda:1

API Comparators#

With comparator_policy="api", the comparator queries a remote model. comparator_api_format selects the request format:

• comparator_api_format="generic" (default): The request includes the prompt, actor index, task item, candidate count, and generation settings; completions are read from comparator_api_response_field.

• comparator_api_format="openai": Chat Completions, also used by compatible endpoints such as DeepSeek’s.

• comparator_api_format="anthropic": Anthropic Messages API (https://api.anthropic.com/v1/messages, ANTHROPIC_API_KEY).

• comparator_api_format="openai_responses" (alias responses or codex): OpenAI Responses API (https://api.openai.com/v1/responses, OPENAI_API_KEY).

The Anthropic and Responses formats have no n parameter, so they send one request per candidate. Requests use the standard centralized or decentralized comparator prompts.

comparator_policy=api
comparator_api_format=anthropic
comparator_api_url=https://api.anthropic.com/v1/messages
comparator_api_model=claude-sonnet-4-20250514
comparator_api_key_env=ANTHROPIC_API_KEY

Replay Buffer#

The replay buffer selects the stored pairs used for training in each round. The current round’s dataset is saved first, so it is eligible alongside earlier rounds. preference_replay_mode sets which rounds are eligible:

• preference_replay_mode="current" (default): All pairs from the current round, without historical sampling.

• preference_replay_mode="nearest_k": The latest preference_replay_k non-empty round datasets, including the current one.

• preference_replay_mode="all_history": Every non-empty round dataset of the current trainer run.

• preference_replay_mode="lambda_decay": The same history, weighted toward recent rounds.

Each row of the diagram shows which round datasets a mode can draw from and how strongly each is weighted, with the newest round on the right.

nearest_k means the most recent rounds, not nearest-neighbor retrieval. In nearest_k and all_history, eligible datasets have equal weight regardless of size, and historical modes sample from them rather than training on every stored pair. Replayed pairs are re-tokenized from the saved round files. Setting preference_replay_dir to an earlier run alone does not import its files, while training checkpoints restore the replay index and its data.

Lambda Decay#

With age \( a=0 \) for the newest of \( K \) eligible datasets, preference_replay_lambda (\( \lambda \)) sets the selection weights:

\( \lambda=0 \) keeps only the newest dataset, smaller values favor recent rounds, and \( \lambda=1 \) weights every dataset equally.

Sampling#

Historical modes assign each of the preference_replay_sample_size draws to a dataset by weight, then sample pairs within each dataset. Sampling proceeds without repetition when the assigned count does not exceed the dataset size; larger counts cycle through shuffled copies. The combined sample is shuffled. The default sample size is the current round’s pair count, or the latest non-empty dataset’s size if the current round is empty.

current rejects preference_replay_sample_size and preference_replay_lambda and accepts preference_replay_k only as None or 1. The following configuration instead samples 256 pairs from all rounds, weighted toward recent ones:

from comlrl.trainers.preference import MADPOIterConfig

args = MADPOIterConfig(
    num_agents=2,
    preference_replay_mode="lambda_decay",
    preference_replay_lambda=0.8,
    preference_replay_sample_size=256,
)

MARLHFIterConfig accepts the same options. To replay only the latest rounds, users can set preference_replay_mode="nearest_k" and the number of rounds through preference_replay_k.

MARLHF#

Multi-Agent Reinforcement Learning from Human Feedback (MARLHF) fits a joint reward model to the preference pairs, then trains the policies with RL on the learned reward. Because the reward model scores the team’s complete response, separate actors learn from one shared reward.

Reward Learning#

The JointRewardModel \( \widetilde{R}_{\psi}(\mathbf{h}_t,\mathbf{a}_t) \) reads all agents’ prompts and responses as one text and applies a scalar head to the causal-LM backbone’s last hidden state at the final non-padding token. It minimizes the Bradley–Terry loss over the training pairs \( \mathcal{B} \):

where \( \sigma \) is the sigmoid function. The loss raises the winner’s score relative to the loser’s; it does not regress onto the task rewards. Training Entities describes how the reward model is loaded.

Policy Learning#

After reward learning, \( \widetilde{R}_{\psi} \) is fixed and takes the place of the task reward. The policies sample new joint responses for each task, \( \widetilde{R}_{\psi} \) scores them without gradients, and a MARL backend updates each agent’s policy to maximize the learned return:

rl_algorithm selects the backend:

• rl_algorithm="magrpo" (default): MAGRPO, with a group-mean baseline.

• rl_algorithm="mareinforce": MAREINFORCE, with raw returns.

• rl_algorithm="marloo": MARLOO, with a leave-one-out baseline.

• rl_algorithm="maremax": MAReMax, with a group-maximum baseline.

• rl_algorithm="maac": MAAC, with a centralized critic.

• rl_algorithm="iac": IAC, with per-actor critics or attached value heads.

Multi-Agent REINFORCE and Multi-Agent Actor-Critic describe these updates, with \( \widetilde{R}_{\psi}(\mathbf{h}_t,\mathbf{a}_t) \) in place of the reward \( r_t \). Optional frozen-reference KL regularization (reference_kl_enabled, reference_kl_coef, reference_devices) applies only to policy learning: each agent’s sampled KL estimate, scaled by reference_kl_coef, is subtracted from its reward signal.

reward_model_name is required when the actors are model objects without a model identifier. MAAC and separate-critic IAC likewise require explicit critic sources or a usable critic model identifier.

MADPO#

Multi-Agent Direct Preference Optimization (MADPO) trains the policies directly from the preference pairs, without a reward model or an RL stage. Its loss combines all agents’ preference margins, so each policy learns from the team’s outcome.

Each agent’s margin is the difference between its log-probabilities of the winning and losing responses, and the joint margin sums these across agents:

MADPO minimizes

where \( \eta>0 \) is dpo_beta (default 0.1). The task reward only orders each pair; the loss uses the policies’ log-probabilities. The gradient for agent \( i \) is

The weight \( \sigma(-\eta\,\Delta_{\boldsymbol{\theta}}) \) depends on the team’s joint margin, while the gradient flows only through agent \( i \)’s own margin. In each batch, the trainer computes all margins once without gradients, then updates each agent with its own differentiable margin plus the others’ detached margins.

The objective is reference-free: the inherited reference_kl_enabled, reference_kl_coef, and reference_devices settings add no reference terms to the MADPO loss. The training metrics turn_1/reward_mean and turn_1/expected_return average the candidate rewards recorded when the pairs were collected, so they do not reflect the updated policies.

The decentralized scorer excludes the last completion token after trimming padding, so each agent’s margin omits that token’s log-probability. The centralized scorer covers the full joint completion.