mirror of
https://github.com/wassname/DeepRL.git
synced 2026-10-29 12:40:58 +08:00
Major refactor to support LSTM layer
This commit is contained in:
Shangtong Zhang committed
1 parent
af17bbbf94
commit
9c73abfbff
8 files changed
+370
-204
No files matched your search
+127
-38
@@ -12,63 +12,57 @@ import numpy as np
|
||||
|
||||
# Base class for all kinds of network
|
||||
class BasicNet:
|
||||
def __init__(self, optimizer_fn, gpu):
|
||||
def __init__(self, optimizer_fn, gpu, LSTM=False):
|
||||
if optimizer_fn is not None:
|
||||
self.optimizer = optimizer_fn(self.parameters())
|
||||
self.gpu = gpu and torch.cuda.is_available()
|
||||
self.LSTM = LSTM
|
||||
if self.gpu:
|
||||
print 'Transferring network to GPU...'
|
||||
self.cuda()
|
||||
print 'Network transferred.'
|
||||
|
||||
def to_torch_variable(self, x, dtype='float32'):
|
||||
x = torch.from_numpy(np.asarray(x, dtype=dtype))
|
||||
if not isinstance(x, torch.FloatTensor):
|
||||
x = torch.from_numpy(np.asarray(x, dtype=dtype))
|
||||
if self.gpu:
|
||||
x = x.cuda()
|
||||
return Variable(x)
|
||||
|
||||
def reset(self, terminal):
|
||||
if not self.LSTM:
|
||||
return
|
||||
if terminal:
|
||||
self.h.data.zero_()
|
||||
self.c.data.zero_()
|
||||
self.h = Variable(self.h.data)
|
||||
self.c = Variable(self.c.data)
|
||||
|
||||
# Base class for value based methods
|
||||
class VanillaNet(BasicNet):
|
||||
def predict(self, x, to_numpy=True):
|
||||
def predict(self, x, to_numpy=False):
|
||||
y = self.forward(x)
|
||||
if to_numpy:
|
||||
y = y.cpu().data.numpy()
|
||||
return y
|
||||
|
||||
def gradient(self, x, actions, targets):
|
||||
y = self.forward(x)
|
||||
y = y.gather(1, actions)
|
||||
loss = self.criterion(y, targets)
|
||||
loss.backward()
|
||||
|
||||
# Base class for actor critic method
|
||||
class ActorCriticNet(BasicNet):
|
||||
def predict(self, x):
|
||||
phi = self.forward(x)
|
||||
return F.softmax(self.fc_actor(phi)).cpu().data.numpy()
|
||||
|
||||
def gradient(self, x, actions, rewards):
|
||||
phi = self.forward(x)
|
||||
logit = self.fc_actor(phi)
|
||||
prob = F.softmax(logit)
|
||||
log_prob_ = F.log_softmax(logit)
|
||||
state_value = self.fc_critic(phi)
|
||||
log_prob = log_prob_.gather(1, actions)
|
||||
advantage = (rewards - state_value).detach()
|
||||
policy_loss = -torch.sum(log_prob * advantage)
|
||||
value_loss = 0.5 * torch.sum(torch.pow(rewards - state_value, 2))
|
||||
entropy = -torch.sum(torch.mul(prob, log_prob_))
|
||||
loss = policy_loss + value_loss - self.xentropy_weight * entropy
|
||||
loss.backward()
|
||||
nn.utils.clip_grad_norm(self.parameters(), self.grad_threshold)
|
||||
phi = self.forward(x, True)
|
||||
pre_prob = self.fc_actor(phi)
|
||||
prob = F.softmax(pre_prob)
|
||||
log_prob = F.log_softmax(pre_prob)
|
||||
value = self.fc_critic(phi)
|
||||
return prob, log_prob, value
|
||||
|
||||
def critic(self, x):
|
||||
phi = self.forward(x)
|
||||
return self.fc_critic(phi).cpu().data.numpy()
|
||||
phi = self.forward(x, False)
|
||||
return self.fc_critic(phi)
|
||||
|
||||
# Base class for dueling architecture
|
||||
class DuelingNet(BasicNet):
|
||||
def predict(self, x, to_numpy=True):
|
||||
def predict(self, x, to_numpy=False):
|
||||
phi = self.forward(x)
|
||||
value = self.fc_value(phi)
|
||||
advantange = self.fc_advantage(phi)
|
||||
@@ -77,6 +71,8 @@ class DuelingNet(BasicNet):
|
||||
return q.cpu().data.numpy()
|
||||
return q
|
||||
|
||||
# Starting of several network instances
|
||||
|
||||
# Network for CartPole with value based methods
|
||||
class FullyConnectedNet(nn.Module, VanillaNet):
|
||||
def __init__(self, dims, optimizer_fn=None, gpu=True):
|
||||
@@ -178,21 +174,30 @@ class DuelingConvNet(nn.Module, DuelingNet):
|
||||
class FCActorCriticNet(nn.Module, ActorCriticNet):
|
||||
def __init__(self,
|
||||
dims,
|
||||
xentropy_weight=0.01,
|
||||
grad_threshold=40,
|
||||
gpu=True):
|
||||
LSTM=False):
|
||||
super(FCActorCriticNet, self).__init__()
|
||||
self.fc1 = nn.Linear(dims[0], dims[1])
|
||||
if LSTM:
|
||||
self.layer1 = nn.LSTMCell(dims[0], dims[1])
|
||||
else:
|
||||
self.layer1 = nn.Linear(dims[0], dims[1])
|
||||
self.fc_actor = nn.Linear(dims[1], dims[2])
|
||||
self.fc_critic = nn.Linear(dims[1], 1)
|
||||
self.xentropy_weight = xentropy_weight
|
||||
self.grad_threshold = grad_threshold
|
||||
BasicNet.__init__(self, optimizer_fn=None, gpu=gpu)
|
||||
BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
|
||||
if LSTM:
|
||||
self.h = self.to_torch_variable(np.zeros((1, dims[1])))
|
||||
self.c = self.to_torch_variable(np.zeros((1, dims[1])))
|
||||
|
||||
def forward(self, x):
|
||||
def forward(self, x, update_LSTM=True):
|
||||
x = self.to_torch_variable(x)
|
||||
x = x.view(x.size(0), -1)
|
||||
phi = self.fc1(x)
|
||||
if self.LSTM:
|
||||
h, c = self.layer1(x, (self.h, self.c))
|
||||
if update_LSTM:
|
||||
self.h = h
|
||||
self.c = c
|
||||
phi = h
|
||||
else:
|
||||
phi = self.layer1(x)
|
||||
return phi
|
||||
|
||||
# Network for pixel Atari game with actor critic
|
||||
@@ -222,3 +227,87 @@ class ConvActorCriticNet(nn.Module, ActorCriticNet):
|
||||
y = y.view(y.size(0), -1)
|
||||
return F.elu(self.fc4(y))
|
||||
|
||||
class OpenAIConvActorCriticNet(nn.Module, ActorCriticNet):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
n_actions,
|
||||
LSTM=False):
|
||||
super(OpenAIConvActorCriticNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
|
||||
self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
|
||||
self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
|
||||
self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
|
||||
|
||||
self.LSTM = LSTM
|
||||
hidden_units = 256
|
||||
|
||||
if LSTM:
|
||||
self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units)
|
||||
else:
|
||||
self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
|
||||
|
||||
self.fc_actor = nn.Linear(hidden_units, n_actions)
|
||||
self.fc_critic = nn.Linear(hidden_units, 1)
|
||||
BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
|
||||
if LSTM:
|
||||
self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
|
||||
self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
|
||||
|
||||
def forward(self, x, update_LSTM=True):
|
||||
x = self.to_torch_variable(x)
|
||||
y = F.elu(self.conv1(x))
|
||||
y = F.elu(self.conv2(y))
|
||||
y = F.elu(self.conv3(y))
|
||||
y = F.elu(self.conv4(y))
|
||||
y = y.view(y.size(0), -1)
|
||||
if self.LSTM:
|
||||
h, c = self.layer5(y, (self.h, self.c))
|
||||
if update_LSTM:
|
||||
self.h = h
|
||||
self.c = c
|
||||
phi = h
|
||||
else:
|
||||
phi = F.elu(self.layer5(y))
|
||||
return phi
|
||||
|
||||
class OpenAIConvNet(nn.Module, VanillaNet):
|
||||
def __init__(self,
|
||||
in_channels,
|
||||
n_actions,
|
||||
LSTM=False):
|
||||
super(OpenAIConvNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(in_channels, 32, 3, stride=2, padding=1)
|
||||
self.conv2 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
|
||||
self.conv3 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
|
||||
self.conv4 = nn.Conv2d(32, 32, 3, stride=2, padding=1)
|
||||
|
||||
self.LSTM = LSTM
|
||||
hidden_units = 256
|
||||
|
||||
if LSTM:
|
||||
self.layer5 = nn.LSTMCell(32 * 3 * 3, hidden_units)
|
||||
else:
|
||||
self.layer5 = nn.Linear(32 * 3 * 3, hidden_units)
|
||||
|
||||
self.fc6 = nn.Linear(hidden_units, n_actions)
|
||||
BasicNet.__init__(self, optimizer_fn=None, gpu=False, LSTM=LSTM)
|
||||
if LSTM:
|
||||
self.h = self.to_torch_variable(np.zeros((1, hidden_units)))
|
||||
self.c = self.to_torch_variable(np.zeros((1, hidden_units)))
|
||||
|
||||
def forward(self, x, update_LSTM=True):
|
||||
x = self.to_torch_variable(x)
|
||||
y = F.elu(self.conv1(x))
|
||||
y = F.elu(self.conv2(y))
|
||||
y = F.elu(self.conv3(y))
|
||||
y = F.elu(self.conv4(y))
|
||||
y = y.view(y.size(0), -1)
|
||||
if self.LSTM:
|
||||
h, c = self.layer5(y, (self.h, self.c))
|
||||
if update_LSTM:
|
||||
self.h = h
|
||||
self.c = c
|
||||
phi = h
|
||||
else:
|
||||
phi = F.elu(self.layer5(y))
|
||||
return self.fc6(phi)
|
||||
Reference in new issue
Block a user