Use classes instead of lambdas for schedules #2125
Conversation
akanto
force-pushed
the
save-load-portability
branch
from
52b6ad1 to
63cfb2e
Compare
…non-portable schedules Previously, using closures (e.g., lambdas) for learning_rate or clip_range caused segmentation faults when loading models across different platforms (e.g., macOS to Linux), because cloudpickle could not safely serialize/deserialize them. This commit rewrites: - `constant_fn` as a `ConstantSchedule` class - `get_schedule_fn` as a `FloatConverterSchedule` class - `get_linear_fn` as a `LinearSchedule` class All schedules are now proper callable classes, making them portable and safely pickleable. Old functions are kept (marked as deprecated) for backward compatibility when loading existing models.
akanto
force-pushed
the
save-load-portability
branch
from
63cfb2e to
a6d8c07
Compare
araffin
changed the title
| @@ -273,7 +273,7 @@ def logger(self) -> Logger: | |||
|
|
|||
| def _setup_lr_schedule(self) -> None: | |||
| """Transform to callable if needed.""" | |||
| self.lr_schedule = get_schedule_fn(self.learning_rate) | |||
| self.lr_schedule = FloatConverterSchedule(self.learning_rate) | |||
There was a problem hiding this comment.
I would rather have a new get_schedule helper here because FloatConverterSchedule does more than what its name suggest
PS: you don't have to force push for every edit, you can simply push new commits to the same branch
There was a problem hiding this comment.
or maybe rename FloatConverterSchedule to something else (the issue is that Schedule type is already taken... maybe FloatSchedule?
There was a problem hiding this comment.
I have renamed it to FloatSchedule, as recommended, but maybe a ScheduleWrapper could be a better name, since it is basically a wrapper that ensures that a constant is transformed to Schedule, and ensures that any callable returning float values.
I wanted to change the original logic as minimal as possible, but maybe the original function is not clean enough, since it is doing multiple things.
stable_baselines3/common/utils.py
Outdated
| @@ -78,6 +78,35 @@ def update_learning_rate(optimizer: th.optim.Optimizer, learning_rate: float) -> | |||
| param_group["lr"] = learning_rate | |||
|
|
|||
|
|
|||
| class FloatConverterSchedule: | |||
There was a problem hiding this comment.
This class actually does more no? it enforces that we have a schedule object that can be pickled, no?
Maybe rename it to FloatSchedule and update the docstring to explain its utility (in addition to casting to float)
There was a problem hiding this comment.
same comment as above
- Renamed FloatConverterSchedule to FloatSchedule to better reflect its purpose. - Moved parameter documentation to the class-level docstring for proper Sphinx support
araffin
changed the title
|
Thanks, I can send a pull request. I guess a commit is not enough, but the stable-baslines3 release of 2.6.1 is required to update the dependencies: https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/blob/master/setup.py#L70, and after that, I can update all |
|
I've released an alpha that you can use: |
…rtable schedules
Previously, using closures (e.g., lambdas) for learning_rate or clip_range caused segmentation faults when loading models across different platforms (e.g., macOS to Linux), because cloudpickle could not safely serialize/deserialize them.
Description
Refactor schedule-related helper functions into proper classes. This ensures full portability and prevents segfaults when
loading models across different operating systems. Introduces
ConstantSchedule,CappedLinearSchedule, andFloatConverterSchedulefor supporting portability across different operating systems.This commit rewrites:
constant_fnas aConstantScheduleclassget_schedule_fnas aFloatConverterScheduleclassget_linear_fnas aLinearScheduleclassAll schedules are now proper callable classes, making them portable and safely pickleable. Old functions are kept (marked as deprecated) for backward compatibility when loading existing models.
Motivation and Context
Fixes cross-platform segmentation faults (#2115) caused by non-portable closures (like lambdas) being serialized into model files. After this change, saved models are robust, portable, and no longer crash at load time if they are moved across different operating systems.
Types of changes
Checklist
make format(required)make check-codestyleandmake lint(required)make pytestandmake typeboth pass. (required)make doc(required)Note: You can run most of the checks using
make commit-checks.Note: we are using a maximum length of 127 characters per line