wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-07-27 11:26:41 +08:00
Code Issues Packages Projects Releases Wiki Activity

Policy gradient example: Support multi-GPU training (#584)

Browse Source 
* add tf metrics

* comments

* fix network scopes

* add doc

* initial work

* try with 3 virtual cpus

* clean up metrics

* use format string

* fix trace level

* back to pong

* always run summary on cpu

* plot intermediate and final sgd stats

* add back a global step

* update

* add timeline

* use staging area and reuse weights properly

* stage at cpu

* whoops, stage only the batch

* clean up a bit

* fix py flake

* wip

* create an optimizer graph per device

* print timeline on 5th batch instead

* print examples per second

* log placement for training ops

* force placement on cpu:0

* try separating weights onto different gpus

* try using nccl

* add cpu fallback

* remove space from date

* check has gpu device

* fix flag config

* checkpoint

* wip

* update

* add some timing

* trace loading

* try cpu

* revert that

* remove expensive test

* lint

* cleanups

* clean up timers

* clean it up a bit

* fix code for non-scalar action spaces

* address some nits

* fix quotes

* efficient shuffling between sgd epochs
This commit is contained in:
Eric Liang
2017-06-13 06:03:25 +00:00
committed by Philipp Moritz
parent 1916475e14
commit 4374ad1453
7 changed files with 378 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/policy_gradient/examples/example.py
+65 -38
View File
@@ -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)
examples/policy_gradient/reinforce/agent.py
+234 -17
View File
@@ -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
examples/policy_gradient/reinforce/filter.py
-2
View File
@@ -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
examples/policy_gradient/reinforce/models/fcnet.py
+5 -5
View File
@@ -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)
examples/policy_gradient/reinforce/models/visionnet.py
+5 -5
View File
@@ -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])
examples/policy_gradient/reinforce/policy.py
+18 -36
View File
@@ -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
examples/policy_gradient/reinforce/utils.py
+51 -11
View File
@@ -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