How captured model calls become training rows: one chain walked from the leaf, or a list cut at every root

Explainer for NVIDIA-NeMo/RL #4124 (paired with NVIDIA-NeMo/Gym #3368). RL links pin c617a2c; Gym links pin b27e0db.

Before
Ordinary token capture builds one row per rollout by walking from the last call up to the first. If the harness rewrites its history, the whole rollout is masked.
This PR
With context compaction (CC), each rewrite starts a new root. RL picks the calls the agent accepted and cuts them into one row per chain.
Status
Proposed ā€” open PR, not merged. The values below come from a real run of the PR's code, with a stand-in model.

Nine pages on this PR, one question each. This one: how the reassembler turns a rollout's captured calls into training rows.

Background. Every model call in a rollout leaves two things behind. Gym's ledger holds one small record per call, including its parent call. TransferQueue holds the call's new tokens only: the observation tokens since the parent (no gradient) plus the tokens the model generated (gradient). A call with no parent is a root; it stores its whole prompt. The reassembler joins the two stores to build training rows. Page 2 shows how the data gets there. Dotted words have a short explanation: hover or tap them.

Without CC, every request is the previous one plus new observations, so the calls form one chain. RL takes the call the agent ended on (the leaf), follows parent links up to the root, and joins the stored tokens in order (Gym _terminal_chain). Your mental model is right for this case. With CC, a rewrite breaks the chain, so the ledger has several roots. RL no longer picks one leaf. It matches the agent's final output to the calls that produced it (select_captured_calls) and cuts that list at every root (_plan_selected_calls). Each piece then goes through the same leaf-to-root walk.

The same 5 calls, without and with compaction

one rollout Ā· 5 model calls Ā· the agent rejects call 4 and asks again (call 5) Ā· each call here generates 1 token Compaction off Ā· Gym's ledger is one chain c1 c2 c3 c4 āœ— rejected c5 leaf cum_len 3 ā†’ 6 ā†’ 9 ā†’ 12 (c4 and c5 both 12) 1. leaf = the response the agent ended on (c5) 2. follow parents up: c5 ā†’ c3 ā†’ c2 ā†’ c1, then reverse nemo_gym.py L1257-1300 (pick leaf) rebuild.py L201-216 (walk) 3. join each call's stored tokens, root first c1 [10,11,1001] + c2 [20,21,1002] + c3 [30,31,1003] + c5 [50,51,1005] rebuild.py L303-331 1 row Ā· group_g0 Ā· 12 tokens Ā· 4 with gradient [10,11,1001,20,21,1002,30,31,1003,50,51,1005] c4 is read, checked and cleaned up, but trains nothing. If the agent had rewritten its history here, Gym could not find a parent, and the whole rollout would be masked (unresolved_parent). Compaction on Ā· the ledger has two roots c1 c2 rewrite c3 c4 āœ— rejected c5parent: ā€“parent: c1parent: ā€“parent: c3 1. order: match the agent's output, answer by answer ā†’ [c1, c2, c3, c5] (c4's answer is not in it) gym_selection.py L175-229 2. cut: read each call's own parent field, no walking c1 none ā†’ new [c1] Ā· c2 ā† c1 ā†’ [c1, c2]c3 none ā†’ new [c3] Ā· c5 ā† c3 ā†’ [c3, c5] rollout_reassembler.py L1149-1163 3. reuse the old row builder per chain (re-walks as a check) group_g0_s0 Ā· 6 tok [10,11,1001,20,21,1002] group_g0_s1 Ā· 6 tok [10,11,1003,50,51,1005] Both rows get the rollout's one reward (1.0) and its one advantage (+0.5). Each generated token is in exactly one row. rollout_reassembler.py L1350-1376 (rows _s0, _s1) root (no parent; stores its whole prompt) continues its parent (stores new tokens only) rejected; never trained training row what RL does, in order Token values come from the PR's test stand-in for vLLM: every root's prompt is [10, 11], call i adds [10Ā·i, 10Ā·i+1] and generates 1000+i. Real prompts and outputs are far longer; the shapes, parents and row split are what the real code produced.

With compaction on, step by step. RL does not search the tree for leaves. It reads the order from the agent's output, cuts by each call's own parent field, and only then walks:

ā‘  match the agent's output to calls, left to right (gym_selection.py L175-229) output items 0-1 [reasoning 1, answer] ā†’ only c1 produced exactly these ā†’ c1 output items 2-3 [reasoning 2, answer] ā†’ c2 output items 4-5 [reasoning 3, answer] ā†’ c3 output items 6-7 [reasoning 5, answer] ā†’ c5 (c4's output is not in the conversation) selected = [c1, c2, c3, c5] ā† conversation order, no tree search ā‘” cut in one pass, looking only at each call's parent (rollout_reassembler.py L1149-1163) c1 parent: none ā†’ start chain 0: [c1] c2 parent: c1 = previous call ā†’ chain 0: [c1, c2] c3 parent: none ā†’ start chain 1: [c3] c5 parent: c3 = previous call ā†’ chain 1: [c3, c5] (a parent that is not the previous selected call is an error) ā‘¢ reuse the old row builder on each chain; it re-walks the chain as a check and joins tokens (leaf set at L1166-1172, walk in Gym rebuild.py L201-216) chain 0: c2 ā†’ c1 ā†’ stop (no parent) ā†’ row group_g0_s0 chain 1: c5 ā†’ c3 ā†’ stop (no parent) ā†’ row group_g0_s1

Step ā‘” already knows every chain; step ā‘¢ only reuses the existing finalizer, which checks hashes and lengths and joins the stored tokens, and RL confirms it finds the same calls (L1263-1273).

Why the output and not just the tree: without compaction, the leaf's ancestors are the whole conversation. With compaction, c5's ancestors stop at c3, and nothing in the tree says c1 and c2 come first or were accepted. The agent's output is the only record of what it kept, and in what order.

What RL checks while it cuts

The cut is strict, so a wrong ledger fails loudly instead of training on the wrong tokens:

  1. A selected call must continue the call just before it, unless it is a root. Otherwise the agent's output skipped a call, and RL raises (L1156-1161).
  2. Each chain is rebuilt with the ordinary code (finalize_rollout, the same hash and length checks), and the rebuilt call order must equal the selected order (L1250-1273).
  3. All or nothing per rollout: if any chain fails, every row of that rollout is masked (L1329-1349).
  4. One advantage per rollout, not per row. GRPO runs on one reward per rollout, then copies the result to all its rows (single_controller.py L5378-5393). A rollout cut into 3 rows still counts once in its group.
ledger, compaction on (real run; ids shortened) call parent mode prev_len delta_len cum_len response c1 a894ā€¦ - text 0 3 3 chatcmpl-1 c2 6e2cā€¦ c1 token_in 3 3 6 chatcmpl-2 c3 b483ā€¦ - text 0 3 3 chatcmpl-3 ā† compaction root c4 9577ā€¦ c3 token_in 3 3 6 chatcmpl-4 ā† rejected, not in output c5 df66ā€¦ c3 token_in 3 3 6 chatcmpl-5 ā† accepted retry
So what. Your mental model holds inside each piece: every row is still built by walking one chain from its leaf up to its root. What CC adds is a step before that walk. The agent's own output decides which calls were accepted, and every root starts a new piece. That is why a rewrite now costs one extra row instead of the whole rollout. It is also why the harness must return its full accepted history: a chain missing from the output is a chain that never trains (review comment on gym_selection.py).