From 9cf289b997320168cc026302d0275c6ad270dd8f Mon Sep 17 00:00:00 2001 From: Antonin Raffin Date: Fri, 6 Sep 2019 10:44:55 +0200 Subject: [PATCH] Bug fixes + add evaluate script --- .coveragerc | 10 ++++++ .gitignore | 1 + scripts/run_tests.sh | 2 ++ setup.cfg | 12 +++++++ tests/test_td3.py | 11 ++++-- torchy_baselines/common/base_class.py | 9 +++-- torchy_baselines/common/evaluation.py | 21 +++++++++++ torchy_baselines/common/policies.py | 11 +++--- torchy_baselines/common/replay_buffer.py | 35 ++----------------- torchy_baselines/td3/policies.py | 6 +++- torchy_baselines/td3/td3.py | 44 +++++++++++++++--------- 11 files changed, 97 insertions(+), 65 deletions(-) create mode 100644 .coveragerc create mode 100755 scripts/run_tests.sh create mode 100644 setup.cfg diff --git a/.coveragerc b/.coveragerc new file mode 100644 index 0000000..4e5d7bd --- /dev/null +++ b/.coveragerc @@ -0,0 +1,10 @@ +[run] +branch = False +omit = + tests/* + setup.py + +[report] +exclude_lines = + pragma: no cover + raise NotImplementedError() diff --git a/.gitignore b/.gitignore index 5a26d08..ad72a2e 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ __pycache__/ _build/ *.npz +*.pth # Setuptools distribution and build folders. /dist/ diff --git a/scripts/run_tests.sh b/scripts/run_tests.sh new file mode 100755 index 0000000..d765f3a --- /dev/null +++ b/scripts/run_tests.sh @@ -0,0 +1,2 @@ +#!/bin/bash +python -m pytest --cov-config .coveragerc --cov-report html --cov-report term --cov=. -v diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..346ddf6 --- /dev/null +++ b/setup.cfg @@ -0,0 +1,12 @@ +[metadata] +# This includes the license file in the wheel. +license_file = LICENSE + +[tool:pytest] +# Deterministic ordering for tests; useful for pytest-xdist. +env = + PYTHONHASHSEED=0 +filterwarnings = + # Gym warnings + ignore:Parameters to load are deprecated.:DeprecationWarning + ignore:the imp module is deprecated in favour of importlib:PendingDeprecationWarning diff --git a/tests/test_td3.py b/tests/test_td3.py index d29d007..ad2a116 100644 --- a/tests/test_td3.py +++ b/tests/test_td3.py @@ -1,8 +1,13 @@ +import os + import gym from torchy_baselines import TD3 -def test_simple_run(): +def test_pendulum(): env = gym.make("Pendulum-v0") - model = TD3('MlpPolicy', env, policy_kwargs=dict(net_arch=[64, 64]), verbose=1) - model.learn(total_timesteps=50000) + model = TD3('MlpPolicy', env, policy_kwargs=dict(net_arch=[64, 64]), start_timesteps=100, verbose=1) + model.learn(total_timesteps=500, eval_freq=100) + model.save("test_save") + model.load("test_save") + os.remove("test_save.pth") diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index c7eaee4..1e78608 100644 --- a/torchy_baselines/common/base_class.py +++ b/torchy_baselines/common/base_class.py @@ -19,11 +19,10 @@ class BaseRLModel(ABC): """ def __init__(self, policy, env, policy_base, policy_kwargs=None, verbose=0): - # if isinstance(policy, str) and policy_base is not None: - # self.policy = get_policy_from_name(policy_base, policy) - # else: - # self.policy = policy - self.policy = None + if isinstance(policy, str) and policy_base is not None: + self.policy = get_policy_from_name(policy_base, policy) + else: + self.policy = policy self.env = env self.verbose = verbose self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs diff --git a/torchy_baselines/common/evaluation.py b/torchy_baselines/common/evaluation.py index e69de29..fb3f37f 100644 --- a/torchy_baselines/common/evaluation.py +++ b/torchy_baselines/common/evaluation.py @@ -0,0 +1,21 @@ +import numpy as np + + +def evaluate_policy(model, env, n_eval_episodes=10, deterministic=True, render=False): + """ + Runs policy for n episodes and returns average reward + """ + mean_reward = 0.0 + for _ in range(n_eval_episodes): + obs = env.reset() + done = False + while not done: + action = model.predict(np.array(obs), deterministic=deterministic) + obs, reward, done, _ = env.step(action) + mean_reward += reward + if render: + env.render() + + mean_reward /= n_eval_episodes + + return mean_reward diff --git a/torchy_baselines/common/policies.py b/torchy_baselines/common/policies.py index 417d3e8..8047c43 100644 --- a/torchy_baselines/common/policies.py +++ b/torchy_baselines/common/policies.py @@ -1,7 +1,7 @@ -from abc import ABC +import torch.nn as nn -class BasePolicy(ABC): +class BasePolicy(nn.Module): """ The base policy object @@ -10,16 +10,13 @@ class BasePolicy(ABC): """ def __init__(self, observation_space, action_space, device='cpu'): + super(BasePolicy, self).__init__() self.observation_space = observation_space self.action_space = action_space self.device = device -_policy_registry = { - # ActorCriticPolicy: { - # "MlpPolicy": MlpPolicy, - # } -} +_policy_registry = dict() def get_policy_from_name(base_policy_type, name): diff --git a/torchy_baselines/common/replay_buffer.py b/torchy_baselines/common/replay_buffer.py index 1d15b4b..b5a953c 100644 --- a/torchy_baselines/common/replay_buffer.py +++ b/torchy_baselines/common/replay_buffer.py @@ -1,40 +1,11 @@ import numpy as np import torch as th -# Code based on: -# https://github.com/openai/baselines/blob/master/baselines/deepq/replay_buffer.py - -# Expects tuples of (state, next_state, action, reward, done) -# class ReplayBuffer(object): -# def __init__(self, max_size=1e6): -# self.storage = [] -# self.max_size = max_size -# self.ptr = 0 -# -# def add(self, data): -# if len(self.storage) == self.max_size: -# self.storage[int(self.ptr)] = data -# self.ptr = (self.ptr + 1) % self.max_size -# else: -# self.storage.append(data) -# -# def sample(self, batch_size): -# ind = np.random.randint(0, len(self.storage), size=batch_size) -# x, y, u, r, d = [], [], [], [], [] -# -# for i in ind: -# X, Y, U, R, D = self.storage[i] -# x.append(np.array(X, copy=False)) -# y.append(np.array(Y, copy=False)) -# u.append(np.array(U, copy=False)) -# r.append(np.array(R, copy=False)) -# d.append(np.array(D, copy=False)) -# -# return np.array(x), np.array(y), np.array(u), np.array(r).reshape(-1, 1), np.array(d).reshape(-1, 1) - class ReplayBuffer(object): - + """ + Taken from https://github.com/apourchot/CEM-RL + """ def __init__(self, buffer_size, state_dim, action_dim, device='cpu'): super(ReplayBuffer, self).__init__() # params diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index b9a8c27..a1dcd42 100644 --- a/torchy_baselines/td3/policies.py +++ b/torchy_baselines/td3/policies.py @@ -1,7 +1,7 @@ import torch as th import torch.nn as nn -from torchy_baselines.common.policies import BasePolicy +from torchy_baselines.common.policies import BasePolicy, register_policy class Actor(nn.Module): @@ -80,3 +80,7 @@ class TD3Policy(BasePolicy): def make_critic(self): return Critic(self.state_dim, self.action_dim, self.net_arch).to(self.device) + +MlpPolicy = TD3Policy + +register_policy("MlpPolicy", MlpPolicy) diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 6656aa2..5702e87 100644 --- a/torchy_baselines/td3/td3.py +++ b/torchy_baselines/td3/td3.py @@ -5,6 +5,7 @@ import numpy as np from torchy_baselines.common.base_class import BaseRLModel from torchy_baselines.common.replay_buffer import ReplayBuffer from torchy_baselines.common.utils import set_random_seed +from torchy_baselines.common.evaluation import evaluate_policy from torchy_baselines.td3.policies import TD3Policy @@ -17,13 +18,12 @@ class TD3(BaseRLModel): def __init__(self, policy, env, policy_kwargs=None, verbose=0, buffer_size=int(1e6), learning_rate=1e-3, seed=0, device='cpu', - action_noise_std=0.1, start_timesteps=10000, _init_setup_model=True): + action_noise_std=0.1, start_timesteps=100, _init_setup_model=True): super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose) self.max_action = float(self.action_space.high) self.replay_buffer = None - self.policy = None self.device = device self.action_noise_std = action_noise_std self.learning_rate = learning_rate @@ -39,8 +39,8 @@ class TD3(BaseRLModel): set_random_seed(self.seed, using_cuda=self.device != 'cpu') self.replay_buffer = ReplayBuffer(self.buffer_size, state_dim, action_dim, self.device) - self.policy = TD3Policy(self.observation_space, self.action_space, - self.learning_rate, device=self.device, **self.policy_kwargs) + self.policy = self.policy(self.observation_space, self.action_space, + self.learning_rate, device=self.device, **self.policy_kwargs) self._create_aliases() def _create_aliases(self): @@ -75,7 +75,7 @@ class TD3(BaseRLModel): state, action, next_state, done, reward = self.replay_buffer.sample(batch_size) # Select action according to policy and add clipped noise - noise = action.data.normal_(0, policy_noise).to(self.device) + noise = action.clone().data.normal_(0, policy_noise) noise = noise.clamp(-noise_clip, noise_clip) next_action = (self.actor_target(next_state) + noise).clamp(-1, 1) @@ -114,24 +114,33 @@ class TD3(BaseRLModel): target_param.data.copy_(tau * param.data + (1 - tau) * target_param.data) def learn(self, total_timesteps, callback=None, seed=None, log_interval=100, - tb_log_name="TD3", reset_num_timesteps=True): - num_timesteps = 0 + eval_freq=-1, n_eval_episodes=5, tb_log_name="TD3", reset_num_timesteps=True): + timesteps_since_eval = 0 episode_num = 0 done = True + evaluations = [] - while num_timesteps < total_timesteps: + while self.num_timesteps < total_timesteps: + + if callback is not None: + # Only stop training if return value is False, not when it is None. + if callback(locals(), globals()) is False: + break if done: - if num_timesteps > 0: - print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format( - num_timesteps, episode_num, episode_timesteps, episode_reward)) + 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.train(episode_timesteps) # Evaluate episode - # if timesteps_since_eval >= args.eval_freq: - # timesteps_since_eval %= args.eval_freq - # evaluations.append(evaluate_policy(policy)) + if eval_freq > 0 and timesteps_since_eval >= eval_freq: + timesteps_since_eval %= eval_freq + evaluations.append(evaluate_policy(self, self.env, n_eval_episodes)) + if self.verbose > 0: + print("Eval num_timesteps={}, mean_reward={:.2f}".format(self.num_timesteps, evaluations[-1])) # Reset environment obs = self.env.reset() @@ -140,10 +149,10 @@ class TD3(BaseRLModel): episode_num += 1 # Select action randomly or according to policy - if num_timesteps < self.start_timesteps: + if self.num_timesteps < self.start_timesteps: action = self.env.action_space.sample() else: - action = self.policy.select_action(np.array(obs)) + action = self.select_action(np.array(obs)) if self.action_noise_std > 0: # NOTE: in the original implementation, the noise is applied to the unscaled action @@ -162,8 +171,9 @@ class TD3(BaseRLModel): obs = new_obs episode_timesteps += 1 - num_timesteps += 1 + self.num_timesteps += 1 timesteps_since_eval += 1 + return self def save(self, path): if not path.endswith('.pth'):