Merge branch 'master' into feat/mps-support

docs/misc/changelog.rst

@@ -4,7 +4,7 @@ Changelog
 ==========
 
 
-Release 1.7.0a0 (WIP)
+Release 1.7.0a1 (WIP)
 --------------------------
 
 Breaking Changes:
@@ -22,6 +22,9 @@ SB3-Contrib
 
 Bug Fixes:
 ^^^^^^^^^^
+- Fix return type of ``evaluate_actions`` in ``ActorCritcPolicy`` to reflect that entropy is an optional tensor (@Rocamonde)
+- Fix type annotation of ``policy`` in ``BaseAlgorithm`` and ``OffPolicyAlgorithm``
+- Allowed model trained with Python 3.7 to be loaded with Python 3.8+ without the ``custom_objects`` workaround
 
 Deprecations:
 ^^^^^^^^^^^^^

stable_baselines3/common/base_class.py

@@ -60,7 +60,7 @@ class BaseAlgorithm(ABC):
     """
     The base of RL algorithms
 
-    :param policy: Policy object
+    :param policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
     :param env: The environment to learn from
                 (if registered in Gym, can be str. Can be None for loading trained models)
     :param learning_rate: learning rate for the optimizer,
@@ -89,7 +89,7 @@ class BaseAlgorithm(ABC):
 
     def __init__(
         self,
-        policy: Type[BasePolicy],
+        policy: Union[str, Type[BasePolicy]],
         env: Union[GymEnv, str, None],
         learning_rate: Union[float, Schedule],
         policy_kwargs: Optional[Dict[str, Any]] = None,

stable_baselines3/common/off_policy_algorithm.py

@@ -28,7 +28,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
     """
     The base for Off-Policy algorithms (ex: SAC/TD3)
 
-    :param policy: Policy object
+    :param policy: The policy model to use (MlpPolicy, CnnPolicy, ...)
     :param env: The environment to learn from
                 (if registered in Gym, can be str. Can be None for loading trained models)
     :param learning_rate: learning rate for the optimizer,
@@ -75,7 +75,7 @@ class OffPolicyAlgorithm(BaseAlgorithm):
 
     def __init__(
         self,
-        policy: Type[BasePolicy],
+        policy: Union[str, Type[BasePolicy]],
         env: Union[GymEnv, str],
         learning_rate: Union[float, Schedule],
         buffer_size: int = 1_000_000,  # 1e6

stable_baselines3/common/policies.py

@@ -2,7 +2,6 @@
 
 import collections
 import copy
-import warnings
 from abc import ABC, abstractmethod
 from functools import partial
 from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union
@@ -613,7 +612,7 @@ class ActorCriticPolicy(BasePolicy):
         """
         return self.get_distribution(observation).get_actions(deterministic=deterministic)
 
-    def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
+    def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor) -> Tuple[th.Tensor, th.Tensor, Optional[th.Tensor]]:
         """
         Evaluate actions according to the current policy,
         given the observations.
@@ -629,7 +628,8 @@ class ActorCriticPolicy(BasePolicy):
         distribution = self._get_action_dist_from_latent(latent_pi)
         log_prob = distribution.log_prob(actions)
         values = self.value_net(latent_vf)
-        return values, log_prob, distribution.entropy()
+        entropy = distribution.entropy()
+        return values, log_prob, entropy
 
     def get_distribution(self, obs: th.Tensor) -> Distribution:
         """

stable_baselines3/common/save_util.py

@@ -162,13 +162,15 @@ def json_to_data(json_string: str, custom_objects: Optional[Dict[str, Any]] = No
             try:
                 base64_object = base64.b64decode(serialization.encode())
                 deserialized_object = cloudpickle.loads(base64_object)
-            except (RuntimeError, TypeError):
+            except (RuntimeError, TypeError, AttributeError) as e:
                 warnings.warn(
                     f"Could not deserialize object {data_key}. "
-                    + "Consider using `custom_objects` argument to replace "
-                    + "this object."
+                    "Consider using `custom_objects` argument to replace "
+                    "this object.\n"
+                    f"Exception: {e}"
                 )
-            return_data[data_key] = deserialized_object
+            else:
+                return_data[data_key] = deserialized_object
         else:
             # Read as it is
             return_data[data_key] = data_item

stable_baselines3/common/vec_env/vec_normalize.py

@@ -1,5 +1,4 @@
 import pickle
-import warnings
 from copy import deepcopy
 from typing import Any, Dict, List, Optional, Union
 

stable_baselines3/sac/policies.py

@@ -1,4 +1,3 @@
-import warnings
 from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import gym

stable_baselines3/version.txt

@@ -1 +1 @@
-1.7.0a0
+1.7.0a1

tests/test_distributions.py

@@ -77,6 +77,8 @@ def test_get_distribution(dummy_model_distribution_obs_and_actions):
         distribution = model.policy.get_distribution(observations)
         log_prob_2 = distribution.log_prob(actions)
         entropy_2 = distribution.entropy()
+        assert entropy_1 is not None
+        assert entropy_2 is not None
         assert th.allclose(log_prob_1, log_prob_2)
         assert th.allclose(entropy_1, entropy_2)
 

tests/test_save_load.py

@@ -1,7 +1,10 @@
+import base64
 import io
+import json
 import os
 import pathlib
 import warnings
+import zipfile
 from collections import OrderedDict
 from copy import deepcopy
 
@@ -690,3 +693,33 @@ def test_save_load_large_model(tmp_path):
 
     # clear file from os
     os.remove(tmp_path / "test_save.zip")
+
+
+def test_load_invalid_object(tmp_path):
+    # See GH Issue #1122 for an example
+    # of invalid object loading
+    path = str(tmp_path / "ppo_pendulum.zip")
+    PPO("MlpPolicy", "Pendulum-v1", learning_rate=lambda _: 1.0).save(path)
+
+    with zipfile.ZipFile(path, mode="r") as archive:
+        json_data = json.loads(archive.read("data").decode())
+
+    # Intentionally corrupt the data
+    serialization = json_data["learning_rate"][":serialized:"]
+    base64_object = base64.b64decode(serialization.encode())
+    new_bytes = base64_object.replace(b"CodeType", b"CodeTyps")
+    base64_encoded = base64.b64encode(new_bytes).decode()
+    json_data["learning_rate"][":serialized:"] = base64_encoded
+    serialized_data = json.dumps(json_data, indent=4)
+
+    with open(tmp_path / "data", "w") as f:
+        f.write(serialized_data)
+    # 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

tests/test_vec_normalize.py

@@ -1,5 +1,4 @@
 import operator
-import warnings
 
 import gym
 import numpy as np