Skip to content

Resume RND exploration schedules from PPO checkpoints - #240

Merged
ClemensSchwarke merged 3 commits into
leggedrobotics:mainfrom
Afloat16:rsl-rnd-schedule-checkpoint
Oct 8, 2026
Merged

ClemensSchwarke merged 3 commits into
leggedrobotics:mainfrom
Afloat16:rsl-rnd-schedule-checkpoint

Conversation

@Afloat16

@Afloat16 Afloat16 commented Oct 7, 2026

Copy link
Copy Markdown
Contributor

PPO checkpoints restore RND parameters and Adam state but omit the Python update_counter. Loading a checkpoint restarts step and linear exploration schedules even when every model and optimizer has been restored. This changes the next intrinsic reward and the subsequent policy objective after resuming.

Persist rnd_update_counter as plain integer PPO checkpoint metadata and restore it when the RND state is selected for loading. Older checkpoints default to the historical initial count of zero. Selective loads that exclude RND retain the current counter. The RND model state remains tensor-only for the existing distributed broadcast path, and the training step gains no new tensor conversion or synchronization.

Six new cases construct actual PPO, MLPModel, RolloutStorage and RND components, perform on-disk weights-only round trips, and compare the next reward with independent step/linear formulas. They also cover legacy checkpoints, selective loads, no-RND checkpoints and the broadcast-state constraint. The complete original PPO and RND test modules are included unchanged in the matched before/after runs.

Genuine CPU validation with Torch 2.7.1, NumPy 1.26.4 and the complete package source: the same 31 collected cases change from 3 failures and 28 passes on the original source to 31 passes with the fix, including all 25 tests in the affected original modules unchanged. No tests were skipped. The final source passes the pinned Ruff 0.14.0 checks and the repository's complete pre-commit run --all-files hooks.

CUDA, simulators, distributed execution and ONNX export are outside this validation scope.

@ClemensSchwarke

Copy link
Copy Markdown
Collaborator

Hi @Afloat16,
thanks for the PRs. Would it be possible to add your real name to the contributors?
Thanks, Clemens

@ClemensSchwarke
ClemensSchwarke merged commit 72c8956 into leggedrobotics:main Oct 8, 2026
7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants