Secure deserialization (.load()) by default

docs/guide/save_format.md

@@ -58,3 +58,80 @@ Cons:
 - More complex implementation.
 - Still relies partly on cloudpickle for complex objects (e.g. custom functions)
   with can lead to [incompatibilities](https://github.com/DLR-RM/stable-baselines3/issues/172) between Python versions.
+
+## Secure Deserialization
+
+:::{warning}
+**Loading untrusted checkpoints can execute arbitrary Python code.**
+Starting with SB3 2.10, all `load()` methods use `deserialization_mode="safe"` by default, which blocks
+arbitrary code execution during deserialization at the cost of skipping non-whitelisted serialized entries.
+:::
+
+The `deserialization_mode` parameter is available on all load methods:
+
+- `stable_baselines3.common.base_class.BaseAlgorithm.load`
+- `stable_baselines3.common.save_util.json_to_data`
+- `stable_baselines3.common.save_util.load_from_pkl`
+- `stable_baselines3.common.save_util.load_from_zip_file`
+- `stable_baselines3.common.off_policy_algorithm.OffPolicyAlgorithm.load_replay_buffer`
+- `stable_baselines3.common.vec_env.vec_normalize.VecNormalize.load`
+
+Each accepts one of two modes:
+
+### Safe mode (default)
+
+In `deserialization_mode="safe"`, SB3 uses a restricted unpickler that only allows a fixed allowlist of
+known-safe types (SB3 classes, gymnasium spaces, numpy types, PyTorch types, cloudpickle internals).
+If a serialized entry references a type outside the allowlist, it throws an error (that is caught), and the
+user must supply a safe replacement via the `custom_objects` argument or extends the whitelist (see below).
+
+```python
+from stable_baselines3 import PPO
+
+# If the checkpoint contains a custom learning-rate schedule that is not
+# in the allowlist, you must provide it via custom_objects:
+loaded = PPO.load(
+    "model.zip",
+    custom_objects={
+        "learning_rate": 0.0003,
+        "lr_schedule": lambda progress: progress * 0.0003,
+    },
+)
+```
+
+### Legacy mode
+
+In `deserialization_mode="legacy"`, SB3 falls back to the standard `cloudpickle` / `pickle` loader.
+This preserves full backward compatibility with models saved before SB3 2.10 but **executes arbitrary
+Python code** embedded in the checkpoint. A `UserWarning` is emitted.
+
+```python
+# Restores the pre-2.10 loading behavior for checkpoints that contain
+# lambda functions, local classes, or custom gym environments:
+loaded = PPO.load("model.zip", deserialization_mode="legacy")
+```
+
+### Extending the Safe Allowlist
+
+If you have custom types (e.g. a custom environment or a custom space) that are safe to deserialize,
+you can register them with the allowlist using `stable_baselines3.common.safe_globals.add_safe_globals`:
+
+```python
+from stable_baselines3.common.safe_globals import add_safe_globals
+from my_module import MyCustomSpace
+
+add_safe_globals(MyCustomSpace)
+loaded = PPO.load("model.zip", deserialization_mode="safe")
+```
+
+For a temporary, scope-limited registration, use the `stable_baselines3.common.safe_globals.safe_globals`
+context manager:
+
+```python
+from stable_baselines3.common.safe_globals import safe_globals
+from my_module import MyCustomSpace
+
+with safe_globals(MyCustomSpace):
+    loaded = PPO.load("model.zip", deserialization_mode="safe")
+# MyCustomSpace is automatically removed from the allowlist on exit
+```

docs/misc/changelog.md

@@ -2,6 +2,47 @@
 
 # Changelog
 
+## Release 2.10.0a0 (TBD)
+
+**Secure deserialization by default**
+
+:::{warning}
+Models saved with cloudpickle that contain arbitrary Python code (e.g. lambda functions, local classes, or custom spaces) will now produce warnings and skip the affected entries when loaded with the new default `deserialization_mode="safe"`. Use `deserialization_mode="legacy"` to restore the old behavior. You can find more information in the [Saving and Loading](https://stable-baselines3.readthedocs.io/en/master/guide/save_reload.html) documentation.
+:::
+
+### Breaking Changes:
+
+- `deserialization_mode` now defaults to `"safe"` for all load methods (`BaseAlgorithm.load`, `json_to_data`, `load_from_pkl`, `load_from_zip_file`, `load_replay_buffer`, `VecNormalize.load`). This blocks arbitrary code execution during deserialization at the cost of skipping non-whitelisted serialized entries.
+- `BaseModel.load()` now uses `torch.load(..., weights_only=True)` with a safe-globals allowlist for policy state-dicts.
+
+### New Features:
+
+- Added `deserialization_mode` parameter to all load methods (`"safe"` or `"legacy"`) to mitigate deserialization of Untrusted Data. `"safe"` mode uses a restricted unpickler that only allows a fixed allowlist of known-safe SB3/gymnasium/numpy types.
+- Added `add_safe_globals()` function to register custom classes as safe for restricted deserialization (a la `torch.serialization.add_safe_globals`).
+- Added `safe_globals` context manager to temporarily extend the safe-globals allowlist for a single load call.
+- Added `DeserializationMode` type alias (`Literal["safe", "legacy"]`) to `stable_baselines3.common.type_aliases`.
+
+### Bug Fixes:
+
+### [SB3-Contrib]
+
+### [RL Zoo]
+
+### [SBX] (SB3 + Jax)
+
+### Deprecations:
+
+### Others:
+
+- Added regression tests in `tests/test_security_deserialization.py`.
+- Improved test coverage for `save_util.py`: added tests for `BadZipFile` error handling in `load_from_zip_file` and `IsADirectoryError`/`FileNotFoundError` handling in `open_path_pathlib`.
+- Improved test coverage for `sb2_compat/rmsprop_tf_like.py`: added `TestRMSpropTFLike` test class covering parameter validation, centered/momentum/weight_decay branches, closure support, and pickle roundtrip.
+
+### Documentation:
+
+- Updated save/reload guide with a dedicated section on secure deserialization, explaining safe vs. legacy mode and how to handle `custom_objects` in safe mode.
+
+
 ## Release 2.9.0 (2026-06-15)
 
 **Updated dependencies (pandas is now optional, gymnasium 1.3.0 support, torch>=2.8)**

stable_baselines3/common/base_class.py

@@ -23,7 +23,13 @@
 from stable_baselines3.common.policies import BasePolicy
 from stable_baselines3.common.preprocessing import check_for_nested_spaces, is_image_space, is_image_space_channels_first
 from stable_baselines3.common.save_util import load_from_zip_file, recursive_getattr, recursive_setattr, save_to_zip_file
-from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, Schedule, TensorDict
+from stable_baselines3.common.type_aliases import (
+    DeserializationMode,
+    GymEnv,
+    MaybeCallback,
+    Schedule,
+    TensorDict,
+)
 from stable_baselines3.common.utils import (
     FloatSchedule,
     check_for_correct_spaces,
@@ -648,6 +654,7 @@ def load(  # noqa: C901
         custom_objects: dict[str, Any] | None = None,
         print_system_info: bool = False,
         force_reset: bool = True,
+        deserialization_mode: DeserializationMode = "safe",
         **kwargs,
     ) -> SelfBaseAlgorithm:
         """
@@ -671,6 +678,16 @@ def load(  # noqa: C901
         :param force_reset: Force call to ``reset()`` before training
             to avoid unexpected behavior.
             See https://github.com/DLR-RM/stable-baselines3/issues/597
+        :param deserialization_mode: How to handle cloudpickle-serialized objects
+            in the checkpoint's ``data`` JSON.
+
+            - ``"legacy"``: Deserialize with cloudpickle for full
+              backward compatibility.  A security warning is emitted because
+              cloudpickle deserialization can execute arbitrary Python code.
+            - ``"safe"`` (default): Attempt restricted deserialization using an
+              allowlist of known-safe globals. Entries that cannot be safely
+              deserialized (and are not provided via ``custom_objects``) are skipped
+              with a warning.
         :param kwargs: extra arguments to change the model when loading
         :return: new model instance with loaded parameters
         """
@@ -683,6 +700,7 @@ def load(  # noqa: C901
             device=device,
             custom_objects=custom_objects,
             print_system_info=print_system_info,
+            deserialization_mode=deserialization_mode,
         )
 
         assert data is not None, "No data found in the saved file"

stable_baselines3/common/off_policy_algorithm.py

@@ -16,7 +16,15 @@
 from stable_baselines3.common.noise import ActionNoise, VectorizedActionNoise
 from stable_baselines3.common.policies import BasePolicy
 from stable_baselines3.common.save_util import load_from_pkl, save_to_pkl
-from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback, RolloutReturn, Schedule, TrainFreq, TrainFrequencyUnit
+from stable_baselines3.common.type_aliases import (
+    DeserializationMode,
+    GymEnv,
+    MaybeCallback,
+    RolloutReturn,
+    Schedule,
+    TrainFreq,
+    TrainFrequencyUnit,
+)
 from stable_baselines3.common.utils import safe_mean, should_collect_more_steps
 from stable_baselines3.common.vec_env import VecEnv
 from stable_baselines3.her.her_replay_buffer import HerReplayBuffer
@@ -228,6 +236,7 @@ def load_replay_buffer(
         self,
         path: str | pathlib.Path | io.BufferedIOBase,
         truncate_last_traj: bool = True,
+        deserialization_mode: DeserializationMode = "safe",
     ) -> None:
         """
         Load a replay buffer from a pickle file.
@@ -237,8 +246,18 @@ def load_replay_buffer(
             If set to ``True``, we assume that the last trajectory in the replay buffer was finished
             (and truncate it).
             If set to ``False``, we assume that we continue the same trajectory (same episode).
+        :param deserialization_mode: How to handle pickle deserialization.
+
+            - ``"safe"`` (default): Deserialize using a restricted unpickler that
+              only allows a fixed allowlist of known-safe SB3/gymnasium/numpy types.
+              Any pickle payload referencing a type outside this allowlist is
+              rejected with a clear error.
+            - ``"legacy"``: Deserialize with ``pickle.load()``.  This preserves
+              backward compatibility but can **execute arbitrary Python code**
+              embedded in the pickle file.
+
         """
-        self.replay_buffer = load_from_pkl(path, self.verbose)
+        self.replay_buffer = load_from_pkl(path, self.verbose, deserialization_mode=deserialization_mode)
         assert isinstance(self.replay_buffer, ReplayBuffer), "The replay buffer must inherit from ReplayBuffer class"
 
         # Backward compatibility with SB3 < 2.1.0 replay buffer

stable_baselines3/common/policies.py

@@ -22,6 +22,7 @@
     make_proba_distribution,
 )
 from stable_baselines3.common.preprocessing import get_action_dim, is_image_space, maybe_transpose, preprocess_obs
+from stable_baselines3.common.safe_globals import register_sb3_safe_globals
 from stable_baselines3.common.torch_layers import (
     BaseFeaturesExtractor,
     CombinedExtractor,
@@ -173,9 +174,9 @@ def load(cls: type[SelfBaseModel], path: str, device: th.device | str = "auto")
         :return:
         """
         device = get_device(device)
-        # Note(antonin): we cannot use `weights_only=True` here because we need to allow
-        # gymnasium imports for the policy to be loaded successfully
-        saved_variables = th.load(path, map_location=device, weights_only=False)
+        # Ensure SB3 policy types are registered for torch.load with weights_only=True
+        register_sb3_safe_globals()
+        saved_variables = th.load(path, map_location=device, weights_only=True)
 
         # Create policy object
         model = cls(**saved_variables["data"])
@@ -985,3 +986,26 @@ def q1_forward(self, obs: th.Tensor, actions: th.Tensor) -> th.Tensor:
         with th.no_grad():
             features = self.extract_features(obs, self.features_extractor)
         return self.q_networks[0](th.cat([features, actions], dim=1))
+
+
+# ---------------------------------------------------------------------------
+# 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]