saymrwulf/stable-baselines3
1
0
Fork 0
mirror of https://github.com/saymrwulf/stable-baselines3.git synced 2026-09-15 22:10:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Reformat and code cleanup

Browse source
This commit is contained in:
Antonin RAFFIN 2020-04-23 15:18:21 +02:00
parent 71df3c7409
commit 7ae54206ce
15 changed files with 63 additions and 59 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
tests/test_cnn.py
View file

@ -6,7 +6,6 @@ import pytest
from torchy_baselines import A2C, PPO, SAC, TD3
from torchy_baselines.common.identity_env import FakeImageEnv
SAVE_PATH = './cnn_model.zip'
@ -16,7 +15,7 @@ def test_cnn(model_class):
# Atari after preprocessing: 84x84x1, here we are using lower resolution
# to check that the network handle it automatically
env = FakeImageEnv(screen_height=40, screen_width=40, n_channels=1,
discrete = model_class not in {SAC, TD3})
discrete=model_class not in {SAC, TD3})
if model_class in {A2C, PPO}:
kwargs = dict(n_steps=100)
else:

1
tests/test_distributions.py
View file

@ -89,7 +89,6 @@ def test_categorical():
# mean negative log likelihood == entropy
dist = CategoricalDistribution(N_ACTIONS)
set_random_seed(1)
state = th.rand(N_SAMPLES, N_FEATURES)
action_logits = th.rand(N_SAMPLES, N_ACTIONS)
dist = dist.proba_distribution(action_logits)

4
tests/test_monitor.py
View file

@ -10,7 +10,7 @@ from torchy_baselines.common.monitor import Monitor, get_monitor_files, load_res
def test_monitor(tmp_path):
"""
test the monitor wrapper
Test the monitor wrapper
"""
env = gym.make("CartPole-v1")
env.seed(0)
@ -22,7 +22,7 @@ def test_monitor(tmp_path):
ep_lengths = []
ep_len, ep_reward = 0, 0
for _ in range(total_steps):
_, reward, done, _ = monitor_env.step(0)
_, reward, done, _ = monitor_env.step(monitor_env.action_space.sample())
ep_len += 1
ep_reward += reward
if done:

4
tests/test_run.py
View file

@ -4,10 +4,10 @@ import pytest
from torchy_baselines import A2C, PPO, SAC, TD3
from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise
action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1))
normal_action_noise = NormalActionNoise(np.zeros(1), 0.1 * np.ones(1))
@pytest.mark.parametrize('action_noise', [action_noise, OrnsteinUhlenbeckActionNoise(np.zeros(1), 0.1 * np.ones(1))])
@pytest.mark.parametrize('action_noise', [normal_action_noise, OrnsteinUhlenbeckActionNoise(np.zeros(1), 0.1 * np.ones(1))])
def test_td3(action_noise):
model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]),
learning_starts=100, verbose=1, create_eval_env=True, action_noise=action_noise)

6
torchy_baselines/common/base_class.py
View file

@ -241,8 +241,8 @@ class BaseRLModel(ABC):
if (observation_space != env.observation_space
# Special cases for images that need to be transposed
and not (is_image_space(env.observation_space)
and observation_space == VecTransposeImage.transpose_space(env.observation_space))
):
and observation_space == VecTransposeImage.transpose_space(env.observation_space)
)):
return False
if action_space != env.action_space:
return False
@ -884,7 +884,7 @@ class OffPolicyRLModel(BaseRLModel):
action_noise.reset()
# Display training infos
if self.verbose >= 1 and log_interval is not None and (self._episode_num) % log_interval == 0:
if self.verbose >= 1 and log_interval is not None and self._episode_num % log_interval == 0:
fps = int(self.num_timesteps / (time.time() - self.start_time))
logger.logkv("episodes", self._episode_num)
if len(self.ep_info_buffer) > 0 and len(self.ep_info_buffer[0]) > 0:

1
torchy_baselines/common/identity_env.py
View file

@ -133,6 +133,7 @@ class FakeImageEnv(Env):
else:
self.action_space = Box(low=-1, high=1, shape=(5,), dtype=np.float32)
self.ep_length = 10
self.current_step = 0
def reset(self) -> np.ndarray:
self.current_step = 0

11
torchy_baselines/common/policies.py
View file

@ -8,7 +8,7 @@ import torch.nn as nn
import numpy as np
from torchy_baselines.common.preprocessing import preprocess_obs, get_flattened_obs_dim, is_image_space
from torchy_baselines.common.utils import get_device, get_schedule_fn
from torchy_baselines.common.utils import get_device
from torchy_baselines.common.vec_env import VecTransposeImage
@ -17,8 +17,9 @@ class BaseFeaturesExtractor(nn.Module):
Base class that represents a features extractor.
:param observation_space: (gym.Space)
:param feature_dim: (int) Number of features extracted.
:param features_dim: (int) Number of features extracted.
"""
def __init__(self, observation_space: gym.Space, features_dim: int = 0):
super(BaseFeaturesExtractor, self).__init__()
assert features_dim > 0
@ -40,6 +41,7 @@ class FlattenExtractor(BaseFeaturesExtractor):
:param observation_space: (gym.Space)
"""
def __init__(self, observation_space: gym.Space):
super(FlattenExtractor, self).__init__(observation_space, get_flattened_obs_dim(observation_space))
self.flatten = nn.Flatten()
@ -53,9 +55,10 @@ class NatureCNN(BaseFeaturesExtractor):
CNN from DQN nature paper: https://arxiv.org/abs/1312.5602
:param observation_space: (gym.Space)
:param feature_dim: (int) Number of features extracted.
:param features_dim: (int) Number of features extracted.
This corresponds to the number of unit for the last layer.
"""
def __init__(self, observation_space: gym.spaces.Box,
features_dim: int = 512):
super(NatureCNN, self).__init__(observation_space, features_dim)
@ -104,6 +107,7 @@ class BasePolicy(nn.Module):
: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,
device: Union[th.device, str] = 'auto',
@ -500,6 +504,7 @@ class MlpExtractor(nn.Module):
:param activation_fn: (Type[nn.Module]) The activation function to use for the networks.
:param device: (th.device)
"""
def __init__(self, feature_dim: int,
net_arch: List[Union[int, Dict[str, List[int]]]],
activation_fn: Type[nn.Module],

2
torchy_baselines/common/preprocessing.py
View file

@ -1,4 +1,4 @@
from typing import Tuple, Union
from typing import Tuple
import numpy as np
import torch as th

2
torchy_baselines/common/type_aliases.py
View file

@ -1,7 +1,7 @@
"""
Common aliases for type hint
"""
from typing import Union, Dict, Any, NamedTuple, Optional, List, Callable, Tuple
from typing import Union, Dict, Any, NamedTuple, List, Callable, Tuple
import numpy as np
import torch as th

4
torchy_baselines/common/vec_env/vec_transpose.py
View file

@ -1,5 +1,3 @@
import warnings
import typing
import numpy as np
from gym import spaces
@ -11,7 +9,6 @@ if typing.TYPE_CHECKING:
from torchy_baselines.common.type_aliases import GymStepReturn
class VecTransposeImage(VecEnvWrapper):
"""
Re-order channels, from WxHxC to CxWxH.
@ -19,6 +16,7 @@ class VecTransposeImage(VecEnvWrapper):
:param venv: (VecEnv)
"""
def __init__(self, venv: VecEnv):
assert is_image_space(venv.observation_space), 'The observation space must be an image'

1
torchy_baselines/ppo/policies.py
View file

@ -286,7 +286,6 @@ class PPOPolicy(BasePolicy):
MlpPolicy = PPOPolicy
class CnnPolicy(PPOPolicy):
"""
CnnPolicy class (with both actor and critic) for A2C and derivates (PPO).

13
torchy_baselines/sac/policies.py
View file

@ -41,6 +41,7 @@ class Actor(BasePolicy):
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],
@ -79,7 +80,6 @@ class Actor(BasePolicy):
self.latent_pi = nn.Sequential(*latent_pi_net)
last_layer_dim = net_arch[-1] if len(net_arch) > 0 else features_dim
if self.use_sde:
latent_sde_dim = last_layer_dim
# Separate feature extractor for SDE
@ -195,6 +195,7 @@ class Critic(BasePolicy):
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],
@ -256,6 +257,7 @@ class SACPolicy(BasePolicy):
: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,
@ -328,7 +330,8 @@ class SACPolicy(BasePolicy):
# Do not optimize the shared feature extractor with the critic loss
# otherwise, there are gradient computation issues
# Another solution: having duplicated features extractor but requires more memory and computation
critic_parameters = [param for name, param in self.critic.named_parameters() if 'features_extractor' not in name]
critic_parameters = [param for name, param in self.critic.named_parameters() if
'features_extractor' not in name]
self.critic.optimizer = self.optimizer_class(critic_parameters, lr=lr_schedule(1),
**self.optimizer_kwargs)
@ -357,8 +360,8 @@ class SACPolicy(BasePolicy):
def make_critic(self) -> Critic:
return Critic(**self.net_args).to(self.device)
def forward(self, obs: th.Tensor) -> th.Tensor:
return self.predict(obs, deterministic=False)
def forward(self, obs: th.Tensor, deterministic: bool = False) -> th.Tensor:
return self._predict(obs, deterministic=deterministic)
def _predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
return self.actor(observation, deterministic)
@ -394,6 +397,7 @@ class CnnPolicy(SACPolicy):
: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,
@ -428,6 +432,5 @@ class CnnPolicy(SACPolicy):
optimizer_kwargs)
register_policy("MlpPolicy", MlpPolicy)
register_policy("CnnPolicy", CnnPolicy)

1
torchy_baselines/sac/sac.py
View file

@ -6,7 +6,6 @@ import numpy as np
from torchy_baselines.common import logger
from torchy_baselines.common.base_class import OffPolicyRLModel
from torchy_baselines.common.buffers import ReplayBuffer
from torchy_baselines.common.type_aliases import GymEnv, MaybeCallback
from torchy_baselines.common.noise import ActionNoise
from torchy_baselines.sac.policies import SACPolicy

5
torchy_baselines/td3/policies.py
View file

@ -8,7 +8,7 @@ from torchy_baselines.common.preprocessing import get_action_dim
from torchy_baselines.common.policies import (BasePolicy, register_policy, create_mlp,
create_sde_features_extractor, NatureCNN,
BaseFeaturesExtractor, FlattenExtractor)
from torchy_baselines.common.distributions import StateDependentNoiseDistribution, Distribution
from torchy_baselines.common.distributions import StateDependentNoiseDistribution
class Actor(BasePolicy):
@ -147,7 +147,7 @@ class Actor(BasePolicy):
given the observations. Only useful when using SDE.
:param obs: (th.Tensor)
:param action: (th.Tensor)
:param actions: (th.Tensor)
:return: (th.Tensor, th.Tensor) log likelihood of taking those actions
and entropy of the action distribution.
"""
@ -485,5 +485,6 @@ class CnnPolicy(TD3Policy):
optimizer_class,
optimizer_kwargs)
register_policy("MlpPolicy", MlpPolicy)
register_policy("CnnPolicy", CnnPolicy)