Refactor policies

tests/test_td3.py

@@ -13,7 +13,7 @@ def test_pendulum():
 
 
 def test_cemrl():
-    model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16, 16]), pop_size=2, n_grad=1,
+    model = CEMRL('MlpPolicy', 'Pendulum-v0', policy_kwargs=dict(net_arch=[16]), pop_size=2, n_grad=1,
                  start_timesteps=100, verbose=1)
     model.learn(total_timesteps=1000, eval_freq=500)
     model.save("test_save")

torchy_baselines/cem_rl/cem_rl.py

@@ -89,10 +89,14 @@ class CEMRL(TD3):
                         self.train_critic(actor_steps // self.n_grad)
                         self.train_actor(actor_steps)
                     else:
-                        # Closer to td3: policy delay and it scales
-                        # with a bigger population
-                        # but less training steps per agent
-                        for it in range(2 * (actor_steps // self.n_grad)):
+                        # Closer to td3: with policy delay
+                        if self.update_style == 'td3_like':
+                            n_training_steps = actor_steps
+                        else:
+                            # scales with a bigger population
+                            # but less training steps per agent
+                            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)
                             self.train_critic(replay_data=replay_data)

torchy_baselines/common/base_class.py

@@ -1,8 +1,9 @@
 from abc import ABCMeta, abstractmethod
 
 
-import numpy as np
 import gym
+import torch as th
+import numpy as np
 
 from torchy_baselines.common.policies import get_policy_from_name
 
@@ -16,14 +17,25 @@ class BaseRLModel(object):
                 (if registered in Gym, can be str. Can be None for loading trained models)
     :param verbose: (int) the verbosity level: 0 none, 1 training information, 2 debug
     :param policy_base: (BasePolicy) the base policy used by this method
+    :param device: (str or th.device) Device on which the code should.
+        By default, it will try to use a Cuda compatible device and fallback to cpu
+        if it is not possible.
     """
     __metaclass__ = ABCMeta
 
-    def __init__(self, policy, env, policy_base, policy_kwargs=None, verbose=0):
+    def __init__(self, policy, env, policy_base, policy_kwargs=None, verbose=0, device='auto'):
         if isinstance(policy, str) and policy_base is not None:
             self.policy = get_policy_from_name(policy_base, policy)
         else:
             self.policy = policy
+
+        if device == 'auto':
+            device = 'cuda' if th.cuda.is_available() else 'cpu'
+
+        self.device = th.device(device)
+        if verbose > 0:
+            print("Using {} device".format(self.device))
+
         self.env = env
         self.verbose = verbose
         self.policy_kwargs = {} if policy_kwargs is None else policy_kwargs

torchy_baselines/td3/policies.py

@@ -4,10 +4,24 @@ import torch.nn as nn
 from torchy_baselines.common.policies import BasePolicy, register_policy
 
 
+def create_mlp(input_dim, output_dim, net_arch,
+               activation_fn=nn.ReLU, squash_out=False):
+    modules = [nn.Linear(input_dim, net_arch[0]), activation_fn()]
+
+    for idx in range(len(net_arch) - 1):
+        modules.append(nn.Linear(net_arch[idx], net_arch[idx + 1]))
+        modules.append(activation_fn())
+
+    modules.append(nn.Linear(net_arch[-1], output_dim))
+    if squash_out:
+        modules.append(nn.Tanh())
+    return modules
+
+
 class BaseNetwork(nn.Module):
     """docstring for BaseNetwork."""
 
-    def __init__(self, device='cpu'):
+    def __init__(self):
         super(BaseNetwork, self).__init__()
 
     def load_from_vector(self, vector):
@@ -36,49 +50,41 @@ class Actor(BaseNetwork):
             net_arch = [400, 300]
 
         # TODO: orthogonal initialization?
-
-        self.actor_net = nn.Sequential(
-            nn.Linear(state_dim, net_arch[0]),
-            activation_fn(),
-            nn.Linear(net_arch[0], net_arch[1]),
-            activation_fn(),
-            nn.Linear(net_arch[1], action_dim),
-            nn.Tanh(),
-        )
+        actor_net = create_mlp(state_dim, action_dim, net_arch, activation_fn, squash_out=True)
+        self.actor_net = nn.Sequential(*actor_net)
 
     def forward(self, x):
         return self.actor_net(x)
 
 
 class Critic(BaseNetwork):
-    def __init__(self, state_dim, action_dim, net_arch=None, activation_fn=nn.ReLU):
+    def __init__(self, state_dim, action_dim,
+                 net_arch=None, activation_fn=nn.ReLU):
         super(Critic, self).__init__()
 
         if net_arch is None:
             net_arch = [400, 300]
 
-        self.q1_net = nn.Sequential(
-            nn.Linear(state_dim + action_dim, net_arch[0]),
-            activation_fn(),
-            nn.Linear(net_arch[0], net_arch[1]),
-            activation_fn(),
-            nn.Linear(net_arch[1], 1),
-        )
+        # TODO: solve  pytorch parameter registration
+        # for _ in range(n_critics):
+        #     q_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn)
+        #     self.q_net = nn.Sequential(*q_net)
+        #     self.q_networks.append(self.q_net)
 
-        self.q2_net = nn.Sequential(
-            nn.Linear(state_dim + action_dim, net_arch[0]),
-            activation_fn(),
-            nn.Linear(net_arch[0], net_arch[1]),
-            activation_fn(),
-            nn.Linear(net_arch[1], 1),
-        )
+        q1_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn)
+        self.q1_net = nn.Sequential(*q1_net)
+
+        q2_net = create_mlp(state_dim + action_dim, 1, net_arch, activation_fn)
+        self.q2_net = nn.Sequential(*q2_net)
+
+        self.q_networks = [self.q1_net, self.q2_net]
 
     def forward(self, obs, action):
         qvalue_input = th.cat([obs, action], dim=1)
-        return self.q1_net(qvalue_input), self.q2_net(qvalue_input)
+        return [q_net(qvalue_input) for q_net in self.q_networks]
 
     def q1_forward(self, obs, action):
-        return self.q1_net(th.cat([obs, action], dim=1))
+        return self.q_networks[0](th.cat([obs, action], dim=1))
 
 
 class TD3Policy(BasePolicy):

torchy_baselines/td3/td3.py

@@ -24,14 +24,10 @@ class TD3(BaseRLModel):
                  batch_size=100,
                 _init_setup_model=True):
 
-        super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose)
-
-        if device == 'auto':
-            device = 'cuda' if th.cuda.is_available() else 'cpu'
+        super(TD3, self).__init__(policy, env, TD3Policy, policy_kwargs, verbose, device)
 
         self.max_action = np.abs(self.action_space.high)
         self.replay_buffer = None
-        self.device = device
         self.action_noise_std = action_noise_std
         self.learning_rate = learning_rate
         self.buffer_size = buffer_size
@@ -45,7 +41,7 @@ class TD3(BaseRLModel):
 
     def _setup_model(self, seed=None):
         state_dim, action_dim = self.observation_space.shape[0], self.action_space.shape[0]
-        set_random_seed(self.seed, using_cuda=self.device != 'cpu')
+        set_random_seed(self.seed, using_cuda=self.device == th.device('cuda'))
 
         if self.env is not None:
             self.env.seed(self.seed)