Spaces:
Running
Running
| # Copyright 2017 The TensorFlow Authors All Rights Reserved. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Unless required by applicable law or agreed to in writing, software | |
| # distributed under the License is distributed on an "AS IS" BASIS, | |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| # See the License for the specific language governing permissions and | |
| # limitations under the License. | |
| # ============================================================================== | |
| """Objectives to compute loss and value targets. | |
| Implements Actor Critic, PCL (vanilla PCL, Unified PCL, Trust PCL), and TRPO. | |
| """ | |
| import tensorflow as tf | |
| import numpy as np | |
| class Objective(object): | |
| def __init__(self, learning_rate, clip_norm): | |
| self.learning_rate = learning_rate | |
| self.clip_norm = clip_norm | |
| def get_optimizer(self, learning_rate): | |
| """Optimizer for gradient descent ops.""" | |
| return tf.train.AdamOptimizer(learning_rate=learning_rate, | |
| epsilon=2e-4) | |
| def training_ops(self, loss, learning_rate=None): | |
| """Gradient ops.""" | |
| opt = self.get_optimizer(learning_rate) | |
| params = tf.trainable_variables() | |
| grads = tf.gradients(loss, params) | |
| if self.clip_norm: | |
| grads, global_norm = tf.clip_by_global_norm(grads, self.clip_norm) | |
| tf.summary.scalar('grad_global_norm', global_norm) | |
| return opt.apply_gradients(zip(grads, params)) | |
| def get(self, rewards, pads, values, final_values, | |
| log_probs, prev_log_probs, target_log_probs, | |
| entropies, logits, | |
| target_values, final_target_values): | |
| """Get objective calculations.""" | |
| raise NotImplementedError() | |
| def discounted_future_sum(values, discount, rollout): | |
| """Discounted future sum of time-major values.""" | |
| discount_filter = tf.reshape( | |
| discount ** tf.range(float(rollout)), [-1, 1, 1]) | |
| expanded_values = tf.concat( | |
| [values, tf.zeros([rollout - 1, tf.shape(values)[1]])], 0) | |
| conv_values = tf.transpose(tf.squeeze(tf.nn.conv1d( | |
| tf.expand_dims(tf.transpose(expanded_values), -1), discount_filter, | |
| stride=1, padding='VALID'), -1)) | |
| return conv_values | |
| def discounted_two_sided_sum(values, discount, rollout): | |
| """Discounted two-sided sum of time-major values.""" | |
| roll = float(rollout) | |
| discount_filter = tf.reshape( | |
| discount ** tf.abs(tf.range(-roll + 1, roll)), [-1, 1, 1]) | |
| expanded_values = tf.concat( | |
| [tf.zeros([rollout - 1, tf.shape(values)[1]]), values, | |
| tf.zeros([rollout - 1, tf.shape(values)[1]])], 0) | |
| conv_values = tf.transpose(tf.squeeze(tf.nn.conv1d( | |
| tf.expand_dims(tf.transpose(expanded_values), -1), discount_filter, | |
| stride=1, padding='VALID'), -1)) | |
| return conv_values | |
| def shift_values(values, discount, rollout, final_values=0.0): | |
| """Shift values up by some amount of time. | |
| Those values that shift from a value beyond the last value | |
| are calculated using final_values. | |
| """ | |
| roll_range = tf.cumsum(tf.ones_like(values[:rollout, :]), 0, | |
| exclusive=True, reverse=True) | |
| final_pad = tf.expand_dims(final_values, 0) * discount ** roll_range | |
| return tf.concat([discount ** rollout * values[rollout:, :], | |
| final_pad], 0) | |
| class ActorCritic(Objective): | |
| """Standard Actor-Critic.""" | |
| def __init__(self, learning_rate, clip_norm=5, | |
| policy_weight=1.0, critic_weight=0.1, | |
| tau=0.1, gamma=1.0, rollout=10, | |
| eps_lambda=0.0, clip_adv=None, | |
| use_target_values=False): | |
| super(ActorCritic, self).__init__(learning_rate, clip_norm=clip_norm) | |
| self.policy_weight = policy_weight | |
| self.critic_weight = critic_weight | |
| self.tau = tau | |
| self.gamma = gamma | |
| self.rollout = rollout | |
| self.clip_adv = clip_adv | |
| self.eps_lambda = tf.get_variable( # TODO: need a better way | |
| 'eps_lambda', [], initializer=tf.constant_initializer(eps_lambda), | |
| trainable=False) | |
| self.new_eps_lambda = tf.placeholder(tf.float32, []) | |
| self.assign_eps_lambda = self.eps_lambda.assign( | |
| 0.99 * self.eps_lambda + 0.01 * self.new_eps_lambda) | |
| self.use_target_values = use_target_values | |
| def get(self, rewards, pads, values, final_values, | |
| log_probs, prev_log_probs, target_log_probs, | |
| entropies, logits, | |
| target_values, final_target_values): | |
| not_pad = 1 - pads | |
| batch_size = tf.shape(rewards)[1] | |
| entropy = not_pad * sum(entropies) | |
| rewards = not_pad * rewards | |
| value_estimates = not_pad * values | |
| log_probs = not_pad * sum(log_probs) | |
| target_values = not_pad * tf.stop_gradient(target_values) | |
| final_target_values = tf.stop_gradient(final_target_values) | |
| sum_rewards = discounted_future_sum(rewards, self.gamma, self.rollout) | |
| if self.use_target_values: | |
| last_values = shift_values( | |
| target_values, self.gamma, self.rollout, | |
| final_target_values) | |
| else: | |
| last_values = shift_values(value_estimates, self.gamma, self.rollout, | |
| final_values) | |
| future_values = sum_rewards + last_values | |
| baseline_values = value_estimates | |
| adv = tf.stop_gradient(-baseline_values + future_values) | |
| if self.clip_adv: | |
| adv = tf.minimum(self.clip_adv, tf.maximum(-self.clip_adv, adv)) | |
| policy_loss = -adv * log_probs | |
| critic_loss = -adv * baseline_values | |
| regularizer = -self.tau * entropy | |
| policy_loss = tf.reduce_mean( | |
| tf.reduce_sum(policy_loss * not_pad, 0)) | |
| critic_loss = tf.reduce_mean( | |
| tf.reduce_sum(critic_loss * not_pad, 0)) | |
| regularizer = tf.reduce_mean( | |
| tf.reduce_sum(regularizer * not_pad, 0)) | |
| # loss for gradient calculation | |
| loss = (self.policy_weight * policy_loss + | |
| self.critic_weight * critic_loss + regularizer) | |
| raw_loss = tf.reduce_mean( # TODO | |
| tf.reduce_sum(not_pad * policy_loss, 0)) | |
| gradient_ops = self.training_ops( | |
| loss, learning_rate=self.learning_rate) | |
| tf.summary.histogram('log_probs', tf.reduce_sum(log_probs, 0)) | |
| tf.summary.histogram('rewards', tf.reduce_sum(rewards, 0)) | |
| tf.summary.scalar('avg_rewards', | |
| tf.reduce_mean(tf.reduce_sum(rewards, 0))) | |
| tf.summary.scalar('policy_loss', | |
| tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss))) | |
| tf.summary.scalar('critic_loss', | |
| tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss))) | |
| tf.summary.scalar('loss', loss) | |
| tf.summary.scalar('raw_loss', raw_loss) | |
| return (loss, raw_loss, future_values, | |
| gradient_ops, tf.summary.merge_all()) | |
| class PCL(ActorCritic): | |
| """PCL implementation. | |
| Implements vanilla PCL, Unified PCL, and Trust PCL depending | |
| on provided inputs. | |
| """ | |
| def get(self, rewards, pads, values, final_values, | |
| log_probs, prev_log_probs, target_log_probs, | |
| entropies, logits, | |
| target_values, final_target_values): | |
| not_pad = 1 - pads | |
| batch_size = tf.shape(rewards)[1] | |
| rewards = not_pad * rewards | |
| value_estimates = not_pad * values | |
| log_probs = not_pad * sum(log_probs) | |
| target_log_probs = not_pad * tf.stop_gradient(sum(target_log_probs)) | |
| relative_log_probs = not_pad * (log_probs - target_log_probs) | |
| target_values = not_pad * tf.stop_gradient(target_values) | |
| final_target_values = tf.stop_gradient(final_target_values) | |
| # Prepend. | |
| not_pad = tf.concat([tf.ones([self.rollout - 1, batch_size]), | |
| not_pad], 0) | |
| rewards = tf.concat([tf.zeros([self.rollout - 1, batch_size]), | |
| rewards], 0) | |
| value_estimates = tf.concat( | |
| [self.gamma ** tf.expand_dims( | |
| tf.range(float(self.rollout - 1), 0, -1), 1) * | |
| tf.ones([self.rollout - 1, batch_size]) * | |
| value_estimates[0:1, :], | |
| value_estimates], 0) | |
| log_probs = tf.concat([tf.zeros([self.rollout - 1, batch_size]), | |
| log_probs], 0) | |
| prev_log_probs = tf.concat([tf.zeros([self.rollout - 1, batch_size]), | |
| prev_log_probs], 0) | |
| relative_log_probs = tf.concat([tf.zeros([self.rollout - 1, batch_size]), | |
| relative_log_probs], 0) | |
| target_values = tf.concat( | |
| [self.gamma ** tf.expand_dims( | |
| tf.range(float(self.rollout - 1), 0, -1), 1) * | |
| tf.ones([self.rollout - 1, batch_size]) * | |
| target_values[0:1, :], | |
| target_values], 0) | |
| sum_rewards = discounted_future_sum(rewards, self.gamma, self.rollout) | |
| sum_log_probs = discounted_future_sum(log_probs, self.gamma, self.rollout) | |
| sum_prev_log_probs = discounted_future_sum(prev_log_probs, self.gamma, self.rollout) | |
| sum_relative_log_probs = discounted_future_sum( | |
| relative_log_probs, self.gamma, self.rollout) | |
| if self.use_target_values: | |
| last_values = shift_values( | |
| target_values, self.gamma, self.rollout, | |
| final_target_values) | |
| else: | |
| last_values = shift_values(value_estimates, self.gamma, self.rollout, | |
| final_values) | |
| future_values = ( | |
| - self.tau * sum_log_probs | |
| - self.eps_lambda * sum_relative_log_probs | |
| + sum_rewards + last_values) | |
| baseline_values = value_estimates | |
| adv = tf.stop_gradient(-baseline_values + future_values) | |
| if self.clip_adv: | |
| adv = tf.minimum(self.clip_adv, tf.maximum(-self.clip_adv, adv)) | |
| policy_loss = -adv * sum_log_probs | |
| critic_loss = -adv * (baseline_values - last_values) | |
| policy_loss = tf.reduce_mean( | |
| tf.reduce_sum(policy_loss * not_pad, 0)) | |
| critic_loss = tf.reduce_mean( | |
| tf.reduce_sum(critic_loss * not_pad, 0)) | |
| # loss for gradient calculation | |
| loss = (self.policy_weight * policy_loss + | |
| self.critic_weight * critic_loss) | |
| # actual quantity we're trying to minimize | |
| raw_loss = tf.reduce_mean( | |
| tf.reduce_sum(not_pad * adv * (-baseline_values + future_values), 0)) | |
| gradient_ops = self.training_ops( | |
| loss, learning_rate=self.learning_rate) | |
| tf.summary.histogram('log_probs', tf.reduce_sum(log_probs, 0)) | |
| tf.summary.histogram('rewards', tf.reduce_sum(rewards, 0)) | |
| tf.summary.histogram('future_values', future_values) | |
| tf.summary.histogram('baseline_values', baseline_values) | |
| tf.summary.histogram('advantages', adv) | |
| tf.summary.scalar('avg_rewards', | |
| tf.reduce_mean(tf.reduce_sum(rewards, 0))) | |
| tf.summary.scalar('policy_loss', | |
| tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss))) | |
| tf.summary.scalar('critic_loss', | |
| tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss))) | |
| tf.summary.scalar('loss', loss) | |
| tf.summary.scalar('raw_loss', tf.reduce_mean(raw_loss)) | |
| tf.summary.scalar('eps_lambda', self.eps_lambda) | |
| return (loss, raw_loss, | |
| future_values[self.rollout - 1:, :], | |
| gradient_ops, tf.summary.merge_all()) | |
| class TRPO(ActorCritic): | |
| """TRPO.""" | |
| def get(self, rewards, pads, values, final_values, | |
| log_probs, prev_log_probs, target_log_probs, | |
| entropies, logits, | |
| target_values, final_target_values): | |
| not_pad = 1 - pads | |
| batch_size = tf.shape(rewards)[1] | |
| rewards = not_pad * rewards | |
| value_estimates = not_pad * values | |
| log_probs = not_pad * sum(log_probs) | |
| prev_log_probs = not_pad * prev_log_probs | |
| target_values = not_pad * tf.stop_gradient(target_values) | |
| final_target_values = tf.stop_gradient(final_target_values) | |
| sum_rewards = discounted_future_sum(rewards, self.gamma, self.rollout) | |
| if self.use_target_values: | |
| last_values = shift_values( | |
| target_values, self.gamma, self.rollout, | |
| final_target_values) | |
| else: | |
| last_values = shift_values(value_estimates, self.gamma, self.rollout, | |
| final_values) | |
| future_values = sum_rewards + last_values | |
| baseline_values = value_estimates | |
| adv = tf.stop_gradient(-baseline_values + future_values) | |
| if self.clip_adv: | |
| adv = tf.minimum(self.clip_adv, tf.maximum(-self.clip_adv, adv)) | |
| policy_loss = -adv * tf.exp(log_probs - prev_log_probs) | |
| critic_loss = -adv * baseline_values | |
| policy_loss = tf.reduce_mean( | |
| tf.reduce_sum(policy_loss * not_pad, 0)) | |
| critic_loss = tf.reduce_mean( | |
| tf.reduce_sum(critic_loss * not_pad, 0)) | |
| raw_loss = policy_loss | |
| # loss for gradient calculation | |
| if self.policy_weight == 0: | |
| policy_loss = 0.0 | |
| elif self.critic_weight == 0: | |
| critic_loss = 0.0 | |
| loss = (self.policy_weight * policy_loss + | |
| self.critic_weight * critic_loss) | |
| gradient_ops = self.training_ops( | |
| loss, learning_rate=self.learning_rate) | |
| tf.summary.histogram('log_probs', tf.reduce_sum(log_probs, 0)) | |
| tf.summary.histogram('rewards', tf.reduce_sum(rewards, 0)) | |
| tf.summary.scalar('avg_rewards', | |
| tf.reduce_mean(tf.reduce_sum(rewards, 0))) | |
| tf.summary.scalar('policy_loss', | |
| tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss))) | |
| tf.summary.scalar('critic_loss', | |
| tf.reduce_mean(tf.reduce_sum(not_pad * policy_loss))) | |
| tf.summary.scalar('loss', loss) | |
| tf.summary.scalar('raw_loss', raw_loss) | |
| return (loss, raw_loss, future_values, | |
| gradient_ops, tf.summary.merge_all()) | |