diff --git a/setup.cfg b/setup.cfg index 037c68f..2420141 100644 --- a/setup.cfg +++ b/setup.cfg @@ -27,6 +27,7 @@ per-file-ignores = ./stable_baselines3/__init__.py:F401 ./stable_baselines3/common/__init__.py:F401 ./stable_baselines3/a2c/__init__.py:F401 + ./stable_baselines3/cmaes/__init__.py:F401 ./stable_baselines3/ddpg/__init__.py:F401 ./stable_baselines3/dqn/__init__.py:F401 ./stable_baselines3/ppo/__init__.py:F401 diff --git a/setup.py b/setup.py index 8401c7c..62ba10e 100644 --- a/setup.py +++ b/setup.py @@ -115,6 +115,8 @@ setup( "tensorboard", # Checking memory taken by replay buffer "psutil", + # Enable CMA + "cma", ], }, description="Pytorch version of Stable Baselines, implementations of reinforcement learning algorithms.", diff --git a/stable_baselines3/__init__.py b/stable_baselines3/__init__.py index 48c88c1..c1c87d0 100644 --- a/stable_baselines3/__init__.py +++ b/stable_baselines3/__init__.py @@ -8,6 +8,11 @@ from stable_baselines3.sac import SAC from stable_baselines3.td3 import TD3 from stable_baselines3.tqc import TQC +try: + from stable_baselines3.cmaes import CMAES +except ImportError: + CMAES = None + # Read version from file version_file = os.path.join(os.path.dirname(__file__), "version.txt") with open(version_file, "r") as file_handler: diff --git a/stable_baselines3/cmaes/__init__.py b/stable_baselines3/cmaes/__init__.py new file mode 100644 index 0000000..f41804a --- /dev/null +++ b/stable_baselines3/cmaes/__init__.py @@ -0,0 +1,2 @@ +from stable_baselines3.cmaes.cmaes import CMAES +from stable_baselines3.cmaes.policies import CnnPolicy, MlpPolicy diff --git a/stable_baselines3/cmaes/cmaes.py b/stable_baselines3/cmaes/cmaes.py new file mode 100644 index 0000000..e6926e0 --- /dev/null +++ b/stable_baselines3/cmaes/cmaes.py @@ -0,0 +1,178 @@ +import time +from typing import Any, Dict, Optional, Type, Union + +import cma +import numpy as np +import torch as th + +from stable_baselines3.cmaes.policies import CMAESPolicy +from stable_baselines3.common import logger +from stable_baselines3.common.base_class import BaseAlgorithm +from stable_baselines3.common.policies import BasePolicy +from stable_baselines3.common.type_aliases import GymEnv, MaybeCallback +from stable_baselines3.common.utils import safe_mean + + +class CMAES(BaseAlgorithm): + def __init__( + self, + policy: Type[BasePolicy], + env: Union[GymEnv, str], + n_steps: int = 200, + n_individuals: int = -1, + std_init: float = 0.5, + best_individual: Optional[np.ndarray] = None, + policy_kwargs: Dict[str, Any] = None, + tensorboard_log: Optional[str] = None, + verbose: int = 0, + device: Union[th.device, str] = "auto", + create_eval_env: bool = False, + monitor_wrapper: bool = True, + seed: Optional[int] = None, + _init_setup_model: bool = True, + ): + + super(CMAES, self).__init__( + policy=policy, + env=env, + policy_base=CMAESPolicy, + learning_rate=0.0, + policy_kwargs=policy_kwargs, + tensorboard_log=tensorboard_log, + verbose=verbose, + device=device, + support_multi_env=True, + create_eval_env=create_eval_env, + monitor_wrapper=monitor_wrapper, + seed=seed, + use_sde=False, + ) + + self.policy_kwargs["device"] = self.device + self.best_individual = best_individual + self.best_ever = None + self.std_init = std_init + self.n_steps = n_steps + self.es = None + if _init_setup_model: + self._setup_model() + + def _setup_model(self) -> None: + self.set_random_seed(self.seed) + + self.policy = self.policy_class( + self.observation_space, self.action_space, **self.policy_kwargs # pytype:disable=not-instantiable + ) + self.policy = self.policy.to(self.device) + self.actor = self.policy.actor + + def learn( + self, + total_timesteps: int, + callback: MaybeCallback = None, + log_interval: int = 10, + eval_env: Optional[GymEnv] = None, + eval_freq: int = -1, + n_eval_episodes: int = 5, + tb_log_name: str = "run", + eval_log_path: Optional[str] = None, + reset_num_timesteps: bool = True, + ) -> "CMAES": + + total_timesteps, callback = self._setup_learn( + total_timesteps, eval_env, callback, eval_freq, n_eval_episodes, eval_log_path, reset_num_timesteps, tb_log_name + ) + + callback.on_training_start(locals(), globals()) + + if self.best_individual is None: + self.best_individual = self.actor.parameters_to_vector() + + if self.best_ever is None: + self.best_ever = cma.optimization_tools.BestSolution() + + if self.es is None: + options = {"seed": self.seed} + if self.env.num_envs > 1: + options['popsize'] = self.env.num_envs + self.es = cma.CMAEvolutionStrategy(self.best_individual, self.std_init, options) + continue_training = True + + while self.num_timesteps < total_timesteps and not self.es.stop() and continue_training: + candidates = self.es.ask() + returns = np.zeros((len(candidates),)) + candidate_idx = 0 + candidate_steps = 0 + + self.actor.load_from_vector(candidates[candidate_idx]) + callback.on_rollout_start() + + while candidate_idx < len(candidates): + # TODO support num_envs > 0 + action, _ = self.actor.predict(self._last_obs, deterministic=True) + + # Rescale and perform action + new_obs, reward, done, infos = self.env.step(action) + + # Only stop training if return value is False, not when it is None. + if callback.on_step() is False: + continue_training = False + break + + returns[candidate_idx] += reward + + # Retrieve reward and episode length if using Monitor wrapper + self._update_info_buffer(infos, done) + + self._last_obs = new_obs + + self.num_timesteps += 1 + candidate_steps += 1 + self._update_current_progress_remaining(self.num_timesteps, self._total_timesteps) + + if candidate_steps > self.n_steps: + candidate_idx += 1 + candidate_steps = 0 + if candidate_idx < len(candidates): + self.actor.load_from_vector(candidates[candidate_idx]) + else: + break + + if done: + self._episode_num += 1 + + # Log training infos + if log_interval is not None and self._episode_num % log_interval == 0: + self._dump_logs() + + callback.on_rollout_end() + + # TODO: inject best solution from time to time? + self.es.tell(candidates, -1 * returns) + # TODO: load best individual when using predict + self.best_ever.update(self.es.best) + self.best_individual = self.best_ever.x + self.policy.best_actor.load_from_vector(self.best_individual) + + callback.on_training_end() + + return self + + def _dump_logs(self) -> None: + """ + Write log. + """ + fps = int(self.num_timesteps / (time.time() - self.start_time)) + logger.record("time/episodes", self._episode_num, exclude="tensorboard") + if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0: + logger.record("rollout/ep_rew_mean", safe_mean([ep_info["r"] for ep_info in self.ep_info_buffer])) + logger.record("rollout/ep_len_mean", safe_mean([ep_info["l"] for ep_info in self.ep_info_buffer])) + logger.record("time/fps", fps) + logger.record("time/time_elapsed", int(time.time() - self.start_time), exclude="tensorboard") + logger.record("time/total timesteps", self.num_timesteps, exclude="tensorboard") + logger.record("rollout/best_ever", -self.best_ever.f) + + if len(self.ep_success_buffer) > 0: + logger.record("rollout/success rate", safe_mean(self.ep_success_buffer)) + # Pass the number of timesteps for tensorboard + logger.dump(step=self.num_timesteps) diff --git a/stable_baselines3/cmaes/policies.py b/stable_baselines3/cmaes/policies.py new file mode 100644 index 0000000..3a5b988 --- /dev/null +++ b/stable_baselines3/cmaes/policies.py @@ -0,0 +1,185 @@ +from typing import Any, Dict, List, Optional, Type, Union + +import gym +import torch as th +from torch import nn as nn + +from stable_baselines3.common.policies import BasePolicy, register_policy +from stable_baselines3.common.preprocessing import get_action_dim +from stable_baselines3.common.torch_layers import BaseFeaturesExtractor, FlattenExtractor, NatureCNN, create_mlp + + +class Actor(BasePolicy): + """ + Actor network (policy) for TD3. + + :param observation_space: (gym.spaces.Space) Obervation space + :param action_space: (gym.spaces.Space) Action space + :param net_arch: ([int]) Network architecture + :param features_extractor: (nn.Module) Network to extract features + (a CNN when using images, a nn.Flatten() layer otherwise) + :param features_dim: (int) Number of features + :param activation_fn: (Type[nn.Module]) Activation function + :param normalize_images: (bool) Whether to normalize images or not, + dividing by 255.0 (True by default) + :param device: (Union[th.device, str]) Device on which the code should run. + """ + + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + net_arch: List[int], + features_extractor: nn.Module, + features_dim: int, + activation_fn: Type[nn.Module] = nn.ReLU, + normalize_images: bool = True, + device: Union[th.device, str] = "auto", + ): + super(Actor, self).__init__( + observation_space, + action_space, + features_extractor=features_extractor, + normalize_images=normalize_images, + device=device, + squash_output=True, + ) + + self.features_extractor = features_extractor + self.normalize_images = normalize_images + self.net_arch = net_arch + self.features_dim = features_dim + self.activation_fn = activation_fn + + action_dim = get_action_dim(self.action_space) + actor_net = create_mlp(features_dim, action_dim, net_arch, activation_fn, squash_output=True) + # Deterministic action + self.mu = nn.Sequential(*actor_net) + + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update( + dict( + net_arch=self.net_arch, + features_dim=self.features_dim, + activation_fn=self.activation_fn, + features_extractor=self.features_extractor, + ) + ) + return data + + def forward(self, obs: th.Tensor, deterministic: bool = True) -> th.Tensor: + # assert deterministic, 'The TD3 actor only outputs deterministic actions' + features = self.extract_features(obs) + return self.mu(features) + + def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: + return self.forward(observation, deterministic=deterministic) + + +class CMAESPolicy(BasePolicy): + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + net_arch: Optional[List[int]] = None, + device: Union[th.device, str] = "auto", + activation_fn: Type[nn.Module] = nn.ReLU, + features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor, + features_extractor_kwargs: Optional[Dict[str, Any]] = None, + normalize_images: bool = True, + optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None, + ): + super(CMAESPolicy, self).__init__( + observation_space, + action_space, + device, + features_extractor_class, + features_extractor_kwargs, + optimizer_class=optimizer_class, + optimizer_kwargs=optimizer_kwargs, + squash_output=True, + ) + + # Default network architecture, from the original paper + if net_arch is None: + if features_extractor_class == FlattenExtractor: + net_arch = [64, 64] + else: + net_arch = [] + + self.features_extractor = features_extractor_class(self.observation_space, **self.features_extractor_kwargs) + self.features_dim = self.features_extractor.features_dim + + self.net_arch = net_arch + self.activation_fn = activation_fn + self.net_args = { + "observation_space": self.observation_space, + "action_space": self.action_space, + "features_extractor": self.features_extractor, + "features_dim": self.features_dim, + "net_arch": self.net_arch, + "activation_fn": self.activation_fn, + "normalize_images": normalize_images, + "device": device, + } + self.actor = Actor(**self.net_args).to(self.device) + self.best_actor = Actor(**self.net_args).to(self.device) + + def _get_data(self) -> Dict[str, Any]: + data = super()._get_data() + + data.update( + dict( + net_arch=self.net_args["net_arch"], + activation_fn=self.net_args["activation_fn"], + optimizer_class=self.optimizer_class, + optimizer_kwargs=self.optimizer_kwargs, + features_extractor_class=self.features_extractor_class, + features_extractor_kwargs=self.features_extractor_kwargs, + ) + ) + return data + + def forward(self, observation: th.Tensor, deterministic: bool = False): + return self._predict(observation, deterministic=deterministic) + + def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor: + return self.best_actor(observation, deterministic=deterministic) + + +MlpPolicy = CMAESPolicy + + +class CnnPolicy(CMAESPolicy): + def __init__( + self, + observation_space: gym.spaces.Space, + action_space: gym.spaces.Space, + net_arch: Optional[List[int]] = None, + device: Union[th.device, str] = "auto", + activation_fn: Type[nn.Module] = nn.ReLU, + features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN, + features_extractor_kwargs: Optional[Dict[str, Any]] = None, + normalize_images: bool = True, + optimizer_class: Type[th.optim.Optimizer] = th.optim.Adam, + optimizer_kwargs: Optional[Dict[str, Any]] = None, + ): + super(CnnPolicy, self).__init__( + observation_space, + action_space, + net_arch, + device, + activation_fn, + features_extractor_class, + features_extractor_kwargs, + normalize_images, + optimizer_class, + optimizer_kwargs, + ) + + +register_policy("MlpPolicy", MlpPolicy) +register_policy("CnnPolicy", CnnPolicy) diff --git a/tests/test_run.py b/tests/test_run.py index 0e1a4f6..7221117 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -1,7 +1,7 @@ import numpy as np import pytest -from stable_baselines3 import A2C, DDPG, DQN, PPO, SAC, TD3, TQC +from stable_baselines3 import A2C, CMAES, DDPG, DQN, PPO, SAC, TD3, TQC from stable_baselines3.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1)) @@ -82,6 +82,15 @@ def test_n_critics(n_critics): ) model.learn(total_timesteps=1000) +# "CartPole-v1" +@pytest.mark.parametrize("env_id", ["MountainCarContinuous-v0"]) +def test_cmaes(env_id): + if CMAES is None: + return + model = CMAES("MlpPolicy", env_id, seed=0, policy_kwargs=dict(net_arch=[64]), verbose=1, create_eval_env=True) + model.learn(total_timesteps=50000, eval_freq=10000) + + # def test_crr(tmp_path): # model = TQC(