diff --git a/docs/guide/save_format.md b/docs/guide/save_format.md index d4e416d2d..f6ffdc2f7 100644 --- a/docs/guide/save_format.md +++ b/docs/guide/save_format.md @@ -58,3 +58,91 @@ 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 **may execute 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") + +# If you only load models from trusted sources, you can silence the warning with: +# import warnings +# warnings.filterwarnings("ignore", message="Loading a model checkpoint that contains cloudpickle-serialized objects", category=UserWarning) +``` + +### 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.SafeGlobals` +context manager: + +```python +from stable_baselines3.common.safe_globals import SafeGlobals +from my_module import MyCustomSpace + +with SafeGlobals(MyCustomSpace): + loaded = PPO.load("model.zip", deserialization_mode="safe") +# MyCustomSpace is automatically removed from the allowlist on exit +``` + +:::{note} +With the default `deserialization_mode="safe"`, you may encounter a `pickle.UnpicklingError` or `Could not deserialize object` warning when loading checkpoints containing custom types that are not in the allowlist. The error message +will indicate the missing type (e.g., `Global 'my_module.MyCustomType' is not in the safe deserialization allowlist`). +You can either use `add_safe_globals()` or `SafeGlobals` to register your custom types, pass `custom_objects=...` at load time, or switch to `deserialization_mode="legacy"` +if you trust the checkpoint source. +::: diff --git a/docs/misc/changelog.md b/docs/misc/changelog.md index 5f291be1d..f9d2a9179 100644 --- a/docs/misc/changelog.md +++ b/docs/misc/changelog.md @@ -2,6 +2,41 @@ # 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 and context manager to register custom classes as safe for restricted deserialization (à la `torch.serialization.add_safe_globals`). + +### Bug Fixes: + +### [SB3-Contrib] + +### [RL Zoo] + +### [SBX] (SB3 + Jax) + +### Deprecations: + +### Others: + +### 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.2a0 (2026-07-18) ### Breaking Changes: @@ -29,6 +64,7 @@ ### Documentation: + ## Release 2.9.0 (2026-06-15) **Updated dependencies (pandas is now optional, gymnasium 1.3.0 support, torch>=2.8)** diff --git a/stable_baselines3/common/base_class.py b/stable_baselines3/common/base_class.py index e1216ab6b..41971516f 100644 --- a/stable_baselines3/common/base_class.py +++ b/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 = 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" diff --git a/stable_baselines3/common/off_policy_algorithm.py b/stable_baselines3/common/off_policy_algorithm.py index 1eaf79d7c..49b92705c 100644 --- a/stable_baselines3/common/off_policy_algorithm.py +++ b/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 = 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 diff --git a/stable_baselines3/common/policies.py b/stable_baselines3/common/policies.py index 93f0a254f..fae8b9fe0 100644 --- a/stable_baselines3/common/policies.py +++ b/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"]) diff --git a/stable_baselines3/common/safe_globals.py b/stable_baselines3/common/safe_globals.py new file mode 100644 index 000000000..80c28627c --- /dev/null +++ b/stable_baselines3/common/safe_globals.py @@ -0,0 +1,714 @@ +""" +Central safe-globals registration for safe deserialization. + +This module collects the allowlist of types that are permitted during +restricted deserialization. The same allowlist is used by two mechanisms: + +1. ``_RestrictedUnpickler`` (pickle / cloudpickle data stored inside zip + checkpoints and in raw .pkl files). +2. ``torch.serialization.add_safe_globals`` so that ``torch.load(..., + weights_only=True)`` accepts the same types when loading standalone + policy pickles (.pkl) and PyTorch state-dicts (.pth) inside zip + checkpoints. + +Both mechanisms share this single source of truth so that adding a new +type to the allowlist is a one-line change. +""" + +from __future__ import annotations + +import collections +import pickle +from collections.abc import Callable +from typing import Any + +import numpy as np +import torch as th + +# --------------------------------------------------------------------------- +# Lazily-populated sets (populated by _register_safe_globals) +# --------------------------------------------------------------------------- + +# String allowlist for _RestrictedUnpickler (module.ClassPath). +_SAFE_GLOBALS_STR: set[str] | None = None + +# Actual objects registered with torch.serialization.add_safe_globals. +# We register them at most once. +_TORCH_REGISTRATION_DONE = False + + +def _collect_numpy_types() -> list[type | Callable[..., Any]]: + """Collect numpy types needed for gymnasium spaces and array reconstruction. + + For the **string allowlist** (restricted unpickler), numpy-internal pickle + helpers are handled by ``_NUMPY_PICKLE_INTERNALS_STR`` so that both + numpy 1.x and 2.x module paths are accepted without crashing on import. + + For **torch.serialization.add_safe_globals** (``weights_only=True``), we + also need the *live objects* of the currently installed numpy version. + Those are imported here inside try/except blocks so the code works on + either numpy 1.x or 2.x. + """ + types: list[type | Callable[..., Any]] = [ + np.ndarray, + np.dtype, + np.float32, + np.float64, + np.int32, + np.int64, + np.int8, + np.int16, + np.uint8, + np.uint16, + np.uint32, + np.uint64, + np.bool_, + ] + + # 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, + ] + + # numpy random internals (gymnasium spaces store internal RNG state) + types += [ + np.random.bit_generator.BitGenerator, + np.random.bit_generator.SeedSequence, + np.random._pcg64.PCG64, + np.random._mt19937.MT19937, + np.random._philox.Philox, + np.random._sfc64.SFC64, + np.random._generator.Generator, + np.random.mtrand.RandomState, + ] + + # numpy pickle reconstruction helpers (functions) — try the numpy 2.x path + # first, then fall back to the numpy 1.x path. If neither works, we skip + # registration with torch (the string allowlist still covers pickle). + try: + import numpy._core.multiarray as np_ma + import numpy._core.numeric as np_numeric + + types += [ + np_ma._reconstruct, # type: ignore[attr-defined] + np_ma.scalar, # type: ignore[attr-defined] + np_numeric._frombuffer, # type: ignore[attr-defined] + ] + except ModuleNotFoundError: # pragma: no cover + try: + import numpy.core.multiarray as np_ma_np1 + import numpy.core.numeric as np_numeric_np1 + + types += [ + np_ma_np1._reconstruct, # type: ignore[attr-defined] + np_ma_np1.scalar, # type: ignore[attr-defined] + np_numeric_np1._frombuffer, # type: ignore[attr-defined] + ] + except ModuleNotFoundError: + pass + + try: + import numpy.random._pickle as np_rp + + types += [ + np_rp.__generator_ctor, + np_rp.__bit_generator_ctor, + np_rp.__randomstate_ctor, + ] + except ModuleNotFoundError: # pragma: no cover + pass + + # Cython unpickle helpers in numpy random + try: + import numpy.random.bit_generator as np_bg + + types += [ + np_bg.__pyx_unpickle_SeedSequence, # type: ignore[attr-defined] + np_bg.__pyx_unpickle_SeedlessSeedSequence, # type: ignore[attr-defined] + ] + except ModuleNotFoundError: # pragma: no cover + pass + + return types + + +# --------------------------------------------------------------------------- +# Explicit string allowlist for numpy internal pickle helpers. +# These cover BOTH numpy 1.x (numpy.core.*) and numpy 2.x (numpy._core.*) +# module paths, because we cannot import them reliably at module load time. +# The restricted unpickler relies on these strings, so checkpoints saved +# with either numpy version load correctly. +# --------------------------------------------------------------------------- + +_NUMPY_PICKLE_INTERNALS_STR: tuple[str, ...] = ( + # numpy 2.x paths + "numpy._core.multiarray._reconstruct", + "numpy._core.multiarray.scalar", + "numpy._core.numeric._frombuffer", + # numpy 1.x paths + "numpy.core.multiarray._reconstruct", + "numpy.core.multiarray.scalar", + "numpy.core.numeric._frombuffer", + # numpy random pickle helpers (shared across versions) + "numpy.random._pickle.__generator_ctor", + "numpy.random._pickle.__bit_generator_ctor", + "numpy.random._pickle.__randomstate_ctor", + # Cython unpickle helpers in numpy.random.bit_generator + "numpy.random.bit_generator.__pyx_unpickle_SeedSequence", + "numpy.random.bit_generator.__pyx_unpickle_SeedlessSeedSequence", +) + + +def _collect_cloudpickle_types() -> list: + """Collect cloudpickle 2.x / 3.x internal helpers needed by torch weights_only.""" + import cloudpickle.cloudpickle as cp + + types = [ + # SECURITY: Function reconstruction helpers are commented out because they + # allow arbitrary code execution by reconstructing malicious functions. + # These should NOT be in the allowlist for safe deserialization. + # cp._make_function, + # cp._make_cell, + # cp._make_empty_cell, + # cp._function_setstate, + cp._make_skeleton_class, + cp._make_skeleton_enum, + cp._builtin_type, + ] + + # cloudpickle 3.x extras (may not exist in older versions) + for name in ( + "_make_dict_items", + "_make_dict_keys", + "_make_dict_values", + "_make_typevar", + ): + if hasattr(cp, name): + types.append(getattr(cp, name)) + + return types + + +def _collect_base_types() -> list[type | Callable[..., Any]]: + """Collect types that don't depend on SB3 internal imports (no circular import risk).""" + from gymnasium import spaces + from torch import nn + + types: list[type | Callable[..., Any]] = [] + + # Gymnasium spaces + types += [ + spaces.Box, + spaces.Discrete, + spaces.MultiBinary, + spaces.MultiDiscrete, + spaces.Dict, + spaces.Tuple, + spaces.Space, + ] + + # Try to also register old gym spaces (for models saved with gym <= 0.26) + try: # pragma: no cover + import gym + + types += [ + gym.spaces.Box, + gym.spaces.Discrete, + gym.spaces.MultiBinary, + gym.spaces.MultiDiscrete, + gym.spaces.Dict, + gym.spaces.Tuple, + gym.spaces.Space, + ] + except ImportError: # pragma: no cover + pass + + # Numpy types + types += _collect_numpy_types() + + # PyTorch types + types += [ + th.optim.Adam, + th.optim.SGD, + th.optim.RMSprop, + th.device, + ] + + # PyTorch nn.Module subclasses used as policy constructor parameters + types += [ + nn.Linear, + nn.Conv2d, + nn.Flatten, + nn.Sequential, + nn.Tanh, + nn.ReLU, + nn.ELU, + nn.LeakyReLU, + nn.SiLU, + nn.GELU, + nn.LayerNorm, + nn.BatchNorm2d, + ] + + # Cloudpickle internals + types += _collect_cloudpickle_types() + + # Standard library containers and builtins + types += [ + collections.deque, + collections.OrderedDict, + getattr, # getattr is somehow needed for save/load policy with weight_only=True + setattr, + ] + + return types + + +def _get_safe_globals_str() -> set[str]: + """Return the string allowlist for _RestrictedUnpickler, populating it lazily.""" + global _SAFE_GLOBALS_STR + if _SAFE_GLOBALS_STR is None: + _register_safe_globals() + assert _SAFE_GLOBALS_STR is not None + return _SAFE_GLOBALS_STR + + +def _register_safe_globals() -> None: + """Register safe types with both torch.serialization and build string allowlist. + + This function is called lazily (first time the allowlist is needed) to avoid + import-time side effects that could trigger circular imports. + + The base set of types (gymnasium, numpy, cloudpickle, etc.) is collected here. + SB3-specific types (policies, buffers, extractors, VecNormalize) must be + registered separately via ``register_sb3_safe_globals()`` because those modules + import from this file and we cannot import them at module load time. + """ + global _SAFE_GLOBALS_STR, _TORCH_REGISTRATION_DONE + + if _TORCH_REGISTRATION_DONE: # pragma: no cover + return + + base_types = _collect_base_types() + + # Build string allowlist for _RestrictedUnpickler + _SAFE_GLOBALS_STR = set() + for safe_type in base_types: + if hasattr(safe_type, "__module__") and hasattr(safe_type, "__qualname__"): + _SAFE_GLOBALS_STR.add(f"{safe_type.__module__}.{safe_type.__qualname__}") + else: # pragma: no cover + # Some objects (like numpy random functions) may not have these + try: + _SAFE_GLOBALS_STR.add(f"{type(safe_type).__module__}.{safe_type.__name__}") + except AttributeError: + pass + + # Numpy-internal pickle helpers: both numpy 1.x (numpy.core.*) and + # numpy 2.x (numpy._core.*) paths, plus random pickle helpers. + for _name in _NUMPY_PICKLE_INTERNALS_STR: + _SAFE_GLOBALS_STR.add(_name) + + # Explicit string entries for builtins and cloudpickle internals (these are + # functions/objects whose module/qualname doesn't match what pickle embeds) + for _name in ( + "builtins.tuple", + "builtins.dict", + "builtins.list", + "builtins.set", + "builtins.frozenset", + "builtins.str", + "builtins.int", + "builtins.float", + "builtins.bool", + "builtins.bytes", + "builtins.bytearray", + "builtins.object", + "builtins.type", + "cloudpickle.cloudpickle", + "cloudpickle.cloudpickle.__newobj__", + # SECURITY: Function reconstruction helpers are commented out because they + # allow arbitrary code execution. Do NOT add these to the allowlist. + # "cloudpickle.cloudpickle._make_skeleton_function", + # "cloudpickle.cloudpickle._make_stepfunc", + # "cloudpickle.cloudpickle._make_fileless_lambda", + # "cloudpickle.cloudpickle.make_function_from_globals", + "cloudpickle.cloudpickle.make_dict_fromnamedtuple", + "cloudpickle.cloudpickle.make_dict_fromnamedtuple_with_defaults", + "cloudpickle.cloudpickle.make_dynamic_classlookup", + "cloudpickle.cloudpickle.make_instance_from_reduce", + "cloudpickle.cloudpickle.make_local_from_global", + "cloudpickle.cloudpickle.make_numpy_array", + "cloudpickle.cloudpickle.make_numpy_scalar", + "cloudpickle.cloudpickle.make_opaque_object", + "cloudpickle.cloudpickle.make_object_from_newargs", + "cloudpickle.cloudpickle.make_object_from_newargsreduce", + "cloudpickle.cloudpickle.make_repr_from_name", + "cloudpickle.cloudpickle.make_seq", + "cloudpickle.cloudpickle.make_set", + "cloudpickle.cloudpickle.make_skeleton_class", + "cloudpickle.cloudpickle.make_skeleton_enum", + "cloudpickle.cloudpickle.make_super", + "cloudpickle.cloudpickle.make_type_var", + "cloudpickle.cloudpickle.make_type_var_tuple", + "cloudpickle.cloudpickle.make_typed_dict", + "cloudpickle.cloudpickle.make_unordered_set", + "cloudpickle.cloudpickle.restore_class", + "cloudpickle.cloudpickle.restore_class_attr_descriptors", + # SECURITY: restore_function allows arbitrary code execution + # "cloudpickle.cloudpickle.restore_function", + "cloudpickle.cloudpickle._class_setstate", + "cloudpickle.cloudpickle._fillvar", + # SECURITY: subimport is intentionally NOT whitelisted as it allows arbitrary + # module imports during deserialization, which would defeat the entire + # allowlist mechanism. Do NOT add this to the allowlist. + # "cloudpickle.cloudpickle.subimport", + "cloudpickle.cloudpickle._lookup_module_and_obj_in_qualname", + "cloudpickle.cloudpickle.whichmodule", + "types.FunctionType", + "types.ModuleType", + "types.CellType", + # PyTorch tensor/storage reconstruction helpers (needed for optimizer state + # and model weight pickles). These are safe: they only reconstruct tensors, + # storages, and parameters from serialised data. + "torch._utils._rebuild_tensor", + "torch._utils._rebuild_tensor_v2", + "torch._utils._rebuild_tensor_v3", + "torch._utils._rebuild_parameter", + "torch._utils._rebuild_parameter_with_state", + "torch._utils._rebuild_sparse_tensor", + "torch._utils._rebuild_nested_tensor", + "torch._utils._rebuild_qtensor", + "torch._utils._rebuild_device_tensor_from_cpu_tensor", + "torch._utils._rebuild_device_tensor_from_numpy", + "torch._utils._rebuild_wrapper_subclass", + "torch._utils._get_restore_location", + "torch._tensor._rebuild_from_type", + "torch._tensor._rebuild_from_type_v2", + "torch.storage._load_from_bytes", + "torch.storage._get_storage_from_sequence", + "copyreg.__newobj_ex__", + "copyreg.__newobj__", + "copyreg._reconstruct", + "copyreg.reconstructor", + "copyreg._reduce_ex", + ): + _SAFE_GLOBALS_STR.add(_name) + + # Register with torch.serialization + th.serialization.add_safe_globals(base_types) # type: ignore[arg-type] + _TORCH_REGISTRATION_DONE = True + + +def register_sb3_safe_globals() -> None: + """Register SB3-specific types (policies, buffers, extractors, VecNormalize). + + This must be called after SB3 modules are fully loaded (no circular import). + It is called lazily from ``policies.py`` when policy loading actually happens. + """ + global _SAFE_GLOBALS_STR + + # Ensure base registration is done first + if not _TORCH_REGISTRATION_DONE: + _register_safe_globals() + + # We can skip if already registered (check for a marker) + if _SAFE_GLOBALS_STR is not None and "stable_baselines3.common.buffers.ReplayBuffer" in _SAFE_GLOBALS_STR: + return + + from stable_baselines3.a2c.policies import ( + ActorCriticPolicy as A2CActorCriticPolicy, + ) + from stable_baselines3.a2c.policies import ( + CnnPolicy as A2CCnnPolicy, + ) + from stable_baselines3.a2c.policies import ( + MlpPolicy as A2CMlpPolicy, + ) + from stable_baselines3.a2c.policies import ( + MultiInputPolicy as A2CMultiInputPolicy, + ) + from stable_baselines3.common.buffers import ( + DictReplayBuffer, + ReplayBuffer, + RolloutBuffer, + ) + from stable_baselines3.common.distributions import ( + BernoulliDistribution, + CategoricalDistribution, + DiagGaussianDistribution, + MultiCategoricalDistribution, + StateDependentNoiseDistribution, + ) + from stable_baselines3.common.noise import ( + ActionNoise, + NormalActionNoise, + OrnsteinUhlenbeckActionNoise, + VectorizedActionNoise, + ) + from stable_baselines3.common.policies import ( + ActorCriticCnnPolicy, + ActorCriticPolicy, + BasePolicy, + ContinuousCritic, + MultiInputActorCriticPolicy, + ) + from stable_baselines3.common.running_mean_std import RunningMeanStd + from stable_baselines3.common.torch_layers import ( + BaseFeaturesExtractor, + CombinedExtractor, + FlattenExtractor, + MlpExtractor, + NatureCNN, + ) + from stable_baselines3.common.type_aliases import TrainFreq, TrainFrequencyUnit + from stable_baselines3.common.utils import ( + ConstantSchedule, + FloatSchedule, + LinearSchedule, + ) + from stable_baselines3.common.vec_env.vec_normalize import VecNormalize + from stable_baselines3.ddpg.policies import ( + CnnPolicy as DdpgCnnPolicy, + ) + from stable_baselines3.ddpg.policies import ( + MlpPolicy as DdpgMlpPolicy, + ) + from stable_baselines3.ddpg.policies import ( + MultiInputPolicy as DdpgMultiInputPolicy, + ) + from stable_baselines3.dqn.policies import ( + CnnPolicy as DqnCnnPolicy, + ) + from stable_baselines3.dqn.policies import ( + MlpPolicy as DqnMlpPolicy, + ) + from stable_baselines3.dqn.policies import ( + MultiInputPolicy as DqnMultiInputPolicy, + ) + from stable_baselines3.dqn.policies import ( + QNetwork, + ) + from stable_baselines3.her.goal_selection_strategy import GoalSelectionStrategy + from stable_baselines3.her.her_replay_buffer import HerReplayBuffer + from stable_baselines3.ppo.policies import ( + ActorCriticPolicy as PpoActorCriticPolicy, + ) + from stable_baselines3.ppo.policies import ( + CnnPolicy as PpoCnnPolicy, + ) + from stable_baselines3.ppo.policies import ( + MlpPolicy as PpoMlpPolicy, + ) + from stable_baselines3.ppo.policies import ( + MultiInputPolicy as PpoMultiInputPolicy, + ) + from stable_baselines3.sac.policies import ( + CnnPolicy as SacCnnPolicy, + ) + from stable_baselines3.sac.policies import ( + MlpPolicy as SacMlpPolicy, + ) + from stable_baselines3.sac.policies import ( + MultiInputPolicy as SacMultiInputPolicy, + ) + from stable_baselines3.td3.policies import ( + CnnPolicy as Td3CnnPolicy, + ) + from stable_baselines3.td3.policies import ( + MlpPolicy as Td3MlpPolicy, + ) + from stable_baselines3.td3.policies import ( + MultiInputPolicy as Td3MultiInputPolicy, + ) + + sb3_types = [ + # Buffers + ReplayBuffer, + DictReplayBuffer, + RolloutBuffer, + HerReplayBuffer, + GoalSelectionStrategy, + # Schedules + FloatSchedule, + ConstantSchedule, + LinearSchedule, + TrainFreq, + TrainFrequencyUnit, + # Action noise + ActionNoise, + NormalActionNoise, + OrnsteinUhlenbeckActionNoise, + VectorizedActionNoise, + # Extractors + FlattenExtractor, + NatureCNN, + CombinedExtractor, + MlpExtractor, + BaseFeaturesExtractor, + # Distributions + BernoulliDistribution, + CategoricalDistribution, + DiagGaussianDistribution, + MultiCategoricalDistribution, + StateDependentNoiseDistribution, + # Base policy classes + BasePolicy, + ActorCriticPolicy, + ActorCriticCnnPolicy, + MultiInputActorCriticPolicy, + ContinuousCritic, + # Algo-specific policies + A2CActorCriticPolicy, + A2CCnnPolicy, + A2CMlpPolicy, + A2CMultiInputPolicy, + PpoActorCriticPolicy, + PpoCnnPolicy, + PpoMlpPolicy, + PpoMultiInputPolicy, + DdpgCnnPolicy, + DdpgMlpPolicy, + DdpgMultiInputPolicy, + DqnCnnPolicy, + DqnMlpPolicy, + DqnMultiInputPolicy, + QNetwork, + SacCnnPolicy, + SacMlpPolicy, + SacMultiInputPolicy, + Td3CnnPolicy, + Td3MlpPolicy, + Td3MultiInputPolicy, + # VecNormalize and running stats + RunningMeanStd, + VecNormalize, + ] + + # Add to string allowlist + for safe_type in sb3_types: + if _SAFE_GLOBALS_STR is not None and hasattr(safe_type, "__module__") and hasattr(safe_type, "__qualname__"): + _SAFE_GLOBALS_STR.add(f"{safe_type.__module__}.{safe_type.__qualname__}") + + # Register with torch.serialization + th.serialization.add_safe_globals(sb3_types) # type: ignore[arg-type] + + +class _RestrictedUnpickler(pickle.Unpickler): + """Unpickler that only allows globals from a predefined allowlist. + + Blocks cloudpickle payloads that contain arbitrary code execution + while still permitting the types that SB3 legitimately serialises. + """ + + def find_class(self, module: str, name: str) -> Any: + global_full = f"{module}.{name}" + if global_full in get_safe_globals(): + return super().find_class(module, name) + raise pickle.UnpicklingError( + f"Global {global_full!r} is not in the safe deserialization allowlist. " + "This is likely an attempt to execute arbitrary code via a crafted " + "checkpoint. Use deserialization_mode='legacy' if you trust this file, " + "or report the type to the SB3 maintainers." + ) + + +# --------------------------------------------------------------------------- +# Public API: extend the safe-globals allowlist (a la torch.serialization) +# --------------------------------------------------------------------------- + +_USER_SAFE_GLOBALS: set[str] = set() + + +def add_safe_globals( + safe_globals: list[type | tuple[type, str]] | type | tuple[type, str], +) -> None: + """Register one or more classes/functions as safe for ``deserialization_mode="safe"``. + + This is the SB3 equivalent of :func:`torch.serialization.add_safe_globals`. + Types registered here will be permitted by the restricted unpickler when + ``deserialization_mode="safe"`` is used, for *all* subsequent loads in the + current process. + + Each item can be: + + * A class or function object. Its fully-qualified name + ``module.qualname`` will be used automatically. + * A ``(class_or_function, "explicit.module.Path")`` tuple when the + pickle payload uses a different module path than the live object + (e.g., the checkpoint was saved on a machine with a different + package name). + + **Only register types you trust.** The security guarantee of + ``deserialization_mode="safe"`` is that *only* allowlisted globals can + be invoked during unpickling. + + **Example** + + .. code-block:: python + + from stable_baselines3.common.safe_globals import add_safe_globals + + class MyCustomSpace(gymnasium.spaces.Space): + ... + + add_safe_globals([MyCustomSpace]) + model = PPO.load("checkpoint.zip", deserialization_mode="safe") + """ + if not isinstance(safe_globals, list): + safe_globals = [safe_globals] + + for item in safe_globals: + if isinstance(item, tuple): + obj, explicit_path = item + _USER_SAFE_GLOBALS.add(explicit_path) + th.serialization.add_safe_globals([item]) # type: ignore[arg-type] + else: + _USER_SAFE_GLOBALS.add(f"{item.__module__}.{item.__qualname__}") + th.serialization.add_safe_globals([item]) # type: ignore[arg-type] + + +class SafeGlobals: + """Context-manager that temporarily adds globals to the safe allowlist. + + The added types are automatically removed when the block exits. + + **Example** + + .. code-block:: python + + from stable_baselines3.common.safe_globals import SafeGlobals + + with SafeGlobals([MyCustomSpace]): + model = PPO.load("checkpoint.zip", deserialization_mode="safe") + """ + + def __init__( + self, + safe_globals: list[type | tuple[type, str]] | type | tuple[type, str], + ) -> None: + self._items = safe_globals if isinstance(safe_globals, list) else [safe_globals] + + def __enter__(self) -> SafeGlobals: + self._backup = _USER_SAFE_GLOBALS.copy() + add_safe_globals(self._items) + return self + + def __exit__(self, *args) -> None: + _USER_SAFE_GLOBALS.clear() + _USER_SAFE_GLOBALS.update(self._backup) + + +def get_safe_globals() -> set[str]: + """Return the complete allowlist (built-in + user-registered) for safe deserialization.""" + base = _get_safe_globals_str() + return base | _USER_SAFE_GLOBALS diff --git a/stable_baselines3/common/save_util.py b/stable_baselines3/common/save_util.py index 3569eb887..f015b5c28 100644 --- a/stable_baselines3/common/save_util.py +++ b/stable_baselines3/common/save_util.py @@ -18,10 +18,18 @@ import torch as th import stable_baselines3 as sb3 -from stable_baselines3.common.type_aliases import TensorDict +from stable_baselines3.common.safe_globals import ( + _RestrictedUnpickler, +) +from stable_baselines3.common.type_aliases import DeserializationMode, TensorDict from stable_baselines3.common.utils import get_device, get_system_info +def _cloudpickle_loads_safe(data: bytes) -> Any: + """Deserialize cloudpickle data using the restricted allowlist unpickler.""" + return _RestrictedUnpickler(io.BytesIO(data)).load() + + def recursive_getattr(obj: Any, attr: str, *args) -> Any: """ Recursive version of getattr @@ -128,7 +136,11 @@ def data_to_json(data: dict[str, Any]) -> str: return json_string -def json_to_data(json_string: str, custom_objects: dict[str, Any] | None = None) -> dict[str, Any]: +def json_to_data( + json_string: str, + custom_objects: dict[str, Any] | None = None, + deserialization_mode: DeserializationMode = DeserializationMode.SAFE, +) -> dict[str, Any]: """ Turn JSON serialization of class-parameters back into dictionary. @@ -140,14 +152,33 @@ def json_to_data(json_string: str, custom_objects: dict[str, Any] | None = None) will be used instead. Similar to custom_objects in ``keras.models.load_model``. Useful when you have an object in file that can not be deserialized. + :param deserialization_mode: How to handle cloudpickle-serialized objects + stored in the checkpoint's ``data`` JSON. + + - ``"safe"`` (default): Deserialize using a restricted unpickler that + only allows a fixed allowlist of known-safe SB3/gymnasium/numpy types. + Any cloudpickle payload referencing a type outside this allowlist is + skipped with a warning. + - ``"legacy"``: Deserialize all ``:serialized:`` entries with the + unrestricted cloudpickle loader. This preserves full backward + compatibility but **executes arbitrary Python code** embedded in the + checkpoint. :return: Loaded class parameters. """ + # Ensure SB3 types are registered for safe deserialization + from stable_baselines3.common.safe_globals import register_sb3_safe_globals + + register_sb3_safe_globals() if custom_objects is not None and not isinstance(custom_objects, dict): raise ValueError("custom_objects argument must be a dict or None") + if deserialization_mode not in (DeserializationMode.SAFE, DeserializationMode.LEGACY): + raise ValueError(f"deserialization_mode must be 'legacy' or 'safe', got {deserialization_mode!r}") + json_dict = json.loads(json_string) # This will be filled with deserialized data return_data = {} + warned_once = False for data_key, data_item in json_dict.items(): if custom_objects is not None and data_key in custom_objects.keys(): # If item is provided in custom_objects, replace @@ -156,14 +187,26 @@ def json_to_data(json_string: str, custom_objects: dict[str, Any] | None = None) elif isinstance(data_item, dict) and ":serialized:" in data_item.keys(): # If item is dictionary with ":serialized:" # key, this means it is serialized with cloudpickle. + if not warned_once: + warned_once = True + if deserialization_mode == DeserializationMode.LEGACY: + 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, + ) + serialization = data_item[":serialized:"] - # Try-except deserialization in case we run into - # errors. If so, we can tell bit more information to - # user. try: base64_object = base64.b64decode(serialization.encode()) - deserialized_object = cloudpickle.loads(base64_object) - except (RuntimeError, TypeError, AttributeError) as e: + if deserialization_mode == DeserializationMode.SAFE: + deserialized_object = _cloudpickle_loads_safe(base64_object) + else: + deserialized_object = cloudpickle.loads(base64_object) + except (RuntimeError, TypeError, AttributeError, pickle.UnpicklingError) as e: warnings.warn( f"Could not deserialize object {data_key}. " "Consider using `custom_objects` argument to replace " @@ -356,7 +399,11 @@ def save_to_pkl(path: str | pathlib.Path | io.BufferedIOBase, obj: Any, verbose: file.close() -def load_from_pkl(path: str | pathlib.Path | io.BufferedIOBase, verbose: int = 0) -> Any: +def load_from_pkl( + path: str | pathlib.Path | io.BufferedIOBase, + verbose: int = 0, + deserialization_mode: DeserializationMode = DeserializationMode.SAFE, +) -> Any: """ Load an object from the path. If a suffix is provided in the path, it will use that suffix. If the path does not exist, it will attempt to load using the .pkl suffix. @@ -365,9 +412,35 @@ def load_from_pkl(path: str | pathlib.Path | io.BufferedIOBase, verbose: int = 0 if save_path is a str or pathlib.Path and mode is "w", single dispatch ensures that the path actually exists. If path is a io.BufferedIOBase the path exists. :param verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages + :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 **executes arbitrary Python code** embedded + in the pickle file. A ``UserWarning`` is emitted. """ + if deserialization_mode not in (DeserializationMode.SAFE, DeserializationMode.LEGACY): + raise ValueError(f"deserialization_mode must be 'legacy' or 'safe', got {deserialization_mode!r}") + file = open_path(path, "r", verbose=verbose, suffix="pkl") - obj = pickle.load(file) + + if deserialization_mode == DeserializationMode.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 @@ -380,6 +453,7 @@ def load_from_zip_file( device: th.device | str = "auto", verbose: int = 0, print_system_info: bool = False, + deserialization_mode: DeserializationMode = DeserializationMode.SAFE, ) -> tuple[dict[str, Any] | None, TensorDict, TensorDict | None]: """ Load model data from a .zip archive @@ -397,6 +471,8 @@ def load_from_zip_file( :param verbose: Verbosity level: 0 for no output, 1 for info messages, 2 for debug messages :param print_system_info: Whether to print or not the system info about the saved model. + :param deserialization_mode: How to handle cloudpickle-serialized objects + in the checkpoint's ``data`` JSON. See :func:`json_to_data` for details. :return: Class parameters, model state_dicts (aka "params", dict of state_dict) and dict of pytorch variables """ @@ -421,7 +497,7 @@ def load_from_zip_file( if "system_info.txt" in namelist: print("== SAVED MODEL SYSTEM INFO ==") print(archive.read("system_info.txt").decode()) - else: + else: # pragma: no cover warnings.warn( "The model was saved with SB3 <= 1.2.0 and thus cannot print system information.", UserWarning, @@ -431,7 +507,11 @@ def load_from_zip_file( # Load class parameters that are stored # with either JSON or pickle (not PyTorch variables). json_data = archive.read("data").decode() - data = json_to_data(json_data, custom_objects=custom_objects) + data = json_to_data( + json_data, + custom_objects=custom_objects, + deserialization_mode=deserialization_mode, + ) # Check for all .pth files and load them using th.load. # "pytorch_variables.pth" stores PyTorch variables, and any other .pth diff --git a/stable_baselines3/common/type_aliases.py b/stable_baselines3/common/type_aliases.py index 3865d6da7..fee24536d 100644 --- a/stable_baselines3/common/type_aliases.py +++ b/stable_baselines3/common/type_aliases.py @@ -24,6 +24,13 @@ MaybeCallback = Union[None, Callable, list["BaseCallback"], "BaseCallback"] PyTorchObs = Union[th.Tensor, TensorDict] # noqa: UP007 + +# Deserialization mode for secure checkpoint loading +class DeserializationMode(str, Enum): + SAFE = "safe" + LEGACY = "legacy" + + # A schedule takes the remaining progress as input # and outputs a scalar (e.g. learning rate, clip range, ...) Schedule = Callable[[float], float] diff --git a/stable_baselines3/common/vec_env/vec_normalize.py b/stable_baselines3/common/vec_env/vec_normalize.py index dea4f4c0b..755339934 100644 --- a/stable_baselines3/common/vec_env/vec_normalize.py +++ b/stable_baselines3/common/vec_env/vec_normalize.py @@ -1,5 +1,6 @@ import inspect import pickle +import warnings from copy import deepcopy from typing import Any @@ -9,6 +10,7 @@ from stable_baselines3.common import utils from stable_baselines3.common.preprocessing import is_image_space from stable_baselines3.common.running_mean_std import RunningMeanStd +from stable_baselines3.common.type_aliases import DeserializationMode from stable_baselines3.common.vec_env.base_vec_env import VecEnv, VecEnvStepReturn, VecEnvWrapper @@ -308,16 +310,48 @@ def reset(self) -> np.ndarray | dict[str, np.ndarray]: return self.normalize_obs(obs) @staticmethod - def load(load_path: str, venv: VecEnv) -> "VecNormalize": + def load( + load_path: str, + venv: VecEnv, + deserialization_mode: DeserializationMode = DeserializationMode.SAFE, + ) -> "VecNormalize": """ Loads a saved VecNormalize object. :param load_path: the path to load from. :param venv: the VecEnv to wrap. + :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 **executes arbitrary Python code** + embedded in the pickle file. A ``UserWarning`` is emitted. :return: """ - with open(load_path, "rb") as file_handler: - vec_normalize = pickle.load(file_handler) + if deserialization_mode not in (DeserializationMode.SAFE, DeserializationMode.LEGACY): + raise ValueError(f"deserialization_mode must be 'legacy' or 'safe', got {deserialization_mode!r}") + + if deserialization_mode == DeserializationMode.SAFE: + from stable_baselines3.common.safe_globals import ( + _RestrictedUnpickler, + register_sb3_safe_globals, + ) + + register_sb3_safe_globals() + with open(load_path, "rb") as file_handler: + vec_normalize = _RestrictedUnpickler(file_handler).load() + else: + warnings.warn( + "Loading a VecNormalize pickle file with pickle deserialization " + "(deserialization_mode='legacy'). This can execute arbitrary Python " + "code from the file. Only load pickle files from trusted sources. ", + UserWarning, + ) + with open(load_path, "rb") as file_handler: + vec_normalize = pickle.load(file_handler) vec_normalize.set_venv(venv) return vec_normalize diff --git a/stable_baselines3/version.txt b/stable_baselines3/version.txt index a36506279..1e8c33284 100644 --- a/stable_baselines3/version.txt +++ b/stable_baselines3/version.txt @@ -1 +1 @@ -2.9.2a0 +2.10.0a0 diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 7a65dcbc5..a1c4f233f 100644 --- a/tests/test_save_load.py +++ b/tests/test_save_load.py @@ -630,7 +630,8 @@ def test_open_file_str_pathlib(tmp_path, pathtype): assert fp1.closed with warnings.catch_warnings(record=True) as record: assert load_from_pkl(pathtype(f"{tmp_path}/t1")) == "foo" - assert not record + # Only the expected security warning from load_from_pkl; no path-related warnings + assert all("deserializer" in str(warning.message).lower() for warning in record) # test custom suffix with open_path(pathtype(f"{tmp_path}/t1.custom_ext"), "w") as fp1: @@ -638,7 +639,7 @@ def test_open_file_str_pathlib(tmp_path, pathtype): assert fp1.closed with warnings.catch_warnings(record=True) as record: assert load_from_pkl(pathtype(f"{tmp_path}/t1.custom_ext")) == "foo" - assert not record + assert all("deserializer" in str(warning.message).lower() for warning in record) # test without suffix with open_path(pathtype(f"{tmp_path}/t1"), "w", suffix="pkl") as fp1: @@ -646,7 +647,7 @@ def test_open_file_str_pathlib(tmp_path, pathtype): assert fp1.closed with warnings.catch_warnings(record=True) as record: assert load_from_pkl(pathtype(f"{tmp_path}/t1.pkl")) == "foo" - assert not record + assert all("deserializer" in str(warning.message).lower() for warning in record) # test that a warning is raised when the path doesn't exist with open_path(pathtype(f"{tmp_path}/t2.pkl"), "w") as fp1: @@ -654,11 +655,14 @@ def test_open_file_str_pathlib(tmp_path, pathtype): assert fp1.closed with warnings.catch_warnings(record=True) as record: assert load_from_pkl(open_path(pathtype(f"{tmp_path}/t2"), "r", suffix="pkl")) == "foo" - assert len(record) == 0 + # Only security warning, no "path not found" warning + assert all("deserializer" in str(warning.message).lower() for warning in record) with warnings.catch_warnings(record=True) as record: assert load_from_pkl(open_path(pathtype(f"{tmp_path}/t2"), "r", suffix="pkl", verbose=2)) == "foo" - assert len(record) == 1 + # Security warning + path-not-found verbose warning + non_security = [warning for warning in record if "deserializer" not in str(warning.message).lower()] + assert len(non_security) == 1 fp = pathlib.Path(f"{tmp_path}/t2").open("w") fp.write("rubbish") @@ -729,6 +733,9 @@ def test_save_load_large_model(tmp_path): def test_load_invalid_object(tmp_path): # See GH Issue #1122 for an example # of invalid object loading + # Note: This test uses deserialization_mode="legacy" because it tests with lambda + # functions which cannot be deserialized in safe mode (they require _make_function + # which is intentionally not in the allowlist for security reasons). path = str(tmp_path / "ppo_pendulum.zip") PPO("MlpPolicy", "Pendulum-v1", learning_rate=lambda _: 1.0).save(path) @@ -748,12 +755,19 @@ def test_load_invalid_object(tmp_path): # Replace with the corrupted file # probably doesn't work on windows os.system(f"cd {tmp_path}; zip ppo_pendulum.zip data") - with pytest.warns(UserWarning, match=r"custom_objects"): - PPO.load(path) - # Load with custom object, no warnings with warnings.catch_warnings(record=True) as record: - PPO.load(path, custom_objects=dict(learning_rate=lambda _: 1.0)) - assert len(record) == 0 + warnings.simplefilter("always") + PPO.load(path, deserialization_mode="legacy") + assert len(record) == 2 + assert any("cloudpickle-serialized" in str(warning.message) for warning in record) + assert any("custom_objects" in str(warning.message) for warning in record) + # Load with custom object: the only warning should be the security warning + # (no "Could not deserialize" or "custom_objects" warnings) + with warnings.catch_warnings(record=True) as record: + PPO.load(path, custom_objects=dict(learning_rate=lambda _: 1.0), deserialization_mode="legacy") + # Filter out the expected security warning + non_security = [warning for warning in record if "cloudpickle-serialized" not in str(warning.message)] + assert len(non_security) == 0 def test_dqn_target_update_interval(tmp_path): @@ -767,8 +781,8 @@ def test_dqn_target_update_interval(tmp_path): assert model.target_update_interval == 100 -# Turn warnings into errors -@pytest.mark.filterwarnings("error") +# Turn ResourceWarnings into errors (not our security warnings) +@pytest.mark.filterwarnings("error::ResourceWarning") def test_no_resource_warning(tmp_path): # Check behavior of save/load # see https://github.com/DLR-RM/stable-baselines3/issues/1751 @@ -811,7 +825,8 @@ def test_cast_lr_schedule(tmp_path): assert type(model.lr_schedule(1.0)) is float assert np.allclose(model.lr_schedule(0.5), 0.5 * np.sin(1.0)) model.save(tmp_path / "ppo.zip") - model = PPO.load(tmp_path / "ppo.zip") + with pytest.warns(UserWarning, match=r"cloudpickle-serialized"): + model = PPO.load(tmp_path / "ppo.zip", deserialization_mode="legacy") assert type(model.lr_schedule(1.0)) is float assert np.allclose(model.lr_schedule(0.5), 0.5 * np.sin(1.0)) @@ -854,7 +869,8 @@ def test_save_load_backward_compatible(tmp_path, model_class): model.save(tmp_path / "test_schedule_safe.zip") - model = model_class.load(tmp_path / "test_schedule_safe.zip", env=env) + with pytest.warns(UserWarning, match=r"cloudpickle-serialized"): + model = model_class.load(tmp_path / "test_schedule_safe.zip", env=env, deserialization_mode="legacy") assert model.learning_rate(0) == 0.001 assert model.learning_rate.__name__ == "" diff --git a/tests/test_security_deserialization.py b/tests/test_security_deserialization.py new file mode 100644 index 000000000..302cf916f --- /dev/null +++ b/tests/test_security_deserialization.py @@ -0,0 +1,579 @@ +""" +Security tests for deserialization of untrusted data remediation. + +These tests verify that the ``deserialization_mode`` parameter correctly gates +unsafe deserialization across all SB3 load paths: + +1. ``json_to_data`` (cloudpickle inside checkpoint .zip ``data`` JSON) +2. ``load_from_pkl`` (raw pickle, e.g. replay buffers) +3. ``VecNormalize.load`` (raw pickle) +4. ``BaseAlgorithm.load`` / ``PPO.load`` etc. (zip checkpoint entry point) +""" + +import base64 +import io +import json +import os +import pathlib +import pickle +import warnings +import zipfile + +import cloudpickle +import gymnasium as gym +import numpy as np +import pytest + +from stable_baselines3 import DDPG, PPO, SAC +from stable_baselines3.common.noise import NormalActionNoise +from stable_baselines3.common.safe_globals import ( + _USER_SAFE_GLOBALS, + SafeGlobals, + add_safe_globals, + get_safe_globals, +) +from stable_baselines3.common.save_util import ( + json_to_data, + load_from_pkl, + save_to_pkl, +) +from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize + +# --------------------------------------------------------------------------- +# Helpers: craft malicious payloads that would execute code on deserialization +# --------------------------------------------------------------------------- + + +def _make_evil_class(sentinel_path: pathlib.Path) -> type: + """ + Return a class whose instances, when pickled and then unpickled, + write the sentinel file. Uses __reduce__ with a top-level function + (not an instance method) so that both ``pickle`` and ``cloudpickle`` + trigger the side effect on load. + """ + + def _write_sentinel() -> None: + sentinel_path.write_text("PWNED") + + class Evil: + def __reduce__(self): + return (_write_sentinel, ()) + + return Evil + + +def _evil_cloudpickle_payload(sentinel_path: pathlib.Path) -> str: + """Return a base64-encoded cloudpickle payload that writes the sentinel.""" + return base64.b64encode(cloudpickle.dumps(_make_evil_class(sentinel_path)())).decode() + + +def _inject_malicious_entry(zip_path: str, sentinel_path: pathlib.Path) -> None: + """ + Open an existing SB3 checkpoint zip, inject a cloudpickle-serialized + malicious object into the ``data`` JSON, and write back. + """ + with zipfile.ZipFile(zip_path) as z: + data = json.loads(z.read("data").decode()) + + payload = _evil_cloudpickle_payload(sentinel_path) + data["_malicious_"] = {":type:": "", ":serialized:": payload} + + with zipfile.ZipFile(zip_path, "r") as zin: + with zipfile.ZipFile(zip_path + ".tmp", "w") as zout: + for name in zin.namelist(): + content = zin.read(name) + if name == "data": + content = json.dumps(data).encode() + zout.writestr(name, content) + os.replace(zip_path + ".tmp", zip_path) + + +# --------------------------------------------------------------------------- +# Test 1: json_to_data blocks malicious cloudpickle in safe mode +# --------------------------------------------------------------------------- + + +def test_json_to_data_safe_blocks_malicious(tmp_path): + """json_to_data with deserialization_mode='safe' must skip :serialized: entries.""" + sentinel = tmp_path / "sentinel" + payload = _evil_cloudpickle_payload(sentinel) + json_str = json.dumps( + { + "learning_rate": 0.001, # plain JSON, should load fine + "_evil_": {":type:": "", ":serialized:": payload}, + "verbose": 0, + } + ) + + # Safe mode: should skip the malicious entry + with warnings.catch_warnings(record=True) as rec: + result = json_to_data(json_str, deserialization_mode="safe") + + assert not sentinel.exists(), "Sentinel file created: RCE was NOT blocked!" + assert "_evil_" not in result, "Malicious entry should have been skipped" + assert result["learning_rate"] == 0.001, "Plain JSON entries should still load" + assert result["verbose"] == 0 + # Check that warnings were emitted (safe-mode warning + deserialization error) + assert len(rec) >= 1, "Expected at least one warning" + + +def test_json_to_data_legacy_allows_malicious(tmp_path): + """json_to_data with deserialization_mode='legacy' must still deserialize (backward compat).""" + sentinel = tmp_path / "sentinel" + payload = _evil_cloudpickle_payload(sentinel) + json_str = json.dumps( + { + "learning_rate": 0.001, + "_evil_": {":type:": "", ":serialized:": payload}, + } + ) + + # Legacy mode: should deserialize and execute the payload + with warnings.catch_warnings(record=True): + _ = json_to_data(json_str, deserialization_mode="legacy") + + # In legacy mode the evil payload runs + assert sentinel.exists(), "Legacy mode should have deserialized the payload" + + +# --------------------------------------------------------------------------- +# Test 2: load_from_pkl blocks in safe mode +# --------------------------------------------------------------------------- + + +def test_load_from_pkl_safe_blocks_evil(tmp_path): + """load_from_pkl with deserialization_mode='safe' blocks evil globals.""" + path = tmp_path / "test.pkl" + sentinel = tmp_path / "sentinel" + evil_class = _make_evil_class(sentinel) + + # cloudpickle can serialize local functions; standard pickle cannot + with open(path, "wb") as f: + cloudpickle.dump(evil_class(), f) + + # Safe mode: should block because __reduce__ invokes a non-allowlisted global + with pytest.raises(pickle.UnpicklingError, match="not in the safe deserialization allowlist"): + load_from_pkl(path, deserialization_mode="safe") + + assert not sentinel.exists(), "Sentinel file created: RCE was NOT blocked!" + + +def test_load_from_pkl_legacy_warns(tmp_path): + """load_from_pkl with deserialization_mode='legacy' must emit a warning.""" + path = tmp_path / "test.pkl" + save_to_pkl(path, {"key": "value"}) + + with warnings.catch_warnings(record=True) as rec: + result = load_from_pkl(path, deserialization_mode="legacy") + assert result == {"key": "value"} + assert any("pickle deserialization" in str(warning.message).lower() for warning in rec) + + +# --------------------------------------------------------------------------- +# Test 3: VecNormalize.load blocks in safe mode +# --------------------------------------------------------------------------- + + +def test_vec_normalize_load_safe_works(tmp_path): + """VecNormalize.load with deserialization_mode='safe' succeeds for trusted data.""" + + venv = DummyVecEnv([lambda: gym.make("CartPole-v1")]) + vec_normalize = VecNormalize(venv) + # Run a few steps so the stats are non-trivial + for _ in range(20): + vec_normalize.reset() + vec_normalize.step([vec_normalize.action_space.sample()]) + pkl_path = tmp_path / "vecnormalize.pkl" + vec_normalize.save(str(pkl_path)) + + # Safe mode: should load successfully because VecNormalize uses only allowlisted types + loaded = VecNormalize.load(str(pkl_path), venv, deserialization_mode="safe") + assert loaded.obs_rms.mean.shape == vec_normalize.obs_rms.mean.shape + # Verify the running stats were preserved + np.testing.assert_allclose(loaded.obs_rms.mean, vec_normalize.obs_rms.mean, atol=1e-6) + + +# --------------------------------------------------------------------------- +# Test 4: BaseAlgorithm.load / PPO.load with malicious checkpoint +# --------------------------------------------------------------------------- + + +def test_ppo_load_safe_blocks_rce(tmp_path): + """ + PPO.load with deserialization_mode='safe' must block a malicious checkpoint + that would normally execute arbitrary code via cloudpickle. + + In safe mode, serialized entries like observation_space, action_space, + learning_rate, etc. are skipped. The user must provide them via + custom_objects or pass an env. Here we pass custom_objects to supply + the required objects. + """ + + sentinel = tmp_path / "sentinel" + + # Create a clean model and save + model = PPO("MlpPolicy", "CartPole-v1", n_steps=64, device="cpu") + zip_path = str(tmp_path / "model.zip") + model.save(zip_path) + + # Inject malicious entry + _inject_malicious_entry(zip_path, sentinel) + + # Safe mode: must provide custom_objects for serialized entries + env = gym.make("CartPole-v1") + custom_objects = { + "observation_space": env.observation_space, + "action_space": env.action_space, + "policy_class": PPO.policy_aliases["MlpPolicy"], + "rollout_buffer_class": model.rollout_buffer_class, + "learning_rate": 0.0, + "lr_schedule": lambda _: 0.0, + "clip_range": lambda _: 0.0, + } + with warnings.catch_warnings(record=True): + loaded = PPO.load( + zip_path, + env=env, + device="cpu", + deserialization_mode="safe", + custom_objects=custom_objects, + ) + + assert not sentinel.exists(), "Sentinel file created after PPO.load(safe): RCE was NOT blocked!" + # Verify the loaded model works + obs = loaded.get_env().reset() + loaded.predict(obs, deterministic=True) + + +def test_ppo_load_legacy_allows_rce(tmp_path): + """ + PPO.load with deserialization_mode='legacy' still deserializes malicious + checkpoints (backward compatibility). + """ + sentinel = tmp_path / "sentinel" + + model = PPO("MlpPolicy", "CartPole-v1", n_steps=64, device="cpu") + zip_path = str(tmp_path / "model.zip") + model.save(zip_path) + + _inject_malicious_entry(zip_path, sentinel) + + # Legacy mode: payload executes + with warnings.catch_warnings(record=True): + PPO.load(zip_path, device="cpu", deserialization_mode="legacy") + + assert sentinel.exists(), "Legacy mode should have executed the payload" + + +# --------------------------------------------------------------------------- +# Test 5: Invalid deserialization_mode raises +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "func_and_args", + [ + # json_to_data + lambda: (json_to_data, (json.dumps({}), None, "bad_mode")), + # load_from_pkl + lambda: (load_from_pkl, (io.BytesIO(pickle.dumps("x")), 0, "bad_mode")), + # VecNormalize.load + lambda: (VecNormalize.load, ("dummy.pkl", None, "bad_mode")), + # OffPolicyAlgorithm.load_replay_buffer (via SAC) + lambda: ( + lambda path, mode: SAC("MlpPolicy", "Pendulum-v1", device="cpu").load_replay_buffer( + path, deserialization_mode=mode + ), + ("dummy.pkl", "bad_mode"), + ), + ], +) +def test_invalid_deserialization_mode_raises(func_and_args): + func, args = func_and_args() + with pytest.raises(ValueError, match="deserialization_mode must be"): + func(*args) + + +# --------------------------------------------------------------------------- +# Test 6: custom_objects overrides still work in safe mode +# --------------------------------------------------------------------------- + + +def test_load_replay_buffer_safe_works(tmp_path): + """load_replay_buffer with deserialization_mode='safe' succeeds for trusted buffers.""" + model = SAC("MlpPolicy", "Pendulum-v1", buffer_size=1000, learning_starts=50, device="cpu") + model.learn(80) + original_size = model.replay_buffer.size() + pkl_path = str(tmp_path / "replay_buffer.pkl") + model.save_replay_buffer(pkl_path) + + # Safe mode: should load successfully + model.load_replay_buffer(pkl_path, deserialization_mode="safe") + + assert model.replay_buffer.size() == original_size + + +def test_model_load_safe_with_action_noise(tmp_path): + """Model checkpoints with action noise load in safe mode.""" + noise = NormalActionNoise(mean=[0], sigma=[0.1]) + model = DDPG( + "MlpPolicy", + "Pendulum-v1", + buffer_size=1000, + learning_starts=50, + action_noise=noise, + device="cpu", + ) + zip_path = str(tmp_path / "model_noise.zip") + model.save(zip_path) + + # Safe mode: should load successfully with action noise intact + loaded = DDPG.load(zip_path, device="cpu", deserialization_mode="safe") + assert loaded.action_noise is not None + assert isinstance(loaded.action_noise, NormalActionNoise) + + +def test_vec_normalize_load_safe_blocks_evil(tmp_path): + """VecNormalize.load with deserialization_mode='safe' blocks evil globals.""" + import cloudpickle + + sentinel = tmp_path / "sentinel" + venv = DummyVecEnv([lambda: gym.make("CartPole-v1")]) + evil_class = _make_evil_class(sentinel) + evil_data = cloudpickle.dumps(evil_class()) + + pkl_path = tmp_path / "vecnormalize_evil.pkl" + with open(pkl_path, "wb") as f: + f.write(evil_data) + + # Safe mode: should block the evil global + with pytest.raises(pickle.UnpicklingError, match="not in the safe deserialization allowlist"): + VecNormalize.load(str(pkl_path), venv, deserialization_mode="safe") + + assert not sentinel.exists(), "Sentinel file created: RCE was NOT blocked!" + + +def test_json_to_data_safe_with_custom_objects(tmp_path): + """In safe mode, keys provided via custom_objects are used instead of skipping.""" + sentinel = tmp_path / "sentinel" + payload = _evil_cloudpickle_payload(sentinel) + json_str = json.dumps( + { + "learning_rate": 0.001, + "lr_schedule": {":type:": "", ":serialized:": payload}, + } + ) + + # Provide a safe replacement via custom_objects + custom = {"lr_schedule": lambda x: 0.0} + result = json_to_data(json_str, custom_objects=custom, deserialization_mode="safe") + + assert not sentinel.exists(), "Sentinel created: custom_objects override failed" + assert callable(result["lr_schedule"]) + assert result["lr_schedule"](0.5) == 0.0 + + +# --------------------------------------------------------------------------- +# Test 7: add_safe_globals and SafeGlobals context manager +# --------------------------------------------------------------------------- + + +def test_add_safe_globals_persists(): + """add_safe_globals() should persist across calls.""" + # Clear user globals to start fresh + _USER_SAFE_GLOBALS.clear() + + class MyCustomType: + pass + + # Before registration: MyCustomType should NOT be in the allowlist + before = get_safe_globals() + qualname = f"{MyCustomType.__module__}.{MyCustomType.__qualname__}" + assert qualname not in before, "Custom type should not be in allowlist yet" + + # Register it + add_safe_globals([MyCustomType]) + + # After registration: MyCustomType should be in the allowlist + after = get_safe_globals() + assert qualname in after, "Custom type should be in allowlist after add_safe_globals" + + # Clean up + _USER_SAFE_GLOBALS.discard(qualname) + + +def test_safe_globals_context_manager(tmp_path): + """safe_globals context manager should restore allowlist on exit.""" + import cloudpickle + + # Clear user globals to start fresh + _USER_SAFE_GLOBALS.clear() + + class ScopedType: + value = 42 + + qualname = f"{ScopedType.__module__}.{ScopedType.__qualname__}" + + # Before context: not in allowlist + assert qualname not in get_safe_globals() + + # Inside context: should be in allowlist + with SafeGlobals([ScopedType]): + assert qualname in get_safe_globals() + # Verify the type can actually be deserialized + payload = cloudpickle.dumps(ScopedType()) + from stable_baselines3.common.save_util import _cloudpickle_loads_safe + + obj = _cloudpickle_loads_safe(payload) + assert obj.value == 42 + + # After context: should be removed + assert qualname not in get_safe_globals(), "safe_globals context manager did not restore allowlist on exit" + + +# --------------------------------------------------------------------------- +# Test 8: KNOWN VULNERABILITY - cloudpickle function reconstruction bypass +# --------------------------------------------------------------------------- + + +def test_safe_mode_function_reconstruction_blocked(tmp_path): + """ + Test that deserialization_mode='safe' blocks arbitrary function + deserialization via cloudpickle's internal helpers. + + After fixing the vulnerability by removing _make_function, _make_cell, etc. + from the allowlist, functions should no longer be deserializable in safe mode. + """ + sentinel = tmp_path / "pwned_by_function" + + # Create a malicious function and serialize it with cloudpickle + def malicious_function(): + sentinel.write_text("PWNED") + + payload = cloudpickle.dumps(malicious_function) + + # Try to deserialize with safe mode + from stable_baselines3.common.save_util import _cloudpickle_loads_safe + + # This should now raise an error + with pytest.raises(pickle.UnpicklingError, match="not in the safe deserialization allowlist"): + _cloudpickle_loads_safe(payload) + + # Verify the sentinel was NOT created + assert not sentinel.exists(), "Function should NOT have been deserialized and executed" + + +def test_load_from_pkl_safe_blocks_function_rce(tmp_path): + """ + Test that load_from_pkl with deserialization_mode='safe' blocks + arbitrary function deserialization, preventing RCE. + """ + sentinel = tmp_path / "sentinel_function" + + # Create a malicious function + def write_sentinel(): + sentinel.write_text("PWNED") + + # Serialize it + path = tmp_path / "malicious_func.pkl" + with open(path, "wb") as f: + cloudpickle.dump(write_sentinel, f) + + # Load with safe mode - this should now block + with pytest.raises(pickle.UnpicklingError, match="not in the safe deserialization allowlist"): + load_from_pkl(path, deserialization_mode="safe") + + # Verify the sentinel was NOT created + assert not sentinel.exists(), "Function should NOT have been loaded and executed" + + +def test_json_to_data_safe_blocks_function_rce(tmp_path): + """ + Test that json_to_data with deserialization_mode='safe' blocks + arbitrary function deserialization from :serialized: entries. + """ + sentinel = tmp_path / "sentinel_json" + + # Create a malicious function that doesn't capture any external objects + # to avoid issues with pathlib not being in the allowlist + def write_sentinel(): + with open(str(sentinel), "w") as f: + f.write("PWNED") + + # Serialize it + payload = base64.b64encode(cloudpickle.dumps(write_sentinel)).decode() + + # Create JSON with serialized function + json_str = json.dumps({"lr_schedule": {":type:": "", ":serialized:": payload}}) + + # Load with safe mode - this should now skip/block the function + with pytest.warns(UserWarning, match=r"Could not deserialize object lr_schedule"): + result = json_to_data(json_str, deserialization_mode="safe") + + # The function should have been skipped (not loaded) + assert "lr_schedule" not in result, "Function should have been skipped" + + # Verify the sentinel was NOT created + assert not sentinel.exists(), "Function should NOT have been loaded from JSON and executed" + + +# --------------------------------------------------------------------------- +# Test 9: Verify the vulnerability with a complete PPO checkpoint +# --------------------------------------------------------------------------- + + +def test_ppo_load_safe_blocks_function_rce(tmp_path): + """ + Test that PPO.load with deserialization_mode='safe' blocks + arbitrary function deserialization from injected checkpoint data. + """ + sentinel = tmp_path / "sentinel_ppo" + + # Create a clean model + model = PPO("MlpPolicy", "CartPole-v1", n_steps=64, device="cpu") + zip_path = str(tmp_path / "model.zip") + model.save(zip_path) + + # Create a malicious function + def write_sentinel(): + sentinel.write_text("PWNED") + + # Inject the malicious function into the checkpoint's data JSON + with zipfile.ZipFile(zip_path) as z: + data = json.loads(z.read("data").decode()) + + payload = base64.b64encode(cloudpickle.dumps(write_sentinel)).decode() + data["_malicious_func"] = {":type:": "", ":serialized:": payload} + + with zipfile.ZipFile(zip_path, "r") as zin: + with zipfile.ZipFile(zip_path + ".tmp", "w") as zout: + for name in zin.namelist(): + content = zin.read(name) + if name == "data": + content = json.dumps(data).encode() + zout.writestr(name, content) + os.replace(zip_path + ".tmp", zip_path) + + # Load with safe mode and custom_objects to satisfy other requirements + env = gym.make("CartPole-v1") + custom_objects = { + "observation_space": env.observation_space, + "action_space": env.action_space, + "policy_class": PPO.policy_aliases["MlpPolicy"], + "rollout_buffer_class": model.rollout_buffer_class, + "learning_rate": 0.0, + "lr_schedule": lambda _: 0.0, + "clip_range": lambda _: 0.0, + } + + with pytest.warns(UserWarning, match=r"Could not deserialize object _malicious_func"): + PPO.load( + zip_path, + env=env, + device="cpu", + deserialization_mode="safe", + custom_objects=custom_objects, + ) + + # Verify the sentinel was NOT created + assert not sentinel.exists(), "Function should NOT have been loaded from checkpoint and executed" diff --git a/tests/test_utils.py b/tests/test_utils.py index 8edf67dd0..7f709698c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -41,7 +41,14 @@ @pytest.mark.parametrize("vec_env_cls", [None, SubprocVecEnv]) @pytest.mark.parametrize("wrapper_class", [None, gym.wrappers.RecordEpisodeStatistics]) def test_make_vec_env(env_id, n_envs, vec_env_cls, wrapper_class): - env = make_vec_env(env_id, n_envs, vec_env_cls=vec_env_cls, wrapper_class=wrapper_class, monitor_dir=None, seed=0) + env = make_vec_env( + env_id, + n_envs, + vec_env_cls=vec_env_cls, + wrapper_class=wrapper_class, + monitor_dir=None, + seed=0, + ) assert env.num_envs == n_envs @@ -76,7 +83,13 @@ def test_make_vec_env_func_checker(): @pytest.mark.parametrize("terminal_on_life_loss", [True, False]) @pytest.mark.parametrize("clip_reward", [True]) def test_make_atari_env( - env_id, noop_max, action_repeat_probability, frame_skip, screen_size, terminal_on_life_loss, clip_reward + env_id, + noop_max, + action_repeat_probability, + frame_skip, + screen_size, + terminal_on_life_loss, + clip_reward, ): n_envs = 2 wrapper_kwargs = { @@ -124,7 +137,13 @@ def test_vec_env_kwargs(): def test_vec_env_wrapper_kwargs(): - env = make_vec_env("MountainCarContinuous-v0", n_envs=1, seed=0, wrapper_class=MaxAndSkipEnv, wrapper_kwargs={"skip": 3}) + env = make_vec_env( + "MountainCarContinuous-v0", + n_envs=1, + seed=0, + wrapper_class=MaxAndSkipEnv, + wrapper_kwargs={"skip": 3}, + ) assert env.get_attr("_skip")[0] == 3