Does Gym return the whole forest, and does RL pick the paths?

NVIDIA-NeMo/RL #4124 (needs Gym #4128 and Gym #3368). Old RL links pin c617a2c, the first reviewed head; new RL links pin 2f3649d, the current head; Gym links pin 507d15c.

Before
At the first reviewed head, RL read the agent's kept calls as one list, in order. Each call had to continue the one before it, unless it started a new tree. A branch was an error.
This PR now
RL walks each end call back to its root, one row per end call. A new setting, token_capture.selection_policy, picks which calls count: returned_history (the default: the calls in the agent's output) or all_captured (every call Gym recorded).
Status
Mostly done. Gym sends the whole forest, and RL now chooses the paths. Still open: the two policies are a fixed list, not a plug-in point, and the stacked #4285 adds its own, different setting for the same choice. A pending review comment asks for one shared entry point.

Nine pages on this PR, one question each. This one: what the author's redesign changed in how calls become rows.

Background. Each time the agent calls the model, Gym saves one call record: an ID and the ID of the call it continues (its parent), or none if the call starts fresh. Compaction makes a fresh start, so one rollout becomes several trees: a forest. The agent also returns its normal output, which lists the answers it kept. A discarded retry is not in that list. RL joins the two and builds training rows. Dotted words have a short explanation: hover or tap them.

Gym already sent the whole forest at the first reviewed head. Its call list holds every call, discarded ones included, with parent links. That part did not change. Two RL steps changed. First, which calls count: the new selection_policy setting (nemo_gym.py L1305). Second, how those calls become rows, _plan_selected_calls. The old code walked the kept list top-down and required one straight line per tree. The new code sorts the end calls and walks each one back to its root. The rows below are the real output of each commit's _plan_selected_calls, run on the same four small graphs.

Same calls, old code vs the two new settings

Gym sends every call + parent Old (c617a2c) walk kept list top-down returned_history default: agent output all_captured every recorded call 1 · no compaction c1 c2 c3 kept by agent: c1 c2 c3 row 0: c1 → c2 → c3 row 0: c1 → c2 → c3 row 0: c1 → c2 → c3 all three the same 2 · compaction + a retried call c1 c2 compaction ↓ c3 c4 c5 c4 discardedc5 = retry, kept kept by agent: c1 c2 c3 c5 row 0: c1 → c2 row 1: c3 → c5 row 0: c1 → c2 row 1: c3 → c5 row 0: c1 → c2 row 1: c3 → c4 row 2: c3 → c5 same as old 3 · branch, agent keeps both c1 c2 c3 e.g. two subagents kept by agent: c1 c2 c3 ValueError: continuation must follow its predecessor row 0: c1 → c2 row 1: c1 → c3 row 0: c1 → c2 row 1: c1 → c3 c1 trains in row 0 only 4 · same branch, agent keeps only c1, c3 c1 c2 c3 not in agent output row 0: c1 → c3 row 0: c1 → c3 row 0: c1 → c2 row 1: c1 → c3 agent output picks Gym's list picks kept by the agent in Gym's list, not kept training row (root first) rollout fails to build

returned_history gives the same rows as the old code everywhere except case 3, where the old code failed. all_captured also trains the calls the agent discarded (c4, c2), with the same reward as the rest of the rollout. One exception: a call that failed, that nothing continues, and that isn't the scored answer is left out (L1324). The script is plan_old_vs_new.py: it pulls _plan_selected_calls out of each commit and runs it on Gym's own record types.

Where the paths meet

With forest_training off, RL still uses the older path, not the planner. All paths end in the same Gym function, verify_and_linearize, which takes one end call and walks back to the root. They differ only in which calls count:

# all: Gym's list of every call + parent (RolloutReceipt.manifest) forest_training off end call = the scored answer → 1 row finalize_group L742 → finalize_rollout → verify_and_linearize returned_history calls = matched from agent output select_captured_calls all_captured calls = every record − failed unused leaves inline in nemo_gym.py then, for both: end calls = counted calls with no counted child L1300 finalize_rollout L1427 → verify_and_linearize → N rows

Your proposal vs what landed

ProposalAt 2f3649d
Gym always returns the whole forest Yes, and it already did at the first reviewed head. The new head needs Gym #4128 (the close route and an admission field). That PR is public, but RL's Gym pin doesn't include it yet.
RL walks from each leaf back to the root Yes. One row per end call, through the same walk as the path without forest training. A shared call trains once (loss off in later rows).
RL options for which paths to train: main path, whole tree, subagents Mostly. returned_history is the main path the agent kept; all_captured is the whole tree. They are a fixed list (a Literal plus an if/elif), not a plug-in point. The stacked #4285 adds main_chain, a different all, and custom functions under its own call_selection setting.
Simple agent and compaction use the same reassembler Partly. With forest_training on, both go through the planner (case 1 gives one row). With it off, the older path still runs.
Subagents Partly. all_captured skips the one-session check, so calls from several sessions train if Gym records them under the rollout. They all share one reward, and no agent roles are used. returned_history still raises on more than one session.
Bottom line. The PR now does what you asked: Gym sends the forest, and an RL setting chooses which paths train. One thing is left. all_captured sits inline in nemo_gym.py, while returned_history is a function in gym_selection.py. The pending review asks for both to sit behind one lookup in gym_selection.py, so #4285's policies land as new entries under this setting instead of a second one (proposed, not in the PR).