def network_fn(X, nenv=1): nbatch = X.shape[0] nsteps = nbatch // nenv h = nature_cnn(X, **conv_kwargs) M = tf.placeholder(tf.float32, [nbatch]) #mask (done t-1) S = tf.placeholder(tf.float32, [nenv, 2 * nlstm]) #states xs = batch_to_seq(h, nenv, nsteps) ms = batch_to_seq(M, nenv, nsteps) if layer_norm: h5, snew = utils.lnlstm(xs, ms, S, scope='lnlstm', nh=nlstm) else: h5, snew = utils.lstm(xs, ms, S, scope='lstm', nh=nlstm) h = seq_to_batch(h5) initial_state = np.zeros(S.shape.as_list(), dtype=float) return h, { 'S': S, 'M': M, 'state': snew, 'initial_state': initial_state }
def q_retrace(R, D, q_i, v, rho_i, nenvs, nsteps, gamma): """ Calculates q_retrace targets :param R: Rewards :param D: Dones :param q_i: Q values for actions taken :param v: V values :param rho_i: Importance weight for each action :return: Q_retrace values """ rho_bar = batch_to_seq(tf.minimum(1.0, rho_i), nenvs, nsteps, True) # list of len steps, shape [nenvs] rs = batch_to_seq(R, nenvs, nsteps, True) # list of len steps, shape [nenvs] ds = batch_to_seq(D, nenvs, nsteps, True) # list of len steps, shape [nenvs] q_is = batch_to_seq(q_i, nenvs, nsteps, True) vs = batch_to_seq(v, nenvs, nsteps + 1, True) v_final = vs[-1] qret = v_final qrets = [] for i in range(nsteps - 1, -1, -1): check_shape([qret, ds[i], rs[i], rho_bar[i], q_is[i], vs[i]], [[nenvs]] * 6) qret = rs[i] + gamma * qret * (1.0 - ds[i]) qrets.append(qret) qret = (rho_bar[i] * (qret - q_is[i])) + vs[i] qrets = qrets[::-1] qret = seq_to_batch(qrets, flat=True) return qret
def __init__(self, sess, ob_space, ac_space, nenv, nsteps, nstack, reuse=False, nlstm=256): nbatch = nenv * nsteps nh, nw, nc = ob_space.shape ob_shape = (nbatch, nh, nw, nc * nstack) nact = ac_space.n X = tf.placeholder(tf.uint8, ob_shape) # obs M = tf.placeholder(tf.float32, [nbatch]) # mask (done t-1) S = tf.placeholder(tf.float32, [nenv, nlstm * 2]) # states with tf.variable_scope("model", reuse=reuse): h = nature_cnn(X) # lstm xs = batch_to_seq(h, nenv, nsteps) ms = batch_to_seq(M, nenv, nsteps) h5, snew = lstm(xs, ms, S, 'lstm1', nh=nlstm) h5 = seq_to_batch(h5) pi_logits = fc(h5, 'pi', nact, init_scale=0.01) pi = tf.nn.softmax(pi_logits) q = fc(h5, 'q', nact) a = sample(pi_logits) # could change this to use self.pi instead self.initial_state = np.zeros((nenv, nlstm * 2), dtype=np.float32) self.X = X self.M = M self.S = S self.pi = pi # actual policy params now self.q = q def step(ob, state, mask, *args, **kwargs): # returns actions, mus, states a0, pi0, s = sess.run([a, pi, snew], {X: ob, S: state, M: mask}) return a0, pi0, s self.step = step
def strip(var, nenvs, nsteps, flat=False): vars = batch_to_seq(var, nenvs, nsteps + 1, flat) return seq_to_batch(vars[:-1], flat)