Bug fixes + add evaluate script

.coveragerc

@@ -0,0 +1,10 @@
+[run]
+branch = False
+omit =
+    tests/*
+    setup.py
+
+[report]
+exclude_lines =
+    pragma: no cover
+    raise NotImplementedError()

.gitignore

@@ -11,6 +11,7 @@
 __pycache__/
 _build/
 *.npz
+*.pth
 
 # Setuptools distribution and build folders.
 /dist/

scripts/run_tests.sh

@@ -0,0 +1,2 @@
+#!/bin/bash
+python -m pytest --cov-config .coveragerc --cov-report html --cov-report term --cov=. -v

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

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

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

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

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

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

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)

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,7 +39,7 @@ 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.policy = self.policy(self.observation_space, self.action_space,
                                   self.learning_rate, device=self.device, **self.policy_kwargs)
         self._create_aliases()
 
@@ -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:
+                if self.num_timesteps > 0:
+                    if self.verbose > 1:
                         print("Total T: {} Episode Num: {} Episode T: {} Reward: {}".format(
-                          num_timesteps, episode_num, episode_timesteps, episode_reward))
+                              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'):