/
nn.py
27 lines (20 loc) · 810 Bytes
/
nn.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
import numpy as np
import tensorflow as tf
from keras.layers import Dense, Dropout, LeakyReLU
from keras.models import Sequential
class FullyConnectedMLP(object):
"""Vanilla two hidden layer multi-layer perceptron"""
def __init__(self, obs_shape, act_shape, h_size=64):
input_dim = np.prod(obs_shape) + np.prod(act_shape)
self.model = Sequential()
self.model.add(Dense(h_size, input_dim=input_dim))
self.model.add(LeakyReLU())
self.model.add(Dropout(0.5))
self.model.add(Dense(h_size))
self.model.add(LeakyReLU())
self.model.add(Dropout(0.5))
self.model.add(Dense(1))
def run(self, obs, act):
flat_obs = tf.contrib.layers.flatten(obs)
x = tf.concat([flat_obs, act], axis=1)
return self.model(x)