Secure deserialization (.load()) by default#2264
Open
araffin wants to merge 16 commits into
Open
Conversation
araffin
changed the title
.load()) by default
There was a problem hiding this comment.
Pull request overview
This PR hardens Stable-Baselines3 model/component loading against unsafe pickle/cloudpickle deserialization by introducing a default "safe" deserialization mode across load paths, backed by a restricted unpickler + allowlist, and updates tests/docs accordingly.
Changes:
- Add
deserialization_mode: Literal["safe","legacy"]across key load APIs and route checkpoint JSON/pickle loading through a restricted unpickler by default. - Introduce a centralized safe-globals allowlist (
common/safe_globals.py) plus user extension APIs (add_safe_globals,safe_globals) and enabletorch.load(..., weights_only=True)for policy loads. - Add regression/security tests and update docs/changelog for the breaking behavior.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 9 comments.
Show a summary per file
| File | Description |
|---|---|
tests/test_vec_envs.py |
Adjusts parametrization input type for pytest collection stability. |
tests/test_utils.py |
Adds tests for RMSpropTFLike; modifies schedule-related assertions. |
tests/test_security_deserialization.py |
New end-to-end tests ensuring safe mode blocks RCE across load entry points. |
tests/test_save_load.py |
Updates warning expectations and adds coverage for bad zip + open_path edge cases. |
stable_baselines3/version.txt |
Bumps development version to 2.10.0a0. |
stable_baselines3/common/vec_env/vec_normalize.py |
Adds deserialization_mode to VecNormalize.load() with safe/legacy behavior. |
stable_baselines3/common/type_aliases.py |
Introduces DeserializationMode = Literal["safe","legacy"]. |
stable_baselines3/common/save_util.py |
Adds safe/legacy handling for JSON cloudpickle payloads and .pkl loads; threads mode into zip loading. |
stable_baselines3/common/safe_globals.py |
New centralized allowlist + restricted unpickler + user extension APIs. |
stable_baselines3/common/policies.py |
Switches policy .load() to torch.load(weights_only=True) and registers safe globals. |
stable_baselines3/common/off_policy_algorithm.py |
Adds deserialization_mode to replay buffer load path. |
stable_baselines3/common/base_class.py |
Adds deserialization_mode parameter to algorithm .load() and forwards it to zip loader. |
docs/misc/changelog.md |
Documents breaking change and new APIs; adds release entry. |
docs/guide/save_format.md |
Documents safe vs legacy deserialization and allowlist extension APIs. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+68
to
+81
| # numpy dtype descriptor classes (needed by torch weights_only) | ||
| types += [ | ||
| np.dtypes.Float32DType, | ||
| np.dtypes.Float64DType, | ||
| np.dtypes.Int32DType, | ||
| np.dtypes.Int64DType, | ||
| np.dtypes.Int8DType, | ||
| np.dtypes.Int16DType, | ||
| np.dtypes.UInt8DType, | ||
| np.dtypes.UInt16DType, | ||
| np.dtypes.UInt32DType, | ||
| np.dtypes.UInt64DType, | ||
| np.dtypes.BoolDType, | ||
| ] |
Comment on lines
+68
to
+81
| # numpy dtype descriptor classes (needed by torch weights_only) | ||
| types += [ | ||
| np.dtypes.Float32DType, | ||
| np.dtypes.Float64DType, | ||
| np.dtypes.Int32DType, | ||
| np.dtypes.Int64DType, | ||
| np.dtypes.Int8DType, | ||
| np.dtypes.Int16DType, | ||
| np.dtypes.UInt8DType, | ||
| np.dtypes.UInt16DType, | ||
| np.dtypes.UInt32DType, | ||
| np.dtypes.UInt64DType, | ||
| np.dtypes.BoolDType, | ||
| ] |
Comment on lines
+363
to
+367
| "cloudpickle.cloudpickle._class_setstate", | ||
| "cloudpickle.cloudpickle._fillvar", | ||
| "cloudpickle.cloudpickle.subimport", | ||
| "cloudpickle.cloudpickle._lookup_module_and_obj_in_qualname", | ||
| "cloudpickle.cloudpickle.whichmodule", |
Comment on lines
+991
to
+1011
| # --------------------------------------------------------------------------- | ||
| # Register base policy classes for torch.load with weights_only=True. | ||
| # These are defined in this module, so registration must happen after | ||
| # their class definitions (i.e. at the end of the file). | ||
| # --------------------------------------------------------------------------- | ||
|
|
||
| # Policy classes defined in this module | ||
| _safe_base_policies = [ | ||
| BasePolicy, | ||
| ActorCriticPolicy, | ||
| ActorCriticCnnPolicy, | ||
| MultiInputActorCriticPolicy, | ||
| ContinuousCritic, | ||
| BernoulliDistribution, | ||
| CategoricalDistribution, | ||
| DiagGaussianDistribution, | ||
| MultiCategoricalDistribution, | ||
| StateDependentNoiseDistribution, | ||
| ] | ||
|
|
||
| th.serialization.add_safe_globals(_safe_base_policies) # type: ignore[arg-type] |
Comment on lines
+192
to
+208
| if deserialization_mode == "safe": | ||
| pass | ||
| # warnings.warn( | ||
| # "Loading a model checkpoint that contains cloudpickle-serialized " | ||
| # "objects with a restricted (safe) deserializer. Only known-safe " | ||
| # "SB3/gymnasium/numpy types are allowed. " | ||
| # UserWarning, | ||
| # ) | ||
| else: | ||
| warnings.warn( | ||
| "Loading a model checkpoint that contains cloudpickle-serialized " | ||
| "objects (deserialization_mode='legacy'). This allows arbitrary " | ||
| "Python code execution from the checkpoint file. Only load " | ||
| "checkpoints from trusted sources. To enable safe deserialization, " | ||
| "use deserialization_mode='safe'. ", | ||
| UserWarning, | ||
| ) |
Comment on lines
+344
to
+349
| # warnings.warn( | ||
| # "Loading a VecNormalize pickle file with a restricted (safe) " | ||
| # "deserializer. Only known-safe SB3/gymnasium/numpy types are " | ||
| # "allowed.", | ||
| # UserWarning, | ||
| # ) |
Comment on lines
436
to
454
| file = open_path(path, "r", verbose=verbose, suffix="pkl") | ||
| obj = pickle.load(file) | ||
|
|
||
| if deserialization_mode == "safe": | ||
| # Ensure SB3 types are registered for safe deserialization | ||
| from stable_baselines3.common.safe_globals import register_sb3_safe_globals | ||
|
|
||
| register_sb3_safe_globals() | ||
| obj = _RestrictedUnpickler(file).load() | ||
| else: | ||
| warnings.warn( | ||
| "Loading a .pkl file with pickle deserialization (deserialization_mode='legacy'). " | ||
| "This can execute arbitrary Python code from the file. Only load pickle files " | ||
| "from trusted sources. ", | ||
| UserWarning, | ||
| ) | ||
| obj = pickle.load(file) | ||
| if isinstance(path, (str, pathlib.Path)): | ||
| file.close() | ||
| return obj |
This file contains hidden or 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
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.
Description
See #1831
closes #1852
Related to DLR-RM/rl-baselines3-zoo#505
Note: using opencode with local Qwen 3.6 27B model
Motivation and Context
Types of changes
Checklist
docs/misc/changelog.md) (required).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