Reformat and code cleanup

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:

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)
 

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:

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)

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:

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,
+                            log_std: th.Tensor,
-                           deterministic: bool = False) -> 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,
+                            log_std: th.Tensor,
-                           latent_sde: th.Tensor,
+                            latent_sde: th.Tensor,
-                           deterministic: bool = False) -> 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)

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

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} '
+                                                   f'only with images not with {observation_space} '
-                '(you are probably using `CnnPolicy` instead of `MlpPolicy`)')
+                                                   '(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],

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

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

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'
 

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

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,
+            net_arch=self.net_arch,
-             features_dim=self.features_dim,
+            features_dim=self.features_dim,
-             activation_fn=self.activation_fn,
+            activation_fn=self.activation_fn,
-             use_sde=self.use_sde,
+            use_sde=self.use_sde,
-             log_std_init=self.log_std_init,
+            log_std_init=self.log_std_init,
-             full_std=self.full_std,
+            full_std=self.full_std,
-             sde_net_arch=self.sde_net_arch,
+            sde_net_arch=self.sde_net_arch,
-             use_expln=self.use_expln,
+            use_expln=self.use_expln,
-             features_extractor=self.features_extractor,
+            features_extractor=self.features_extractor,
-             clip_mean=self.clip_mean
+            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'],
+            net_arch=self.net_args['net_arch'],
-             activation_fn=self.net_args['activation_fn'],
+            activation_fn=self.net_args['activation_fn'],
-             use_sde=self.actor_kwargs['use_sde'],
+            use_sde=self.actor_kwargs['use_sde'],
-             log_std_init=self.actor_kwargs['log_std_init'],
+            log_std_init=self.actor_kwargs['log_std_init'],
-             sde_net_arch=self.actor_kwargs['sde_net_arch'],
+            sde_net_arch=self.actor_kwargs['sde_net_arch'],
-             use_expln=self.actor_kwargs['use_expln'],
+            use_expln=self.actor_kwargs['use_expln'],
-             clip_mean=self.actor_kwargs['clip_mean'],
+            clip_mean=self.actor_kwargs['clip_mean'],
-             lr_schedule=self._dummy_schedule,  # dummy lr schedule, not needed for loading policy alone
+            lr_schedule=self._dummy_schedule,  # dummy lr schedule, not needed for loading policy alone
-             optimizer_class=self.optimizer_class,
+            optimizer_class=self.optimizer_class,
-             optimizer_kwargs=self.optimizer_kwargs,
+            optimizer_kwargs=self.optimizer_kwargs,
-             features_extractor_class=self.features_extractor_class,
+            features_extractor_class=self.features_extractor_class,
-             features_extractor_kwargs=self.features_extractor_kwargs
+            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:
+    def forward(self, obs: th.Tensor, deterministic: bool = False) -> th.Tensor:
-        return self.predict(obs, deterministic=False)
+        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)

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

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)