saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-26 19:52:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

ADD CMAES

Browse source
This commit is contained in:
Antonin RAFFIN 2020-07-30 11:17:32 +02:00
parent e0f70531aa
commit 03e9a96d2c
7 changed files with 383 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
setup.cfg
View file

@ -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

2
setup.py
View file

@ -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.",

5
stable_baselines3/__init__.py
View file

@ -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:

2
stable_baselines3/cmaes/__init__.py Normal file
View file

@ -0,0 +1,2 @@
from stable_baselines3.cmaes.cmaes import CMAES
from stable_baselines3.cmaes.policies import CnnPolicy, MlpPolicy

178
stable_baselines3/cmaes/cmaes.py Normal file
View file

@ -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)

185
stable_baselines3/cmaes/policies.py Normal file
View file

@ -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)

11
tests/test_run.py
View file

@ -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(