Reformat

tests/test_distributions.py

@@ -22,6 +22,7 @@ def test_bijector():
     # Check the inverse method
     assert th.isclose(TanhBijector.inverse(squashed_actions), actions).all()
 
+
 @pytest.mark.parametrize("model_class", [A2C, PPO])
 def test_squashed_gaussian(model_class):
     """

tests/test_logger.py

@@ -5,7 +5,8 @@ import pytest
 import numpy as np
 
 from torchy_baselines.common.logger import (make_output_format, read_csv, read_json, DEBUG, ScopedConfigure,
-    info, debug, set_level, configure, logkv, logkvs, dumpkvs, logkv_mean, warn, error, reset)
+                                            info, debug, set_level, configure, logkv, logkvs, dumpkvs, logkv_mean, warn,
+                                            error, reset)
 
 KEY_VALUES = {
     "test": 1,

tests/test_predict.py

@@ -12,6 +12,7 @@ MODEL_LIST = [
     SAC,
 ]
 
+
 @pytest.mark.parametrize("model_class", MODEL_LIST)
 def test_auto_wrap(model_class):
     # test auto wrapping of env into a VecEnv

tests/test_run.py

@@ -1,7 +1,5 @@
-import os
-
-import pytest
 import numpy as np
+import pytest
 
 from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3
 from torchy_baselines.common.noise import NormalActionNoise, OrnsteinUhlenbeckActionNoise

tests/test_save_load.py

@@ -16,6 +16,7 @@ MODEL_LIST = [
     SAC,
 ]
 
+
 @pytest.mark.parametrize("model_class", MODEL_LIST)
 def test_save_load(model_class):
     """

tests/test_vec_normalize.py

@@ -8,9 +8,11 @@ from torchy_baselines import CEMRL, SAC, TD3
 
 ENV_ID = 'Pendulum-v0'
 
+
 def make_env():
     return gym.make(ENV_ID)
 
+
 def check_rms_equal(rmsa, rmsb):
     assert np.all(rmsa.mean == rmsb.mean)
     assert np.all(rmsa.var == rmsb.var)
@@ -34,6 +36,7 @@ def check_vec_norm_equal(norma, normb):
     assert norma.epsilon == normb.epsilon
     assert norma.training == normb.training
 
+
 def _make_warmstart_cartpole():
     """Warm-start VecNormalize by stepping through CartPole"""
     venv = DummyVecEnv([lambda: gym.make("CartPole-v1")])

torchy_baselines/a2c/a2c.py

@@ -1,15 +1,14 @@
-from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
-
-from gym import spaces
 import torch as th
 import torch.nn.functional as F
+from gym import spaces
+from typing import Type, Union, Callable, Optional, Dict, Any
 
-from torchy_baselines.common.utils import explained_variance
 from torchy_baselines.common import logger
-from torchy_baselines.common.type_aliases import GymEnv
 from torchy_baselines.common.callbacks import BaseCallback
-from torchy_baselines.ppo.ppo import PPO
+from torchy_baselines.common.type_aliases import GymEnv
+from torchy_baselines.common.utils import explained_variance
 from torchy_baselines.ppo.policies import PPOPolicy
+from torchy_baselines.ppo.ppo import PPO
 
 
 class A2C(PPO):

torchy_baselines/cem_rl/cem.py

@@ -1,6 +1,5 @@
-from typing import Type, Tuple, Optional, List
-
 import numpy as np
+from typing import Tuple, Optional, List
 
 
 # TODO: add more from https://github.com/hardmaru/estool/blob/master/es.py

torchy_baselines/common/base_class.py

@@ -16,7 +16,7 @@ from torchy_baselines.common.policies import BasePolicy, get_policy_from_name
 from torchy_baselines.common.utils import set_random_seed, get_schedule_fn, update_learning_rate
 from torchy_baselines.common.vec_env import DummyVecEnv, VecEnv, unwrap_vec_normalize, VecNormalize
 from torchy_baselines.common.save_util import data_to_json, json_to_data, recursive_getattr, recursive_setattr
-from torchy_baselines.common.type_aliases import GymEnv, TensorDict, OptimizerStateDict, RolloutReturn
+from torchy_baselines.common.type_aliases import GymEnv, TensorDict, RolloutReturn
 from torchy_baselines.common.callbacks import BaseCallback, CallbackList, ConvertCallback, EvalCallback
 from torchy_baselines.common.monitor import Monitor
 from torchy_baselines.common.noise import ActionNoise
@@ -494,7 +494,7 @@ class BaseRLModel(ABC):
                 if "data" in namelist and load_data:
                     # Load class parameters and convert to string
                     json_data = archive.read("data").decode()
-                    data = json_to_data(json_data, device)
+                    data = json_to_data(json_data)
 
                 if "tensors.pth" in namelist and load_data:
                     # Load extra tensors

torchy_baselines/common/buffers.py

@@ -18,6 +18,7 @@ class BaseBuffer(object):
         to which the values will be converted
     :param n_envs: (int) Number of parallel environments
     """
+
     def __init__(self,
                  buffer_size: int,
                  obs_dim: int,
@@ -123,8 +124,8 @@ class BaseBuffer(object):
             return env.normalize_obs(obs).astype(np.float32)
         return obs
 
-    def _normalize_reward(self,
-                          reward: np.ndarray,
+    @staticmethod
+    def _normalize_reward(reward: np.ndarray,
                           env: Optional[VecNormalize] = None) -> np.ndarray:
         if env is not None:
             return env.normalize_reward(reward).astype(np.float32)
@@ -141,13 +142,13 @@ class ReplayBuffer(BaseBuffer):
     :param device: (th.device)
     :param n_envs: (int) Number of parallel environments
     """
+
     def __init__(self,
                  buffer_size: int,
                  obs_dim: int,
                  action_dim: int,
                  device: Union[th.device, str] = 'cpu',
                  n_envs: int = 1):
-
         super(ReplayBuffer, self).__init__(buffer_size, obs_dim, action_dim, device, n_envs=n_envs)
 
         assert n_envs == 1, "Replay buffer only support single environment for now"
@@ -201,6 +202,7 @@ class RolloutBuffer(BaseBuffer):
     :param gamma: (float) Discount factor
     :param n_envs: (int) Number of parallel environments
     """
+
     def __init__(self,
                  buffer_size: int,
                  obs_dim: int,

torchy_baselines/common/distributions.py

@@ -317,7 +317,7 @@ class StateDependentNoiseDistribution(Distribution):
         self.exploration_matrices = self.weights_dist.rsample((batch_size,))
 
     def proba_distribution_net(self, latent_dim: int, log_std_init: float = -2.0,
-                               latent_sde_dim: Optional[th.Tensor] = None) -> Tuple[nn.Module, nn.Parameter]:
+                               latent_sde_dim: Optional[int] = None) -> Tuple[nn.Module, nn.Parameter]:
         """
         Create the layers and parameter that represent the distribution:
         one output will be the deterministic action, the other parameter will be the
@@ -325,7 +325,7 @@ class StateDependentNoiseDistribution(Distribution):
 
         :param latent_dim: (int) Dimension of the last layer of the policy (before the action layer)
         :param log_std_init: (float) Initial value for the log standard deviation
-        :param latent_sde_dim: (int) Dimension of the last layer of the feature extractor
+        :param latent_sde_dim: (Optional[int]) Dimension of the last layer of the feature extractor
             for SDE. By default, it is shared with the policy network.
         :return: (nn.Linear, nn.Parameter)
         """

torchy_baselines/common/monitor.py

@@ -27,8 +27,9 @@ class Monitor(gym.Wrapper):
         :param env: (gym.Env) The environment
         :param filename: (Optional[str]) the location to save a log file, can be None for no log
         :param allow_early_resets: (bool) allows the reset of the environment before it is done
-        :param reset_keywords: (Tuple[str, ...]) extra keywords for the reset call, if extra parameters are needed at reset
-        :param info_keywords: (Tuple[str, ...]) extra information to log, from the information return of environment.step
+        :param reset_keywords: (Tuple[str, ...]) extra keywords for the reset call,
+            if extra parameters are needed at reset
+        :param info_keywords: (Tuple[str, ...]) extra information to log, from the information return of env.step()
         """
         super(Monitor, self).__init__(env=env)
         self.t_start = time.time()

torchy_baselines/common/noise.py

@@ -9,6 +9,7 @@ class ActionNoise(ABC):
     """
     The action noise base class
     """
+
     def __init__(self):
         super(ActionNoise, self).__init__()
 
@@ -22,6 +23,7 @@ class ActionNoise(ABC):
     def __call__(self):
         pass
 
+
 class NormalActionNoise(ActionNoise):
     """
     A Gaussian action noise
@@ -29,6 +31,7 @@ class NormalActionNoise(ActionNoise):
     :param mean: (float) the mean value of the noise
     :param sigma: (float) the scale of the noise (std here)
     """
+
     def __init__(self, mean, sigma):
         self._mu = mean
         self._sigma = sigma

torchy_baselines/common/save_util.py

@@ -122,14 +122,12 @@ def data_to_json(data: Dict[str, Any]) -> str:
 
 
 def json_to_data(json_string: str,
-                 device: Union[th.device, str] = 'cpu',
                  custom_objects: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
     """
     Turn JSON serialization of class-parameters back into dictionary.
 
     :param json_string: (str) JSON serialization of the class-parameters
         that should be loaded.
-    :param device: torch.device device to which the data should be mapped if errors occur
     :param custom_objects: (dict) Dictionary of objects to replace
         upon loading. If a variable is present in this dictionary as a
         key, it will not be deserialized and the corresponding item

torchy_baselines/common/type_aliases.py

@@ -1,8 +1,7 @@
 """
-Common aliases for type hing
+Common aliases for type hint
 """
-from typing import Union, Type, Optional, Dict, Any, List, NamedTuple
-from collections import namedtuple
+from typing import Union, Dict, Any, NamedTuple, Optional
 
 import numpy as np
 import torch as th

torchy_baselines/common/vec_env/subproc_vec_env.py

@@ -61,11 +61,11 @@ def tile_images(img_nhwc):
     new_width = int(np.ceil(float(n_images) / new_height))
     img_nhwc = np.array(list(img_nhwc) + [img_nhwc[0] * 0 for _ in range(n_images, new_height * new_width)])
     # img_HWhwc
-    out_image = img_nhwc.reshape(new_height, new_width, height, width, n_channels)
+    out_image = img_nhwc.reshape((new_height, new_width, height, width, n_channels))
     # img_HhWwc
     out_image = out_image.transpose(0, 2, 1, 3, 4)
     # img_Hh_Ww_c
-    out_image = out_image.reshape(new_height * height, new_width * width, n_channels)
+    out_image = out_image.reshape((new_height * height, new_width * width, n_channels))
     return out_image
 
 

torchy_baselines/common/vec_env/vec_normalize.py

@@ -86,7 +86,7 @@ class VecNormalize(VecEnvWrapper):
         """
         obs, rews, news, infos = self.venv.step_wait()
         self.old_obs = obs
-        self.old_rews = rews
+        self.old_reward = rews
 
         if self.training:
             self.obs_rms.update(obs)
@@ -146,7 +146,7 @@ class VecNormalize(VecEnvWrapper):
         """
         Returns an unnormalized version of the rewards from the most recent step.
         """
-        return self.old_rews.copy()
+        return self.old_reward.copy()
 
     def reset(self):
         """

torchy_baselines/ppo/policies.py

@@ -9,8 +9,8 @@ import numpy as np
 from torchy_baselines.common.policies import (BasePolicy, register_policy, MlpExtractor,
                                               create_sde_feature_extractor)
 from torchy_baselines.common.distributions import (make_proba_distribution, Distribution,
-    DiagGaussianDistribution, CategoricalDistribution, StateDependentNoiseDistribution)
-
+                                                   DiagGaussianDistribution, CategoricalDistribution,
+                                                   StateDependentNoiseDistribution)
 
 
 class PPOPolicy(BasePolicy):
@@ -183,13 +183,15 @@ class PPOPolicy(BasePolicy):
         action, _ = self._get_action_dist_from_latent(latent_pi, latent_sde, deterministic=deterministic)
         return action
 
-    def evaluate_actions(self, obs: th.Tensor, actions: th.Tensor, deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
+    def evaluate_actions(self, obs: th.Tensor,
+                         actions: th.Tensor,
+                         deterministic: bool = False) -> Tuple[th.Tensor, th.Tensor, th.Tensor]:
         """
         Evaluate actions according to the current policy,
         given the observations.
 
         :param obs: (th.Tensor)
-        :param action: (th.Tensor)
+        :param actions: (th.Tensor)
         :param deterministic: (bool)
         :return: (th.Tensor, th.Tensor, th.Tensor) estimated value, log likelihood of taking those actions
             and entropy of the action distribution.

torchy_baselines/ppo/ppo.py

@@ -167,7 +167,6 @@ class PPO(BaseRLModel):
                 continue_training = False
                 return None, continue_training
 
-
             if self.use_sde and self.sde_sample_freq > 0 and n_steps % self.sde_sample_freq == 0:
                 # Sample a new noise matrix
                 self.policy.reset_noise(env.num_envs)
@@ -227,7 +226,8 @@ class PPO(BaseRLModel):
                 values, log_prob, entropy = self.policy.evaluate_actions(rollout_data.observations, actions)
                 values = values.flatten()
                 # Normalize advantage
-                advantages = (rollout_data.advantages - rollout_data.advantages.mean()) / (rollout_data.advantages.std() + 1e-8)
+                advantages = (rollout_data.advantages - rollout_data.advantages.mean()) / (
+                            rollout_data.advantages.std() + 1e-8)
 
                 # ratio between old and new policy, should be one at the first iteration
                 ratio = th.exp(log_prob - rollout_data.old_log_prob)
@@ -242,7 +242,8 @@ class PPO(BaseRLModel):
                 else:
                     # Clip the different between old and new value
                     # NOTE: this depends on the reward scaling
-                    values_pred = rollout_data.old_values + th.clamp(values - rollout_data.old_values, -clip_range_vf, clip_range_vf)
+                    values_pred = rollout_data.old_values + th.clamp(values - rollout_data.old_values, -clip_range_vf,
+                                                                     clip_range_vf)
                 # Value loss using the TD(gae_lambda) target
                 value_loss = F.mse_loss(rollout_data.returns, values_pred)
 
@@ -275,7 +276,6 @@ class PPO(BaseRLModel):
         if self.clip_range_vf is not None:
             logger.logkv("clip_range_vf", clip_range_vf)
 
-
         logger.logkv("explained_variance", explained_var)
         # TODO: gather stats for the entropy and other losses?
         logger.logkv("entropy_loss", entropy_loss.item())

torchy_baselines/sac/policies.py

@@ -244,6 +244,7 @@ class SACPolicy(BasePolicy):
     def predict(self, observation: th.Tensor, deterministic: bool = False) -> th.Tensor:
         return self.actor.forward(observation, deterministic)
 
+
 MlpPolicy = SACPolicy
 
 register_policy("MlpPolicy", MlpPolicy)

torchy_baselines/td3/policies.py

@@ -1,12 +1,11 @@
-from typing import List, Tuple, Callable, Optional
-
 import torch
 import torch as th
 import torch.nn as nn
+from typing import List, Tuple, Optional
 
+from torchy_baselines.common.distributions import StateDependentNoiseDistribution
 from torchy_baselines.common.policies import BasePolicy, register_policy, create_mlp, BaseNetwork, \
     create_sde_feature_extractor
-from torchy_baselines.common.distributions import StateDependentNoiseDistribution
 
 
 class Actor(BaseNetwork):

torchy_baselines/td3/td3.py

@@ -1,14 +1,12 @@
-from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
-
 import torch as th
 import torch.nn.functional as F
-import numpy as np
+from typing import List, Tuple, Type, Union, Callable, Optional, Dict, Any
 
 from torchy_baselines.common.base_class import OffPolicyRLModel
 from torchy_baselines.common.buffers import ReplayBuffer
-from torchy_baselines.common.type_aliases import ReplayBufferSamples, GymEnv
-from torchy_baselines.common.noise import ActionNoise
 from torchy_baselines.common.callbacks import BaseCallback
+from torchy_baselines.common.noise import ActionNoise
+from torchy_baselines.common.type_aliases import ReplayBufferSamples, GymEnv
 from torchy_baselines.td3.policies import TD3Policy