mirror of
https://github.com/wassname/DeepRL.git
synced 2026-08-22 11:40:47 +08:00
58 lines
2.2 KiB
Python
58 lines
2.2 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 ConvNet(nn.Module, VanillaNet):
|
|
def __init__(self, in_channels, action_dim, gpu=-1):
|
|
super(ConvNet, self).__init__()
|
|
self.body = NatureConvNet(in_channels)
|
|
VanillaNet.__init__(self, self.body.feature_dim, action_dim, gpu)
|
|
|
|
def feature(self, x):
|
|
x = self.variable(x)
|
|
return self.body(x)
|
|
|
|
class DuelingConvNet(nn.Module, DuelingNet):
|
|
def __init__(self, in_channels, action_dim, gpu=-1):
|
|
super(DuelingConvNet, self).__init__()
|
|
self.body = NatureConvNet(in_channels)
|
|
DuelingNet.__init__(self, self.body.feature_dim, action_dim, gpu)
|
|
|
|
def feature(self, x):
|
|
x = self.variable(x)
|
|
return self.body(x)
|
|
|
|
class ActorCriticConvNet(nn.Module, ActorCriticNet):
|
|
def __init__(self, in_channels, action_dim, gpu=-1):
|
|
super(ActorCriticConvNet, self).__init__()
|
|
self.body = NatureConvNet(in_channels)
|
|
ActorCriticNet.__init__(self, self.body.feature_dim, action_dim, gpu)
|
|
|
|
def feature(self, x):
|
|
x = self.variable(x)
|
|
return self.body(x)
|
|
|
|
class CategoricalConvNet(nn.Module, CategoricalNet):
|
|
def __init__(self, in_channels, n_actions, n_atoms, gpu=-1):
|
|
super(CategoricalConvNet, self).__init__()
|
|
self.body = NatureConvNet(in_channels)
|
|
CategoricalNet.__init__(self, self.body.feature_dim, n_actions, n_atoms, gpu)
|
|
|
|
def feature(self, x):
|
|
x = self.variable(x)
|
|
return self.body(x)
|
|
|
|
class QuantileConvNet(nn.Module, QuantileNet):
|
|
def __init__(self, in_channels, n_actions, n_quantiles, gpu=-1):
|
|
super(QuantileConvNet, self).__init__()
|
|
self.body = NatureConvNet(in_channels)
|
|
QuantileNet.__init__(self, self.body.feature_dim, n_actions, n_quantiles, gpu)
|
|
|
|
def feature(self, x):
|
|
x = self.variable(x)
|
|
return self.body(x)
|