Noisy Cross Entropy Method (CEM)

.github/workflows/ci.yml

@@ -28,7 +28,7 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         # cpu version of pytorch
-        pip install torch==1.8.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
+        pip install torch==1.11.0+cpu -f https://download.pytorch.org/whl/torch_stable.html
         # Install dependencies for docs and tests
         pip install stable_baselines3[extra,tests,docs]
         # Install master version

docs/misc/changelog.rst

@@ -3,18 +3,19 @@
 Changelog
 ==========
 
-Release 1.4.1a1 (WIP)
+Release 1.4.1a3 (WIP)
 -------------------------------
 
-
 Breaking Changes:
 ^^^^^^^^^^^^^^^^^
 - Switched minimum Gym version to 0.21.0.
 - Upgraded to Stable-Baselines3 >= 1.4.1a1
+- Changed default policy architecture for ARS/CEM to ``[32]`` instead of ``[64, 64]``
 
 New Features:
 ^^^^^^^^^^^^^
 - Allow PPO to turn of advantage normalization (see `PR #61 <https://github.com/Stable-Baselines-Team/stable-baselines3-contrib/pull/61>`_) @vwxyzjn
+- Added noisy Cross Entropy Method (CEM)
 
 Bug Fixes:
 ^^^^^^^^^^

sb3_contrib/__init__.py

@@ -1,6 +1,7 @@
 import os
 
 from sb3_contrib.ars import ARS
+from sb3_contrib.cem import CEM
 from sb3_contrib.ppo_mask import MaskablePPO
 from sb3_contrib.qrdqn import QRDQN
 from sb3_contrib.tqc import TQC

sb3_contrib/ars/policies.py

@@ -36,7 +36,7 @@ def __init__(
         )
 
         if net_arch is None:
-            net_arch = [64, 64]
+            net_arch = [32]
 
         self.net_arch = net_arch
         self.features_extractor = self.make_features_extractor()

sb3_contrib/cem/__init__.py

@@ -0,0 +1,2 @@
+from sb3_contrib.cem.cem import CEM
+from sb3_contrib.cem.policies import LinearPolicy, MlpPolicy

sb3_contrib/cem/policies.py

@@ -0,0 +1,10 @@
+from stable_baselines3.common.policies import register_policy
+
+from sb3_contrib.ars.policies import ARSLinearPolicy, ARSPolicy
+
+MlpPolicy = ARSPolicy
+LinearPolicy = ARSLinearPolicy
+
+
+register_policy("LinearPolicy", LinearPolicy)
+register_policy("MlpPolicy", MlpPolicy)

sb3_contrib/version.txt

@@ -1 +1 @@
-1.4.1a1
+1.4.1a3

setup.cfg

@@ -23,6 +23,7 @@ ignore = W503,W504,E203,E231  # line breaks before and after binary operators
 per-file-ignores =
 	./sb3_contrib/__init__.py:F401
 	./sb3_contrib/ars/__init__.py:F401
+	./sb3_contrib/cem/__init__.py:F401
 	./sb3_contrib/ppo_mask/__init__.py:F401
 	./sb3_contrib/qrdqn/__init__.py:F401
 	./sb3_contrib/tqc/__init__.py:F401

setup.py

@@ -64,6 +64,7 @@
     package_data={"sb3_contrib": ["py.typed", "version.txt"]},
     install_requires=[
         "stable_baselines3>=1.4.1a1",
+        "torch>=1.11",
     ],
     description="Contrib package of Stable Baselines3, experimental code.",
     author="Antonin Raffin",

tests/test_run.py

@@ -3,7 +3,7 @@
 from stable_baselines3.common.env_util import make_vec_env
 from stable_baselines3.common.vec_env import VecNormalize
 
-from sb3_contrib import ARS, QRDQN, TQC, TRPO, MaskablePPO
+from sb3_contrib import ARS, CEM, QRDQN, TQC, TRPO, MaskablePPO
 from sb3_contrib.common.envs import InvalidActionEnvDiscrete
 from sb3_contrib.common.vec_env import AsyncEval
 
@@ -92,13 +92,25 @@ def test_ars(policy_str, env_id):
     model.learn(total_timesteps=500, log_interval=1, eval_freq=250)
 
 
-def test_ars_multi_env():
+@pytest.mark.parametrize("env_id", ["CartPole-v1", "Pendulum-v1"])
+@pytest.mark.parametrize("policy_str", ["LinearPolicy", "MlpPolicy"])
+def test_cem(policy_str, env_id):
+    model = CEM(policy_str, env_id, pop_size=2, verbose=1, seed=0)
+    model.learn(total_timesteps=500, log_interval=1, eval_freq=250)
+
+
+@pytest.mark.parametrize("model_class", [ARS, CEM])
+def test_es_multi_env(model_class):
     env = make_vec_env("Pendulum-v1", n_envs=2)
-    model = ARS("MlpPolicy", env, n_delta=1)
+    kwargs = dict(n_delta=1) if model_class == ARS else dict(pop_size=2)
+
+    model = model_class("MlpPolicy", env, **kwargs)
     model.learn(total_timesteps=250)
 
+    kwargs = dict(n_delta=2) if model_class == ARS else dict(pop_size=3)
+
     env = VecNormalize(make_vec_env("Pendulum-v1", n_envs=1))
-    model = ARS("MlpPolicy", env, n_delta=2, seed=0)
+    model = model_class("MlpPolicy", env, seed=0, **kwargs)
     # with parallelism
     async_eval = AsyncEval([lambda: VecNormalize(make_vec_env("Pendulum-v1", n_envs=1)) for _ in range(2)], model.policy)
     async_eval.seed(0)