diff --git a/tests/test_cnn.py b/tests/test_cnn.py index 92292ec..1aec9d4 100644 --- a/tests/test_cnn.py +++ b/tests/test_cnn.py @@ -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: diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 7d8f228..de97b13 100644 --- a/tests/test_distributions.py +++ b/tests/test_distributions.py @@ -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) diff --git a/tests/test_monitor.py b/tests/test_monitor.py index d783f25..b21c33b 100644 --- a/tests/test_monitor.py +++ b/tests/test_monitor.py @@ -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: diff --git a/tests/test_run.py b/tests/test_run.py index 5fd1fe1..0c6ea81 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -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) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 48936e1..9541098 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -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: diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index 93b1930..9a67b93 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -149,8 +149,8 @@ class DiagGaussianDistribution(Distribution): return sum_independent_dims(self.distribution.entropy()) def actions_from_params(self, mean_actions: th.Tensor, - log_std: th.Tensor, - deterministic: bool = False) -> th.Tensor: + log_std: th.Tensor, + deterministic: bool = False) -> th.Tensor: # Update the proba distribution self.proba_distribution(mean_actions, log_std) return self.get_actions(deterministic=deterministic) @@ -278,7 +278,7 @@ class CategoricalDistribution(Distribution): return self.distribution.entropy() def actions_from_params(self, action_logits: th.Tensor, - deterministic: bool = False) -> th.Tensor: + deterministic: bool = False) -> th.Tensor: # Update the proba distribution self.proba_distribution(action_logits) return self.get_actions(deterministic=deterministic) @@ -453,9 +453,9 @@ class StateDependentNoiseDistribution(Distribution): return sum_independent_dims(self.distribution.entropy()) def actions_from_params(self, mean_actions: th.Tensor, - log_std: th.Tensor, - latent_sde: th.Tensor, - deterministic: bool = False) -> th.Tensor: + log_std: th.Tensor, + latent_sde: th.Tensor, + deterministic: bool = False) -> th.Tensor: # Update the proba distribution self.proba_distribution(mean_actions, log_std, latent_sde) return self.get_actions(deterministic=deterministic) diff --git a/torchy_baselines/common/identity_env.py b/torchy_baselines/common/identity_env.py index 3975d7e..8a76d3b 100644 --- a/torchy_baselines/common/identity_env.py +++ b/torchy_baselines/common/identity_env.py @@ -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 diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 6dd787c..f6c6644 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -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,17 +55,18 @@ 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) # We assume CxWxH images (channels first) # Re-ordering will be done by pre-preprocessing or wrapper assert is_image_space(observation_space), ('You should use NatureCNN ' - f'only with images not with {observation_space} ' - '(you are probably using `CnnPolicy` instead of `MlpPolicy`)') + f'only with images not with {observation_space} ' + '(you are probably using `CnnPolicy` instead of `MlpPolicy`)') n_input_channels = observation_space.shape[0] self.cnn = nn.Sequential(nn.Conv2d(n_input_channels, 32, kernel_size=8, stride=4, padding=0), nn.ReLU(), @@ -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', @@ -203,13 +207,13 @@ class BasePolicy(nn.Module): # as PyTorch use channel first format if is_image_space(self.observation_space): if (observation.shape == self.observation_space.shape or - observation.shape[1:] == self.observation_space.shape): + observation.shape[1:] == self.observation_space.shape): pass else: # Try to re-order the channels transpose_obs = VecTransposeImage.transpose_image(observation) if (transpose_obs.shape == self.observation_space.shape - or transpose_obs.shape[1:] == self.observation_space.shape): + or transpose_obs.shape[1:] == self.observation_space.shape): observation = transpose_obs vectorized_env = self._is_vectorized_observation(observation, self.observation_space) @@ -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], diff --git a/torchy_baselines/common/preprocessing.py b/torchy_baselines/common/preprocessing.py index 332bea6..434355d 100644 --- a/torchy_baselines/common/preprocessing.py +++ b/torchy_baselines/common/preprocessing.py @@ -1,4 +1,4 @@ -from typing import Tuple, Union +from typing import Tuple import numpy as np import torch as th diff --git a/torchy_baselines/common/type_aliases.py b/torchy_baselines/common/type_aliases.py index 1026361..70b63a8 100644 --- a/torchy_baselines/common/type_aliases.py +++ b/torchy_baselines/common/type_aliases.py @@ -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 diff --git a/torchy_baselines/common/vec_env/vec_transpose.py b/torchy_baselines/common/vec_env/vec_transpose.py index 3b13c39..3a514e9 100644 --- a/torchy_baselines/common/vec_env/vec_transpose.py +++ b/torchy_baselines/common/vec_env/vec_transpose.py @@ -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' diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 12fe737..2b9164a 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -286,7 +286,6 @@ class PPOPolicy(BasePolicy): MlpPolicy = PPOPolicy - class CnnPolicy(PPOPolicy): """ CnnPolicy class (with both actor and critic) for A2C and derivates (PPO). diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 4205e81..ecd6faf 100644 --- a/torchy_baselines/sac/policies.py +++ b/torchy_baselines/sac/policies.py @@ -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 @@ -105,16 +105,16 @@ class Actor(BasePolicy): data = super()._get_data() data.update(dict( - net_arch=self.net_arch, - features_dim=self.features_dim, - activation_fn=self.activation_fn, - use_sde=self.use_sde, - log_std_init=self.log_std_init, - full_std=self.full_std, - sde_net_arch=self.sde_net_arch, - use_expln=self.use_expln, - features_extractor=self.features_extractor, - clip_mean=self.clip_mean + net_arch=self.net_arch, + features_dim=self.features_dim, + activation_fn=self.activation_fn, + use_sde=self.use_sde, + log_std_init=self.log_std_init, + full_std=self.full_std, + sde_net_arch=self.sde_net_arch, + use_expln=self.use_expln, + features_extractor=self.features_extractor, + clip_mean=self.clip_mean )) return data @@ -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) @@ -336,18 +339,18 @@ class SACPolicy(BasePolicy): data = super()._get_data() data.update(dict( - net_arch=self.net_args['net_arch'], - activation_fn=self.net_args['activation_fn'], - use_sde=self.actor_kwargs['use_sde'], - log_std_init=self.actor_kwargs['log_std_init'], - sde_net_arch=self.actor_kwargs['sde_net_arch'], - use_expln=self.actor_kwargs['use_expln'], - clip_mean=self.actor_kwargs['clip_mean'], - lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone - optimizer_class=self.optimizer_class, - optimizer_kwargs=self.optimizer_kwargs, - features_extractor_class=self.features_extractor_class, - features_extractor_kwargs=self.features_extractor_kwargs + net_arch=self.net_args['net_arch'], + activation_fn=self.net_args['activation_fn'], + use_sde=self.actor_kwargs['use_sde'], + log_std_init=self.actor_kwargs['log_std_init'], + sde_net_arch=self.actor_kwargs['sde_net_arch'], + use_expln=self.actor_kwargs['use_expln'], + clip_mean=self.actor_kwargs['clip_mean'], + lr_schedule=self._dummy_schedule, # dummy lr schedule, not needed for loading policy alone + 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 @@ -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) diff --git a/torchy_baselines/sac/sac.py b/torchy_baselines/sac/sac.py index b37658e..c1cbc64 100644 --- a/torchy_baselines/sac/sac.py +++ b/torchy_baselines/sac/sac.py @@ -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 diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 23aa670..b43adaf 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -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)