class blr_model: def __init__(self, x_dim, y_dim, state_dim, action_dim, observation_space_low, observation_space_high, action_bound_low, action_bound_high, unroll_steps, no_samples, no_basis, discount_factor, train_policy_batch_size, train_policy_iterations, hyperparameters, debugging_plot): assert x_dim == state_dim + action_dim assert len(hyperparameters) == y_dim self.x_dim = x_dim self.y_dim = y_dim self.state_dim = state_dim self.action_dim = action_dim self.observation_space_low = observation_space_low self.observation_space_high = observation_space_high self.action_bound_low = action_bound_low self.action_bound_high = action_bound_high self.unroll_steps = unroll_steps self.no_samples = no_samples self.no_basis = no_basis self.discount_factor = discount_factor self.train_policy_batch_size = train_policy_batch_size self.train_policy_iterations = train_policy_iterations self.hyperparameters = hyperparameters self.debugging_plot = debugging_plot self.policy_scope = 'policy_scope' self.policy_reuse_vars = None self.models = [ bayesian_model(self.x_dim, self.observation_space_low, self.observation_space_high, self.action_bound_low, self.action_bound_high, self.no_basis, *self.hyperparameters[i]) for i in range(self.y_dim) ] self.states = tf.placeholder(shape=[None, self.state_dim], dtype=tf.float64) self.batch_size = tf.shape(self.states)[0] #self.batch_size = 3 self.actions = self.build_policy(self.states) self.cum_xx = [ tf.tile(tf.expand_dims(model.cum_xx_pl, axis=0), [self.batch_size * self.no_samples, 1, 1]) for model in self.models ] self.cum_xy = [ tf.tile(tf.expand_dims(model.cum_xy_pl, axis=0), [self.batch_size * self.no_samples, 1, 1]) for model in self.models ] self.unroll(self.states) #self.unroll2(self.states) #TODO: for debugging purposes def unroll2(self, seed_states): assert seed_states.shape.as_list() == [None, self.state_dim] no_samples = self.no_samples unroll_steps = self.unroll_steps #self.reward_model = real_env_pendulum_reward()#Use true model. self.reward_model = ANN(self.state_dim + self.action_dim, 1) self.placeholders_reward = [ tf.placeholder(shape=v.shape, dtype=tf.float64) for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.reward_model.scope) ] self.assign_ops = [ v.assign(pl) for v, pl in zip( tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.reward_model.scope), self.placeholders_reward) ] states = tf.expand_dims(seed_states, axis=1) states = tf.tile(states, [1, no_samples, 1]) states = tf.reshape(states, shape=[-1, self.state_dim]) costs = [] self.next_states = [] for unroll_step in range(unroll_steps): actions = self.build_policy(states) rewards = (self.discount_factor** unroll_step) * self.reward_model.build(states, actions) rewards = tf.reshape(tf.squeeze(rewards, axis=-1), shape=[-1, no_samples]) costs.append(-rewards) states_actions = tf.concat([states, actions], axis=-1) next_states = self.get_next_states2(states_actions) self.next_states.append(next_states) states = next_states costs = tf.stack(costs, axis=-1) self.loss = tf.reduce_mean( tf.reduce_sum(tf.reduce_mean(costs, axis=1), axis=-1)) self.opt = tf.train.AdamOptimizer().minimize( self.loss, var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'policy_scope')) #TODO: for debugging purposes def get_next_states(self, states_actions): self.string = 'unroll2_gns' mu, sigma = [ tf.concat(e, axis=-1) for e in zip(*[ model.posterior_predictive_distribution(states_actions, None) for model in self.models ]) ] self.mus1.append(mu) self.sigmas1.append(sigma) #print mu.shape #print sigma.shape next_state = tfd.MultivariateNormalDiag( loc=mu, scale_diag=tf.sqrt(sigma)).sample() return next_state #TODO: for debugging purposes def get_next_states2(self, states_actions): self.string = 'unroll2_gns2' mus = [] sigmas = [] for model in self.models: mu, sigma = model.mu_sigma(model.cum_xx_pl, model.cum_xy_pl) post_pred_mu, post_pred_sigma = model.post_pred2( states_actions, mu, sigma) mus.append(post_pred_mu) sigmas.append(post_pred_sigma) mus = tf.concat(mus, axis=-1) sigmas = tf.concat(sigmas, axis=-1) self.mus2.append(mus) self.sigmas2.append(sigmas) #print mus.shape #print sigmas.shape next_state = tfd.MultivariateNormalDiag( loc=mus, scale_diag=tf.sqrt(sigmas)).sample() return next_state def unroll(self, seed_states): assert seed_states.shape.as_list() == [None, self.state_dim] no_samples = self.no_samples unroll_steps = self.unroll_steps #self.reward_model = real_env_pendulum_reward()#Use true model. self.reward_model = ANN(self.state_dim + self.action_dim, 1) self.placeholders_reward = [ tf.placeholder(shape=v.shape, dtype=tf.float64) for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.reward_model.scope) ] self.assign_ops = [ v.assign(pl) for v, pl in zip( tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.reward_model.scope), self.placeholders_reward) ] states = tf.expand_dims(seed_states, axis=1) states = tf.tile(states, [1, no_samples, 1]) states = tf.reshape(states, shape=[-1, self.state_dim]) self.mus0 = [] self.sigmas0 = [] self.mus1 = [] self.sigmas1 = [] self.mus2 = [] self.sigmas2 = [] costs = [] self.next_states = [] #ns = [] #bs = [] for unroll_step in range(unroll_steps): print 'unrolling:', unroll_step if self.debugging_plot == True: actions = self.build_policy2(states) else: actions = self.build_policy(states) # Reward rewards = (self.discount_factor** unroll_step) * self.reward_model.build(states, actions) rewards = tf.reshape(tf.squeeze(rewards, axis=-1), shape=[-1, no_samples]) costs.append(-rewards) states_actions = tf.concat([states, actions], axis=-1) mus, sigmas = zip(*[ self.mu_sigma(self.cum_xx[y], self.cum_xy[y], self.models[y].s, self.models[y].noise_sd) for y in range(self.y_dim) ]) bases = [ model.approx_rbf_kern_basis(states_actions) for model in self.models ] #bs.append(bases) mu_pred, sigma_pred = [ tf.concat(e, axis=-1) for e in zip(*[ self.prediction(mu, sigma, basis, model.noise_sd) for mu, sigma, basis, model in zip(mus, sigmas, bases, self.models) ]) ] self.mus0.append(mu_pred) self.sigmas0.append(sigma_pred) self.get_next_states(states_actions) self.get_next_states2(states_actions) next_states = tfd.MultivariateNormalDiag( loc=mu_pred, scale_diag=tf.sqrt(sigma_pred)).sample() #ns.append(tf.split(next_states, self.y_dim, axis=-1)) self.next_states.append( tf.reshape(next_states, shape=[-1, no_samples, self.state_dim])) for y in range(self.y_dim): self.update_posterior(bases[y], next_states[..., y:y + 1], y) states = next_states if self.debugging_plot == False: print 'here1' costs = tf.stack(costs, axis=-1) print 'here2' self.loss = tf.reduce_mean( tf.reduce_sum(tf.reduce_mean(costs, axis=1), axis=-1)) print 'here3' self.opt = tf.train.AdamOptimizer().minimize( self.loss, var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, 'policy_scope')) print 'here4' self.string = 'unroll' def update_posterior(self, X, y, i): X_expanded_dims = tf.expand_dims(X, axis=-1) y_expanded_dims = tf.expand_dims(y, axis=-1) self.cum_xx[i] += tf.matmul( X_expanded_dims, tf.transpose(X_expanded_dims, perm=[0, 2, 1])) self.cum_xy[i] += tf.matmul(X_expanded_dims, y_expanded_dims) def prediction(self, mu, sigma, basis, noise_sd): basis_expanded_dims = tf.expand_dims(basis, axis=-1) mu_pred = tf.matmul(tf.transpose(mu, perm=[0, 2, 1]), basis_expanded_dims) sigma_pred = tf.square(noise_sd) + tf.matmul( tf.matmul(tf.transpose(basis_expanded_dims, perm=[0, 2, 1]), sigma), basis_expanded_dims) return tf.squeeze(mu_pred, axis=-1), tf.squeeze(sigma_pred, axis=-1) def mu_sigma(self, xx, xy, s, noise_sd): noise_sd_sq = tf.square(noise_sd) prior_sigma_inv = tf.matrix_inverse( tf.tile( tf.expand_dims(s * tf.eye(self.no_basis, dtype=tf.float64), axis=0), [self.batch_size * self.no_samples, 1, 1])) A = tf.matrix_inverse(tf.multiply(noise_sd_sq, prior_sigma_inv) + xx) sigma = tf.multiply(noise_sd_sq, A) # Assuming that prior mean is zero vector mu = tf.matmul(A, xy) return mu, sigma def mu_sigma2(self, xx, xy, s, noise_sd, bs, ns, idx): if bs and ns: assert len(zip(*bs)) == self.y_dim assert len(zip(*ns)) == self.y_dim X = zip(*bs)[idx] y = zip(*ns)[idx] X = tf.expand_dims(tf.stack(X, axis=0), axis=-1) XX = tf.matmul(X, tf.transpose(X, perm=[0, 1, 3, 2])) y = tf.expand_dims(tf.stack(y, axis=0), axis=-1) Xy = tf.matmul(X, y) XX_ = tf.reduce_sum(XX, axis=0) Xy_ = tf.reduce_sum(Xy, axis=0) else: XX_ = 0. Xy_ = 0. noise_sd_sq = tf.square(noise_sd) prior_sigma_inv = tf.matrix_inverse( tf.tile( tf.expand_dims(s * tf.eye(self.no_basis, dtype=tf.float64), axis=0), [self.batch_size * self.no_samples, 1, 1])) A = tf.matrix_inverse( tf.multiply(noise_sd_sq, prior_sigma_inv) + xx + XX_) sigma = tf.multiply(noise_sd_sq, A) # Assuming that prior mean is zero vector mu = tf.matmul(A, xy + Xy_) return mu, sigma def update(self, sess, X=None, y=None, memory=None): if memory is not None: states = np.stack([e[0] for e in memory], axis=0) actions = np.stack([e[1] for e in memory], axis=0) y = np.stack([e[3] for e in memory], axis=0) X = np.concatenate([states, actions], axis=-1) for i in range(self.y_dim): self.models[i].update(sess, X, y[..., i]) def act(self, sess, state): state = np.atleast_2d(state) action = sess.run(self.actions, feed_dict={self.states: state}) return action[0] def train(self, sess, memory): feed_dict = {} #TODO: for debugging purposes if self.string == 'unroll': for model in self.models: feed_dict[model.cum_xx_pl] = model.cum_xx feed_dict[model.cum_xy_pl] = model.cum_xy feed_dict[model.mu_placeholder] = model.mu #for testing feed_dict[model.sigma_placeholder] = model.sigma #for testing feed_dict[ model.sigma_prior_pl] = model.sigma_prior #for testing feed_dict[model.mu_prior_pl] = model.mu_prior #for testing elif self.string == 'unroll2_gns': for model in self.models: feed_dict[model.mu_placeholder] = model.mu feed_dict[model.sigma_placeholder] = model.sigma elif self.string == 'unroll2_gns2': for model in self.models: feed_dict[model.cum_xx_pl] = model.cum_xx feed_dict[model.cum_xy_pl] = model.cum_xy feed_dict[model.sigma_prior_pl] = model.sigma_prior feed_dict[model.mu_prior_pl] = model.mu_prior for it in range(self.train_policy_iterations): batch = memory.sample(self.train_policy_batch_size) states = np.stack([b[0] for b in batch], axis=0) feed_dict[self.states] = states mus0, sigmas0, mus1, sigmas1, mus2, sigmas2, next_states, loss, _ = sess.run( [ self.mus0, self.sigmas0, self.mus1, self.sigmas1, self.mus2, self.sigmas2, self.next_states, self.loss, self.opt ], feed_dict=feed_dict) if loss > 1000.: print next_states ''' assert len(mus0) == len(sigmas0) assert len(mus0) == len(mus1) assert len(mus0) == len(sigmas1) assert len(mus0) == len(mus2) assert len(mus0) == len(sigmas2) ''' ''' for mu0, sigma0, mu1, sigma1, mu2, sigma2, ii in zip(mus0, sigmas0, mus1, sigmas1, mus2, sigmas2, range(len(mus0))): try: np.testing.assert_almost_equal(sigma1, sigma2, decimal=4) except: print ii, 'here0' for i in range(len(sigma1)): for j in range(len(sigma1[i])): print sigma1[i, j], sigma2[i, j] exit() try: np.testing.assert_almost_equal(mu1, mu2, decimal=4) except: print ii, 'here3', for i in range(len(mu1)): print mu1[i], mu2[i] exit() try: np.testing.assert_almost_equal(mu0, mu1, decimal=4) except: print ii, 'here1', for i in range(len(mu0)): print mu0[i], mu1[i] exit() try: np.testing.assert_almost_equal(mu0, mu2, decimal=4) except: print ii, 'here2', for i in range(len(m0)): print m0[i], m2[i] exit() try: np.testing.assert_almost_equal(sigma0, sigma1, decimal=4) except: print ii, 'here4', for i in range(len(sigma0)): for j in range(len(sigma0[i])): print sigma0[i, j], sigma1[i, j] exit() try: np.testing.assert_almost_equal(sigma0, sigma2, decimal=4) except: print ii, 'here5', for i in range(len(sigma0)): for j in range(len(sigma0[i])): print sigma0[i, j], sigma2[i, j] exit() ''' print 'iteration:', it, 'loss:', loss, self.string, len(mus0) ''' try: mus0, sigmas0, mus1, sigmas1, mus2, sigmas2, next_states, loss, _ = sess.run([self.mus0, self.sigmas0, self.mus1, self.sigmas1, self.mus2, self.sigmas2, self.next_states, self.loss, self.opt], feed_dict=feed_dict) assert len(mus0) == len(sigmas0) assert len(mus0) == len(mus1) assert len(mus0) == len(sigmas1) assert len(mus0) == len(mus2) assert len(mus0) == len(sigmas2) for mu0, sigma0, mu1, sigma1, mu2, sigma2 in zip(mus0, sigmas0, mus1, sigmas1, mus2, sigmas2): np.testing.assert_almost_equal(mu0, mu1) np.testing.assert_almost_equal(mu0, mu2) np.testing.assert_almost_equal(mu1, mu2) np.testing.assert_almost_equal(sigma0, sigma1) np.testing.assert_almost_equal(sigma0, sigma2) np.testing.assert_almost_equal(sigma1, sigma2) if loss > 1000.: print next_states print 'iteration:', it, 'loss:', loss, self.string except: print 'training step failed.' ''' def build_policy(self, states): assert states.shape.as_list() == [None, self.state_dim] #Fully connected layer 1 fc1 = slim.fully_connected(states, 256, activation_fn=tf.nn.relu, scope=self.policy_scope + '/fc1', reuse=self.policy_reuse_vars) #Fully connected layer 2 fc2 = slim.fully_connected(fc1, 256, activation_fn=tf.nn.relu, scope=self.policy_scope + '/fc2', reuse=self.policy_reuse_vars) #Output layer output = slim.fully_connected(fc2, self.action_dim, activation_fn=tf.nn.tanh, scope=self.policy_scope + '/output', reuse=self.policy_reuse_vars) #Apply action bounds np.testing.assert_array_equal(-self.action_bound_low, self.action_bound_high) action_bound = tf.constant(self.action_bound_high, dtype=tf.float64) policy = tf.multiply(output, action_bound) #Change flag self.policy_reuse_vars = True return policy def build_policy2(self, states): try: self.policy except: self.idx = 0 self.policy = tf.placeholder(shape=[self.unroll_steps, 1], dtype=tf.float64) action = self.policy[self.idx:self.idx + 1, ...] tile_size = tf.shape(states)[0] action_tiled = tf.tile(action, [tile_size, 1]) self.idx += 1 return action_tiled
class direct_policy_search: def __init__(self, state_dim, action_dim, action_bound_high, \ action_bound_low, unroll_length, discount_factor, \ gradient_descent_steps, scope): self.state_dim = state_dim self.action_dim = action_dim self.action_bound_high = action_bound_high self.action_bound_low = action_bound_low self.unroll_length = unroll_length self.discount_factor = discount_factor self.gradient_descent_steps = gradient_descent_steps self.scope = scope #Make sure bounds are same (assumption can be relaxed later) np.testing.assert_array_equal(-self.action_bound_low, self.action_bound_high) #Flags self.policy_reuse_vars = None ''' self.reward_model = ANN(self.state_dim+self.action_dim, 1) self.placeholders_reward = [tf.placeholder(shape=v.shape, dtype=tf.float64) for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.reward_model.scope)] self.assign_ops0 = [v.assign(pl) for v, pl in zip(tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.reward_model.scope), self.placeholders_reward)] ''' #self.reward_model = real_env_pendulum_reward() self.reward_model = mountain_car_continuous_reward_function() #self.state_model = real_env_pendulum_state() #self.state_model = mountain_car_continuous_state_function() self.state_model = ANN(self.state_dim + self.action_dim, self.state_dim) self.placeholders_state = [ tf.placeholder(shape=v.shape, dtype=tf.float64) for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.state_model.scope) ] self.assign_ops1 = [ v.assign(pl) for v, pl in zip( tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self. state_model.scope), self.placeholders_state) ] #Build computational graph (i.e., unroll policy) #self.states = tf.placeholder(shape=[None, self.state_dim], dtype=tf.float32) self.states = tf.placeholder(shape=[None, self.state_dim], dtype=tf.float64) self.action = self.build_policy(self.states) state = self.states action = self.build_policy(state) rewards = [] for i in range(self.unroll_length): print i #reward = pow(self.discount_factor, i) * self.reward_model.build(state, action) #reward = pow(self.discount_factor, i) * self.reward_model.step_tf(state, action) reward = pow(self.discount_factor, i) * self.reward_model.sigmoid_approx(state, action) rewards.append(reward) state = self.state_model.build(state, action) #state = self.state_model.step_tf(state, action) action = self.build_policy(state) rewards = tf.reduce_sum(tf.stack(rewards, axis=-1), axis=-1) print 'here0' self.loss = -tf.reduce_mean(tf.reduce_sum(rewards, axis=-1)) print 'here1' self.opt = tf.train.AdamOptimizer().minimize( self.loss, var_list=tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope)) print 'here2' def act(self, sess, states): states = np.atleast_2d(states) #print sess.run(tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)) action = sess.run(self.action, feed_dict={self.states: states}) return action[0] def train(self, sess, states): for _ in range(self.gradient_descent_steps): loss, _ = sess.run([self.loss, self.opt], feed_dict={self.states: states}) #asin1, asin2, loss, _ = sess.run([self.asin1, self.asin2, self.loss, self.opt], feed_dict={self.states:states}) def build_policy(self, states): assert states.shape.as_list() == [None, self.state_dim] #Fully connected layer 1 fc1 = slim.fully_connected(states, 256, activation_fn=tf.nn.relu, scope=self.scope + '/fc1', reuse=self.policy_reuse_vars) fc2 = slim.fully_connected(fc1, 256, activation_fn=tf.nn.relu, scope=self.scope + '/fc2', reuse=self.policy_reuse_vars) #Output layer output = slim.fully_connected(fc2, self.action_dim, activation_fn=tf.nn.tanh, scope=self.scope + '/output', reuse=self.policy_reuse_vars) #Apply action bounds #action_bound = tf.constant(self.action_bound_high, dtype=tf.float32) action_bound = tf.constant(self.action_bound_high, dtype=tf.float64) policy = tf.multiply(output, action_bound) #Change flag self.policy_reuse_vars = True return policy