From 2469ff3859c7fbb6468ff9900eeadc27d018af21 Mon Sep 17 00:00:00 2001 From: Antonin RAFFIN Date: Sat, 21 Sep 2019 17:17:09 +0200 Subject: [PATCH] Reformat --- tests/test_run.py | 7 ++++--- tests/test_vec_normalize.py | 4 ++-- torchy_baselines/cem_rl/cem_rl.py | 10 ++++------ torchy_baselines/common/base_class.py | 2 -- torchy_baselines/common/buffers.py | 4 +++- torchy_baselines/common/distributions.py | 1 + torchy_baselines/common/utils.py | 3 +-- torchy_baselines/common/vec_env/__init__.py | 4 ++-- torchy_baselines/ppo/policies.py | 2 +- torchy_baselines/td3/policies.py | 1 + torchy_baselines/td3/td3.py | 7 +++---- 11 files changed, 22 insertions(+), 23 deletions(-) diff --git a/tests/test_run.py b/tests/test_run.py index 0583d12..546a3c0 100644 --- a/tests/test_run.py +++ b/tests/test_run.py @@ -1,9 +1,8 @@ import os -import gym - from torchy_baselines import TD3, CEMRL, PPO + def test_td3(): model = TD3('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[64, 64]), start_timesteps=100, verbose=1, create_eval_env=True) @@ -12,14 +11,16 @@ def test_td3(): model.load("test_save") os.remove("test_save.pth") + def test_cemrl(): model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1, - start_timesteps=100, verbose=1, create_eval_env=True) + start_timesteps=100, verbose=1, create_eval_env=True) model.learn(total_timesteps=20000, eval_freq=1000) model.save("test_save") model.load("test_save") os.remove("test_save.pth") + def test_ppo(): model = PPO('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), verbose=1, create_eval_env=True) model.learn(total_timesteps=1000, eval_freq=500) diff --git a/tests/test_vec_normalize.py b/tests/test_vec_normalize.py index fbcee54..78437a3 100644 --- a/tests/test_vec_normalize.py +++ b/tests/test_vec_normalize.py @@ -11,8 +11,8 @@ ENV_ID = 'Pendulum-v0' 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/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index aec935d..be3629c 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/torchy_baselines/cem_rl/cem_rl.py @@ -1,12 +1,10 @@ import time import torch as th -import torch.nn.functional as F -import numpy as np -from torchy_baselines.td3.td3 import TD3 -from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.cem_rl.cem import CEM +from torchy_baselines.common.evaluation import evaluate_policy +from torchy_baselines.td3.td3 import TD3 class CEMRL(TD3): @@ -16,6 +14,7 @@ class CEMRL(TD3): Paper: https://arxiv.org/abs/1810.01222 Code: https://github.com/apourchot/CEM-RL """ + def __init__(self, policy, env, policy_kwargs=None, verbose=0, sigma_init=1e-3, pop_size=10, damp=1e-3, damp_limit=1e-5, elitism=False, n_grad=5, policy_freq=2, batch_size=100, @@ -100,7 +99,7 @@ class CEMRL(TD3): else: # scales with a bigger population # but less training steps per agent - n_training_steps == 2 * (actor_steps // self.n_grad) + n_training_steps = 2 * (actor_steps // self.n_grad) for it in range(n_training_steps): # Sample replay buffer replay_data = self.replay_buffer.sample(self.batch_size) @@ -148,7 +147,6 @@ class CEMRL(TD3): print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( self.num_timesteps, episode_num, episode_timesteps, episode_reward)) - self.es.tell(self.es_params, self.fitnesses) timesteps_since_eval += actor_steps return self diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index bebd824..08fcce7 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -1,6 +1,5 @@ from abc import ABCMeta, abstractmethod - import gym import torch as th import numpy as np @@ -216,7 +215,6 @@ class BaseRLModel(object): if self.eval_env is not None: self.eval_env.seed(seed) - def collect_rollouts(self, env, n_episodes=1, action_noise_std=0.0, deterministic=False, callback=None, remove_timelimits=True, start_timesteps=0, num_timesteps=0, replay_buffer=None): diff --git a/torchy_baselines/common/buffers.py b/torchy_baselines/common/buffers.py index 32c37b3..599338a 100644 --- a/torchy_baselines/common/buffers.py +++ b/torchy_baselines/common/buffers.py @@ -52,10 +52,12 @@ class BaseBuffer(object): def _get_samples(self, batch_inds): raise NotImplementedError() + class ReplayBuffer(BaseBuffer): """ Taken from https://github.com/apourchot/CEM-RL """ + def __init__(self, buffer_size, state_dim, action_dim, device='cpu', n_envs=1): super(ReplayBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs) @@ -89,7 +91,7 @@ class ReplayBuffer(BaseBuffer): class RolloutBuffer(BaseBuffer): def __init__(self, buffer_size, state_dim, action_dim, device='cpu', - lambda_=1, gamma=0.99, n_envs=1): + lambda_=1, gamma=0.99, n_envs=1): super(RolloutBuffer, self).__init__(buffer_size, state_dim, action_dim, device, n_envs=n_envs) # TODO: try the buffer on the gpu? self.lambda_ = lambda_ diff --git a/torchy_baselines/common/distributions.py b/torchy_baselines/common/distributions.py index e0a73a8..83e4fde 100644 --- a/torchy_baselines/common/distributions.py +++ b/torchy_baselines/common/distributions.py @@ -1,6 +1,7 @@ import torch as th from torch.distributions import Normal + class Distribution(object): def __init__(self): super(Distribution, self).__init__() diff --git a/torchy_baselines/common/utils.py b/torchy_baselines/common/utils.py index da3a454..b4887a7 100644 --- a/torchy_baselines/common/utils.py +++ b/torchy_baselines/common/utils.py @@ -1,8 +1,7 @@ import random -import scipy.signal -import torch as th import numpy as np +import torch as th def set_random_seed(seed, using_cuda=False): diff --git a/torchy_baselines/common/vec_env/__init__.py b/torchy_baselines/common/vec_env/__init__.py index 3a58d12..97f6022 100644 --- a/torchy_baselines/common/vec_env/__init__.py +++ b/torchy_baselines/common/vec_env/__init__.py @@ -1,6 +1,6 @@ # flake8: noqa F401 -from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError, VecEnv, VecEnvWrapper, \ - CloudpickleWrapper +from torchy_baselines.common.vec_env.base_vec_env import AlreadySteppingError, NotSteppingError,\ + VecEnv, VecEnvWrapper, CloudpickleWrapper from torchy_baselines.common.vec_env.dummy_vec_env import DummyVecEnv from torchy_baselines.common.vec_env.subproc_vec_env import SubprocVecEnv from torchy_baselines.common.vec_env.vec_frame_stack import VecFrameStack diff --git a/torchy_baselines/ppo/policies.py b/torchy_baselines/ppo/policies.py index 0844d78..e7ae512 100644 --- a/torchy_baselines/ppo/policies.py +++ b/torchy_baselines/ppo/policies.py @@ -44,7 +44,6 @@ class PPOPolicy(BasePolicy): self.log_std = nn.Parameter(th.zeros(self.action_dim)) # Init weights: use orthogonal initialization for module in [self.shared_net, self.actor_net, self.value_net]: - gain = 0.01 if module == self.actor_net else 1.0 # Values from stable-baselines check why gain = { self.shared_net: np.sqrt(2), @@ -98,6 +97,7 @@ class PPOPolicy(BasePolicy): def value_forward(self): pass + MlpPolicy = PPOPolicy register_policy("MlpPolicy", MlpPolicy) diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 99eca0c..41a4371 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -79,6 +79,7 @@ class TD3Policy(BasePolicy): def make_critic(self): return Critic(**self.net_args).to(self.device) + MlpPolicy = TD3Policy register_policy("MlpPolicy", MlpPolicy) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index e47fc79..7a527e8 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -21,7 +21,7 @@ class TD3(BaseRLModel): buffer_size=int(1e6), learning_rate=1e-3, seed=0, device='auto', action_noise_std=0.1, start_timesteps=100, policy_freq=2, batch_size=100, create_eval_env=False, - _init_setup_model=True): + _init_setup_model=True): super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device, create_eval_env=create_eval_env) @@ -135,8 +135,7 @@ class TD3(BaseRLModel): for param, target_param in zip(self.actor.parameters(), self.actor_target.parameters()): target_param.data.copy_(tau_actor * param.data + (1 - tau_actor) * target_param.data) - def train(self, n_iterations, batch_size=100, discount=0.99, - tau=0.005, policy_noise=0.2, noise_clip=0.5, policy_freq=2): + def train(self, n_iterations, batch_size=100, policy_freq=2): for it in range(n_iterations): @@ -177,7 +176,7 @@ class TD3(BaseRLModel): if self.num_timesteps > 0: if self.verbose > 1: print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( - self.num_timesteps, episode_num, episode_timesteps, episode_reward)) + self.num_timesteps, episode_num, episode_timesteps, episode_reward)) self.train(episode_timesteps, batch_size=self.batch_size, policy_freq=self.policy_freq) # Evaluate episode