mirror of
https://github.com/wassname/ray.git
synced 2026-08-17 11:25:34 +08:00
Robert Nishihara
3c5375345f
* Initial commit of evolution strategies example. * Some small simplifications. * Update example to use new API. * Add example to documentation.
252 lines
9.9 KiB
Python
252 lines
9.9 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 logging
|
|
import pickle
|
|
|
|
import h5py
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
|
|
import tf_util as U
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class Policy:
|
|
def __init__(self, *args, **kwargs):
|
|
self.args, self.kwargs = args, kwargs
|
|
self.scope = self._initialize(*args, **kwargs)
|
|
self.all_variables = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, self.scope.name)
|
|
|
|
self.trainable_variables = tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES, self.scope.name)
|
|
self.num_params = sum(int(np.prod(v.get_shape().as_list())) for v in self.trainable_variables)
|
|
self._setfromflat = U.SetFromFlat(self.trainable_variables)
|
|
self._getflat = U.GetFlat(self.trainable_variables)
|
|
|
|
logger.info('Trainable variables ({} parameters)'.format(self.num_params))
|
|
for v in self.trainable_variables:
|
|
shp = v.get_shape().as_list()
|
|
logger.info('- {} shape:{} size:{}'.format(v.name, shp, np.prod(shp)))
|
|
logger.info('All variables')
|
|
for v in self.all_variables:
|
|
shp = v.get_shape().as_list()
|
|
logger.info('- {} shape:{} size:{}'.format(v.name, shp, np.prod(shp)))
|
|
|
|
placeholders = [tf.placeholder(v.value().dtype, v.get_shape().as_list()) for v in self.all_variables]
|
|
self.set_all_vars = U.function(
|
|
inputs=placeholders,
|
|
outputs=[],
|
|
updates=[tf.group(*[v.assign(p) for v, p in zip(self.all_variables, placeholders)])]
|
|
)
|
|
|
|
def _initialize(self, *args, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
def save(self, filename):
|
|
assert filename.endswith('.h5')
|
|
with h5py.File(filename, 'w') as f:
|
|
for v in self.all_variables:
|
|
f[v.name] = v.eval()
|
|
# TODO: it would be nice to avoid pickle, but it's convenient to pass Python objects to _initialize
|
|
# (like Gym spaces or numpy arrays)
|
|
f.attrs['name'] = type(self).__name__
|
|
f.attrs['args_and_kwargs'] = np.void(pickle.dumps((self.args, self.kwargs), protocol=-1))
|
|
|
|
@classmethod
|
|
def Load(cls, filename, extra_kwargs=None):
|
|
with h5py.File(filename, 'r') as f:
|
|
args, kwargs = pickle.loads(f.attrs['args_and_kwargs'].tostring())
|
|
if extra_kwargs:
|
|
kwargs.update(extra_kwargs)
|
|
policy = cls(*args, **kwargs)
|
|
policy.set_all_vars(*[f[v.name][...] for v in policy.all_variables])
|
|
return policy
|
|
|
|
# === Rollouts/training ===
|
|
|
|
def rollout(self, env, *, render=False, timestep_limit=None, save_obs=False, random_stream=None):
|
|
"""
|
|
If random_stream is provided, the rollout will take noisy actions with noise drawn from that stream.
|
|
Otherwise, no action noise will be added.
|
|
"""
|
|
env_timestep_limit = env.spec.tags.get('wrapper_config.TimeLimit.max_episode_steps')
|
|
timestep_limit = env_timestep_limit if timestep_limit is None else min(timestep_limit, env_timestep_limit)
|
|
rews = []
|
|
t = 0
|
|
if save_obs:
|
|
obs = []
|
|
ob = env.reset()
|
|
for _ in range(timestep_limit):
|
|
ac = self.act(ob[None], random_stream=random_stream)[0]
|
|
if save_obs:
|
|
obs.append(ob)
|
|
ob, rew, done, _ = env.step(ac)
|
|
rews.append(rew)
|
|
t += 1
|
|
if render:
|
|
env.render()
|
|
if done:
|
|
break
|
|
rews = np.array(rews, dtype=np.float32)
|
|
if save_obs:
|
|
return rews, t, np.array(obs)
|
|
return rews, t
|
|
|
|
def act(self, ob, random_stream=None):
|
|
raise NotImplementedError
|
|
|
|
def set_trainable_flat(self, x):
|
|
self._setfromflat(x)
|
|
|
|
def get_trainable_flat(self):
|
|
return self._getflat()
|
|
|
|
@property
|
|
def needs_ob_stat(self):
|
|
raise NotImplementedError
|
|
|
|
def set_ob_stat(self, ob_mean, ob_std):
|
|
raise NotImplementedError
|
|
|
|
|
|
def bins(x, dim, num_bins, name):
|
|
scores = U.dense(x, dim * num_bins, name, U.normc_initializer(0.01))
|
|
scores_nab = tf.reshape(scores, [-1, dim, num_bins])
|
|
return tf.argmax(scores_nab, 2) # 0 ... num_bins-1
|
|
|
|
|
|
class MujocoPolicy(Policy):
|
|
def _initialize(self, ob_space, ac_space, ac_bins, ac_noise_std, nonlin_type, hidden_dims, connection_type):
|
|
self.ac_space = ac_space
|
|
self.ac_bins = ac_bins
|
|
self.ac_noise_std = ac_noise_std
|
|
self.hidden_dims = hidden_dims
|
|
self.connection_type = connection_type
|
|
|
|
assert len(ob_space.shape) == len(self.ac_space.shape) == 1
|
|
assert np.all(np.isfinite(self.ac_space.low)) and np.all(np.isfinite(self.ac_space.high)), \
|
|
'Action bounds required'
|
|
|
|
self.nonlin = {'tanh': tf.tanh, 'relu': tf.nn.relu, 'lrelu': U.lrelu, 'elu': tf.nn.elu}[nonlin_type]
|
|
|
|
with tf.variable_scope(type(self).__name__) as scope:
|
|
# Observation normalization
|
|
ob_mean = tf.get_variable(
|
|
'ob_mean', ob_space.shape, tf.float32, tf.constant_initializer(np.nan), trainable=False)
|
|
ob_std = tf.get_variable(
|
|
'ob_std', ob_space.shape, tf.float32, tf.constant_initializer(np.nan), trainable=False)
|
|
in_mean = tf.placeholder(tf.float32, ob_space.shape)
|
|
in_std = tf.placeholder(tf.float32, ob_space.shape)
|
|
self._set_ob_mean_std = U.function([in_mean, in_std], [], updates=[
|
|
tf.assign(ob_mean, in_mean),
|
|
tf.assign(ob_std, in_std),
|
|
])
|
|
|
|
# Policy network
|
|
o = tf.placeholder(tf.float32, [None] + list(ob_space.shape))
|
|
a = self._make_net(tf.clip_by_value((o - ob_mean) / ob_std, -5.0, 5.0))
|
|
self._act = U.function([o], a)
|
|
return scope
|
|
|
|
def _make_net(self, o):
|
|
# Process observation
|
|
if self.connection_type == 'ff':
|
|
x = o
|
|
for ilayer, hd in enumerate(self.hidden_dims):
|
|
x = self.nonlin(U.dense(x, hd, 'l{}'.format(ilayer), U.normc_initializer(1.0)))
|
|
else:
|
|
raise NotImplementedError(self.connection_type)
|
|
|
|
# Map to action
|
|
adim, ahigh, alow = self.ac_space.shape[0], self.ac_space.high, self.ac_space.low
|
|
assert isinstance(self.ac_bins, str)
|
|
ac_bin_mode, ac_bin_arg = self.ac_bins.split(':')
|
|
|
|
if ac_bin_mode == 'uniform':
|
|
# Uniformly spaced bins, from ac_space.low to ac_space.high
|
|
num_ac_bins = int(ac_bin_arg)
|
|
aidx_na = bins(x, adim, num_ac_bins, 'out') # 0 ... num_ac_bins-1
|
|
ac_range_1a = (ahigh - alow)[None, :]
|
|
a = 1. / (num_ac_bins - 1.) * tf.to_float(aidx_na) * ac_range_1a + alow[None, :]
|
|
|
|
elif ac_bin_mode == 'custom':
|
|
# Custom bins specified as a list of values from -1 to 1
|
|
# The bins are rescaled to ac_space.low to ac_space.high
|
|
acvals_k = np.array(list(map(float, ac_bin_arg.split(','))), dtype=np.float32)
|
|
logger.info('Custom action values: ' + ' '.join('{:.3f}'.format(x) for x in acvals_k))
|
|
assert acvals_k.ndim == 1 and acvals_k[0] == -1 and acvals_k[-1] == 1
|
|
acvals_ak = (
|
|
(ahigh - alow)[:, None] / (acvals_k[-1] - acvals_k[0]) * (acvals_k - acvals_k[0])[None, :]
|
|
+ alow[:, None]
|
|
)
|
|
|
|
aidx_na = bins(x, adim, len(acvals_k), 'out') # values in [0, k-1]
|
|
a = tf.gather_nd(
|
|
acvals_ak,
|
|
tf.concat([
|
|
tf.tile(np.arange(adim)[None, :, None], [tf.shape(aidx_na)[0], 1, 1]),
|
|
2,
|
|
tf.expand_dims(aidx_na, -1)
|
|
]) # (n,a,2)
|
|
) # (n,a)
|
|
elif ac_bin_mode == 'continuous':
|
|
a = U.dense(x, adim, 'out', U.normc_initializer(0.01))
|
|
else:
|
|
raise NotImplementedError(ac_bin_mode)
|
|
|
|
return a
|
|
|
|
def act(self, ob, random_stream=None):
|
|
a = self._act(ob)
|
|
if random_stream is not None and self.ac_noise_std != 0:
|
|
a += random_stream.randn(*a.shape) * self.ac_noise_std
|
|
return a
|
|
|
|
@property
|
|
def needs_ob_stat(self):
|
|
return True
|
|
|
|
@property
|
|
def needs_ref_batch(self):
|
|
return False
|
|
|
|
def set_ob_stat(self, ob_mean, ob_std):
|
|
self._set_ob_mean_std(ob_mean, ob_std)
|
|
|
|
def initialize_from(self, filename, ob_stat=None):
|
|
"""
|
|
Initializes weights from another policy, which must have the same architecture (variable names),
|
|
but the weight arrays can be smaller than the current policy.
|
|
"""
|
|
with h5py.File(filename, 'r') as f:
|
|
f_var_names = []
|
|
f.visititems(lambda name, obj: f_var_names.append(name) if isinstance(obj, h5py.Dataset) else None)
|
|
assert set(v.name for v in self.all_variables) == set(f_var_names), 'Variable names do not match'
|
|
|
|
init_vals = []
|
|
for v in self.all_variables:
|
|
shp = v.get_shape().as_list()
|
|
f_shp = f[v.name].shape
|
|
assert len(shp) == len(f_shp) and all(a >= b for a, b in zip(shp, f_shp)), \
|
|
'This policy must have more weights than the policy to load'
|
|
init_val = v.eval()
|
|
# ob_mean and ob_std are initialized with nan, so set them manually
|
|
if 'ob_mean' in v.name:
|
|
init_val[:] = 0
|
|
init_mean = init_val
|
|
elif 'ob_std' in v.name:
|
|
init_val[:] = 0.001
|
|
init_std = init_val
|
|
# Fill in subarray from the loaded policy
|
|
init_val[tuple([np.s_[:s] for s in f_shp])] = f[v.name]
|
|
init_vals.append(init_val)
|
|
self.set_all_vars(*init_vals)
|
|
|
|
if ob_stat is not None:
|
|
ob_stat.set_from_init(init_mean, init_std, init_count=1e5)
|