[feat][checkpoint] Allow selective state restore on resume

[bvolpato]

Summary

Allow checkpoint weights to seed a new training phase while independently controlling global-step, dataloader, optimizer, and learning-rate scheduler restoration.

Exact resume remains unchanged because every option defaults to true.

Changes

• Add trainer.resume_load_global_step, trainer.resume_load_dataloader_state, trainer.resume_load_optimizer_states, and trainer.resume_load_lr_scheduler_states.
• Reset training to step zero when global-step restoration is disabled.
• Skip stale fully-async UID and epoch state for a new training phase.
• Preserve logical epoch progress and epoch-end saves in sync and fully-async trainers when global-step restoration skips the dataloader cursor.
• Pass optimizer and scheduler restore options to policy and critic checkpoint loading.
• Document selective-state resume and separate checkpoint output paths.
• Add focused sync and fully-async regression coverage.

Testing

• pytest tests/train/test_trainer.py -q (21 passed)
• focused selective-resume regressions (6 passed)
• pytest tests/train/test_config.py -q (145 passed)
• pytest tests/train/test_fully_async_trainer.py -q (12 passed)
• pre-commit run --all-files
• documentation type check and production build (66 pages)
• git diff --check upstream/main

[gemini-code-assist]

Code Review

This pull request introduces four new configuration parameters (resume_load_global_step, resume_load_dataloader_state, resume_load_optimizer_states, and resume_load_lr_scheduler_states) to allow users to selectively disable restoring specific states when resuming from a checkpoint. The trainer logic and documentation have been updated accordingly, and new unit tests have been added. Feedback suggests explicitly checking if self.train_dataloader is initialized before attempting to restore its state to avoid potential AttributeError warnings when no dataloader exists.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{Reviewed by Cursor Bugbot for commit b2cde16. Configure here.}