def __init__(self, state_space, base_channels, out_features):
super().__init__()
self.embedding = nn.Embedding(9, base_channels * 2**0)
self.conv = nn.Sequential(
nn.Conv2d(base_channels * 2**0,
base_channels * 2**1,
3,
bias=False),
BatchRenorm2d(base_channels * 2**1),
Activation(),
nn.Conv2d(base_channels * 2**1,
base_channels * 2**2,
3,
bias=False),
BatchRenorm2d(base_channels * 2**2),
Activation(),
nn.Conv2d(base_channels * 2**2,
base_channels * 2**3,
3,
bias=False),
BatchRenorm2d(base_channels * 2**3),
Activation(),
)
self.output = nn.Sequential(
nn.Linear(base_channels * 2**3, out_features), Activation())
def __init__(self, state_space, out_features):
assert len(state_space.shape) == 1
super().__init__()
self.layers = nn.Sequential(
nn.Linear(state_space.shape[0], out_features), Activation())
def __init__(self, in_features, action_space):
super().__init__()
self.layers = nn.Sequential(
nn.Linear(in_features, in_features),
Activation(),
nn.Linear(in_features, action_space.n),
)
def __init__(self, in_features, action_space):
super().__init__()
assert np.array_equal(action_space.low,
np.zeros_like(action_space.low))
assert np.array_equal(action_space.high,
np.ones_like(action_space.high))
self.layers = nn.Sequential(
nn.Linear(in_features, in_features),
Activation(),
nn.Linear(in_features,
np.prod(action_space.shape) * 2),
)
def __init__(self, state_space, base_channels, out_features):
assert len(state_space.shape) == 3
super().__init__()
self.layers = nn.Sequential(
# nn.BatchNorm2d(state_space.shape[2]),
nn.Conv2d(state_space.shape[2],
base_channels * 2**2,
7,
stride=2,
padding=3),
Activation(),
nn.MaxPool2d(3, 2),
nn.Conv2d(base_channels * 2**2,
base_channels * 2**3,
3,
stride=2,
padding=1),
Activation(),
nn.Conv2d(base_channels * 2**3,
base_channels * 2**4,
3,
stride=2,
padding=1),
Activation(),
nn.Conv2d(base_channels * 2**4,
base_channels * 2**5,
3,
stride=2,
padding=1),
Activation(),
)
self.pool = nn.AdaptiveMaxPool2d(1)
self.output = nn.Sequential(
nn.Linear(base_channels * 2**5, out_features), Activation())