mirror of
https://github.com/wassname/ray.git
synced 2026-07-15 11:25:40 +08:00
Eric LiangandPhilipp Moritz
4374ad1453
* 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
269 lines
10 KiB
Python
269 lines
10 KiB
Python
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):
|
|
"""
|
|
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
|
|
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()
|
|
|
|
def load_weights(self, weights):
|
|
self.variables.set_weights(weights)
|
|
|
|
def compute_trajectory(self, gamma, lam, horizon):
|
|
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
|
|
|
|
|
|
RemoteAgent = ray.remote(Agent)
|