mirror of
https://github.com/wassname/DeepRL.git
synced 2026-08-21 11:09:46 +08:00
Mike Clark
ed943a6356
A really minor change. Should we input logits instead of probs into `Categorical`? It will do less converting from logits and back, leading to increased stability during training (theoretically).
197 lines
7.8 KiB
Python
197 lines
7.8 KiB
Python
#######################################################################
|
|
# Copyright (C) 2017 Shangtong Zhang(zhangshangtong.cpp@gmail.com) #
|
|
# Permission given to modify the code as long as you keep this #
|
|
# declaration at the top #
|
|
#######################################################################
|
|
|
|
from .network_utils import *
|
|
from .network_bodies import *
|
|
|
|
class VanillaNet(nn.Module, BaseNet):
|
|
def __init__(self, output_dim, body, gpu=-1):
|
|
super(VanillaNet, self).__init__()
|
|
self.fc_head = layer_init(nn.Linear(body.feature_dim, output_dim))
|
|
self.body = body
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, x, to_numpy=False):
|
|
phi = self.body(self.tensor(x))
|
|
y = self.fc_head(phi)
|
|
if to_numpy:
|
|
y = y.cpu().detach().numpy()
|
|
return y
|
|
|
|
class DuelingNet(nn.Module, BaseNet):
|
|
def __init__(self, action_dim, body, gpu=-1):
|
|
super(DuelingNet, self).__init__()
|
|
self.fc_value = layer_init(nn.Linear(body.feature_dim, 1))
|
|
self.fc_advantage = layer_init(nn.Linear(body.feature_dim, action_dim))
|
|
self.body = body
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, x, to_numpy=False):
|
|
phi = self.body(self.tensor(x))
|
|
value = self.fc_value(phi)
|
|
advantange = self.fc_advantage(phi)
|
|
q = value.expand_as(advantange) + (advantange - advantange.mean(1, keepdim=True).expand_as(advantange))
|
|
if to_numpy:
|
|
return q.cpu().detach().numpy()
|
|
return q
|
|
|
|
class CategoricalNet(nn.Module, BaseNet):
|
|
def __init__(self, action_dim, num_atoms, body, gpu=-1):
|
|
super(CategoricalNet, self).__init__()
|
|
self.fc_categorical = layer_init(nn.Linear(body.feature_dim, action_dim * num_atoms))
|
|
self.action_dim = action_dim
|
|
self.num_atoms = num_atoms
|
|
self.body = body
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, x, to_numpy=False):
|
|
phi = self.body(self.tensor(x))
|
|
pre_prob = self.fc_categorical(phi).view((-1, self.action_dim, self.num_atoms))
|
|
prob = F.softmax(pre_prob, dim=-1)
|
|
if to_numpy:
|
|
return prob.cpu().detach().numpy()
|
|
return prob
|
|
|
|
class QuantileNet(nn.Module, BaseNet):
|
|
def __init__(self, action_dim, num_quantiles, body, gpu=-1):
|
|
super(QuantileNet, self).__init__()
|
|
self.fc_quantiles = layer_init(nn.Linear(body.feature_dim, action_dim * num_quantiles))
|
|
self.action_dim = action_dim
|
|
self.num_quantiles = num_quantiles
|
|
self.body = body
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, x, to_numpy=False):
|
|
phi = self.body(self.tensor(x))
|
|
quantiles = self.fc_quantiles(phi)
|
|
quantiles = quantiles.view((-1, self.action_dim, self.num_quantiles))
|
|
if to_numpy:
|
|
quantiles = quantiles.cpu().detach().numpy()
|
|
return quantiles
|
|
|
|
class OptionCriticNet(nn.Module, BaseNet):
|
|
def __init__(self, body, action_dim, num_options, gpu=-1):
|
|
super(OptionCriticNet, self).__init__()
|
|
self.fc_q = layer_init(nn.Linear(body.feature_dim, num_options))
|
|
self.fc_pi = layer_init(nn.Linear(body.feature_dim, num_options * action_dim))
|
|
self.fc_beta = layer_init(nn.Linear(body.feature_dim, num_options))
|
|
self.num_options = num_options
|
|
self.action_dim = action_dim
|
|
self.body = body
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, x):
|
|
phi = self.body(self.tensor(x))
|
|
q = self.fc_q(phi)
|
|
beta = F.sigmoid(self.fc_beta(phi))
|
|
pi = self.fc_pi(phi)
|
|
pi = pi.view(-1, self.num_options, self.action_dim)
|
|
log_pi = F.log_softmax(pi, dim=-1)
|
|
return q, beta, log_pi
|
|
|
|
class ActorCriticNet(nn.Module):
|
|
def __init__(self, state_dim, action_dim, phi_body, actor_body, critic_body):
|
|
super(ActorCriticNet, self).__init__()
|
|
if phi_body is None: phi_body = DummyBody(state_dim)
|
|
if actor_body is None: actor_body = DummyBody(phi_body.feature_dim)
|
|
if critic_body is None: critic_body = DummyBody(phi_body.feature_dim)
|
|
self.phi_body = phi_body
|
|
self.actor_body = actor_body
|
|
self.critic_body = critic_body
|
|
self.fc_action = layer_init(nn.Linear(actor_body.feature_dim, action_dim), 1e-3)
|
|
self.fc_critic = layer_init(nn.Linear(critic_body.feature_dim, 1), 1e-3)
|
|
|
|
self.actor_params = list(self.actor_body.parameters()) + list(self.fc_action.parameters())
|
|
self.critic_params = list(self.critic_body.parameters()) + list(self.fc_critic.parameters())
|
|
self.phi_params = list(self.phi_body.parameters())
|
|
|
|
class DeterministicActorCriticNet(nn.Module, BaseNet):
|
|
def __init__(self,
|
|
state_dim,
|
|
action_dim,
|
|
actor_opt_fn,
|
|
critic_opt_fn,
|
|
phi_body=None,
|
|
actor_body=None,
|
|
critic_body=None,
|
|
gpu=-1):
|
|
super(DeterministicActorCriticNet, self).__init__()
|
|
self.network = ActorCriticNet(state_dim, action_dim, phi_body, actor_body, critic_body)
|
|
self.actor_opt = actor_opt_fn(self.network.actor_params + self.network.phi_params)
|
|
self.critic_opt = critic_opt_fn(self.network.critic_params + self.network.phi_params)
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, obs, to_numpy=False):
|
|
phi = self.feature(obs)
|
|
action = self.actor(phi)
|
|
if to_numpy:
|
|
return action.cpu().detach().numpy()
|
|
return action
|
|
|
|
def feature(self, obs):
|
|
obs = self.tensor(obs)
|
|
return self.network.phi_body(obs)
|
|
|
|
def actor(self, phi):
|
|
return F.tanh(self.network.fc_action(self.network.actor_body(phi)))
|
|
|
|
def critic(self, phi, a):
|
|
return self.network.fc_critic(self.network.critic_body(phi, a))
|
|
|
|
class GaussianActorCriticNet(nn.Module, BaseNet):
|
|
def __init__(self,
|
|
state_dim,
|
|
action_dim,
|
|
phi_body=None,
|
|
actor_body=None,
|
|
critic_body=None,
|
|
gpu=-1):
|
|
super(GaussianActorCriticNet, self).__init__()
|
|
self.network = ActorCriticNet(state_dim, action_dim, phi_body, actor_body, critic_body)
|
|
self.std = nn.Parameter(torch.ones(1, action_dim))
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, obs, action=None, to_numpy=False):
|
|
obs = self.tensor(obs)
|
|
phi = self.network.phi_body(obs)
|
|
phi_a = self.network.actor_body(phi)
|
|
phi_v = self.network.critic_body(phi)
|
|
mean = F.tanh(self.network.fc_action(phi_a))
|
|
if to_numpy:
|
|
return mean.cpu().detach().numpy()
|
|
v = self.network.fc_critic(phi_v)
|
|
dist = torch.distributions.Normal(mean, self.std)
|
|
if action is None:
|
|
action = dist.sample()
|
|
log_prob = dist.log_prob(action)
|
|
log_prob = torch.sum(log_prob, dim=1, keepdim=True)
|
|
return action, log_prob, self.tensor(np.zeros((log_prob.size(0), 1))), v
|
|
|
|
class CategoricalActorCriticNet(nn.Module, BaseNet):
|
|
def __init__(self,
|
|
state_dim,
|
|
action_dim,
|
|
phi_body=None,
|
|
actor_body=None,
|
|
critic_body=None,
|
|
gpu=-1):
|
|
super(CategoricalActorCriticNet, self).__init__()
|
|
self.network = ActorCriticNet(state_dim, action_dim, phi_body, actor_body, critic_body)
|
|
self.set_gpu(gpu)
|
|
|
|
def predict(self, obs, action=None):
|
|
obs = self.tensor(obs)
|
|
phi = self.network.phi_body(obs)
|
|
phi_a = self.network.actor_body(phi)
|
|
phi_v = self.network.critic_body(phi)
|
|
logits = self.network.fc_action(phi_a)
|
|
v = self.network.fc_critic(phi_v)
|
|
dist = torch.distributions.Categorical(logits=logits)
|
|
if action is None:
|
|
action = dist.sample()
|
|
log_prob = dist.log_prob(action).unsqueeze(-1)
|
|
return action, log_prob, dist.entropy().unsqueeze(-1), v
|