mirror of
https://github.com/wassname/ray.git
synced 2026-07-08 04:50:24 +08:00
Wapaul1
289 lines
7.6 KiB
Python
289 lines
7.6 KiB
Python
# Code in this file is copied and adapted from
|
|
# https://github.com/openai/evolution-strategies-starter.
|
|
|
|
from __future__ import absolute_import
|
|
from __future__ import division
|
|
from __future__ import print_function
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
import functools
|
|
import os
|
|
|
|
# Tensorflow must be at least version 1.0.0 for the example to work.
|
|
if int(tf.__version__.split(".")[0]) < 1:
|
|
raise Exception("Your Tensorflow version is less than 1.0.0. Please update "
|
|
"Tensorflow to the latest version.")
|
|
|
|
# ================================================================
|
|
# Import all names into common namespace
|
|
# ================================================================
|
|
|
|
clip = tf.clip_by_value
|
|
|
|
# Make consistent with numpy
|
|
|
|
|
|
def sum(x, axis=None, keepdims=False):
|
|
return tf.reduce_sum(x, reduction_indices=None if axis is None else [axis],
|
|
keep_dims=keepdims)
|
|
|
|
|
|
def mean(x, axis=None, keepdims=False):
|
|
return tf.reduce_mean(x, reduction_indices=None if axis is None else [axis],
|
|
keep_dims=keepdims)
|
|
|
|
|
|
def var(x, axis=None, keepdims=False):
|
|
meanx = mean(x, axis=axis, keepdims=keepdims)
|
|
return mean(tf.square(x - meanx), axis=axis, keepdims=keepdims)
|
|
|
|
|
|
def std(x, axis=None, keepdims=False):
|
|
return tf.sqrt(var(x, axis=axis, keepdims=keepdims))
|
|
|
|
|
|
def max(x, axis=None, keepdims=False):
|
|
return tf.reduce_max(x, reduction_indices=None if axis is None else [axis],
|
|
keep_dims=keepdims)
|
|
|
|
|
|
def min(x, axis=None, keepdims=False):
|
|
return tf.reduce_min(x, reduction_indices=None if axis is None else [axis],
|
|
keep_dims=keepdims)
|
|
|
|
|
|
def concatenate(arrs, axis=0):
|
|
return tf.concat(arrs, axis)
|
|
|
|
|
|
def argmax(x, axis=None):
|
|
return tf.argmax(x, dimension=axis)
|
|
|
|
# Extras
|
|
|
|
|
|
def l2loss(params):
|
|
if len(params) == 0:
|
|
return tf.constant(0.0)
|
|
else:
|
|
return tf.add_n([sum(tf.square(p)) for p in params])
|
|
|
|
|
|
def lrelu(x, leak=0.2):
|
|
f1 = 0.5 * (1 + leak)
|
|
f2 = 0.5 * (1 - leak)
|
|
return f1 * x + f2 * abs(x)
|
|
|
|
|
|
def categorical_sample_logits(X):
|
|
# https://github.com/tensorflow/tensorflow/issues/456
|
|
U = tf.random_uniform(tf.shape(X))
|
|
return argmax(X - tf.log(-tf.log(U)), axis=1)
|
|
|
|
# Global session
|
|
|
|
|
|
def get_session():
|
|
return tf.get_default_session()
|
|
|
|
|
|
def single_threaded_session():
|
|
tf_config = tf.ConfigProto(inter_op_parallelism_threads=1,
|
|
intra_op_parallelism_threads=1)
|
|
return tf.Session(config=tf_config)
|
|
|
|
|
|
ALREADY_INITIALIZED = set()
|
|
|
|
|
|
def initialize():
|
|
new_variables = set(tf.global_variables()) - ALREADY_INITIALIZED
|
|
get_session().run(tf.variables_initializer(new_variables))
|
|
ALREADY_INITIALIZED.update(new_variables)
|
|
|
|
|
|
def eval(expr, feed_dict=None):
|
|
if feed_dict is None:
|
|
feed_dict = {}
|
|
return get_session().run(expr, feed_dict=feed_dict)
|
|
|
|
|
|
def set_value(v, val):
|
|
get_session().run(v.assign(val))
|
|
|
|
|
|
def load_state(fname):
|
|
saver = tf.train.Saver()
|
|
saver.restore(get_session(), fname)
|
|
|
|
|
|
def save_state(fname):
|
|
os.makedirs(os.path.dirname(fname), exist_ok=True)
|
|
saver = tf.train.Saver()
|
|
saver.save(get_session(), fname)
|
|
|
|
# Model components
|
|
|
|
|
|
def normc_initializer(std=1.0):
|
|
def _initializer(shape, dtype=None, partition_info=None):
|
|
out = np.random.randn(*shape).astype(np.float32)
|
|
out *= std / np.sqrt(np.square(out).sum(axis=0, keepdims=True))
|
|
return tf.constant(out)
|
|
return _initializer
|
|
|
|
|
|
def dense(x, size, name, weight_init=None, bias=True):
|
|
w = tf.get_variable(name + "/w", [x.get_shape()[1], size],
|
|
initializer=weight_init)
|
|
ret = tf.matmul(x, w)
|
|
if bias:
|
|
b = tf.get_variable(name + "/b", [size],
|
|
initializer=tf.zeros_initializer())
|
|
return ret + b
|
|
else:
|
|
return ret
|
|
|
|
# Basic Stuff
|
|
|
|
|
|
def function(inputs, outputs, updates=None, givens=None):
|
|
if isinstance(outputs, list):
|
|
return _Function(inputs, outputs, updates, givens=givens)
|
|
elif isinstance(outputs, dict):
|
|
f = _Function(inputs, outputs.values(), updates, givens=givens)
|
|
return lambda *inputs: dict(zip(outputs.keys(), f(*inputs)))
|
|
else:
|
|
f = _Function(inputs, [outputs], updates, givens=givens)
|
|
return lambda *inputs: f(*inputs)[0]
|
|
|
|
|
|
class _Function(object):
|
|
def __init__(self, inputs, outputs, updates, givens, check_nan=False):
|
|
assert all(len(i.op.inputs) == 0 for i in inputs), ("inputs should all be "
|
|
"placeholders")
|
|
self.inputs = inputs
|
|
updates = updates or []
|
|
self.update_group = tf.group(*updates)
|
|
self.outputs_update = list(outputs) + [self.update_group]
|
|
self.givens = {} if givens is None else givens
|
|
self.check_nan = check_nan
|
|
|
|
def __call__(self, *inputvals):
|
|
assert len(inputvals) == len(self.inputs)
|
|
feed_dict = dict(zip(self.inputs, inputvals))
|
|
feed_dict.update(self.givens)
|
|
results = get_session().run(self.outputs_update, feed_dict=feed_dict)[:-1]
|
|
if self.check_nan:
|
|
if any(np.isnan(r).any() for r in results):
|
|
raise RuntimeError("Nan detected")
|
|
return results
|
|
|
|
|
|
# Graph traversal
|
|
|
|
VARIABLES = {}
|
|
|
|
# Flat vectors
|
|
|
|
|
|
def var_shape(x):
|
|
out = [k.value for k in x.get_shape()]
|
|
assert all(isinstance(a, int) for a in out), ("shape function assumes that "
|
|
"shape is fully known")
|
|
return out
|
|
|
|
|
|
def numel(x):
|
|
return intprod(var_shape(x))
|
|
|
|
|
|
def intprod(x):
|
|
return int(np.prod(x))
|
|
|
|
|
|
def flatgrad(loss, var_list):
|
|
grads = tf.gradients(loss, var_list)
|
|
return tf.concat([tf.reshape(grad, [numel(v)], 0)
|
|
for (v, grad) in zip(var_list, grads)])
|
|
|
|
|
|
class SetFromFlat(object):
|
|
def __init__(self, var_list, dtype=tf.float32):
|
|
assigns = []
|
|
shapes = list(map(var_shape, var_list))
|
|
total_size = np.sum([intprod(shape) for shape in shapes])
|
|
|
|
self.theta = theta = tf.placeholder(dtype, [total_size])
|
|
start = 0
|
|
assigns = []
|
|
for (shape, v) in zip(shapes, var_list):
|
|
size = intprod(shape)
|
|
assigns.append(tf.assign(v, tf.reshape(theta[start:start + size],
|
|
shape)))
|
|
start += size
|
|
assert start == total_size
|
|
self.op = tf.group(*assigns)
|
|
|
|
def __call__(self, theta):
|
|
get_session().run(self.op, feed_dict={self.theta: theta})
|
|
|
|
|
|
class GetFlat(object):
|
|
def __init__(self, var_list):
|
|
self.op = tf.concat([tf.reshape(v, [numel(v)]) for v in var_list], 0)
|
|
|
|
def __call__(self):
|
|
return get_session().run(self.op)
|
|
|
|
# Misc
|
|
|
|
|
|
def scope_vars(scope, trainable_only):
|
|
"""Get variables inside a scope. The scope can be specified as a string."""
|
|
return tf.get_collection((tf.GraphKeys.TRAINABLE_VARIABLES if trainable_only
|
|
else tf.GraphKeys.GLOBAL_VARIABLES),
|
|
scope=(scope if isinstance(scope, str)
|
|
else scope.name))
|
|
|
|
|
|
def in_session(f):
|
|
@functools.wraps(f)
|
|
def newfunc(*args, **kwargs):
|
|
with tf.Session():
|
|
f(*args, **kwargs)
|
|
return newfunc
|
|
|
|
|
|
# A mapping from name -> (placeholder, dtype, shape).
|
|
_PLACEHOLDER_CACHE = {}
|
|
|
|
|
|
def get_placeholder(name, dtype, shape):
|
|
print("calling get_placeholder", name)
|
|
if name in _PLACEHOLDER_CACHE:
|
|
out, dtype1, shape1 = _PLACEHOLDER_CACHE[name]
|
|
assert dtype1 == dtype and shape1 == shape
|
|
return out
|
|
else:
|
|
out = tf.placeholder(dtype=dtype, shape=shape, name=name)
|
|
_PLACEHOLDER_CACHE[name] = (out, dtype, shape)
|
|
return out
|
|
|
|
|
|
def get_placeholder_cached(name):
|
|
return _PLACEHOLDER_CACHE[name][0]
|
|
|
|
|
|
def flattenallbut0(x):
|
|
return tf.reshape(x, [-1, intprod(x.get_shape().as_list()[1:])])
|
|
|
|
|
|
def reset():
|
|
global _PLACEHOLDER_CACHE
|
|
global VARIABLES
|
|
_PLACEHOLDER_CACHE = {}
|
|
VARIABLES = {}
|
|
tf.reset_default_graph()
|