forked from keon/deep-q-learning
-
Notifications
You must be signed in to change notification settings - Fork 1
/
models.py
34 lines (25 loc) · 939 Bytes
/
models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import numpy as np
from keras.models import Sequential
from keras.layers import Dense
from keras.optimizers import Adam
class Model:
def build(self):
raise NotImplementedError
def predict(self, x):
raise NotImplementedError
def train(self, x, y, config):
raise NotImplementedError
class Cartpole_nn(Model):
def __init__(self, config):
self.model = self.build(config)
def build(self, config):
model = Sequential()
model.add(Dense(24, input_dim=config.state_size, activation="relu"))
model.add(Dense(24, activation="relu"))
model.add(Dense(config.action_size, activation="linear"))
model.compile(loss=config.loss, optimizer=config.optimizer)
return model
def predict(self, x):
return self.model.predict(x)
def train(self, x, y, **kwargs):
return self.model.fit(np.array(x), np.array(y), epochs=1, verbose=0)