Support shared-policy multi-agent RL training on the context-compaction capture path #4400

Need

Agent harnesses can delegate work to subagents with independent contexts. Training only the parent omits the workers’ generated actions. Training both requires reconstructing each agent’s actual conditioning, assigning loss to the intended actions, and keeping every training row associated with the task attempt that produced its reward.

Support both parent-only and parent-plus-worker training with one shared policy. Each agent may compact its context independently.

Foundation

Build on:

• NeMo-RL #4124 — context compaction
• NeMo-Gym #3368 — context compaction

These changes already separate one logical rollout from its physical training rows, capture actual generation tokens and model inputs, verify continuation versus context rewrite, identify accepted generations, and preserve rollout-level reward/advantage semantics.

Multi-agent support should generalize that path. Harnesses retain responsibility for delegation, communication, scheduling, and compaction; NeMo-RL owns training reconstruction.

Proposed approach

1. Identify accepted and training-eligible calls across agents. Extend shared ingestion to handle multiple sessions through ordinary outputs and stable call references. Distinguish accepted actions from retries, discarded branches, and auxiliary calls.
2. Reconstruct a forest of contexts. Replace sequential segment grouping with traversal of verified parent links. Find endpoints among selected calls and reconstruct their ancestry. Support independent roots, interleaved continuations, compaction, and shared-prefix branches.
3. Preserve unique action ownership. Shared ancestors may appear as conditioning in several rows, but each eligible generation receives loss exactly once. Worker answers inserted into the parent’s context remain masked observations.
4. Enable concurrent capture. Replace the attempt-wide sequential admission restriction with independent context reservations and per-call accounting. Preserve checks for ambiguous ancestry, stale continuations, and unresolved acknowledgements.
5. Close the whole episode safely. Reconcile admitted calls before publishing rows or cleaning up staging. A parent response finishing must not silently discard outstanding worker calls.
6. Reuse logical-owner training. All physical rows retain their task’s ownership. Compute one reward contribution per logical rollout in GRPO and distribute its advantage to eligible actions.

Initial scope and milestones

• Establish a parent-only baseline using an existing harness.
• Add deterministic reconstruction tests for multiple contexts and shared ancestors.
• Qualify concurrent capture, delayed results, cancellation, and episode closure.
• Demonstrate parent-plus-worker training with OpenCode.
• Qualify a second harness, such as LangGraph, to test the generality of the integration.
• Exercise compaction within a worker and preserve existing single-agent/CC behavior.

Initial scope is one shared trainable policy and one task-level reward. Separate policies, agent-specific credit assignment, and decoder prefix-sharing optimizations can follow independently.

Acceptance criteria

A controlled episode should include a parent, two overlapping workers, parent resumption, reversed worker completion order, and compaction in one worker. Verify:

• Exact conditioning, captured tokens, masks, and behavior logprobs.
• Complete accounting of admitted calls.
• Once-only loss for eligible actions.
• One logical-rollout contribution to the GRPO baseline.
• Loss/gradient agreement with an independently constructed per-call reference.
• A real learner update containing worker actions.

Related work

RL #4285 provides relevant ancestry traversal and shared-ancestor masking. Adapt those ideas to CC’s accepted-call selection and logical-owner machinery.

RL #4286 addresses shared-prefix computation and can remain a separate optimization.