saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-07-30 20:18:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Start CNN support (failing for SAC)

Browse source
This commit is contained in:
Antonin RAFFIN 2020-04-21 16:22:46 +02:00
parent 6d2fc328c9
commit 93c2a01f91
9 changed files with 273 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
docs/misc/changelog.rst
View file

@ -3,7 +3,7 @@
Changelog
==========
Pre-Release 0.5.0a1 (WIP)
Pre-Release 0.5.0a2 (WIP)
------------------------------
Breaking Changes:
@ -15,6 +15,7 @@ New Features:
- Added ``optimizer`` and ``optimizer_kwargs`` to ``policy_kwargs`` in order to easily
customizer optimizers
- Complete independent save/load for policies
- Add ``CnnPolicies`` to support images as input (caveat: only support Atari resolution for now)
Bug Fixes:

16
tests/test_cnn.py Normal file
View file

@ -0,0 +1,16 @@
import pytest
from torchy_baselines import A2C, PPO, SAC, TD3
from torchy_baselines.common.identity_env import FakeImageEnv
@pytest.mark.parametrize('model_class', [A2C, PPO, SAC])
def test_cnn(model_class):
# Fake grayscale with frameskip
env = FakeImageEnv(screen_height=84, screen_width=84, n_channels=1,
discrete = model_class not in {SAC, TD3})
if model_class in {A2C, PPO}:
kwargs = dict(n_steps=100)
else:
kwargs = dict(buffer_size=500)
_ = model_class('CnnPolicy', env, **kwargs).learn(500)

1
tests/test_custom_policy.py
View file

@ -7,6 +7,7 @@ from torchy_baselines import A2C, PPO, SAC, TD3
@pytest.mark.parametrize('net_arch', [
[12, dict(vf=[16], pi=[8])],
[4],
[],
[4, 4],
[12, dict(vf=[8, 4], pi=[8])],
[12, dict(vf=[8], pi=[8, 4])],

37
torchy_baselines/common/identity_env.py
View file

@ -1,5 +1,6 @@
import numpy as np
from typing import List
import numpy as np
from gym import Env
from gym.spaces import Discrete, MultiDiscrete, MultiBinary, Box
@ -102,3 +103,37 @@ class IdentityEnvMultiBinary(IdentityEnv):
self.action_space = MultiBinary(dim)
self.observation_space = self.action_space
self.reset()
class FakeImageEnv(Env):
def __init__(self, action_dim: int = 6,
screen_height: int = 210,
screen_width: int = 160,
n_channels: int = 3,
discrete: bool = True):
"""
Fake atari environment for testing purposes.
"""
self.observation_space = Box(low=0, high=255, shape=(screen_height, screen_width, n_channels), dtype=np.uint8)
if discrete:
self.action_space = Discrete(action_dim)
else:
self.action_space = Box(low=-1, high=1, shape=(5,), dtype=np.float32)
self.ep_length = 10
def reset(self):
self.current_step = 0
return self.observation_space.sample()
def step(self, action: int):
reward = 0.0
self.current_step += 1
done = self.current_step >= self.ep_length
return self.observation_space.sample(), reward, done, {}
def render(self, mode='human'):
pass
def get_action_meanings(self) -> List[str]:
return ['NOOP']

46
torchy_baselines/common/policies.py
View file

@ -7,7 +7,7 @@ import torch as th
import torch.nn as nn
import numpy as np
from torchy_baselines.common.preprocessing import preprocess_obs
from torchy_baselines.common.preprocessing import preprocess_obs, get_obs_dim
from torchy_baselines.common.utils import get_device, get_schedule_fn
@ -289,7 +289,8 @@ def create_mlp(input_dim: int,
modules.append(activation_fn())
if output_dim > 0:
modules.append(nn.Linear(net_arch[-1], output_dim))
last_layer_dim = net_arch[-1] if len(net_arch) > 0 else input_dim
modules.append(nn.Linear(last_layer_dim, output_dim))
if squash_output:
modules.append(nn.Tanh())
return modules
@ -453,3 +454,44 @@ class MlpExtractor(nn.Module):
"""
shared_latent = self.shared_net(features)
return self.policy_net(shared_latent), self.value_net(shared_latent)
class BaseFeaturesExtractor(nn.Module):
def __init__(self, observation_space: gym.Space, features_dim: int = 0):
super(BaseFeaturesExtractor, self).__init__()
assert features_dim > 0
self._observation_space = observation_space
self._features_dim = features_dim
@property
def features_dim(self) -> int:
return self._features_dim
def forward(self, observations: th.Tensor) -> th.Tensor:
raise NotImplementedError()
class FlattenExtractor(BaseFeaturesExtractor):
def __init__(self, observation_space: gym.Space):
super(FlattenExtractor, self).__init__(observation_space, get_obs_dim(observation_space))
self.flatten = nn.Flatten()
def forward(self, observations: th.Tensor) -> th.Tensor:
return self.flatten(observations)
class NatureCNN(BaseFeaturesExtractor):
def __init__(self, observation_space: gym.Space,
features_dim: int = 512):
super(NatureCNN, self).__init__(observation_space, features_dim)
# TODO: custom init?
# we assume WxHxC images
# TODO: compute shape before flatten
n_input_channels = observation_space.shape[-1]
self.cnn = nn.Sequential(nn.Conv2d(n_input_channels, 32, 8, stride=4), nn.ReLU(),
nn.Conv2d(32, 64, 4, stride=2), nn.ReLU(),
nn.Conv2d(64, 32, 3, stride=1), nn.ReLU(), nn.Flatten(),
nn.Linear(32 * 7 * 7, features_dim), nn.ReLU())
def forward(self, observations: th.Tensor) -> th.Tensor:
return self.cnn(observations)

8
torchy_baselines/common/preprocessing.py
View file

@ -47,8 +47,12 @@ def preprocess_obs(obs: th.Tensor, observation_space: spaces.Space,
:return: (th.Tensor)
"""
if isinstance(observation_space, spaces.Box):
if is_image_space(observation_space) and normalize_images:
return obs.float() / 255.0
if is_image_space(observation_space):
# Re-order from BxWxHxC to BxCxWxH
obs = obs.permute(0, 3, 1, 2) # .contiguous()?
if normalize_images:
return obs.float() / 255.0
return obs.float()
return obs.float()
elif isinstance(observation_space, spaces.Discrete):
# One hot encoding and convert to float to avoid errors

93
torchy_baselines/ppo/policies.py
View file

@ -6,9 +6,9 @@ import torch as th
import torch.nn as nn
import numpy as np
from torchy_baselines.common.preprocessing import get_obs_dim
from torchy_baselines.common.policies import (BasePolicy, register_policy, MlpExtractor,
create_sde_features_extractor)
create_sde_features_extractor, NatureCNN,
BaseFeaturesExtractor, FlattenExtractor)
from torchy_baselines.common.distributions import (make_proba_distribution, Distribution,
DiagGaussianDistribution, CategoricalDistribution,
StateDependentNoiseDistribution)
@ -37,6 +37,7 @@ class PPOPolicy(BasePolicy):
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
:param squash_output: (bool) Whether to squash the output using a tanh function,
this allows to ensure boundaries when using SDE.
:param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
@ -58,6 +59,7 @@ class PPOPolicy(BasePolicy):
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
squash_output: bool = False,
features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
@ -65,7 +67,10 @@ class PPOPolicy(BasePolicy):
# Default network architecture, from stable-baselines
if net_arch is None:
net_arch = [dict(pi=[64, 64], vf=[64, 64])]
if features_extractor_class == FlattenExtractor:
net_arch = [dict(pi=[64, 64], vf=[64, 64])]
else:
net_arch = []
self.net_arch = net_arch
self.activation_fn = activation_fn
@ -79,9 +84,11 @@ class PPOPolicy(BasePolicy):
self.optimizer_class = optimizer
self.optimizer_kwargs = optimizer_kwargs
self.ortho_init = ortho_init
# In the future, feature_extractor will be replaced with a CNN
self.features_extractor = nn.Flatten()
self.features_dim = get_obs_dim(self.observation_space)
self.features_extractor_class = features_extractor_class
self.features_extractor = features_extractor_class(self.observation_space)
self.features_dim = self.features_extractor.features_dim
self.normalize_images = normalize_images
self.log_std_init = log_std_init
dist_kwargs = None
@ -119,7 +126,8 @@ class PPOPolicy(BasePolicy):
lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone
optimizer=self.optimizer_class,
optimizer_kwargs=self.optimizer_kwargs,
ortho_init=self.ortho_init
ortho_init=self.ortho_init,
features_extractor_class=self.features_extractor_class
))
return data
@ -266,4 +274,75 @@ class PPOPolicy(BasePolicy):
MlpPolicy = PPOPolicy
class CnnPolicy(PPOPolicy):
"""
CnnPolicy class (with both actor and critic) for A2C and derivates (PPO).
:param observation_space: (gym.spaces.Space) Observation space
:param action_space: (gym.spaces.Space) Action space
:param lr_schedule: (Callable) Learning rate schedule (could be constant)
:param net_arch: ([int or dict]) The specification of the policy and value networks.
:param device: (str or th.device) Device on which the code should run.
:param activation_fn: (Type[nn.Module]) Activation function
:param ortho_init: (bool) Whether to use or not orthogonal initialization
:param use_sde: (bool) Whether to use State Dependent Exploration or not
:param log_std_init: (float) Initial value for the log standard deviation
:param full_std: (bool) Whether to use (n_features x n_actions) parameters
for the std instead of only (n_features,) when using SDE
:param sde_net_arch: ([int]) Network architecture for extracting features
when using SDE. If None, the latent features from the policy will be used.
Pass an empty list to use the states as features.
:param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` to ensure
a positive standard deviation (cf paper). It allows to keep variance
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
:param squash_output: (bool) Whether to squash the output using a tanh function,
this allows to ensure boundaries when using SDE.
:param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
``th.optim.Adam`` by default
:param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments,
excluding the learning rate, to pass to the optimizer
"""
def __init__(self,
observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[Union[int, Dict[str, List[int]]]]] = None,
device: Union[th.device, str] = 'auto',
activation_fn: Type[nn.Module] = nn.Tanh,
ortho_init: bool = True,
use_sde: bool = False,
log_std_init: float = 0.0,
full_std: bool = True,
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
squash_output: bool = False,
features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN,
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
super(CnnPolicy, self).__init__(observation_space,
action_space,
lr_schedule,
net_arch,
device,
activation_fn,
ortho_init,
use_sde,
log_std_init,
full_std,
sde_net_arch,
use_expln,
squash_output,
features_extractor_class,
normalize_images,
optimizer,
optimizer_kwargs)
register_policy("MlpPolicy", MlpPolicy)
register_policy("CnnPolicy", CnnPolicy)

91
torchy_baselines/sac/policies.py
View file

@ -4,9 +4,10 @@ import gym
import torch as th
import torch.nn as nn
from torchy_baselines.common.preprocessing import get_action_dim, get_obs_dim
from torchy_baselines.common.preprocessing import get_action_dim
from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp,
create_sde_features_extractor)
create_sde_features_extractor, NatureCNN,
BaseFeaturesExtractor, FlattenExtractor)
from torchy_baselines.common.distributions import SquashedDiagGaussianDistribution, StateDependentNoiseDistribution
# CAP the standard deviation of the actor
@ -75,9 +76,11 @@ class Actor(BasePolicy):
self.use_expln = use_expln
self.full_std = full_std
self.clip_mean = clip_mean
last_layer_dim = net_arch[-1] if len(net_arch) > 0 else features_dim
if self.use_sde:
latent_sde_dim = net_arch[-1]
latent_sde_dim = last_layer_dim
# Separate feature extractor for SDE
if sde_net_arch is not None:
self.sde_features_extractor, latent_sde_dim = create_sde_features_extractor(features_dim, sde_net_arch,
@ -94,8 +97,8 @@ class Actor(BasePolicy):
self.mu = nn.Sequential(self.mu, nn.Hardtanh(min_val=-clip_mean, max_val=clip_mean))
else:
self.action_dist = SquashedDiagGaussianDistribution(action_dim)
self.mu = nn.Linear(net_arch[-1], action_dim)
self.log_std = nn.Linear(net_arch[-1], action_dim)
self.mu = nn.Linear(last_layer_dim, action_dim)
self.log_std = nn.Linear(last_layer_dim, action_dim)
def _get_data(self) -> Dict[str, Any]:
data = super()._get_data()
@ -239,6 +242,7 @@ class SACPolicy(BasePolicy):
a positive standard deviation (cf paper). It allows to keep variance
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
:param clip_mean: (float) Clip the mean output when using SDE to avoid numerical instability.
:param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
@ -257,13 +261,17 @@ class SACPolicy(BasePolicy):
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
clip_mean: float = 2.0,
features_extractor_class: Type[BaseFeaturesExtractor] = FlattenExtractor,
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
super(SACPolicy, self).__init__(observation_space, action_space, device, squash_output=True)
if net_arch is None:
net_arch = [256, 256]
if features_extractor_class == FlattenExtractor:
net_arch = [256, 256]
else:
net_arch = []
if optimizer_kwargs is None:
optimizer_kwargs = {}
@ -271,9 +279,9 @@ class SACPolicy(BasePolicy):
self.optimizer_class = optimizer
self.optimizer_kwargs = optimizer_kwargs
# In the future, features_extractor will be replaced with a CNN
self.features_extractor = nn.Flatten()
self.features_dim = get_obs_dim(self.observation_space)
self.features_extractor_class = features_extractor_class
self.features_extractor = features_extractor_class(self.observation_space)
self.features_dim = self.features_extractor.features_dim
self.net_arch = net_arch
self.activation_fn = activation_fn
@ -325,7 +333,8 @@ class SACPolicy(BasePolicy):
clip_mean=self.actor_kwargs['clip_mean'],
lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone
optimizer=self.optimizer_class,
optimizer_kwargs=self.optimizer_kwargs
optimizer_kwargs=self.optimizer_kwargs,
features_extractor_class=self.features_extractor_class
))
return data
@ -344,4 +353,66 @@ class SACPolicy(BasePolicy):
MlpPolicy = SACPolicy
class CnnPolicy(SACPolicy):
"""
Policy class (with both actor and critic) for SAC.
:param observation_space: (gym.spaces.Space) Observation space
:param action_space: (gym.spaces.Space) Action space
:param lr_schedule: (callable) Learning rate schedule (could be constant)
:param net_arch: (Optional[List[int]]) The specification of the policy and value networks.
:param device: (str or th.device) Device on which the code should run.
:param activation_fn: (Type[nn.Module]) Activation function
:param use_sde: (bool) Whether to use State Dependent Exploration or not
:param log_std_init: (float) Initial value for the log standard deviation
:param sde_net_arch: ([int]) Network architecture for extracting features
when using SDE. If None, the latent features from the policy will be used.
Pass an empty list to use the states as features.
:param use_expln: (bool) Use ``expln()`` function instead of ``exp()`` when using SDE to ensure
a positive standard deviation (cf paper). It allows to keep variance
above zero and prevent it from growing too fast. In practice, ``exp()`` is usually enough.
:param clip_mean: (float) Clip the mean output when using SDE to avoid numerical instability.
:param features_extractor_class: (Type[BaseFeaturesExtractor]) Features extractor to use.
:param normalize_images: (bool) Whether to normalize images or not,
dividing by 255.0 (True by default)
:param optimizer: (Type[th.optim.Optimizer]) The optimizer to use,
``th.optim.Adam`` by default
:param optimizer_kwargs: (Optional[Dict[str, Any]]) Additional keyword arguments,
excluding the learning rate, to pass to the optimizer
"""
def __init__(self, observation_space: gym.spaces.Space,
action_space: gym.spaces.Space,
lr_schedule: Callable,
net_arch: Optional[List[int]] = None,
device: Union[th.device, str] = 'auto',
activation_fn: Type[nn.Module] = nn.ReLU,
use_sde: bool = False,
log_std_init: float = -3,
sde_net_arch: Optional[List[int]] = None,
use_expln: bool = False,
clip_mean: float = 2.0,
features_extractor_class: Type[BaseFeaturesExtractor] = NatureCNN,
normalize_images: bool = True,
optimizer: Type[th.optim.Optimizer] = th.optim.Adam,
optimizer_kwargs: Optional[Dict[str, Any]] = None):
super(CnnPolicy, self).__init__(observation_space,
action_space,
lr_schedule,
net_arch,
device,
activation_fn,
use_sde,
log_std_init,
sde_net_arch,
use_expln,
clip_mean,
features_extractor_class,
normalize_images,
optimizer,
optimizer_kwargs)
register_policy("MlpPolicy", MlpPolicy)
register_policy("CnnPolicy", CnnPolicy)

2
torchy_baselines/version.txt
View file

@ -1 +1 @@
0.5.0a1
0.5.0a2