Context compaction, redesigned: the row cut moves from "after the rollout" to "inside the vLLM worker"

Explainer for NVIDIA-NeMo/RL #4124, paired with NVIDIA-NeMo/Gym #3368. RL links pin c617a2c; Gym links pin the Gym PR head b27e0db, or 17abf2d (Gym main) for code the Gym PR deletes.

Before: NeMo-RL main has no context compaction. An earlier design (never merged) had Gym send every call's full prompt tokens, and RL found the cuts afterwards.
This PR: The agent still decides when to compact, and it tells RL nothing about it. RL's vLLM worker detects the effect on each call: if the new request no longer extends the previous call's history, that call starts a new row. RL rebuilds the rows from ordinary Gym output.
Status: Proposed — open PR, not merged. It needs Gym #3368 merged and the Gym pin bumped first. Review findings are in the PR's review comments.

Nine pages on this PR, one question each. This one: how the new design differs from the earlier one.

Background. A harness is the agent loop in Gym that talks to the model. On a long multi-turn task its prompt grows every turn, mostly from screenshots. Context compaction means the harness sometimes rewrites its prompt, for example by replacing old images with the text "[Earlier image omitted]". After a rewrite the new prompt no longer starts with the old one, so the two cannot share one training sequence, and one rollout becomes several training rows. TransferQueue is where the new design keeps the tokens. Dotted words have a short explanation: hover or tap them.

Both designs end with the same rows. The example below is one rollout of 5 turns, with a rewrite after every 2 actions that keeps only the newest image. It gives 3 rows: turns {1,2}, {3,4} and {5}. In each row, only the tokens the model sampled carry gradient. In both designs the agent decides when to compact. The difference is who turns that rewrite into a row boundary, when, and from what evidence. In the new design that decision is one short function, decide_capture_input. It runs inside the vLLM worker before generation starts.

Same 5 turns, two ways to find the cut

The new check runs before any tokens exist, so it cannot compare tokens. Instead it compares the harness's own messages. Gym hashes the messages of every call it served and stores that hash as a ReplaySummary. The worker hashes the start of the new request the same way. If the hashes match and the request only adds user or tool messages, it is a continuation: the worker reuses the parent's exact sampled tokens and stores only the new part. Otherwise, including a request shape the worker cannot splice, it is a new root: the request is rendered from scratch and stored in full. The harness sends only messages: no tokens, no segment IDs and no parent hint. Gym proposes the earlier call whose history the new request extends (Gym sink.py L352-388), and the worker checks that proposal.

What changed, hop by hop

	Earlier design	This PR
Who cuts	The agent compacts; RL finds the cut after the rollout by scanning every call	The agent compacts; RL's vLLM worker finds the cut per call, at request time, by checking whether the request still extends its parent (vllm_worker_async.py). The reassembler only follows the parent links.
Continuation test	"Is the old prompt a token prefix of the new one?" plus the same test on image IDs. It was also cross-checked against Gym's own list of rewrites.	"Do the source messages and render options hash the same as the parent's summary?" There is no second check against Gym's rewrite list.
Gym → RL	One large object per rollout (`ContextCompactedResponse`, deleted by Gym #3368): every call's full prompt tokens, rewrite events, and images by ID	The ordinary rollout output plus a small call list with no tokens (nemo_gym.py). Tokens and pixels go straight from the worker to TransferQueue.
Storage	The full prompt again on every call	A continuation stores only the new part; only a root stores the full prompt
Images	Stored once by ID in a `MediaArena`; RL processed them again for training	The harness re-sends the images it kept. The worker keeps vLLM's own processed pixels. A kept image whose pixels changed is an error, not a new row (captured_media.py).
Harness	Had to return the special response object	Any harness whose normal output can be matched back to the calls it made (`select_captured_calls`). It may also reject a response and resample.
Training math	Advantage once per rollout, then copied to rows; one global token count as the loss denominator; a special scheduler step size	The same advantage and loss rules (`build_logical_owner_batch`). The scheduler trick is no longer needed: SingleController takes one optimizer step per training step whatever the row count (megatron_policy_worker.py).
Scope	Classic GRPO loop, Megatron, no data store	SingleController GRPO with TransferQueue. Megatron fixed-batch only: no sequence packing, no dynamic batching, no interleaved pipeline (config.py). No checkpoint or resume (config.py). One model call at a time per rollout.

Which harnesses work once this merges

Harness (Gym agent)	How it compacts	Evidence in the PRs	Status
`simple_agent_with_compaction` (config)	Gym's shared compaction client: drops old images and old reasoning every K actions	The PR's CPU tests and the manufactured image runs in the PR description	Supported; the reference harness. Now requires training token capture, so it no longer runs in plain Gym evaluation (L148-151).
`aviary_agent`	Its own option: hides old environment-state messages (`collapse_old_env_states`)	One CPU test with the real agent (test L268)	Tested on CPU only
`browsecomp_agent`	Its own option: keeps only the last N tool rounds (`keep_rounds`)	CPU test (L316); its progress-board recovery mode is tested to fail, because the recovered answer was never generated by the model (L345)	Tested on CPU only; recovery mode not supported
`opencode_agent`	OpenCode's own summarization	Only a private fork, in the live runs of the PR description; the stock agent just records that compaction happened	Not supported as shipped (guide L21)
Any other harness	Its own policy	None	Untested. It must make one model call at a time, not stream (or use buffered Responses streaming), return its whole accepted conversation, and never present text the model did not generate as model output (guide L19-21).

Added in place, or a copy? A copy, now a thin one. Before this PR, simple_agent_with_compaction was already a forked copy of simple_agent's loop, with its own compaction package and its own model server (vllm_model_with_compaction). The Gym PR moves the compaction logic into shared Gym code (nemo_gym/context_management/), deletes the old package (about 1,700 lines) and the special model server, and shrinks the agent to 233 lines. It is still a separate agent: it shares one helper with simple_agent (import L30), and a new parity test runs both loops side by side. simple_agent itself only had that helper pulled out; it has no compaction support.

What you gain, and what you give up

Gain — rows are exactly what was served. A continuation reuses the parent's sampled tokens instead of rendering them again, so a template quirk can never create a false cut or a row the model did not see.
Gain — much less data moves. The earlier design repeated the full prompt on every call. Here only a rewrite repeats it.
Gain — any harness. Gym no longer needs a special response type, so a harness keeps its own compaction policy and its normal output.
Give up — the second opinion. The cut now depends on one hash comparison, not a token test checked against Gym's own rewrite list. Fields the hash ignores (item IDs, the JSON spacing of tool arguments) count as "unchanged". An edit to only those fields never reaches the model. Training is still exact, because it uses the tokens that were served.
Give up — tolerance. A kept image whose processed pixels changed on a continuation is now an error, where the earlier design opened a new row. The harness must also return its full accepted history, or the earlier rows cannot be found.
Open issue — a dropped prompt group ends the whole run. A failed model call now masks only its own rollout, and its siblings still train. But when no agent result comes back (a timeout, a transport error, or an error raised inside the agent such as a context-guard rejection), or every rollout in the group fails, the whole group is dropped. Compaction requires every group in a step (config.py L1437-1438) and never replaces a dropped one, so the run stops.
🔴 review comment — context-compaction.md:47

What to review. The training math did not change. What changed is the evidence behind the cut and what data moves. Correctness now rests on three places: the worker's decision (decide_capture_input), matching Gym output back to calls (select_captured_calls), and the chain planning in the reassembler (_plan_selected_calls). The PR cannot merge before Gym #3368 lands and the Gym pin is bumped.