Files
DeepRL/network/shallow_network.py
T

60 lines
2.3 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 .base_network import *
class FCNet(nn.Module, VanillaNet):
def __init__(self, state_dim, hidden_size, action_dim, gpu=-1):
super(FCNet, self).__init__()
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
VanillaNet.__init__(self, hidden_size, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class DuelingFCNet(nn.Module, DuelingNet):
def __init__(self, state_dim, hidden_size, action_dim, gpu=-1):
super(DuelingFCNet, self).__init__()
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
DuelingNet.__init__(self, hidden_size, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class ActorCriticFCNet(nn.Module, ActorCriticNet):
def __init__(self, state_dim, hidden_size, action_dim, gpu=-1):
super(ActorCriticFCNet, self).__init__()
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
ActorCriticNet.__init__(self, hidden_size, action_dim, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class CategoricalFCNet(nn.Module, CategoricalNet):
def __init__(self, state_dim, n_actions, n_atoms, gpu=-1):
super(CategoricalFCNet, self).__init__()
hidden_size = 64
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
CategoricalNet.__init__(self, hidden_size, n_actions, n_atoms, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)
class QuantileFCNet(nn.Module, QuantileNet):
def __init__(self, state_dim, n_actions, n_quantiles, gpu=-1):
super(QuantileFCNet, self).__init__()
hidden_size = 64
self.fc_body = TwoLayerFCNet(state_dim, hidden_size)
QuantileNet.__init__(self, hidden_size, n_actions, n_quantiles, gpu)
def feature(self, x):
x = self.variable(x)
return self.fc_body(x)