Mosaic: Credit Assignment for Multi-Session Agentic RL — When Training Assumptions Break from Inference Reality

A Overlooked Disconnect

Imagine asking an AI assistant to research a complex question online. It runs a first search round, finds the clues insufficient, compresses what it has into a summary, and opens a new search round. Or it dispatches several sub-agents in parallel and aggregates the results.

In the agent world this is called multi-session operation — a single task trajectory split into multiple segments.

Here's the problem: mainstream Agentic RL training methods (GRPO, PPO, etc.) assume the agent completes everything in one continuous context. During training, the whole trajectory receives one reward that is "broadcast" to all tokens. But at inference, the trajectory is split — some segments made key decisions, some just organized information, some wandered off track.

Broadcasting the same reward to all segments is like a company paying everyone the same bonus — whether you closed the big deal or slacked off.

The Mosaic paper (arXiv: 2610.22174) targets exactly this disconnect.

The Core Problem: Three Flaws of Broadcasting

When agent trajectories span multiple sessions, traditional "broadcast-style" advantage estimation has three fatal issues:

1. No contribution distinction: A successful trajectory may contain sessions with critical contributions and sessions that just wasted tokens. Broadcasting cannot tell them apart.

2. No guidance on session allocation: If one trajectory succeeded with 5 sessions and another succeeded with 3, which is better? Broadcast-style methods provide no signal to guide the agent toward finishing with fewer sessions.

3. Train-inference mismatch: Training assumes a single session; inference uses multiple. The model's learned ability to "make decisions within one context" gets disrupted by context compression and sub-agent forking at inference time.

Mosaic's Solution: Hierarchical Credit Assignment

Mosaic's core idea is critic-free hierarchical optimization — no extra value network, with advantage estimation at two levels.

Trajectory level: encoding "session usage" into advantage

At the trajectory level, Mosaic looks not only at "did the task succeed" but also "how many sessions were used." Among successful trajectories, the one with fewer sessions gets higher advantage — like rewarding a 3-person team over a 5-person team for the same project.

Two ranking rules are designed:

• C1 (among successful trajectories): fewer sessions ranks higher
• C2 (among failed trajectories): among equal failures, more thorough exploration (more sessions) ranks higher — because failure may be bad luck, and thorough exploration shows the agent was genuinely trying

Session level: marginal contribution shaping

At the session level, Mosaic estimates each session's marginal information gain — how much would the trajectory's success probability drop if this session were removed?

It's like evaluating each member of a team: not by how much they did, but by whether the project would collapse without them. Key decision-makers get high marginal contribution; information-organizing sessions get low.

A gating mechanism handles a subtle case: trajectories where all sessions fail. You can't simply rank by marginal contribution (all negative), so the method preserves directional consistency — more exploratory sessions receive relatively higher advantage.

Experimental Results: Fewer Sessions, Higher Success Rates

Experiments use the BrowseComp benchmark with Qwen3.5-4B, testing two multi-session scenarios:

Context compression scenario:

• Mosaic achieves higher success rates than broadcasting
• Uses fewer sessions on average
• Per-session utility (marginal contribution per session) improves significantly

Sub-agent forking scenario:

• Also outperforms broadcasting
• Better session allocation — key decisions get more "budget," ineffective exploration is suppressed

The paper also includes a mechanistic analysis: comparing Mosaic's advantage estimates against "true" per-session contribution (via leave-one-out estimation), Mosaic's estimates correlate with true contribution far better than broadcasting. This shows hierarchical advantage estimation genuinely captures each session's real contribution.

Why This Paper Matters

1. It exposes an overlooked disconnect

Most of the Agentic RL community focuses on "better reward design" or "scaling rollouts," but few noticed the session-structure mismatch between training and inference. Mosaic puts it on the table.

2. It offers a simple, effective fix

No extra critic network, no complex training pipeline changes — just two adjustments to advantage estimation. "Small change, big effect" methods like this tend to be adopted quickly by the community.

3. It points to key future directions

The paper's discussion raises several forward-looking directions:

• Cross-session state transitions as learning objectives: not just assigning credit, but learning "how to pass key information from one session to the next"
• Asynchronous training: multi-session rollouts are naturally suited to asynchrony — different sessions can proceed in parallel
• RL inside black-box agent frameworks: how to do RL when the agent framework (e.g., Claude Code, Cursor) is a black box?

A Deeper Insight: The Granularity of Credit Assignment

Mosaic raises a more general question: how fine-grained should credit assignment in RL be?

• Coarsest: trajectory-level reward broadcast to all tokens (traditional GRPO)
• Medium: session-level advantage (Mosaic)
• Finer: chunk-level attribution (e.g., Hy-Memory's chunk-level attribution)
• Finest: token-level credit assignment (theoretical ideal, computationally infeasible)

Each level down gives more accurate credit but higher estimation variance. Mosaic's session-level choice is a sweet spot — fine enough to distinguish contributions, coarse enough to keep estimation stable.

This parallels human organizational management strikingly well: a CEO doesn't score every employee (too fine, too noisy) or pay everyone identically (too coarse, unfair). Department-level evaluation is the common compromise. Mosaic's session-level advantage estimation is the agent world's "department-level KPI."

Limitations and Reflections

The paper has some debatable points:

1. Marginal contribution estimation: Estimating marginal contribution by measuring success-probability change after removing a session requires extra rollouts — non-trivial compute overhead. The paper doesn't discuss this trade-off in detail.

2. Tested only on BrowseComp: BrowseComp is a web search task, naturally suited to multi-session work. But code generation, math reasoning, etc. may have different multi-session structures — code generation may need fewer session switches.

3. Representativeness of Qwen3.5-4B: A 4B model's behavior may differ from larger models. At 70B or beyond, broadcasting may be good enough — large models may learn to distinguish session contributions on their own.

Conclusion

Mosaic reveals a simple but overlooked fact: when agents shift from "doing it all in one go" to "doing it in segments," training methods must evolve too. Broadcast-style advantage estimation is a product of the single-session era; the multi-session era needs session-aware credit assignment.

This is not just a technical issue but a philosophy of engineering: when your system's structure changes, your evaluation method must change with it. Otherwise you're measuring a new world with an old ruler.

---

Paper: https://arxiv.org/abs/2610.22174 Project page: https://lark-ai-lab.github.io/Mosaic/ Code: https://github.com/LARK-AI-Lab/Mosaic