diff --git a/tests/test_td3.py b/tests/test_td3.py index 4bc8e6e..089e79f 100644 --- a/tests/test_td3.py +++ b/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") diff --git a/torchy_baselines/cem_rl/cem_rl.py b/torchy_baselines/cem_rl/cem_rl.py index 1dc594b..ef27124 100644 --- a/torchy_baselines/cem_rl/cem_rl.py +++ b/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) diff --git a/torchy_baselines/common/base_class.py b/torchy_baselines/common/base_class.py index 2f5bc78..a7f6bdb 100644 --- a/torchy_baselines/common/base_class.py +++ b/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 diff --git a/torchy_baselines/td3/policies.py b/torchy_baselines/td3/policies.py index 8279778..a886af7 100644 --- a/torchy_baselines/td3/policies.py +++ b/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): diff --git a/torchy_baselines/td3/td3.py b/torchy_baselines/td3/td3.py index 4a037cb..6cee111 100644 --- a/torchy_baselines/td3/td3.py +++ b/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)