| 22 | class CNN: |
| 23 | |
| 24 | def __init__(self, sess, ob_space, ac_space, nenv, nsteps, nstack, reuse=False): |
| 25 | gain = np.sqrt(2) |
| 26 | nbatch = nenv * nsteps |
| 27 | nh, nw, nc = ob_space.shape |
| 28 | ob_shape = (nbatch, nh, nw, nc * nstack) |
| 29 | X = tf.placeholder(tf.uint8, ob_shape) # obs |
| 30 | X_normal = tf.cast(X, tf.float32) / 255.0 |
| 31 | with tf.variable_scope("model", reuse=reuse): |
| 32 | h1 = conv(X_normal, 32, 8, 4, gain) |
| 33 | h2 = conv(h1, 64, 4, 2, gain) |
| 34 | h3 = conv(h2, 64, 3, 1, gain) |
| 35 | h3 = tf.layers.flatten(h3) |
| 36 | h4 = dense(h3, 512, gain=gain) |
| 37 | pi = dense(h4, ac_space.n, act=None) |
| 38 | vf = dense(h4, 1, act=None) |
| 39 | |
| 40 | v0 = vf[:, 0] |
| 41 | a0 = sample(pi) |
| 42 | # self.initial_state = [] # State reserved for LSTM |
| 43 | |
| 44 | def step(ob): |
| 45 | a, v = sess.run([a0, v0], {X: ob}) |
| 46 | return a, v#, [] # dummy state |
| 47 | |
| 48 | def value(ob): |
| 49 | return sess.run(v0, {X: ob}) |
| 50 | |
| 51 | self.X = X |
| 52 | self.pi = pi |
| 53 | self.vf = vf |
| 54 | self.step = step |
| 55 | self.value = value |