[RLlib] Fix torch scheduler stepping and reporting. #48125
Merged
Conversation
This file contains 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
…t no duplicates or fragments are added to the replay buffer b/c it cannot handle these. Furthermore, refined tests for 'OfflinePreLearner'. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
…mplete_episodes' when recording episodes. This ensures that episodes can be read in again for training. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
…dium'. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
…ent read batch size in case 'EpisodeType' or 'BatchType' data is stored in offline datasets. Added some docstrings. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
…d RLlib schedulers in case a user set torch schedulers. These learning rates get now also correctly reported via this method. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
simonsays1980
added
rllib
Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
sven1977
changed the title
sven1977
reviewed
Oct 23, 2024
|
|
||
| # If the module uses learning rate schedulers, step them here. | ||
| if module_id in self._lr_schedulers: | ||
| @OverrideToImplementCustomLogic_CallToSuperRecommended |
There was a problem hiding this comment.
Great! This is so much better to do this in this method here (instead of in apply_gradients).
sven1977
reviewed
Oct 23, 2024
sven1977
reviewed
Oct 23, 2024
sven1977
reviewed
Oct 23, 2024
There was a problem hiding this comment.
LGTM. Awesome fix, thanks @simonsays1980 . Just one important request about adding the assertion into the example script such that this script really acts as a unit test for this functionality.
…er and switched the if-else condition to save us an indent. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
…rates. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
…he number of timesteps for testing in the example because CI tests were failing. Signed-off-by: simonsays1980 <simon.zehnder@gmail.com>
sven1977
reviewed
Nov 4, 2024
| from ray.rllib.utils.framework import try_import_torch | ||
| from ray.rllib.utils.test_utils import add_rllib_example_script_args | ||
|
|
||
| torch, _ = try_import_torch() | ||
|
|
||
| parser = add_rllib_example_script_args(default_reward=450.0, default_timesteps=200000) | ||
|
|
||
| class LRChecker(DefaultCallbacks): |
sven1977
reviewed
Nov 4, 2024
| from ray.rllib.utils.framework import try_import_torch | ||
| from ray.rllib.utils.test_utils import add_rllib_example_script_args | ||
|
|
||
| torch, _ = try_import_torch() | ||
|
|
||
| parser = add_rllib_example_script_args(default_reward=450.0, default_timesteps=200000) | ||
|
|
||
| class LRChecker(DefaultCallbacks): |
There was a problem hiding this comment.
Suggested change
| class LRChecker(DefaultCallbacks): | |
| class LRChecker(DefaultCallbacks): | |
| """A custom callback that asserts the functionality of the lr-scheduler used.""" |
sven1977
reviewed
Nov 4, 2024
| for optimizer_name in optimizer_names: | ||
| # If learning rate schedulers are provided step them here. Note, | ||
| # stepping them in `TorchLearner.apply_gradients` updates the | ||
| # learning rates during minibatch updates; we want to update |
There was a problem hiding this comment.
thanks for clarifying this.
sven1977
approved these changes
Nov 4, 2024
Jay-ju
pushed a commit
to Jay-ju/ray
that referenced
this pull request
Nov 5, 2024
JP-sDEV
pushed a commit
to JP-sDEV/ray
that referenced
this pull request
Nov 14, 2024
Signed-off-by: JP-sDEV <jon.pablo80@gmail.com>
mohitjain2504
pushed a commit
to mohitjain2504/ray
that referenced
this pull request
Nov 15, 2024
Signed-off-by: mohitjain2504 <mohit.jain@dream11.com>
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.
Why are these changes needed?
When using
torch.optim.lr_scheduler.LRSchedulerschedulers learning rates were wrongly updated and reported b/c they were stepped inTorchLearner.apply_gradientsand reported inLearner.after_gradient_based_update. The latter function also updated any RLlib specific learning rate schedules and therefore reset any learning rate schedules of the torch-specific schedulers.This PR fixes this error and instead overrides the
after_gradient_based_updateto step learning rate schedulers in there and report them correctly. It also avoids the RLlib schedulers to override updated learning rates. In case no torch-specific learning rates are used, thesuper()is called an RLlib schedules do their work.This setup also updates the torch-specific schedulers at the right point in regard to SGD: after an SGD epoch and not in each minibatch step.
Related issue number
Checks
git commit -s) in this PR.scripts/format.shto lint the changes in this PR.method in Tune, I've added it in
doc/source/tune/api/under thecorresponding
.rstfile.