Mosaic trains LLM agents whose trajectories span multiple sessions (from context compaction or subagent forking) by splitting credit two ways: a trajectory-level advantage that rewards finishing with fewer sessions, plus a session-level shaping signal based on how much each session raises the likelihood of the correct answer.
Suppose you’ve built a web-research agent on top of a small open LLM. Hard questions take dozens of tool calls, so you add two tricks that production agents already use: when the context window fills, the agent summarizes its history into a shorter prefix and starts fresh (context compaction); when a subtask is independent, it spawns a subagent that returns a result (subagent forking). Both tricks break the trajectory into multiple separate token streams, called sessions.
Now you want to improve the agent with reinforcement learning. The standard recipe, Group Relative Policy Optimization (GRPO), assumes one continuous trajectory per rollout and assigns a single advantage from the final reward. The usual multi-session workaround, called broadcasting, just copies that trajectory-level advantage to every token in every session. The paper’s complaint: broadcasting cannot tell which session actually helped, and gives no pressure to use sessions efficiently. The policy may learn to spawn more and more sessions without getting better answers.
Mosaic adds two signals on top of broadcasting, both without training a separate value network.
First, a trajectory-level nudge based on session count. Within a group of rollouts for the same prompt, successful trajectories that used fewer sessions get their advantage scaled up; successful ones that used many sessions get it scaled down. For failed trajectories the direction flips: giving up early is penalized more than exploring longer. Crucially, the scaling is bounded (a multiplicative factor in roughly 0.8 to 1.2 in their setup), so a failed-but-long trajectory can never outrank a succeeded one. The ranking is done on centered ranks of session count within the success and failure subgroups separately, so raw outliers don’t dominate. When all rollouts in a group have the same outcome, a fallback adds a small session-count-based signal so the gradient isn’t zero.
Second, a per-session shaping reward based on marginal information gain. Treat each session as a transition from a prefix state v_{j-1} to a new prefix state v_j. Score each prefix by the average log-probability the current policy assigns to the ground-truth answer when conditioned on that prefix. The session’s reward is the increase in that score. Intuitively: did this session make the right answer easier to generate? These per-session rewards are normalized separately for each session type (normal, compaction, subagent) because their raw score distributions differ a lot. The last session is excluded because it usually contains the answer itself, which would inflate the signal.
The two signals are combined additively into the session’s advantage, with gates that prevent the shaping term from flipping the sign of the outcome-based advantage.
for group in prompt_groups:
for traj in group:
A_traj = grpo_advantage(traj, group) # outcome-based
u = centered_rank(traj.num_sessions, same_outcome_subgroup)
A_sat = (1 - rho * u) * A_traj # session-aware
for j, session in enumerate(traj.sessions[:-1]):
r_delta = logprob(answer | prefix_after[j]) \
- logprob(answer | prefix_before[j])
r_hat = normalize(r_delta, by=session.type, group=group)
session.advantage = A_sat + lambda_r * r_hat # final per-session
All experiments use Qwen3.5-4B as the base agent on BrowseComp, a hard web-navigation benchmark, with ReAct-style tool use (search and visit).
•
Under context compaction, Mosaic reaches 38.86% accuracy versus 35.31% for broadcasting and 34.68% for single-session Group Relative Policy Optimization (GRPO) run with compaction at test time. Mosaic also uses fewer sessions per trajectory on average (2.35 vs 3.17).
•
Under subagent forking, Mosaic reaches 32.46% versus 31.44% for broadcasting, again with fewer sessions (3.77 vs 4.43).
•
An ablation splits the two components. Both the session-aware trajectory advantage and the session-level shaping help on their own. Combining them helps clearly under compaction but only modestly under subagent forking. The authors speculate that subagent trajectories are graph-shaped rather than linear, so a single session’s marginal information gain is a weaker proxy for its true contribution.
•
Compared to IGPO (Information Gain Policy Optimization), which computes information gain at every turn rather than per session, Mosaic achieves higher accuracy and spends substantially less wall time computing the shaping signal.
•
Behavioral analysis: successful multi-session training (both broadcasting and Mosaic) raises the fraction of sessions with positive information gain and the average positive gain per session. Mosaic strengthens this pattern without pushing session counts higher.
All numbers are on a single 1,266-question benchmark with one base model, so they describe a specific setting, not a universal result.
•
If you train agents with Group Relative Policy Optimization (GRPO) and your inference workflow already breaks trajectories across sessions (summarization, forking, harness-managed restarts), broadcasting is leaving signal on the table. Worth testing a session-count-aware advantage scaling: it’s a few lines on top of your existing GRPO code and does not need a critic.
•
The marginal-information-gain shaping needs two extra forward passes per non-terminal session to score the ground-truth answer under the current policy. That requires training-time access to the model’s log-probabilities and labeled answers, so it fits open-weight training setups, not hosted-API fine-tuning.
•
If you only care about one of the two mechanisms, the ablation suggests the trajectory-level session-count signal is the safer bet, especially for non-linear workflows like subagent forking where per-session credit assignment is noisier.
•
For a research reader: the paper reframes the question from “how many sessions should an agent use” to “how much useful progress per session,” and operationalizes “useful” as a change in answer log-likelihood. That framing is reusable even if you don’t adopt Mosaic itself.
•
Single benchmark (BrowseComp), single base model (Qwen3.5-4B, 4B parameters), one training set of 1,000 tasks. Generalization to other agent tasks or larger models is untested here.
•
The information-gain signal depends on having a ground-truth answer to score against. It applies to verifiable QA-style tasks, not open-ended generation.
•
Rewards are binary success/failure in the training setup; the ordering guarantees for the session-aware advantage assume this binary structure.
•
Subagent gains are smaller and the authors attribute this to the non-linear trajectory structure; per-session credit assignment via marginal gain is a weaker fit when one session’s output feeds many downstream sessions.
•
The training setup is synchronous (colocated rollout and update). The paper flags, but does not solve, how multi-session optimization behaves under asynchronous rollout sampling where trajectories of very different lengths finish at different times.