diff --git a/examples/policy_gradient/examples/example.py b/examples/policy_gradient/examples/example.py index ab5fd87a7..b262cbc0d 100644 --- a/examples/policy_gradient/examples/example.py +++ b/examples/policy_gradient/examples/example.py @@ -5,27 +5,40 @@ from __future__ import print_function from datetime import datetime import argparse +import time + import ray +import numpy as np import tensorflow as tf from reinforce.env import (NoPreprocessor, AtariRamPreprocessor, AtariPixelPreprocessor) from reinforce.agent import Agent, RemoteAgent from reinforce.rollout import collect_samples -from reinforce.utils import iterate, shuffle +from reinforce.utils import shuffle + config = {"kl_coeff": 0.2, "num_sgd_iter": 30, "max_iterations": 1000, "sgd_stepsize": 5e-5, - "sgd_batchsize": 128, + # TODO(pcm): Expose the choice between gpus and cpus + # as a command line argument. + "devices": ["/cpu:%d" % i for i in range(4)], + "tf_session_args": { + "device_count": {"CPU": 4}, + "log_device_placement": False, + "allow_soft_placement": True, + }, + "sgd_batchsize": 128, # total size across all devices "entropy_coeff": 0.0, "clip_param": 0.3, "kl_target": 0.01, "timesteps_per_batch": 40000, "num_agents": 5, "tensorboard_log_dir": "/tmp/ray", - "trace_level": tf.RunOptions.NO_TRACE} + "full_trace_nth_sgd_batch": -1, + "full_trace_data_load": False} if __name__ == "__main__": @@ -47,24 +60,28 @@ if __name__ == "__main__": preprocessor = AtariRamPreprocessor() elif mdp_name == "CartPole-v0": preprocessor = NoPreprocessor() + elif mdp_name == "Walker2d-v1": + preprocessor = NoPreprocessor() else: print("No environment was chosen, so defaulting to Pong-v0.") mdp_name = "Pong-v0" preprocessor = AtariPixelPreprocessor() print("Using the environment {}.".format(mdp_name)) - agents = [RemoteAgent.remote(mdp_name, 1, preprocessor, config, False) + agents = [RemoteAgent.remote(mdp_name, 1, preprocessor, config, True) for _ in range(config["num_agents"])] - agent = Agent(mdp_name, 1, preprocessor, config, True) + agent = Agent(mdp_name, 1, preprocessor, config, False) kl_coeff = config["kl_coeff"] file_writer = tf.summary.FileWriter( - '{}/trpo_{}_{}'.format( - config["tensorboard_log_dir"], mdp_name, datetime.today()), + "{}/trpo_{}_{}".format( + config["tensorboard_log_dir"], mdp_name, + str(datetime.today()).replace(" ", "_")), agent.sess.graph) global_step = 0 for j in range(config["max_iterations"]): + iter_start = time.time() print("== iteration", j) weights = ray.put(agent.get_weights()) [a.load_weights.remote(weights) for a in agents] @@ -72,7 +89,7 @@ if __name__ == "__main__": agents, config["timesteps_per_batch"], 0.995, 1.0, 2000) print("total reward is ", total_reward) print("trajectory length mean is ", traj_len_mean) - print("timesteps: ", trajectory["dones"].shape[0]) + print("timesteps:", trajectory["dones"].shape[0]) traj_stats = tf.Summary(value=[ tf.Summary.Value( tag="policy_gradient/rollouts/mean_reward", @@ -85,40 +102,44 @@ if __name__ == "__main__": trajectory["advantages"] = ((trajectory["advantages"] - trajectory["advantages"].mean()) / trajectory["advantages"].std()) - print("Computing policy (optimizer='" + agent.optimizer.get_name() + - "', iterations=" + str(config["num_sgd_iter"]) + + rollouts_end = time.time() + print("Computing policy (iterations=" + str(config["num_sgd_iter"]) + ", stepsize=" + str(config["sgd_stepsize"]) + "):") names = ["iter", "loss", "kl", "entropy"] print(("{:>15}" * len(names)).format(*names)) + num_devices = len(config["devices"]) trajectory = shuffle(trajectory) - ppo = agent.ppo + shuffle_end = time.time() + tuples_per_device = agent.load_data( + trajectory, j == 0 and config["full_trace_data_load"]) + load_end = time.time() + rollouts_time = rollouts_end - iter_start + shuffle_time = shuffle_end - rollouts_end + load_time = load_end - shuffle_end + sgd_time = 0 for i in range(config["num_sgd_iter"]): - # Test on current set of rollouts. - run_options = tf.RunOptions(trace_level=config["trace_level"]) - run_metadata = tf.RunMetadata() - loss, kl, entropy = agent.sess.run( - [ppo.loss, ppo.mean_kl, ppo.mean_entropy], - feed_dict={ppo.observations: trajectory["observations"], - ppo.advantages: trajectory["advantages"], - ppo.actions: trajectory["actions"].squeeze(), - ppo.prev_logits: trajectory["logprobs"], - ppo.kl_coeff: kl_coeff}, - options=run_options, - run_metadata=run_metadata) + sgd_start = time.time() + batch_index = 0 + num_batches = int(tuples_per_device) // int(agent.per_device_batch_size) + loss, kl, entropy = [], [], [] + permutation = np.random.permutation(num_batches) + while batch_index < num_batches: + full_trace = ( + i == 0 and j == 0 and + batch_index == config["full_trace_nth_sgd_batch"]) + batch_loss, batch_kl, batch_entropy = agent.run_sgd_minibatch( + permutation[batch_index] * agent.per_device_batch_size, + kl_coeff, full_trace, file_writer) + loss.append(batch_loss) + kl.append(batch_kl) + entropy.append(batch_entropy) + batch_index += 1 + loss = np.mean(loss) + kl = np.mean(kl) + entropy = np.mean(entropy) + sgd_end = time.time() print("{:>15}{:15.5e}{:15.5e}{:15.5e}".format(i, loss, kl, entropy)) - # Run SGD for training on current set of rollouts. - for batch in iterate(trajectory, config["sgd_batchsize"]): - run_options = tf.RunOptions(trace_level=config["trace_level"]) - run_metadata = tf.RunMetadata() - agent.sess.run( - [agent.train_op], - feed_dict={ppo.observations: batch["observations"], - ppo.advantages: batch["advantages"], - ppo.actions: batch["actions"].squeeze(), - ppo.prev_logits: batch["logprobs"], - ppo.kl_coeff: kl_coeff}, - options=run_options, - run_metadata=run_metadata) + values = [] if i == config["num_sgd_iter"] - 1: metric_prefix = "policy_gradient/sgd/final_iter/" @@ -140,9 +161,15 @@ if __name__ == "__main__": sgd_stats = tf.Summary(value=values) file_writer.add_summary(sgd_stats, global_step) global_step += 1 + sgd_time += sgd_end - sgd_start if kl > 2.0 * config["kl_target"]: kl_coeff *= 1.5 elif kl < 0.5 * config["kl_target"]: kl_coeff *= 0.5 - print("kl div = ", kl) - print("kl coeff = ", kl_coeff) + print("kl div:", kl) + print("kl coeff:", kl_coeff) + print("rollouts time:", rollouts_time) + print("shuffle time:", shuffle_time) + print("load time:", load_time) + print("sgd time:", sgd_time) + print("sgd examples/s:", len(trajectory["observations"]) / sgd_time) diff --git a/examples/policy_gradient/reinforce/agent.py b/examples/policy_gradient/reinforce/agent.py index d2a797362..b2037edc9 100644 --- a/examples/policy_gradient/reinforce/agent.py +++ b/examples/policy_gradient/reinforce/agent.py @@ -2,39 +2,255 @@ from __future__ import absolute_import from __future__ import division from __future__ import print_function +from collections import namedtuple + +import gym.spaces import tensorflow as tf import os +from tensorflow.python.client import timeline + import ray +from reinforce.distributions import Categorical, DiagGaussian from reinforce.env import BatchedEnv from reinforce.policy import ProximalPolicyLoss from reinforce.filter import MeanStdFilter from reinforce.rollout import rollouts, add_advantage_values +from reinforce.utils import make_divisible_by, average_gradients + + +# Each tower is a copy of the policy graph pinned to a specific device +Tower = namedtuple("Tower", ["init_op", "grads", "policy"]) class Agent(object): - def __init__(self, name, batchsize, preprocessor, config, use_gpu): - if not use_gpu: + """ + Agent class that holds the simulator environment and the policy. + + Initializes the tensorflow graphs for both training and evaluation. + One common policy graph is initialized on '/cpu:0' and holds all the shared + network weights. When run as a remote agent, only this graph is used. + + When the agent is initialized locally with multiple GPU devices, copies of + the policy graph are also placed on each GPU. These per-GPU graphs share the + common policy network weights but take device-local input tensors. + + The idea here is that training data can be bulk-loaded onto these + device-local variables. Synchronous SGD can then be run in parallel over + this GPU-local data. + """ + + def __init__(self, name, batchsize, preprocessor, config, is_remote): + if is_remote: os.environ["CUDA_VISIBLE_DEVICES"] = "" + devices = ["/cpu:0"] + else: + devices = config["devices"] + self.devices = devices + self.config = config self.env = BatchedEnv(name, batchsize, preprocessor=preprocessor) if preprocessor.shape is None: preprocessor.shape = self.env.observation_space.shape - self.sess = tf.Session() - with tf.name_scope("policy_gradient/train"): - with tf.name_scope("proximal_policy_loss"): - self.ppo = ProximalPolicyLoss(self.env.observation_space, - self.env.action_space, preprocessor, - config, self.sess) - with tf.name_scope("adam_optimizer"): - self.optimizer = tf.train.AdamOptimizer(config["sgd_stepsize"]) - self.train_op = self.optimizer.minimize(self.ppo.loss) - self.variables = ray.experimental.TensorFlowVariables(self.ppo.loss, - self.sess) - self.observation_filter = MeanStdFilter(preprocessor.shape, clip=None) - self.reward_filter = MeanStdFilter((), clip=5.0) + if is_remote: + config_proto = tf.ConfigProto() + else: + config_proto = tf.ConfigProto(**config["tf_session_args"]) + self.preprocessor = preprocessor + self.sess = tf.Session(config=config_proto) + + # Defines the training inputs. + self.kl_coeff = tf.placeholder(name="newkl", shape=(), dtype=tf.float32) + self.observations = tf.placeholder(tf.float32, + shape=(None,) + preprocessor.shape) + self.advantages = tf.placeholder(tf.float32, shape=(None,)) + + action_space = self.env.action_space + if isinstance(action_space, gym.spaces.Box): + # The first half of the dimensions are the means, the second half are the + # standard deviations. + self.action_dim = action_space.shape[0] + self.action_shape = (self.action_dim,) + self.logit_dim = 2 * self.action_dim + self.actions = tf.placeholder(tf.float32, shape=(None, self.action_dim)) + self.distribution_class = DiagGaussian + elif isinstance(action_space, gym.spaces.Discrete): + self.action_dim = action_space.n + self.action_shape = () + self.logit_dim = self.action_dim + self.actions = tf.placeholder(tf.int64, shape=(None,)) + self.distribution_class = Categorical + else: + raise NotImplemented("action space" + str(type(action_space)) + + "currently not supported") + self.prev_logits = tf.placeholder(tf.float32, shape=(None, self.logit_dim)) + + data_splits = zip( + tf.split(self.observations, len(devices)), + tf.split(self.advantages, len(devices)), + tf.split(self.actions, len(devices)), + tf.split(self.prev_logits, len(devices))) + + # Parallel SGD ops + self.towers = [] + self.batch_index = tf.placeholder(tf.int32) + assert config["sgd_batchsize"] % len(devices) == 0, \ + "Batch size must be evenly divisible by devices" + if is_remote: + self.batch_size = 1 + self.per_device_batch_size = 1 + else: + self.batch_size = config["sgd_batchsize"] + self.per_device_batch_size = int(self.batch_size / len(devices)) + self.optimizer = tf.train.AdamOptimizer(self.config["sgd_stepsize"]) + self.setup_common_policy( + self.observations, self.advantages, self.actions, self.prev_logits) + for device, (obs, adv, acts, plog) in zip(devices, data_splits): + self.towers.append( + self.setup_per_device_policy(device, obs, adv, acts, plog)) + + avg = average_gradients([t.grads for t in self.towers]) + self.train_op = self.optimizer.apply_gradients(avg) + + # Metric ops + with tf.name_scope("test_outputs"): + self.mean_loss = tf.reduce_mean( + tf.stack(values=[t.policy.loss for t in self.towers]), 0) + self.mean_kl = tf.reduce_mean( + tf.stack(values=[t.policy.mean_kl for t in self.towers]), 0) + self.mean_entropy = tf.reduce_mean( + tf.stack(values=[t.policy.mean_entropy for t in self.towers]), 0) + + # References to the model weights + self.variables = ray.experimental.TensorFlowVariables( + self.common_policy.loss, + self.sess) + self.observation_filter = MeanStdFilter(preprocessor.shape, clip=None) + self.reward_filter = MeanStdFilter((), clip=5.0) self.sess.run(tf.global_variables_initializer()) + def setup_common_policy(self, observations, advantages, actions, prev_log): + with tf.variable_scope("tower"): + self.common_policy = ProximalPolicyLoss( + self.env.observation_space, self.env.action_space, + observations, advantages, actions, prev_log, self.logit_dim, + self.kl_coeff, self.distribution_class, self.config, self.sess) + + def setup_per_device_policy( + self, device, observations, advantages, actions, prev_log): + with tf.device(device): + with tf.variable_scope("tower", reuse=True): + all_obs = tf.Variable( + observations, trainable=False, validate_shape=False, + collections=[]) + all_adv = tf.Variable( + advantages, trainable=False, validate_shape=False, collections=[]) + all_acts = tf.Variable( + actions, trainable=False, validate_shape=False, collections=[]) + all_plog = tf.Variable( + prev_log, trainable=False, validate_shape=False, collections=[]) + obs_slice = tf.slice( + all_obs, + [self.batch_index] + [0] * len(self.preprocessor.shape), + [self.per_device_batch_size] + [-1] * len(self.preprocessor.shape)) + obs_slice.set_shape(observations.shape) + adv_slice = tf.slice( + all_adv, [self.batch_index], [self.per_device_batch_size]) + acts_slice = tf.slice( + all_acts, + [self.batch_index] + [0] * len(self.action_shape), + [self.per_device_batch_size] + [-1] * len(self.action_shape)) + plog_slice = tf.slice( + all_plog, [self.batch_index, 0], [self.per_device_batch_size, -1]) + policy = ProximalPolicyLoss( + self.env.observation_space, self.env.action_space, + obs_slice, adv_slice, acts_slice, plog_slice, self.logit_dim, + self.kl_coeff, self.distribution_class, self.config, self.sess) + grads = self.optimizer.compute_gradients( + policy.loss, colocate_gradients_with_ops=True) + + return Tower( + tf.group( + *[all_obs.initializer, + all_adv.initializer, + all_acts.initializer, + all_plog.initializer]), + grads, + policy) + + def load_data(self, trajectories, full_trace): + """ + Bulk loads the specified trajectories into device memory. + + The data is split equally across all the devices. + + Returns: + The number of tuples loaded per device. + """ + + truncated_obs = make_divisible_by( + trajectories["observations"], self.batch_size) + if full_trace: + run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) + else: + run_options = tf.RunOptions(trace_level=tf.RunOptions.NO_TRACE) + run_metadata = tf.RunMetadata() + self.sess.run( + [t.init_op for t in self.towers], + feed_dict={ + self.observations: truncated_obs, + self.advantages: make_divisible_by( + trajectories["advantages"], self.batch_size), + self.actions: make_divisible_by( + trajectories["actions"].squeeze(), self.batch_size), + self.prev_logits: make_divisible_by( + trajectories["logprobs"], self.batch_size), + }, + options=run_options, + run_metadata=run_metadata) + if full_trace: + trace = timeline.Timeline(step_stats=run_metadata.step_stats) + trace_file = open("/tmp/ray/timeline-load.json", "w") + trace_file.write(trace.generate_chrome_trace_format()) + + tuples_per_device = len(truncated_obs) / len(self.devices) + assert tuples_per_device % self.per_device_batch_size == 0 + return tuples_per_device + + def run_sgd_minibatch(self, batch_index, kl_coeff, full_trace, file_writer): + """ + Run a single step of SGD. + + Runs a SGD step over the batch with index batch_index as created by + load_rollouts_data(), updating local weights. + + Returns: + (mean_loss, mean_kl, mean_entropy) evaluated over the batch. + """ + + if full_trace: + run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE) + else: + run_options = tf.RunOptions(trace_level=tf.RunOptions.NO_TRACE) + run_metadata = tf.RunMetadata() + + _, loss, kl, entropy = self.sess.run( + [self.train_op, self.mean_loss, self.mean_kl, self.mean_entropy], + feed_dict={ + self.batch_index: batch_index, + self.kl_coeff: kl_coeff}, + options=run_options, + run_metadata=run_metadata) + + if full_trace: + trace = timeline.Timeline(step_stats=run_metadata.step_stats) + trace_file = open("/tmp/ray/timeline-sgd.json", "w") + trace_file.write(trace.generate_chrome_trace_format()) + file_writer.add_run_metadata( + run_metadata, "sgd_train_{}".format(batch_index)) + + return loss, kl, entropy + def get_weights(self): return self.variables.get_weights() @@ -42,8 +258,9 @@ class Agent(object): self.variables.set_weights(weights) def compute_trajectory(self, gamma, lam, horizon): - trajectory = rollouts(self.ppo, self.env, horizon, self.observation_filter, - self.reward_filter) + trajectory = rollouts( + self.common_policy, + self.env, horizon, self.observation_filter, self.reward_filter) add_advantage_values(trajectory, gamma, lam, self.reward_filter) return trajectory diff --git a/examples/policy_gradient/reinforce/filter.py b/examples/policy_gradient/reinforce/filter.py index de3f782df..692490f58 100644 --- a/examples/policy_gradient/reinforce/filter.py +++ b/examples/policy_gradient/reinforce/filter.py @@ -91,8 +91,6 @@ class MeanStdFilter(object): if self.destd: x = x / (self.rs.std + 1e-8) if self.clip: - if np.amin(x) < -self.clip or np.amax(x) > self.clip: - print("Clipping value to " + str(self.clip)) x = np.clip(x, -self.clip, self.clip) return x diff --git a/examples/policy_gradient/reinforce/models/fcnet.py b/examples/policy_gradient/reinforce/models/fcnet.py index ac2db1802..e431ea44c 100644 --- a/examples/policy_gradient/reinforce/models/fcnet.py +++ b/examples/policy_gradient/reinforce/models/fcnet.py @@ -17,19 +17,19 @@ def normc_initializer(std=1.0): def fc_net(inputs, num_classes=10, logstd=False): - with tf.name_scope("fc_net") as net: + with tf.name_scope("fc_net"): fc1 = slim.fully_connected(inputs, 128, weights_initializer=normc_initializer(1.0), - scope=net + "fc1") + scope="fc1") fc2 = slim.fully_connected(fc1, 128, weights_initializer=normc_initializer(1.0), - scope=net + "fc2") + scope="fc2") fc3 = slim.fully_connected(fc2, 128, weights_initializer=normc_initializer(1.0), - scope=net + "fc3") + scope="fc3") fc4 = slim.fully_connected(fc3, num_classes, weights_initializer=normc_initializer(0.01), - activation_fn=None, scope=net + "fc4") + activation_fn=None, scope="fc4") if logstd: logstd = tf.get_variable(name="logstd", shape=[num_classes], initializer=tf.zeros_initializer) diff --git a/examples/policy_gradient/reinforce/models/visionnet.py b/examples/policy_gradient/reinforce/models/visionnet.py index 9226d3ec2..75adc751e 100644 --- a/examples/policy_gradient/reinforce/models/visionnet.py +++ b/examples/policy_gradient/reinforce/models/visionnet.py @@ -7,10 +7,10 @@ import tensorflow.contrib.slim as slim def vision_net(inputs, num_classes=10): - with tf.name_scope("vision_net") as net: - conv1 = slim.conv2d(inputs, 16, [8, 8], 4, scope=net + "conv1") - conv2 = slim.conv2d(conv1, 32, [4, 4], 2, scope=net + "conv2") - fc1 = slim.conv2d(conv2, 512, [10, 10], padding="VALID", scope=net + "fc1") + with tf.name_scope("vision_net"): + conv1 = slim.conv2d(inputs, 16, [8, 8], 4, scope="conv1") + conv2 = slim.conv2d(conv1, 32, [4, 4], 2, scope="conv2") + fc1 = slim.conv2d(conv2, 512, [10, 10], padding="VALID", scope="fc1") fc2 = slim.conv2d(fc1, num_classes, [1, 1], activation_fn=None, - normalizer_fn=None, scope=net + "fc2") + normalizer_fn=None, scope="fc2") return tf.squeeze(fc2, [1, 2]) diff --git a/examples/policy_gradient/reinforce/policy.py b/examples/policy_gradient/reinforce/policy.py index 2e73a90ec..db378113d 100644 --- a/examples/policy_gradient/reinforce/policy.py +++ b/examples/policy_gradient/reinforce/policy.py @@ -6,59 +6,41 @@ import gym.spaces import tensorflow as tf from reinforce.models.visionnet import vision_net from reinforce.models.fcnet import fc_net -from reinforce.distributions import Categorical, DiagGaussian class ProximalPolicyLoss(object): - def __init__(self, observation_space, action_space, preprocessor, config, - sess): + def __init__( + self, observation_space, action_space, + observations, advantages, actions, prev_logits, logit_dim, + kl_coeff, distribution_class, config, sess): assert (isinstance(action_space, gym.spaces.Discrete) or isinstance(action_space, gym.spaces.Box)) - # Adapting the kl divergence. - self.kl_coeff = tf.placeholder(name="newkl", shape=(), dtype=tf.float32) - self.observations = tf.placeholder(tf.float32, - shape=(None,) + preprocessor.shape) - self.advantages = tf.placeholder(tf.float32, shape=(None,)) + self.prev_dist = distribution_class(prev_logits) + + # Saved so that we can compute actions given different observations + self.observations = observations - if isinstance(action_space, gym.spaces.Box): - # The first half of the dimensions are the means, the second half are the - # standard deviations. - self.action_dim = action_space.shape[0] - self.logit_dim = 2 * self.action_dim - self.actions = tf.placeholder(tf.float32, - shape=(None, action_space.shape[0])) - Distribution = DiagGaussian - elif isinstance(action_space, gym.spaces.Discrete): - self.action_dim = action_space.n - self.logit_dim = self.action_dim - self.actions = tf.placeholder(tf.int64, shape=(None,)) - Distribution = Categorical - else: - raise NotImplemented("action space" + str(type(action_space)) + - "currently not supported") - self.prev_logits = tf.placeholder(tf.float32, shape=(None, self.logit_dim)) - self.prev_dist = Distribution(self.prev_logits) if len(observation_space.shape) > 1: - self.curr_logits = vision_net(self.observations, - num_classes=self.logit_dim) + self.curr_logits = vision_net(observations, num_classes=logit_dim) else: assert len(observation_space.shape) == 1 - self.curr_logits = fc_net(self.observations, num_classes=self.logit_dim) - self.curr_dist = Distribution(self.curr_logits) + self.curr_logits = fc_net(observations, num_classes=logit_dim) + self.curr_dist = distribution_class(self.curr_logits) self.sampler = self.curr_dist.sample() - self.entropy = self.curr_dist.entropy() + # Make loss functions. - self.ratio = tf.exp(self.curr_dist.logp(self.actions) - - self.prev_dist.logp(self.actions)) + self.ratio = tf.exp(self.curr_dist.logp(actions) - + self.prev_dist.logp(actions)) self.kl = self.prev_dist.kl(self.curr_dist) self.mean_kl = tf.reduce_mean(self.kl) + self.entropy = self.curr_dist.entropy() self.mean_entropy = tf.reduce_mean(self.entropy) - self.surr1 = self.ratio * self.advantages + self.surr1 = self.ratio * advantages self.surr2 = tf.clip_by_value(self.ratio, 1 - config["clip_param"], - 1 + config["clip_param"]) * self.advantages + 1 + config["clip_param"]) * advantages self.surr = tf.minimum(self.surr1, self.surr2) - self.loss = tf.reduce_mean(-self.surr + self.kl_coeff * self.kl - + self.loss = tf.reduce_mean(-self.surr + kl_coeff * self.kl - config["entropy_coeff"] * self.entropy) self.sess = sess diff --git a/examples/policy_gradient/reinforce/utils.py b/examples/policy_gradient/reinforce/utils.py index dbc2deca6..8890fa6dd 100644 --- a/examples/policy_gradient/reinforce/utils.py +++ b/examples/policy_gradient/reinforce/utils.py @@ -3,6 +3,7 @@ from __future__ import division from __future__ import print_function import numpy as np +import tensorflow as tf def flatten(weights, start=0, stop=2): @@ -32,17 +33,56 @@ def concatenate(weights_list): def shuffle(trajectory): permutation = np.random.permutation(trajectory["dones"].shape[0]) for key, val in trajectory.items(): - trajectory[key] = val[permutation][permutation] + trajectory[key] = val[permutation] return trajectory -def iterate(trajectory, batchsize): - trajectory = shuffle(trajectory) - curr_index = 0 - # TODO(pcm): This drops some data at the end of the batch. - while curr_index + batchsize < trajectory["dones"].shape[0]: - batch = dict() - for key in trajectory: - batch[key] = trajectory[key][curr_index:(curr_index + batchsize)] - curr_index += batchsize - yield batch +def make_divisible_by(array, n): + return array[0:array.shape[0] - array.shape[0] % n] + + +def average_gradients(tower_grads): + """ + Average gradients across towers. + + Calculate the average gradient for each shared variable across all towers. + Note that this function provides a synchronization point across all towers. + + Args: + tower_grads: List of lists of (gradient, variable) tuples. The outer list + is over individual gradients. The inner list is over the gradient + calculation for each tower. + + Returns: + List of pairs of (gradient, variable) where the gradient has been averaged + across all towers. + + TODO(ekl): We could use NCCL if this becomes a bottleneck. + """ + + average_grads = [] + for grad_and_vars in zip(*tower_grads): + + # Note that each grad_and_vars looks like the following: + # ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN)) + grads = [] + for g, _ in grad_and_vars: + if g is not None: + # Add 0 dimension to the gradients to represent the tower. + expanded_g = tf.expand_dims(g, 0) + + # Append on a 'tower' dimension which we will average over below. + grads.append(expanded_g) + + # Average over the 'tower' dimension. + grad = tf.concat(axis=0, values=grads) + grad = tf.reduce_mean(grad, 0) + + # Keep in mind that the Variables are redundant because they are shared + # across towers. So .. we will just return the first tower's pointer to + # the Variable. + v = grad_and_vars[0][1] + grad_and_var = (grad, v) + average_grads.append(grad_and_var) + + return average_grads