From 18f38f8cf5b969403d5416e62eb9a4461538c0db Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Thu, 12 Mar 2020 11:12:10 +0100 Subject: [PATCH] Reformat --- tests/test_callbacks.py | 4 ++-- tests/test_distributions.py | 1 + tests/test_logger.py | 3 ++- tests/test_predict.py | 1 + tests/test_run.py | 4 +--- tests/test_save_load.py | 1 + tests/test_vec_normalize.py | 7 +++++-- torchy_baselines/a2c/a2c.py | 11 +++++------ torchy_baselines/cem_rl/cem.py | 3 +-- torchy_baselines/common/base_class.py | 4 ++-- torchy_baselines/common/buffers.py | 10 ++++++---- torchy_baselines/common/distributions.py | 4 ++-- torchy_baselines/common/monitor.py | 5 +++-- torchy_baselines/common/noise.py | 3 +++ torchy_baselines/common/save_util.py | 2 -- torchy_baselines/common/type_aliases.py | 5 ++--- .../common/vec_env/base_vec_env.py | 2 +- .../common/vec_env/subproc_vec_env.py | 4 ++-- .../common/vec_env/vec_normalize.py | 6 +++--- torchy_baselines/ppo/policies.py | 10 ++++++---- torchy_baselines/ppo/ppo.py | 18 +++++++++--------- torchy_baselines/sac/policies.py | 1 + torchy_baselines/td3/policies.py | 5 ++--- torchy_baselines/td3/td3.py | 14 ++++++-------- 24 files changed, 67 insertions(+), 61 deletions(-) diff --git a/tests/test_callbacks.py b/tests/test_callbacks.py index 96db45a..a4ffede 100644 --- a/tests/test_callbacks.py +++ b/tests/test_callbacks.py @@ -6,7 +6,7 @@ import gym from torchy_baselines import A2C, CEMRL, PPO, SAC, TD3 from torchy_baselines.common.callbacks import (CallbackList, CheckpointCallback, EvalCallback, - EveryNTimesteps, StopTrainingOnRewardThreshold) + EveryNTimesteps, StopTrainingOnRewardThreshold) @pytest.mark.parametrize("model_class", [A2C, CEMRL, PPO, SAC, TD3]) @@ -44,6 +44,6 @@ def test_callbacks(model_class): # Transform callback into a callback list automatically model.learn(500, callback=[checkpoint_callback, eval_callback]) # Automatic wrapping, old way of doing callbacks - model.learn(500, callback=lambda _locals, _globals : True) + model.learn(500, callback=lambda _locals, _globals: True) if os.path.exists(log_folder): shutil.rmtree(log_folder) diff --git a/tests/test_distributions.py b/tests/test_distributions.py index 4b5d792..7d28ad7 100644 --- a/tests/test_distributions.py +++ b/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): """ diff --git a/tests/test_logger.py b/tests/test_logger.py index b55a616..5ca0437 100644 --- a/tests/test_logger.py +++ b/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, diff --git a/tests/test_predict.py b/tests/test_predict.py index 6f2245c..e68954f 100644 --- a/tests/test_predict.py +++ b/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 diff --git a/tests/test_run.py b/tests/test_run.py index fdfcff6..1a3f991 100644 --- a/tests/test_run.py +++ b/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 diff --git a/tests/test_save_load.py b/tests/test_save_load.py index 9d73ddf..bdec3ed 100644 --- a/tests/test_save_load.py +++ b/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): """ diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index 3c21f69..75fd3a8 100644 --- a/tests/test_vec_normalize.py +++ b/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")]) @@ -50,8 +53,8 @@ def _make_warmstart_cartpole(): def test_runningmeanstd(): """Test RunningMeanStd object""" for (x_1, x_2, x_3) in [ - (np.random.randn(3), np.random.randn(4), np.random.randn(5)), - (np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2))]: + (np.random.randn(3), np.random.randn(4), np.random.randn(5)), + (np.random.randn(3, 2), np.random.randn(4, 2), np.random.randn(5, 2))]: rms = RunningMeanStd(epsilon=0.0, shape=x_1.shape[1:]) x_cat = np.concatenate([x_1, x_2, x_3], axis=0) diff --git a/torchy_baselines/a2c/a2c.py b/torchy_baselines/a2c/a2c.py index 6e0a1f2..5f9730d 100644 --- a/torchy_baselines/a2c/a2c.py +++ b/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): diff --git a/torchy_baselines/cem_rl/cem.py b/torchy_baselines/cem_rl/cem.py index d1d221c..7527b99 100644 --- a/torchy_baselines/cem_rl/cem.py +++ b/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 diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index b77237d..f1d205a 100644 --- a/torchy_baselines/common/base_class.py +++ b/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 diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 6ffe479..5ddc29c 100644 --- a/torchy_baselines/common/buffers.py +++ b/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, @@ -118,13 +119,13 @@ class BaseBuffer(object): @staticmethod def _normalize_obs(obs: np.ndarray, - env: Optional[VecNormalize] = None) -> np.ndarray: + env: Optional[VecNormalize] = None) -> np.ndarray: if env is not None: 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, diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index a7f219a..6a0ab2f 100644 --- a/torchy_baselines/common/distributions.py +++ b/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) """ diff --git a/torchy_baselines/common/monitor.py b/torchy_baselines/common/monitor.py index 3d84b9b..5a87169 100644 --- a/torchy_baselines/common/monitor.py +++ b/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() diff --git a/torchy_baselines/common/noise.py b/torchy_baselines/common/noise.py index fa25f42..8511010 100644 --- a/torchy_baselines/common/noise.py +++ b/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 diff --git a/torchy_baselines/common/save_util.py b/torchy_baselines/common/save_util.py index a9dedbc..85fb6fd 100644 --- a/torchy_baselines/common/save_util.py +++ b/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 diff --git a/torchy_baselines/common/type_aliases.py b/torchy_baselines/common/type_aliases.py index 12c220f..60042be 100644 --- a/torchy_baselines/common/type_aliases.py +++ b/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 diff --git a/torchy_baselines/common/vec_env/base_vec_env.py b/torchy_baselines/common/vec_env/base_vec_env.py index 18e8a59..4a81e6e 100644 --- a/torchy_baselines/common/vec_env/base_vec_env.py +++ b/torchy_baselines/common/vec_env/base_vec_env.py @@ -244,7 +244,7 @@ class VecEnvWrapper(VecEnv): if blocked_class is not None: own_class = f"{type(self).__module__}.{type(self).__name__}" error_str = (f"Error: Recursive attribute lookup for {name} from {own_class} is " - "ambiguous and hides attribute from {blocked_class}") + "ambiguous and hides attribute from {blocked_class}") raise AttributeError(error_str) return self.getattr_recursive(name) diff --git a/torchy_baselines/common/vec_env/subproc_vec_env.py b/torchy_baselines/common/vec_env/subproc_vec_env.py index 9205683..5e6ee85 100644 --- a/torchy_baselines/common/vec_env/subproc_vec_env.py +++ b/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 diff --git a/torchy_baselines/common/vec_env/vec_normalize.py b/torchy_baselines/common/vec_env/vec_normalize.py index ea94bc0..87fb70a 100644 --- a/torchy_baselines/common/vec_env/vec_normalize.py +++ b/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) @@ -122,7 +122,7 @@ class VecNormalize(VecEnvWrapper): """ if self.norm_reward: reward = np.clip(reward / np.sqrt(self.ret_rms.var + self.epsilon), - -self.clip_reward, self.clip_reward) + -self.clip_reward, self.clip_reward) return reward def unnormalize_obs(self, 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): """ diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 491737a..b989776 100644 --- a/torchy_baselines/ppo/policies.py +++ b/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. diff --git a/torchy_baselines/ppo/ppo.py b/torchy_baselines/ppo/ppo.py index a7dcf0b..5ac1a96 100644 --- a/torchy_baselines/ppo/ppo.py +++ b/torchy_baselines/ppo/ppo.py @@ -146,11 +146,11 @@ class PPO(BaseRLModel): self.clip_range_vf = get_schedule_fn(self.clip_range_vf) def collect_rollouts(self, - env: VecEnv, - callback: BaseCallback, - rollout_buffer: RolloutBuffer, - n_rollout_steps: int = 256, - obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]: + env: VecEnv, + callback: BaseCallback, + rollout_buffer: RolloutBuffer, + n_rollout_steps: int = 256, + obs: Optional[np.ndarray] = None) -> Tuple[Optional[np.ndarray], bool]: n_steps = 0 continue_training = True @@ -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()) diff --git a/torchy_baselines/sac/policies.py b/torchy_baselines/sac/policies.py index 15a28f2..1a1dc33 100644 --- a/torchy_baselines/sac/policies.py +++ b/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) diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index f75cdea..3020d26 100644 --- a/torchy_baselines/td3/policies.py +++ b/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): diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index ccb835e..9de88c1 100644 --- a/torchy_baselines/td3/td3.py +++ b/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 @@ -138,9 +136,9 @@ class TD3(OffPolicyRLModel): self.vf_net = self.policy.vf_net def train_critic(self, gradient_steps: int = 1, - batch_size: int = 100, - replay_data: Optional[ReplayBufferSamples] = None, - tau: float = 0.0) -> None: + batch_size: int = 100, + replay_data: Optional[ReplayBufferSamples] = None, + tau: float = 0.0) -> None: # Update optimizer learning rate self._update_learning_rate(self.critic.optimizer)