Repository navigation
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
ananthsub
changed the title
This was referenced Sep 1, 2026
ananthsub
force-pushed
the
ananthsub/checkpoint-gymnasium-agent
branch
2 times, most recently
from
September 3, 2026 11:12
afcc45f to
80101d2
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-gymnasium-agent
branch
from
September 10, 2026 20:48
80101d2 to
3e739ef
Compare
ananthsub
force-pushed
the
ananthsub/checkpoint-gymnasium-agent
branch
from
September 22, 2026 17:29
3e739ef to
97718e3
Compare
Preserve accumulated episode state at completed environment steps so replacement attempts skip reset and continue with consistent reward, usage, and model history. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Retry refused model and resource calls from the committed boundary without adding an extra turn at the step limit. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
ananthsub
force-pushed
the
ananthsub/checkpoint-gymnasium-agent
branch
from
September 22, 2026 17:38
97718e3 to
ca73662
Compare
OliviaViessmann
pushed a commit
to OliviaViessmann/Gym
that referenced
this pull request
Oct 16, 2026
…continuation (NVIDIA-NeMo#3888) ## What changed and why This PR builds on the agent session lifecycle that NVIDIA-NeMo#3784 added to `main`. The base server's seed and close also open and close the session in the checkpoint participant, and Simple Agent exports, restores and retires the sessions the base server holds. Every agent takes part in checkpoints once checkpointing is on. - Agents that set `checkpoint_sessions_supported = True` and implement the session hooks get an `AgentSessionParticipant`. A session is idle, running, or at a boundary. A running activation parks at its next boundary when a checkpoint closes admission, and one waiting on a model call counts as parked at its last boundary. If that call's reply arrives while the checkpoint is open, the session stays parked at the boundary before the call and holds the reply until resume, so its readiness never regresses between prepare and commit. The model may already have recorded the call, because it delivered the reply before its own prepare. After a crash the restored agent calls again from that boundary, and the token-capture builder treats the two rows like a client retry: it keeps the call the agent continued from. Boundaries are lazy snapshots, so a long turn loop does not serialize its history at every step. Restored sessions are installed under the replacement attempt. - The hooks are coroutines, and export takes every session at once: `export_agent_sessions(session_keys)`, `restore_agent_sessions(sessions)`, and `retire_agent_session(session_key)`. An agent whose state lives outside the process, such as a sandbox it runs tools or a harness in, can checkpoint all of it concurrently within the commit. - A legacy `/run` is an episode of its own: seed and the turn loop replay, and verification uses the mode the resources server reported on its seed reply. - Agents without the hooks get a restart-only participant. The server framework installs it on whatever app `setup_webserver` returns, so agents that build their own app, such as `harbor_agent` and the `mini_swe_agent` agents, take part too; it is installed once per app. Their rollouts are restarts: they never block prepare, keep running through the checkpoint, and are reported in prepare's `restarts`, so one opaque harness no longer stalls checkpointing for every environment. The seed reply carries `x-ng-checkpoint-restart`, and only a new seed waits while a checkpoint is open. After a crash, the controller retires the rollout and starts it over from its input. - With checkpointing on, agent work without a rollout ID is refused with `rollout_id_required`, and every agent refuses `num_workers > 1`. - Simple Agent implements the hooks for its environment server sessions and for its legacy `/run` loop. A session's export carries how many activations it has started, because each activation reports its observations under its own invocation ID (`root`, then `activation-2`, and so on). A restored session's next activation therefore gets the next ID, and an activation continued from its restored boundary keeps the ID it already had. - Simple Agent's session export carries its resources cookie jar as it is. Since NVIDIA-NeMo#4449 on `main`, a session without a direct tool grant holds no resources cookie jar (`resources_cookies` is `None`), so the environment server keeps the session cookie from its own seed for `/verify` and `/close_session`. Export and restore keep that `None` instead of turning it into an empty jar, so a restored session's close still reports no jar and the environment server keeps the cookie from its resources seed. - A commit exports only the live sessions whose episode is inside its scope, as it already did for restored sessions. The controller leaves every participant's restarts out of the scope, so a session another server reports as a restart is not exported here either. - An unclaimed restored session, one whose replacement has not started yet, keeps its restored boundary in the next checkpoint while the controller's commit scope still continues it, so a second crash still continues it. Once the scope leaves it out, it is not exported, and it is retired. - **Activations are idempotent.** The environment server sends each activation's index, its 1-based count within the session, in the `x-ng-agent-activation` header. With checkpointing on, the agent exports the index and the reply of its last completed activation with the session. An invocation of an activation that already completed gets the stored reply back without running the loop: no model or tool call, and no change to observations or the activation count. This keeps the environment server and the agent consistent when the environment server's checkpoint still has the invocation ahead of it while the agent's checkpoint has the activation done. An index out of order is refused with 409 `activation_out_of_order`. Without the header, every invocation is a new activation, as before. - A request for an activation that is still running waits for it to end and then gets its reply, instead of being refused. - With checkpointing on, every activation of an agent session must carry the `/ng-rollout/<id>` prefix. One without it is refused with 400 `rollout_id_required` before any replay, claim or activation, since a restore rewrites the session's episode to the replacement attempt and the session can no longer stand in for a missing path. This holds for restored sessions and for sessions seeded in the running process; every caller that activates an agent session is a Gym environment server, and each sends the prefix. - Before it replays a reply or starts an activation, the agent checks the attempt the request's `/ng-rollout/<id>` path names against the session's, and refuses a mismatch with `stale_attempt`, so a crashed attempt's harness cannot use a restored session. - A replayed reply claims a restored session, like a new activation does, so a later commit that leaves the episode out does not free a session the live episode still uses. - The reply is kept only when the request carried an activation index. - **A legacy `/run` sends activation index 1 on its self-dispatch.** Its turn loop records its reply like an environment server session's activation, and the reply is exported with the legacy session. A checkpoint taken after the loop replied but before `/run` recorded its `verify` boundary therefore restores the reply instead of running the loop again; the reply leaves when the `/run` ends or a retire frees it. - **Direct `/v1/responses` calls are restarts.** With checkpointing on, a direct call with a rollout path is reported as a restart on every agent, and one without a rollout path is refused with `rollout_id_required`, as restart-only agents already did. - **Restart-only seeds respect retires.** A restart-only seed checks the retired mark after it waits for resume, a retire waits for seeds in progress, and the restart-only retire frees seeded sessions through `retire_agent_session`. The base `retire_agent_session` now releases a session through the agent's close hook, so a restart-only agent stops a sandbox it started. - **Restarts from earlier stages do not hold up agent prepare.** The agent participant leaves the restarts carried from earlier prepare stages (NVIDIA-NeMo#3886) out of `readiness`, `ready` and its export. An activation of such an episode still parks at its next boundary until resume, but it no longer holds up prepare, so an episode `single_agent_turn` marked as a restart, from a restart-only resources seed or an MCP grant, does not stop the checkpoint while it is inside a slow tool call. - **Retire cost is per rollout.** Sessions, seeds in progress and direct restart calls are indexed by rollout, so a retire touches only its own rollout's, and each index entry leaves with the last item of its rollout. - **Simple Agent boundaries leave out each turn's `question`**, which repeats the request's input and the loop's earlier outputs. They store `question_outputs`, how many outputs each question held, and a restore rebuilds each question from the request's input and the restored outputs. The boundary snapshot is built directly, without a validate-and-dump round trip. - **A completed activation's reply is kept once.** A session-scoped reply to a request that carries an activation index leaves out `_ng_trajectory`, in the original and in the copy kept to replay it, since `single_agent_turn` takes the session's evidence from the `AgentObservationBundle` at close; callers that send no index, such as the UserSim environment server, still get it. A legacy `/run` drops its turn loop's kept reply in the same step that records its `verify` or `return` boundary, which holds the reply. Kept replies stay unserialized until a replay or an export needs them. - **An episode that grants MCP access is a restart.** `single_agent_turn` marks an episode as a restart when it grants an `MCPToolAccess`, so no checkpoint continues an agent that holds an MCP token past its seed; with NVIDIA-NeMo#3886's token fence, a replaced attempt's token stops working after a restore. - **Direct HTTP tool calls name their attempt.** With checkpointing on, `single_agent_turn`'s direct HTTP grant carries the seeding attempt's `/ng-rollout/<capture key>` prefix in its base URL, since NVIDIA-NeMo#3886 refuses a call on a restored session that names no attempt. Simple Agent sends each direct HTTP tool call under the rollout prefix of the activation it is serving, in place of the prefix in the grant, so after a restore the continued session's calls name the replacement attempt; its legacy `/run` tool calls carry the same prefix. No other agent with checkpoint support uses direct HTTP grants, and restart-only agents seed again after a crash and get a fresh grant. - **A legacy `/run` closes its resources session after a replayable verification.** It sends `x-ng-checkpoint-close-follows: 1` on `/verify`, so NVIDIA-NeMo#3886 keeps the session tracked and exported until the close, and then closes the session, also when a replacement is restored at its `return` step. A failed close is logged and keeps the run's verdict. - **A `skip_verification` legacy `/run` records its `return` boundary.** It records the same boundary as a verified run, with the configured reward as its result, so a checkpoint that finishes while the turn loop's reply is on its way parks the run at `return` until resume instead of letting it finish uncaptured. A run restored at `return` returns its recorded result. - **Closing a session stops its activation before cleanup.** A close first stops an activation still running for it, so an activation the environment server gave up on stops calling the model and tools, and only then runs the agent's close hook. A close cancelled while it waits for the activation has not run the hook, so a retried close runs it exactly once. The session stays tracked, and holds up prepare, until the activation has stopped. A close cancels only an activation nothing has cancelled yet, so it never interrupts a retire's wait for the activation's cleanup. - **A close waits for resume while agent admission is closed**, as a seed does, and checks admission before it marks the session as closing. So a commit that has checked readiness always exports an open session, never a closing one. This cannot stall prepare through the environment server's final cleanup: coordination prepares environment servers before agents, and an episode in final cleanup keeps its environment server from being ready, so its agent close finishes before agent admission closes. A close that a retire took over is refused with `stale_attempt`, and checks after every await that its session record is still current, so it never writes a receipt for a session a retire freed. - A retire cancels the attempt's activations, or a restart-only agent's `/run` calls, and waits until they have stopped before the retire hook runs. The hook is where a sandboxed agent stops its sandbox. A retired session stays tracked until its activation has stopped and the legacy episode's retire and the retire hook have finished, so a retire cut short by its deadline at any of those points leaves it for a retry, which waits again without cancelling a second time. Meanwhile the session blocks the next checkpoint and is not exported, and a legacy `/run` blocks whether or not its session is still tracked. - A retired attempt's late activation is refused with `stale_attempt` until the controller forgets the rollout. - A session seed that arrives while a checkpoint is open waits for resume, and one already in progress holds up prepare and is exported. The environment server seeds inside a replay step that a checkpoint does not wait for, so the agent's session must match whichever boundary the environment server records. Refusing the seed failed the episode, and admitting it but leaving it out of the checkpoint let a replacement invoke a session nothing restored; a random checkpoint sequence against a real `single_agent_turn` deployment found both. A new session that is not a seed, such as a legacy `/run`, is still refused. Seeds are counted per session key, so a retried seed overlapping the first is tracked until both finish. A retire also waits for seeds of its attempt that are in progress, and a seed refused because its attempt was retired frees the state it built before the retire replies. A restore is refused while a seed is in progress. - `SimpleResponsesAPIAgent._free_session` frees a retired session that has no close receipt; it was `_discard_session`. - The participant answers the readiness wait's `ready()` without building its blocker list, and notifies only while a checkpoint is open. ## How it works ### Where this PR sits in the overall flow The highlighted part is what this PR adds. ```mermaid flowchart LR C["Controller<br/>NeMo RL, or rollout collection"] CO["Coordination<br/>prepare, commit, restore, resume,<br/>retire, forget"] K["Control plane on every server<br/>phases, lease, storage,<br/>retire: refuse, stop, free;<br/>forget"] subgraph G["One participant per Gym server"] E["Environment server<br/>episode steps"] M["Policy model<br/>held responses, generation cuts"] A["Agent<br/>sessions parked at boundaries"] R["Resources server<br/>session state"] end W["Inference worker<br/>stages cut prefixes"] D[("Checkpoint directory<br/>records, then manifest")] L[("Capture ledger<br/>retire and delete from NVIDIA-NeMo#3938, NVIDIA-NeMo#3939")] C --> CO --> K K --> E & M & A & R M --> W M --> L G --> D classDef this fill:#fde68a,stroke:#b45309,stroke-width:2px,color:#1f2937 class A this ``` One checkpoint, a crash, and the restore, end to end: ```mermaid sequenceDiagram participant C as Controller participant G as Gym participants participant D as Checkpoint directory C->>G: prepare, in order environment, model, agent, resources Note over G: admission closes, in-flight work parks at a boundary,<br/>undelivered model responses are held G-->>C: prepared, or blockers at the deadline C->>G: commit with the episodes the controller continues G->>D: each participant writes its records, then its manifest Note over G: restored state the commit's scope leaves out is deleted C->>C: publish the checkpoint with the controller's own state C->>G: resume, in order resources, agent, model, environment Note over C,G: crash - every Gym process dies C->>G: restore the checkpoint in fresh processes, all or nothing D-->>G: records installed under attempt N+1, attempt N's capture ledger retired C->>G: resume C->>G: /run as attempt N+1 continues each episode from its boundary Note over C,G: retiring an episode, only while no checkpoint is open C->>G: retire - environment, then agent, then model and resources Note over G: each server marks the attempt, stops its work, waits, frees its state, then replies C->>G: forget, once the rollout is done for good - same order Note over G: the mark refusing the retired attempt goes away ``` ### This PR An agent session's states. Prepare is ready once no activation is between boundaries. ```mermaid flowchart LR I(["idle"]) -->|an activation starts| R(["running"]) R -->|before a model call, or after a tool round| B(["at boundary"]) B -->|next step, when no checkpoint is open| R B -->|parked while a checkpoint is open| B R -->|the activation ends| I ``` Which participant an agent gets. ```mermaid flowchart TD Q{"checkpoint_sessions_supported"} -->|yes| P["AgentSessionParticipant<br/>export_agent_sessions, restore_agent_sessions,<br/>retire_agent_session"] Q -->|no| O["Restart-only: reported in prepare, never blocks,<br/>retired and started over after a crash"] ``` Closing or retiring a session stops its activation first: close cancels it, and retire also waits for it before the retire hook, which is where a sandboxed agent stops its sandbox. ```mermaid flowchart LR E["close from the environment server,<br/>or retire"] --> X["cancel the running activation"] --> W["retire waits for it"] --> H["retire_agent_session hook<br/>stops a sandbox the agent owns"] ``` ## Where this sits in the stack This is one PR in a stack of draft PRs that re-cut partial-rollout checkpointing onto environment servers. Each PR is one commit, and its base is the branch of the PR before it, so each diff shows only that PR's change. The stack builds on the token-capture cleanup PRs NVIDIA-NeMo#3938 (capture ledger `retire` and `delete`) and NVIDIA-NeMo#3939 (complete-record `retire` and `delete`), which have merged. NVIDIA-NeMo#3882 to NVIDIA-NeMo#3887 have merged too, so NVIDIA-NeMo#3888's base is `main`. The checkpoint stack uses those operations to free the ledgers of retired attempts and to clear a dead execution's capture files before a restore. The capture-ledger cache bound (NVIDIA-NeMo#4268) and the striped lock files (NVIDIA-NeMo#3937) are independent of the stack. 1. NVIDIA-NeMo#3882 (`ananthsub/partial-ckpt-core`): feat(checkpoint): add the participant control plane, episode steps, and coordination (merged) 2. NVIDIA-NeMo#3883 (`ananthsub/partial-ckpt-policy-model`): feat(checkpoint): make policy model servers checkpoint participants with generation cuts (merged) 3. NVIDIA-NeMo#3884 (`ananthsub/partial-ckpt-model-worker-cuts`): feat(token-capture): add worker staging helpers for generation cuts (merged) 4. NVIDIA-NeMo#3885 (`ananthsub/partial-ckpt-environment`): feat(checkpoint): continue environment server episodes from their boundaries (merged) 5. NVIDIA-NeMo#3886 (`ananthsub/partial-ckpt-resources`): feat(checkpoint): add the resources server participant with asynchronous session hooks (merged) 6. NVIDIA-NeMo#3887 (`ananthsub/partial-ckpt-verifier-declarations`): feat(checkpoint): declare twenty training verifiers stateless with replayable verification (merged) 7. NVIDIA-NeMo#3888 (`ananthsub/partial-ckpt-agent`): feat(checkpoint): add the agent session participant and Simple Agent continuation **(this PR)** 8. NVIDIA-NeMo#3889 (`ananthsub/partial-ckpt-e2e`): test(checkpoint): add a process-level end-to-end suite driven by coordination 9. NVIDIA-NeMo#3893 (`ananthsub/partial-ckpt-rollout-collection`): feat(checkpoint): checkpoint evaluation runs from rollout collection 10. NVIDIA-NeMo#3903 (`ananthsub/partial-ckpt-telemetry`): feat(checkpoint): spans and metrics for partial-rollout checkpoints 11. NVIDIA-NeMo#3909 (`ananthsub/partial-ckpt-multi-worker`): feat(checkpoint): resources, agent, and environment servers with several workers 12. NVIDIA-NeMo#4197 (`ananthsub/partial-ckpt-fault-tests`): test(checkpoint): check checkpoint invariants over random sequences of operations Partial-rollout checkpointing lets a training controller, such as NeMo RL, checkpoint Gym while rollouts are in flight and, after a crash, continue those rollouts from their last safe point instead of starting them over. Checkpointing is off by default; with the `checkpoint:` block unset, no server installs any checkpoint routes or behavior. ## Relationship to the old stack (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) - Supersedes NVIDIA-NeMo#2942 (agent participant) and NVIDIA-NeMo#2943 (Simple Agent turns). Completed-result acknowledgements, prepare polling, and the refusal retry loop are gone. - NVIDIA-NeMo#2944 (Gymnasium agent turns) is not carried; Gymnasium continuation moves to a separate PR, and until then Gymnasium agents are restart-only. ## Issue No tracking issue exists for this re-cut. The design and the mapping from the old stack are described in this PR series, and the old stack's PRs (NVIDIA-NeMo#2939 to NVIDIA-NeMo#2946) carry the original discussion. ## Validation Run on this branch, on the current `main`: - Each of these behaviors has a test that fails without its change: - an agent session activation without the rollout prefix is refused, for a restored session and for one seeded in the running process, with no model call, replay or claim, while checkpointing off still accepts it; - a legacy `/run` that skips verification parks at its return boundary until resume; - an episode the environment server reports as a restart, from a restart-only resources seed or an MCP grant, holds up neither the agent stage nor the resources stage of prepare, checked through coordination with a `single_agent_turn` server and real environment, agent and resources participants, and the agent neither waits for nor exports such an episode's activation until resume; - an indexed session reply and its replay carry no trajectory, nor does the export, and a kept reply stays unserialized until a replay or an export needs it; - a checkpoint of a legacy `/run` parked at verify holds its reply once, and a restore still replies; - a restore after the legacy turn loop replied returns that reply without running the loop again; - an episode that grants MCP access is a restart; - a close that arrives after prepare waits for resume and the commit exports its session, a cleanup close of a parked episode after the agent stage waits without holding up prepare, and a retire that frees a session while its close waits leaves no receipt behind; - a close stops the activation before its cleanup starts, and a close cancelled while it stops the activation runs the cleanup once when retried; - a request for an activation still running waits for it and gets its reply; - a request on the replaced attempt's path is refused for the restored session; - a restart-only seed that arrives after its rollout was retired is refused, and a restart-only retire frees its rollout's sessions and waits for a seed in progress; - a direct `/v1/responses` call is a restart, and one without a rollout path is refused; - a boundary leaves out the turn questions, and a resumed loop rebuilds them; - an activation without an index keeps no reply; - a restored session answered with its stored reply survives a commit that leaves its episode out; - with checkpointing on, a direct HTTP grant names the seeding attempt, and a restored session's direct HTTP tool call reaches the restored resources session under the replacement's prefix, through Simple Agent and a real resources server; - a legacy `/run` closes its resources session after a verification a checkpoint does not wait for. - A separate ordering check confirms that an environment episode's cleanup close finishes before the agent stage of prepare; it passes with or without the change. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_checkpoint_*.py`: passed, including: - closing a session stops its running activation; - a retire stops the activation before it replies, and a straggler is retired after resume; - a restart-only agent's `/run` has stopped when its retire replies; - a commit whose scope leaves out a restored session retires it; - a session whose retire hook was cut short stays blocking, is not exported, and is freed by a retry; - a seed arriving during an open checkpoint waits for resume, and one in progress when admission closes holds up prepare and is exported; - overlapping seeds of one session each hold up prepare, a retire waits for a seed in progress and frees what it built, and a restore is refused while a seed is in progress; - a closed session holds up prepare until its activation has stopped, and neither a close nor a retire cancels an activation twice; - a restored legacy `/run` claimed at verify is no longer pending and survives a commit that leaves it out; - an agent that builds its own app still reports its runs as restarts, and the participant is installed once per app; - `ready()` does not build the readiness report, and nothing notifies outside a checkpoint; - a commit does not export live sessions outside its scope. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider responses_api_agents/simple_agent/tests`: passed, including that a restored session's next activation does not reuse an invocation ID and that an activation continued from its restored boundary keeps its own; that a restored activation keeps its boundary through a second checkpoint; that a model reply arriving after prepare is held at the boundary before the call; that an activation already completed is answered with its stored reply without running again; that a legacy `/run` restored at verify or return still reports its trajectory; and that a session without a resources cookie jar restores without one. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider environment_servers/single_agent_turn/tests`: passed, including the timeline where the agent's reply to an invocation is lost to a restore: the environment server invokes again and the restored agent answers with the stored reply while the model sees no new call. - `RAY_TMPDIR=/tmp .venv/bin/python -m pytest -q -p no:cacheprovider tests/unit_tests/test_terminal_attribution.py`: passed, including that a call delivered before a checkpoint and issued again after a restore keeps only the re-issued call. - The full core unit suite (`tests/unit_tests`, eight processes, without `test_opensandbox_compose.py`, which needs the optional `opensandbox` package): passed except a sandbox retry test, which fails the same way on `main` in this development environment. - `pre-commit run --files <files changed by this PR>`: all hooks passed, and no hook modified a file. - Every commit carries a `Signed-off-by` line. ## Rollout evidence - Process-level evidence is in the top PR (episodes on the typed `single_agent_turn` request path and legacy `/run` episodes continue from their last boundary). - **Real model:** see the rollout evidence in NVIDIA-NeMo#3889, whose opt-in tests drove this PR's code with `Qwen/Qwen3-4B-Instruct-2507` on vLLM: saved rollouts continued after a kill and restore, restarts started over, and every rollout finished with reward 1.0. Every rollout ran Simple Agent. ## Compatibility - New agent hooks and class attribute: `checkpoint_sessions_supported`, `export_agent_sessions`, `restore_agent_sessions`, `retire_agent_session`. Existing agents need no change and are restart-only under checkpointing. - With checkpointing on, an agent session activation without the `/ng-rollout/<id>` prefix gets 400 `rollout_id_required`. A caller outside Gym that activates Simple Agent sessions with only the session cookie must add the prefix; no caller in the repository is affected. - With checkpointing on, agent calls without a rollout ID get 400 `rollout_id_required`, and agents with `num_workers > 1` fail at startup. - With checkpointing on, a direct HTTP grant's base URL ends in the seeding attempt's rollout prefix. An agent that continues a session after a restore must send its tool calls under its own activation's prefix; Simple Agent does. - With `skip_verification` and checkpointing on, a legacy `/run` caught by a checkpoint waits at its `return` boundary until resume. Nothing changes with checkpointing off. - New request header `x-ng-agent-activation` on agent session invocations, and a 409 `activation_out_of_order` refusal. An environment server that can invoke an activation again after a restore must send the header. - `AgentSessionParticipant.close_session` is now a coroutine, and the agent implements `SimpleServer.setup_checkpoint_participant(app)`, the hook NVIDIA-NeMo#3886 adds that the framework calls after `setup_webserver`. - A commit exports only the live sessions inside its `episode_ids` scope. A commit request without `episode_ids` still exports every live session. - Nothing changes unless the global `checkpoint:` block sets `enabled: true`. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>
Contributor
Author
|
Closing: this first prototype is replaced by the environment server checkpoint stack tracked in #3024. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
This checkpoint implementation is experimental. Its Python modules live under
nemo_gym._checkpoint; they are not a supported public import surface yet.Summary
How a Gymnasium episode resumes
The Gymnasium agent owns conversation and reward bookkeeping. The resources server owns the environment object. This PR records the agent half of a completed nonterminal step; #2945 and #2946 add the matching environment-state snapshot.
sequenceDiagram participant A as Gymnasium agent participant E as Gymnasium resources server participant M as Policy model participant P as Agent checkpoint participant A->>E: reset only for a fresh attempt E-->>A: observation and session cookie A->>M: generate action M-->>A: action and model-call ID A->>E: step(action) E-->>A: observation, reward, revision, terminal flag alt nonterminal step A->>P: commit boundary with history, reward, usage, cookies, revision P-->>A: continue or park for checkpoint else model or step refused for checkpoint A->>P: park without recording an error P-->>A: retry the same operation after explicit resume else terminal or truncated A->>E: close when supported A-->>A: return result without a resumable boundary endNot writing a terminal boundary matters because the environment session may already have been retired. Restoring that point and issuing another
stepwould continue an episode that no longer exists. The step-budget check happens before another model call, so a legacy boundary at the limit returns a truncated result without creating an extra turn.Stack
Test plan
Generation-safe checkpoint update
Layer validation: 16 tests passed.