Skip to content

[Doc] Document checkpointing to resume RL training - #4230

Open
YeonwooSung wants to merge 3 commits into
pytorch:mainfrom
YeonwooSung:doc/3044-save-load-agents
Open

[Doc] Document checkpointing to resume RL training#4230
YeonwooSung wants to merge 3 commits into
pytorch:mainfrom
YeonwooSung:doc/3044-save-load-agents

Conversation

@YeonwooSung

Copy link
Copy Markdown
Contributor

Description

Adds a Resume training guide to docs/source/reference/checkpoint.rst so users can save and restore an RL agent with the existing :class:~torchrl.checkpoint.Checkpoint`` API.

The new section:

  • Lists the components a resume checkpoint should contain (training model, optimizer, target networks; collector policy / env-transform state / frame count; replay-buffer writer, sampler, transforms, and storage; RNG) and why each one matters in RL.
  • Shows constructing Checkpoint(policy=..., optimizer=..., replay_buffer=..., collector=..., loss_module=..., exploration=..., rng=GlobalRNGState()), then save / load with components={...} and map_location.
  • Clarifies that the payload is the TensorDictModule / state_dict (or the buffer's dump / load), not a pickled mystery container.
  • Cross-links the shared-module case from [Feature Request] Documentation about weight update propagation #3032: if the loss holds the same module, save the policy or the loss once; do not persist two independent copies of the online weights.
  • Contrasts checkpoint (interrupt and resume) with export (deploy for inference) and links the export tutorial.

Includes a compact CartPole DQN snippet a reader can copy.

Motivation and Context

close #3044

The export tutorial is inference-only. Users asked for a copy-paste guide to interrupt training and continue it later. Checkpoint already exists; this PR documents how to use it rather than inventing a parallel API.

  • I have raised an issue to propose this change (required for new features and bug fixes)

Types of changes

  • Documentation (update in the documentation)

Checklist

  • I have read the CONTRIBUTION guide (required)
  • My change requires a change to the documentation.
  • I have updated the tests accordingly (required for a bug fix or a new feature).
  • I have updated the documentation accordingly.

@pytorch-bot

pytorch-bot Bot commented Sep 5, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/4230

Note: Links to docs will display an error until the docs builds have been completed.

❌ 3 New Failures

As of commit f9a17e7 with merge base 1d3de3d (image):

NEW FAILURES - The following jobs have failed:

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@github-actions github-actions Bot added the Documentation Improvements or additions to documentation label Sep 5, 2026
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 5, 2026
Populate the replay buffer from the collector before saving so LazyTensorStorage is initialized and the copy-paste example can run end to end.
Retain and checkpoint SoftUpdate alongside the loss, then demonstrate the optimizer-step and target-step order after restoration so resumed DQN training continues updating its target network.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Documentation Improvements or additions to documentation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature Request] Documentation explaining the recommended way to save and load agents

2 participants