def policy_fn(noptions=64,
nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None,
action_placeholder=None,
option_z_placeholder=None,
dones_placeholder=None):
ob_space = env.observation_space
ac_n = get_action_dim(env)
X = observ_placeholder if observ_placeholder is not None else \
observation_placeholder(ob_space, batch_size=nbatch, name='ob')
X_next = observation_placeholder(ob_space,
batch_size=nbatch,
name='ob_next')
ac = action_placeholder if action_placeholder is not None else \
tf.placeholder(tf.float32, shape=(nbatch, ac_n), name='ac')
op = option_z_placeholder if option_z_placeholder is not None else \
tf.placeholder(tf.float32, shape=(nbatch, noptions), name='op_z')
dones = dones_placeholder if dones_placeholder is not None else \
tf.placeholder(tf.float32, shape=(nbatch), name='dones')
extra_tensors = {}
cnn_fm, policy_latent, vf = pi_vf_fn(X, extra_tensors, nbatch, nsteps)
next_cnn_fm, _, next_vf = pi_vf_fn(X_next,
extra_tensors,
nbatch,
nsteps,
recurrent_subname='next')
assert noptions == cnn_fm.get_shape().as_list()[-1], \
'number of options for VFO should equal to channels of last conv layer'
tf.assert_rank(policy_latent, 2)
option_latent = tf.concat([policy_latent, op], 1)
q_latent = tf.concat([policy_latent, op, ac], 1)
policy = PolicyWithValue(env=env,
observations=X,
next_observations=X_next,
actions=ac,
option_z=op,
dones=dones,
feature_map=cnn_fm,
next_feature_map=next_cnn_fm,
latent=policy_latent,
option_latent=tf.stop_gradient(option_latent),
q_latent=tf.stop_gradient(q_latent),
vf=vf,
next_vf=next_vf,
sess=sess,
**extra_tensors)
return policy
def policy_fn(obs_ph=None,
normalize_observations=False,
vf_latent=None,
sess=None):
# preprocess input
ob_space = env.observation_space
X = obs_ph if obs_ph is not None else observation_placeholder(ob_space)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent = policy_network(encoded_x)
policy = PolicyModel(env=env,
observations=X,
policy_latent=policy_latent,
vf_latent=vf_latent,
estimate_q=estimate_q,
q_network=q_network,
sess=sess,
**extra_tensors)
return policy
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None):
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent_outputs = policy_network(encoded_x)
policy_latent = policy_latent_outputs[-1]
if isinstance(policy_latent, tuple):
policy_latent, recurrent_tensors = policy_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
# TODO recurrent architectures are not supported with value_network=copy yet
vf_latent = _v_net(encoded_x)
policy = PolicyWithValue(env=env,
observations=X,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
policy_latent_outputs.append(policy.pi)
feature_maps = tf.concat(policy_latent_outputs, -1)
return (policy, feature_maps)
def __init__(self, checkpoint_path):
player_base.PlayerBase.__init__(self)
self._action_set = 'default'
self._player_prefix = 'player_0'
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
self._sess = tf.Session(config=config)
stacking = 4
self._stacker = ObservationStacker(stacking)
with tf.variable_scope(self._player_prefix):
with tf.variable_scope('ppo2_model'):
env = DummyEnv(self._action_set, stacking)
ob_space = env.observation_space
X = observation_placeholder(ob_space, batch_size=1)
extra_tensors = {}
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent = gfootball_impala_cnn_network_fn(encoded_x)
self._policy = PolicyWithValue(
env=env,
observations=X,
latent=policy_latent,
vf_latent=policy_latent,
sess=self._sess,
estimate_q=False,
**extra_tensors
)
_load_variables(checkpoint_path, self._sess, prefix=self._player_prefix + '/')
saver = tf.train.Saver()
saver.save(self._sess, "/home/alex/Dropbox/projects/python/kaggle/football/saved_models/11_vs_11_easy_stochastic_v2/11_vs_11_easy_stochastic_v2")
def actor_fn(obs_ph=None,
normalize_observations=False,
vf_latent=None,
sess=None):
# preprocess input
ob_space = observation_space
X = obs_ph if obs_ph is not None else observation_placeholder(ob_space)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
actor_latent = actor_network(encoded_x)
if is_determined is False:
actor_model = StochasticActorModel(action_space=action_space,
observations=X,
actor_latent=actor_latent,
vf_latent=vf_latent,
sess=sess,
**extra_tensors)
else:
actor_model = DeterminedActorModel(action_space=action_space,
observations=X,
actor_latent=actor_latent,
vf_latent=vf_latent,
sess=sess,
**extra_tensors)
return actor_model
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None):
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent = policy_network(encoded_x)
if isinstance(policy_latent, tuple):
policy_latent, recurrent_tensors = policy_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
# TODO recurrent architectures are not supported with value_network=copy yet
vf_latent = _v_net(encoded_x)
# Daniel: doing this for action range.
if limit_act_range:
policy_latent = tf.nn.tanh(
fc(policy_latent,
'pi',
env.action_space.shape[0],
init_scale=0.01,
init_bias=0.0))
policy = PolicyWithValue(env=env,
observations=X,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None,
scope_name="pi"):
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
#encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent, recurrent_tensors = policy_network(encoded_x)
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
vf_latent, _ = _v_net(encoded_x)
policy = PolicyWithValue(env=env,
observations=X,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
trainable_variance=trainable_variance,
trainable_bias=trainable_bias,
init_logstd=init_logstd,
clip=clip,
scope_name=scope_name,
**extra_tensors)
return policy
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None):
print("IN HERE JIJIJIJI cdcddd", str(normalize_observations),
"hdiuwdhwiuh")
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent, recurrent_tensors = policy_network(encoded_x)
print(" IN POLICY LATENT ", str(policy_latent), " and recurrent ",
str(recurrent_tensors))
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
print("policy_network is " + str(policy_network))
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
vf_latent, _ = _v_net(encoded_x)
policy = PolicyWithValue(env=env,
observations=X,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def __init__(self, env, nbatch, network, **policy_kwargs):
ob_space = env.observation_space
self.X = observation_placeholder(ob_space, batch_size=nbatch)
encoded_x = encode_observation(ob_space, self.X)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
self.net = get_network_builder(network)(**policy_kwargs)
self.h1 = self.net(encoded_x)
self.h2 = fc(self.h1, 'vf', 1)
self.out = self.h2[:, 0]
def train_copos(args):
import baselines.common.tf_util as U
sess = U.single_threaded_session()
sess.__enter__()
if MPI is None or MPI.COMM_WORLD.Get_rank() == 0:
rank = 0
configure_logger(args.log_path)
else:
rank = MPI.COMM_WORLD.Get_rank()
configure_logger(args.log_path, format_strs=[])
workerseed = args.seed + 10000 * MPI.COMM_WORLD.Get_rank()
#def policy_fn(name, ob_space, ac_space):
# return CompatibleMlpPolicy(name=name, ob_space=ob_space, ac_space=ac_space,
# hid_size=32, num_hid_layers=2)
set_global_seeds(workerseed)
env = build_env(args, normalize_ob=True)
#env = gym.make(args.env)
#env.seed(workerseed)
timesteps_per_batch = 10000
#timesteps_per_batch=2048
beta = -1
if beta < 0:
nr_episodes = int(args.num_timesteps) // timesteps_per_batch
# Automatically compute beta based on initial entropy and number of iterations
policy = build_policy(env, "mlp", value_network='copy', copos=True)
ob = observation_placeholder(env.observation_space)
with tf.variable_scope("tmp_pi"):
tmp_pi = policy(observ_placeholder=ob)
sess.run(tf.global_variables_initializer())
tmp_ob = np.zeros((1, ) + env.observation_space.shape)
entropy = sess.run(tmp_pi.pd.entropy(), feed_dict={tmp_pi.X: tmp_ob})
beta = 2 * entropy / nr_episodes
print("Initial entropy: " + str(entropy) + ", episodes: " +
str(nr_episodes))
print("Automatically set beta: " + str(beta))
copos_mpi.learn(network='mlp',
env=env,
seed=args.seed,
timesteps_per_batch=timesteps_per_batch,
epsilon=0.01,
beta=beta,
cg_iters=10,
cg_damping=0.1,
max_timesteps=int(args.num_timesteps),
gamma=0.99,
lam=0.98,
vf_iters=5,
vf_stepsize=1e-3)
env.close()
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None):
ob_space = env.observation_space
observation_plh = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and observation_plh.dtype == tf.float32:
ob_plh_normalize_clip, rms = _normalize_clip_observation(
observation_plh)
extra_tensors['rms'] = rms
else:
ob_plh_normalize_clip = observation_plh
ob_plh_normalize_clip = encode_observation(ob_space,
ob_plh_normalize_clip)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent, recurrent_tensors = policy_network(
ob_plh_normalize_clip)
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
ob_plh_normalize_clip, nenv)
extra_tensors.update(recurrent_tensors)
val_net = value_network
if val_net is None or val_net == 'shared':
vf_latent = policy_latent
else:
if val_net == 'copy':
val_net = policy_network
else:
assert callable(val_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
vf_latent, _ = val_net(ob_plh_normalize_clip)
policy = PolicyWithValue(env=env,
observations=observation_plh,
policy_latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def policy_fn(nbatch=None, nsteps=None, sess=None, observ_placeholder=None):
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent = policy_network(encoded_x)
if isinstance(policy_latent, tuple):
policy_latent, recurrent_tensors = policy_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
# TODO recurrent architectures are not supported with value_network=copy yet
vf_latent = _v_net(encoded_x)
policy = PolicyWithValue(
env=env,
observations=X,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors
)
return policy
def train(args,extra_args):
env_type, env_id = get_env_type(args)
print('env_type: {}'.format(env_type))
total_timesteps = int(args.num_timesteps)
seed = args.seed
set_global_seeds(seed)
#workerseed = seed + 10000 * MPI.COMM_WORLD.Get_rank()
#set_global_seeds(workerseed)
learn = get_learn_function(args.alg)
alg_kwargs = get_learn_function_defaults(args.alg, env_type)
alg_kwargs.update(extra_args)
env = build_env(args,normalize_ob=False,normalize_ret=False)
if args.save_video_interval != 0:
env = VecVideoRecorder(env, osp.join(logger.get_dir(), "videos"), record_video_trigger=lambda x: x % args.save_video_interval == 0, video_length=args.save_video_length)
if args.network:
alg_kwargs['network'] = args.network
else:
if alg_kwargs.get('network') is None:
alg_kwargs['network'] = get_default_network(env_type)
#timesteps_per_batch=1024
#timesteps_per_batch=2048
beta = -1
if beta < 0:
#print(alg_kwargs)
nr_episodes = total_timesteps // alg_kwargs['timesteps_per_batch']
# Automatically compute beta based on initial entropy and number of iterations
policy = build_policy(env, alg_kwargs['network'], value_network='copy', normalize_observations=alg_kwargs['normalize_observations'], copos=True)
ob = observation_placeholder(env.observation_space)
sess = U.single_threaded_session()
sess.__enter__()
with tf.variable_scope("tmp_pi"):
tmp_pi = policy(observ_placeholder=ob)
sess.run(tf.global_variables_initializer())
tmp_ob = np.zeros((1,) + env.observation_space.shape)
entropy = sess.run(tmp_pi.pd.entropy(), feed_dict={tmp_pi.X: tmp_ob})
#beta = 2 * entropy / nr_episodes
beta = 0
print("Initial entropy: " + str(entropy) + ", episodes: " + str(nr_episodes))
print("Constantly set beta: " + str(beta))
print('Training {} on {}:{} with arguments \n{}'.format(args.alg, env_type, env_id, alg_kwargs))
model=learn(env=env, seed=seed, beta=beta,
total_timesteps=total_timesteps,
**alg_kwargs)
return model, env
def policy_fn(scope_name="pi",
nbatch=None,
nsteps=None,
sess=sess,
observ_placeholder=None):
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent, recurrent_tensors = policy_network(encoded_x)
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
policy = Policy(observations=X,
action_space=ac_space,
latent=policy_latent,
sess=sess,
train=train,
beta=beta,
l2=l2,
lr=lr,
init_scale=init_scale,
init_bias=init_bias,
trainable_variance=trainable_variance,
trainable_bias=trainable_bias,
init_logstd=init_logstd,
scope_name=scope_name,
clip=clip,
class_weights=class_weights,
**extra_tensors)
return policy
def __init__(self,
*,
network,
env,
ppo_lr=3e-4,
cliprange=0.2,
nsteps=128,
nminibatches=4,
ent_coef=0.0,
vf_coef=0.25,
max_grad_norm=0.5,
gamma=0.99,
lam=0.95,
mpi_rank_weight=1,
comm=None,
**network_kwargs):
self.env = env
self.obs_ph = observation_placeholder(env.observation_space)
self.nsteps = nsteps
self.nminibatches = nminibatches
self.nenvs = self.env.num_envs
self.nsteps = nsteps
self.nbatch = self.nenvs * self.nsteps
self.nbatch_train = self.nbatch // nminibatches
self.ppo_model = PPO(network=network,
env=env,
obs_ph=self.obs_ph,
lr=ppo_lr,
cliprange=cliprange,
nsteps=nsteps,
nminibatches=nminibatches,
ent_coef=ent_coef,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm,
gamma=gamma,
lam=lam,
mpi_rank_weight=mpi_rank_weight,
comm=comm,
load_path=None,
**network_kwargs)
def __init__(self, ob_space, sess=None):
logger.info("Using RNDReward")
self.sess = sess or tf.get_default_session()
# RND.
num_inp = ob_space.shape[0]
num_hid1 = 64 #num_inp // 2
num_hid_pred = 32
rep_size = 10
proportion_of_exp_used_for_predictor_update = 1.
self.X = observation_placeholder(ob_space)
# Random target network.
logger.info("CnnTarget: using shape %s as observation input" % (str(ob_space.shape)))
xr = self.X
xr = tf.nn.leaky_relu(fc(xr, 'fc1r', nh=num_hid1, init_scale=np.sqrt(2)))
X_r = fc(xr, 'fc2r', nh=rep_size, init_scale=np.sqrt(2))
# Predictor network.
xrp = self.X
xrp = tf.nn.leaky_relu(fc(xr, 'fc1r_pred', nh=num_hid1, init_scale=np.sqrt(2)))
X_r_hat = tf.nn.relu(fc(xrp, 'fc1r_hat1_pred', nh=num_hid_pred, init_scale=np.sqrt(2)))
X_r_hat = tf.nn.relu(fc(X_r_hat, 'fc1r_hat2_pred', nh=num_hid_pred, init_scale=np.sqrt(2)))
X_r_hat = fc(X_r_hat, 'fc1r_hat3_pred', nh=rep_size, init_scale=np.sqrt(2))
self.feat_var = tf.reduce_mean(tf.nn.moments(X_r, axes=[0])[1])
self.max_feat = tf.reduce_max(tf.abs(X_r))
self.int_rew = tf.reduce_mean(tf.square(tf.stop_gradient(X_r) - X_r_hat), axis=1)
# self.int_rew = tf.reshape(self.int_rew, (self.sy_nenvs, self.sy_nsteps - 1))
noisy_targets = tf.stop_gradient(X_r)
self.aux_loss = tf.reduce_mean(tf.square(noisy_targets - X_r_hat), -1)
mask = tf.random_uniform(shape=tf.shape(self.aux_loss), minval=0., maxval=1., dtype=tf.float32)
mask = tf.cast(mask < proportion_of_exp_used_for_predictor_update, tf.float32)
self.aux_loss = tf.reduce_sum(mask * self.aux_loss) / tf.maximum(tf.reduce_sum(mask), 1.)
optimizer = tf.train.AdamOptimizer(0.01)
self.train = optimizer.minimize(self.aux_loss)
tf.global_variables_initializer().run(session=sess)
def __init__(self, ob_space, sess=None):
logger.info("Using AEReward")
self.sess = sess or tf.get_default_session()
self.X = observation_placeholder(ob_space)
initializer = tf.variance_scaling_initializer()
num_inp = ob_space.shape[0]
num_hid1 = num_inp // 2
# num_hid2 = 16
# num_hid3 = num_hid1
w1 = tf.Variable(initializer([num_inp, num_hid1]), dtype=tf.float32)
# w2 = tf.Variable(initializer([num_hid1, num_hid2]), dtype=tf.float32)
# w3 = tf.Variable(initializer([num_hid2, num_hid3]), dtype=tf.float32)
# w4 = tf.Variable(initializer([num_hid3, num_inp]), dtype=tf.float32)
w4 = tf.Variable(initializer([num_hid1, num_inp]), dtype=tf.float32)
# h = tf.layers.flatten(X)
# h = fc(h, 'mlp_fc{}'.format(i), nh=num_hid1, init_scale=np.sqrt(2))
# h = fc(h, 'mlp_fc{}'.format(i), nh=num_hid2, init_scale=np.sqrt(2))
# h = fc(h, 'mlp_fc{}'.format(i), nh=num_hid1, init_scale=np.sqrt(2))
b1 = tf.Variable(tf.zeros(num_hid1))
# b2 = tf.Variable(tf.zeros(num_hid2))
# b3 = tf.Variable(tf.zeros(num_hid3))
b4 = tf.Variable(tf.zeros(num_inp))
hid_layer1 = tf.nn.relu(tf.matmul(self.X, w1) + b1)
# hid_layer2 = tf.nn.relu(tf.matmul(hid_layer1, w2) + b2)
# hid_layer3 = tf.nn.relu(tf.matmul(hid_layer2, w3) + b3)
output_layer = tf.nn.relu(tf.matmul(hid_layer1, w4) + b4)
self.loss = tf.reduce_mean(tf.square(output_layer - self.X))
self.bonus = tf.reduce_mean(tf.square(output_layer - self.X), 1) # novelty reward as AE loss per observation
optimizer = tf.train.AdamOptimizer(0.01)
self.train = optimizer.minimize(self.loss)
tf.global_variables_initializer().run(session=sess)
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None):
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(ob_space, encoded_x)
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent, policy_latent_mean, info_loss = policy_network(
encoded_x)
if isinstance(policy_latent, tuple):
raise NotImplementedError()
policy = PolicyWithValue(
env=env,
observations=X,
arch=arch,
latent=policy_latent,
latent_mean=policy_latent_mean,
info_loss=info_loss,
# vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def policy_fn(nbatch=None, nsteps=None, sess=None):
ob_space = env.observation_space
Xs = tf.stack([observation_placeholder(ob_space, batch_size=nbatch)] *
3)
extra_tensors = {}
encoded_x_0 = encode_observation(ob_space, Xs[0])
encoded_x_1 = encode_observation(ob_space, Xs[1])
encoded_x_2 = encode_observation(ob_space, Xs[2])
with tf.variable_scope('pi'):
_, f_features_0 = policy_network(encoded_x_0)
with tf.variable_scope('pi', reuse=True):
policy_latent, f_features_1 = policy_network(encoded_x_1)
with tf.variable_scope('pi', reuse=True):
_, f_features_2 = policy_network(encoded_x_2)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
raise NotImplementedError
policy = ModifiedPolicyWithValue(
env=env,
observations=Xs,
latent=policy_latent,
f_features=[f_features_0, f_features_1, f_features_2],
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def learn(
*,
network,
env,
total_timesteps,
timesteps_per_batch=1024, # what to train on
max_kl=0.001,
cg_iters=10,
gamma=0.99,
lam=1.0, # advantage estimation
seed=None,
ent_coef=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=3,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
load_path=None,
**network_kwargs):
'''
learn a policy function with TRPO algorithm
Parameters:
----------
network neural network to learn. Can be either string ('mlp', 'cnn', 'lstm', 'lnlstm' for basic types)
or function that takes input placeholder and returns tuple (output, None) for feedforward nets
or (output, (state_placeholder, state_output, mask_placeholder)) for recurrent nets
env environment (one of the gym environments or wrapped via baselines.common.vec_env.VecEnv-type class
timesteps_per_batch timesteps per gradient estimation batch
max_kl max KL divergence between old policy and new policy ( KL(pi_old || pi) )
ent_coef coefficient of policy entropy term in the optimization objective
cg_iters number of iterations of conjugate gradient algorithm
cg_damping conjugate gradient damping
vf_stepsize learning rate for adam optimizer used to optimie value function loss
vf_iters number of iterations of value function optimization iterations per each policy optimization step
total_timesteps max number of timesteps
max_episodes max number of episodes
max_iters maximum number of policy optimization iterations
callback function to be called with (locals(), globals()) each policy optimization step
load_path str, path to load the model from (default: None, i.e. no model is loaded)
**network_kwargs keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
Returns:
-------
learnt model
'''
if MPI is not None:
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
else:
nworkers = 1
rank = 0
cpus_per_worker = 1
U.get_session(
config=tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=cpus_per_worker,
intra_op_parallelism_threads=cpus_per_worker))
policy = build_policy(env, network, value_network='copy', **network_kwargs)
set_global_seeds(seed)
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
ob = observation_placeholder(ob_space)
with tf.variable_scope("pi"):
pi = policy(observ_placeholder=ob)
with tf.variable_scope("oldpi"):
oldpi = policy(observ_placeholder=ob)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = ent_coef * meanent
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
dist = meankl
all_var_list = get_trainable_variables("pi")
# var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
# vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
vfadam = MpiAdam(vf_var_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
start += sz
gvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(klgrads, tangents)
]) #pylint: disable=E1111
fvp = U.flatgrad(gvp, var_list)
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(get_variables("oldpi"), get_variables("pi"))
])
compute_losses = U.function([ob, ac, atarg], losses)
compute_lossandgrad = U.function([ob, ac, atarg], losses +
[U.flatgrad(optimgain, var_list)])
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob, ret],
U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - tstart),
color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
if MPI is not None:
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
else:
out = np.copy(x)
return out
U.initialize()
if load_path is not None:
pi.load(load_path)
th_init = get_flat()
if MPI is not None:
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
timesteps_per_batch,
stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
if sum([max_iters > 0, total_timesteps > 0, max_episodes > 0]) == 0:
# noththing to be done
return pi
assert sum([max_iters>0, total_timesteps>0, max_episodes>0]) < 2, \
'out of max_iters, total_timesteps, and max_episodes only one should be specified'
while True:
if callback: callback(locals(), globals())
if total_timesteps and timesteps_so_far >= total_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
logger.log("********** Iteration %i ************" % iters_so_far)
with timed("sampling"):
seg = seg_gen.__next__()
add_vtarg_and_adv(seg, gamma, lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()
) / atarg.std() # standardized advantage function estimate
if hasattr(pi, "ret_rms"): pi.ret_rms.update(tdlamret)
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
assign_old_eq_new() # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > max_kl * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(), vfadam.getflat().sum())) # list of tuples
assert all(
np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
if MPI is not None:
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
else:
listoflrpairs = [lrlocal]
lens, rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if rank == 0:
logger.dump_tabular()
return pi
def __init__(self, timed, policy, ob_space, ac_space, max_kl=0.001, cg_iters=10,
ent_coef=0.0,cg_damping=1e-2,vf_stepsize=3e-4,vf_iters =3,load_path=None, num_reward=1, index=1):
if MPI is not None:
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
else:
nworkers = 1
rank = 0
cpus_per_worker = 1
U.get_session(config=tf.ConfigProto(
allow_soft_placement=True,
inter_op_parallelism_threads=cpus_per_worker,
intra_op_parallelism_threads=cpus_per_worker
))
#################################################################
# ob ac ret atarg 都是 placeholder
# ret atarg 此处应该是向量形式
ob = observation_placeholder(ob_space)
# 创建pi和oldpi
with tf.variable_scope(str(index)+"pi"):
pi = policy(observ_placeholder=ob)
with tf.variable_scope(str(index)+"oldpi"):
oldpi = policy(observ_placeholder=ob)
# 每个reward都可以算一个atarg
atarg = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
#此处的KL div和entropy与reward无关
##################################
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
# entbonus 是entropy loss
entbonus = ent_coef * meanent
#################################
###########################################################
# vferr 用来更新 v 网络
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac))
# advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
# optimgain 用来更新 policy 网络, 应该每个reward有一个
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
###########################################################
dist = meankl
# 定义要优化的变量和 V 网络 adam 优化器
all_var_list = get_trainable_variables(str(index)+"pi")
# var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
# vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
var_list = get_pi_trainable_variables(str(index)+"pi")
vf_var_list = get_vf_trainable_variables(str(index)+"pi")
vfadam = MpiAdam(vf_var_list)
# 把变量展开成一个向量的类
get_flat = U.GetFlat(var_list)
# 这个类可以把一个向量分片赋值给var_list里的变量
set_from_flat = U.SetFromFlat(var_list)
# kl散度的梯度
klgrads = tf.gradients(dist, var_list)
####################################################################
# 拉直的向量
flat_tangent = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan")
# 把拉直的向量重新分成很多向量
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start+sz], shape))
start += sz
####################################################################
####################################################################
# 把kl散度梯度与变量乘积相加
gvp = tf.add_n([tf.reduce_sum(g*tangent) for (g, tangent) in zipsame(klgrads, tangents)]) #pylint: disable=E1111
# 把gvp的梯度展成向量
fvp = U.flatgrad(gvp, var_list)
####################################################################
# 用学习后的策略更新old策略
assign_old_eq_new = U.function([],[], updates=[tf.assign(oldv, newv)
for (oldv, newv) in zipsame(get_variables(str(index)+"oldpi"), get_variables(str(index)+"pi"))])
# 计算loss
compute_losses = U.function([ob, ac, atarg], losses)
# 计算loss和梯度
compute_lossandgrad = U.function([ob, ac, atarg], losses + [U.flatgrad(optimgain, var_list)])
# 计算fvp
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
# 计算值网络的梯度
compute_vflossandgrad = U.function([ob, ret], U.flatgrad(vferr, vf_var_list))
# 初始化variable
U.initialize()
if load_path is not None:
pi.load(load_path)
# 得到初始化的参数向量
th_init = get_flat()
if MPI is not None:
MPI.COMM_WORLD.Bcast(th_init, root=0)
# 把向量the_init的值分片赋值给var_list
set_from_flat(th_init)
#同步
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
self.MPI = MPI
self.pi = pi
self.oldpi = oldpi
self.compute_losses = compute_losses
self.compute_lossandgrad = compute_lossandgrad
self.compute_fvp = compute_fvp
self.compute_vflossandgrad = compute_vflossandgrad
self.assign_old_eq_new = assign_old_eq_new
self.get_flat = get_flat
self.set_from_flat = set_from_flat
self.vfadam = vfadam
# params
self.max_kl = max_kl
self.cg_iters = cg_iters
self.ent_coef = ent_coef
self.cg_damping = cg_damping
self.vf_stepsize = vf_stepsize
self.vf_iters = vf_iters
self.rank = rank
self.index = index
self.timed = timed
def learn(
*,
network,
env,
eval_env,
make_eval_env,
env_id,
total_timesteps,
timesteps_per_batch,
sil_update,
sil_loss, # what to train on
max_kl=0.001,
cg_iters=10,
gamma=0.99,
lam=1.0, # advantage estimation
seed=None,
ent_coef=0.0,
lr=3e-4,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=5,
sil_value=0.01,
sil_alpha=0.6,
sil_beta=0.1,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
save_interval=0,
load_path=None,
# MBL
# For train mbl
mbl_train_freq=5,
# For eval
num_eval_episodes=5,
eval_freq=5,
vis_eval=False,
eval_targs=('mbmf', ),
#eval_targs=('mf',),
quant=2,
# For mbl.step
#num_samples=(1500,),
num_samples=(1, ),
horizon=(2, ),
#horizon=(2,1),
#num_elites=(10,),
num_elites=(1, ),
mbl_lamb=(1.0, ),
mbl_gamma=0.99,
#mbl_sh=1, # Number of step for stochastic sampling
mbl_sh=10000,
#vf_lookahead=-1,
#use_max_vf=False,
reset_per_step=(0, ),
# For get_model
num_fc=2,
num_fwd_hidden=500,
use_layer_norm=False,
# For MBL
num_warm_start=int(1e4),
init_epochs=10,
update_epochs=5,
batch_size=512,
update_with_validation=False,
use_mean_elites=1,
use_ent_adjust=0,
adj_std_scale=0.5,
# For data loading
validation_set_path=None,
# For data collect
collect_val_data=False,
# For traj collect
traj_collect='mf',
# For profile
measure_time=True,
eval_val_err=False,
measure_rew=True,
model_fn=None,
update_fn=None,
init_fn=None,
mpi_rank_weight=1,
comm=None,
vf_coef=0.5,
max_grad_norm=0.5,
log_interval=1,
nminibatches=4,
noptepochs=4,
cliprange=0.2,
**network_kwargs):
'''
learn a policy function with TRPO algorithm
Parameters:
----------
network neural network to learn. Can be either string ('mlp', 'cnn', 'lstm', 'lnlstm' for basic types)
or function that takes input placeholder and returns tuple (output, None) for feedforward nets
or (output, (state_placeholder, state_output, mask_placeholder)) for recurrent nets
env environment (one of the gym environments or wrapped via baselines.common.vec_env.VecEnv-type class
timesteps_per_batch timesteps per gradient estimation batch
max_kl max KL divergence between old policy and new policy ( KL(pi_old || pi) )
ent_coef coefficient of policy entropy term in the optimization objective
cg_iters number of iterations of conjugate gradient algorithm
cg_damping conjugate gradient damping
vf_stepsize learning rate for adam optimizer used to optimie value function loss
vf_iters number of iterations of value function optimization iterations per each policy optimization step
total_timesteps max number of timesteps
max_episodes max number of episodes
max_iters maximum number of policy optimization iterations
callback function to be called with (locals(), globals()) each policy optimization step
load_path str, path to load the model from (default: None, i.e. no model is loaded)
**network_kwargs keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
Returns:
-------
learnt model
'''
if not isinstance(num_samples, tuple): num_samples = (num_samples, )
if not isinstance(horizon, tuple): horizon = (horizon, )
if not isinstance(num_elites, tuple): num_elites = (num_elites, )
if not isinstance(mbl_lamb, tuple): mbl_lamb = (mbl_lamb, )
if not isinstance(reset_per_step, tuple):
reset_per_step = (reset_per_step, )
if validation_set_path is None:
if collect_val_data:
validation_set_path = os.path.join(logger.get_dir(), 'val.pkl')
else:
validation_set_path = os.path.join('dataset',
'{}-val.pkl'.format(env_id))
if eval_val_err:
eval_val_err_path = os.path.join('dataset',
'{}-combine-val.pkl'.format(env_id))
logger.log(locals())
logger.log('MBL_SH', mbl_sh)
logger.log('Traj_collect', traj_collect)
if MPI is not None:
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
else:
nworkers = 1
rank = 0
cpus_per_worker = 1
U.get_session(
config=tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=cpus_per_worker,
intra_op_parallelism_threads=cpus_per_worker))
set_global_seeds(seed)
if isinstance(lr, float): lr = constfn(lr)
else: assert callable(lr)
if isinstance(cliprange, float): cliprange = constfn(cliprange)
else: assert callable(cliprange)
policy = build_policy(env, network, value_network='copy', **network_kwargs)
nenvs = env.num_envs
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
nbatch = nenvs * timesteps_per_batch
nbatch_train = nbatch // nminibatches
is_mpi_root = (MPI is None or MPI.COMM_WORLD.Get_rank() == 0)
ob = observation_placeholder(ob_space)
with tf.variable_scope("pi"):
pi = policy(observ_placeholder=ob)
make_model = lambda: Model(
policy=policy,
ob_space=ob_space,
ac_space=ac_space,
nbatch_act=nenvs,
nbatch_train=nbatch_train,
nsteps=timesteps_per_batch,
ent_coef=ent_coef,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm,
sil_update=sil_update,
sil_value=sil_value,
sil_alpha=sil_alpha,
sil_beta=sil_beta,
sil_loss=sil_loss,
# fn_reward=env.process_reward,
fn_reward=None,
# fn_obs=env.process_obs,
fn_obs=None,
ppo=False,
prev_pi='pi',
silm=pi)
model = make_model()
with tf.variable_scope("oldpi"):
oldpi = policy(observ_placeholder=ob)
make_old_model = lambda: Model(
policy=policy,
ob_space=ob_space,
ac_space=ac_space,
nbatch_act=nenvs,
nbatch_train=nbatch_train,
nsteps=timesteps_per_batch,
ent_coef=ent_coef,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm,
sil_update=sil_update,
sil_value=sil_value,
sil_alpha=sil_alpha,
sil_beta=sil_beta,
sil_loss=sil_loss,
# fn_reward=env.process_reward,
fn_reward=None,
# fn_obs=env.process_obs,
fn_obs=None,
ppo=False,
prev_pi='oldpi',
silm=oldpi)
old_model = make_old_model()
# MBL
# ---------------------------------------
#viz = Visdom(env=env_id)
win = None
eval_targs = list(eval_targs)
logger.log(eval_targs)
make_model_f = get_make_mlp_model(num_fc=num_fc,
num_fwd_hidden=num_fwd_hidden,
layer_norm=use_layer_norm)
mbl = MBL(env=eval_env,
env_id=env_id,
make_model=make_model_f,
num_warm_start=num_warm_start,
init_epochs=init_epochs,
update_epochs=update_epochs,
batch_size=batch_size,
**network_kwargs)
val_dataset = {'ob': None, 'ac': None, 'ob_next': None}
if update_with_validation:
logger.log('Update with validation')
val_dataset = load_val_data(validation_set_path)
if eval_val_err:
logger.log('Log val error')
eval_val_dataset = load_val_data(eval_val_err_path)
if collect_val_data:
logger.log('Collect validation data')
val_dataset_collect = []
def _mf_pi(ob, t=None):
stochastic = True
ac, vpred, _, _ = pi.step(ob, stochastic=stochastic)
return ac, vpred
def _mf_det_pi(ob, t=None):
#ac, vpred, _, _ = pi.step(ob, stochastic=False)
ac, vpred = pi._evaluate([pi.pd.mode(), pi.vf], ob)
return ac, vpred
def _mf_ent_pi(ob, t=None):
mean, std, vpred = pi._evaluate([pi.pd.mode(), pi.pd.std, pi.vf], ob)
ac = np.random.normal(mean, std * adj_std_scale, size=mean.shape)
return ac, vpred
################### use_ent_adjust======> adj_std_scale????????pi action sample
def _mbmf_inner_pi(ob, t=0):
if use_ent_adjust:
return _mf_ent_pi(ob)
else:
#return _mf_pi(ob)
if t < mbl_sh: return _mf_pi(ob)
else: return _mf_det_pi(ob)
# ---------------------------------------
# Run multiple configuration once
all_eval_descs = []
def make_mbmf_pi(n, h, e, l):
def _mbmf_pi(ob):
ac, rew = mbl.step(ob=ob,
pi=_mbmf_inner_pi,
horizon=h,
num_samples=n,
num_elites=e,
gamma=mbl_gamma,
lamb=l,
use_mean_elites=use_mean_elites)
return ac[None], rew
return Policy(step=_mbmf_pi, reset=None)
for n in num_samples:
for h in horizon:
for l in mbl_lamb:
for e in num_elites:
if 'mbmf' in eval_targs:
all_eval_descs.append(('MeanRew', 'MBL_TRPO_SIL',
make_mbmf_pi(n, h, e, l)))
#if 'mbmf' in eval_targs: all_eval_descs.append(('MeanRew-n-{}-h-{}-e-{}-l-{}-sh-{}-me-{}'.format(n, h, e, l, mbl_sh, use_mean_elites), 'MBL_TRPO-n-{}-h-{}-e-{}-l-{}-sh-{}-me-{}'.format(n, h, e, l, mbl_sh, use_mean_elites), make_mbmf_pi(n, h, e, l)))
if 'mf' in eval_targs:
all_eval_descs.append(
('MeanRew', 'TRPO_SIL', Policy(step=_mf_pi, reset=None)))
logger.log('List of evaluation targets')
for it in all_eval_descs:
logger.log(it[0])
pool = Pool(mp.cpu_count())
warm_start_done = False
# ----------------------------------------
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = ent_coef * meanent
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
dist = meankl
all_var_list = get_trainable_variables("pi")
# var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
# vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
vfadam = MpiAdam(vf_var_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
start += sz
gvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(klgrads, tangents)
]) #pylint: disable=E1111
fvp = U.flatgrad(gvp, var_list)
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(get_variables("oldpi"), get_variables("pi"))
])
compute_losses = U.function([ob, ac, atarg], losses)
compute_lossandgrad = U.function([ob, ac, atarg], losses +
[U.flatgrad(optimgain, var_list)])
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob, ret],
U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - tstart),
color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
return out
U.initialize()
if load_path is not None:
pi.load(load_path)
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
if traj_collect == 'mf':
seg_gen = traj_segment_generator(env,
timesteps_per_batch,
model,
stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
if sum([max_iters > 0, total_timesteps > 0, max_episodes > 0]) == 0:
# noththing to be done
return pi
assert sum([max_iters>0, total_timesteps>0, max_episodes>0]) < 2, \
'out of max_iters, total_timesteps, and max_episodes only one should be specified'
while True:
if callback: callback(locals(), globals())
if total_timesteps and timesteps_so_far >= total_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
logger.log("********** Iteration %i ************" % iters_so_far)
with timed("sampling"):
seg = seg_gen.__next__()
if traj_collect == 'mf-random' or traj_collect == 'mf-mb':
seg_mbl = seg_gen_mbl.__next__()
else:
seg_mbl = seg
add_vtarg_and_adv(seg, gamma, lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
# Val data collection
if collect_val_data:
for ob_, ac_, ob_next_ in zip(ob[:-1, 0, ...], ac[:-1, ...],
ob[1:, 0, ...]):
val_dataset_collect.append(
(copy.copy(ob_), copy.copy(ac_), copy.copy(ob_next_)))
# -----------------------------
# MBL update
else:
ob_mbl, ac_mbl = seg_mbl["ob"], seg_mbl["ac"]
mbl.add_data_batch(ob_mbl[:-1, 0, ...], ac_mbl[:-1, ...],
ob_mbl[1:, 0, ...])
mbl.update_forward_dynamic(require_update=iters_so_far %
mbl_train_freq == 0,
ob_val=val_dataset['ob'],
ac_val=val_dataset['ac'],
ob_next_val=val_dataset['ob_next'])
# -----------------------------
if traj_collect == 'mf':
#if traj_collect == 'mf' or traj_collect == 'mf-random' or traj_collect == 'mf-mb':
vpredbefore = seg[
"vpred"] # predicted value function before udpate
model = seg["model"]
atarg = (atarg - atarg.mean()) / atarg.std(
) # standardized advantage function estimate
if hasattr(pi, "ret_rms"): pi.ret_rms.update(tdlamret)
if hasattr(pi, "rms"):
pi.rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
assign_old_eq_new(
) # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > max_kl * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(),
vfadam.getflat().sum())) # list of tuples
assert all(
np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
with timed("SIL"):
lrnow = lr(1.0 - timesteps_so_far / total_timesteps)
l_loss, sil_adv, sil_samples, sil_nlogp = model.sil_train(
lrnow)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
if MPI is not None:
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
else:
listoflrpairs = [lrlocal]
lens, rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if sil_update > 0:
logger.record_tabular("SilSamples", sil_samples)
if rank == 0:
# MBL evaluation
if not collect_val_data:
#set_global_seeds(seed)
default_sess = tf.get_default_session()
def multithread_eval_policy(env_, pi_, num_episodes_,
vis_eval_, seed):
with default_sess.as_default():
if hasattr(env, 'ob_rms') and hasattr(env_, 'ob_rms'):
env_.ob_rms = env.ob_rms
res = eval_policy(env_, pi_, num_episodes_, vis_eval_,
seed, measure_time, measure_rew)
try:
env_.close()
except:
pass
return res
if mbl.is_warm_start_done() and iters_so_far % eval_freq == 0:
warm_start_done = mbl.is_warm_start_done()
if num_eval_episodes > 0:
targs_names = {}
with timed('eval'):
num_descs = len(all_eval_descs)
list_field_names = [e[0] for e in all_eval_descs]
list_legend_names = [e[1] for e in all_eval_descs]
list_pis = [e[2] for e in all_eval_descs]
list_eval_envs = [
make_eval_env() for _ in range(num_descs)
]
list_seed = [seed for _ in range(num_descs)]
list_num_eval_episodes = [
num_eval_episodes for _ in range(num_descs)
]
print(list_field_names)
print(list_legend_names)
list_vis_eval = [
vis_eval for _ in range(num_descs)
]
for i in range(num_descs):
field_name, legend_name = list_field_names[
i], list_legend_names[i],
res = multithread_eval_policy(
list_eval_envs[i], list_pis[i],
list_num_eval_episodes[i],
list_vis_eval[i], seed)
#eval_results = pool.starmap(multithread_eval_policy, zip(list_eval_envs, list_pis, list_num_eval_episodes, list_vis_eval,list_seed))
#for field_name, legend_name, res in zip(list_field_names, list_legend_names, eval_results):
perf, elapsed_time, eval_rew = res
logger.record_tabular(field_name, perf)
if measure_time:
logger.record_tabular(
'Time-%s' % (field_name), elapsed_time)
if measure_rew:
logger.record_tabular(
'SimRew-%s' % (field_name), eval_rew)
targs_names[field_name] = legend_name
if eval_val_err:
fwd_dynamics_err = mbl.eval_forward_dynamic(
obs=eval_val_dataset['ob'],
acs=eval_val_dataset['ac'],
obs_next=eval_val_dataset['ob_next'])
logger.record_tabular('FwdValError', fwd_dynamics_err)
logger.dump_tabular()
#print(logger.get_dir())
#print(targs_names)
#if num_eval_episodes > 0:
# win = plot(viz, win, logger.get_dir(), targs_names=targs_names, quant=quant, opt='best')
# -----------
#logger.dump_tabular()
yield pi
if collect_val_data:
with open(validation_set_path, 'wb') as f:
pickle.dump(val_dataset_collect, f)
logger.log('Save {} validation data'.format(len(val_dataset_collect)))
os.remove(f)
def make_env():
env = gym.make(defaults['env_name'])
env.set_episode_size(defaults['timesteps_per_batch'])
env = bench.Monitor(env, logger.get_dir() and os.path.join(logger.get_dir()), allow_early_resets=True)
return env
env = DummyVecEnv([make_env])
network = mlp(num_layers=defaults['num_layers'], num_hidden=defaults['num_hidden'], layer_norm=defaults['layer_norm'])
policy = build_policy(env, network, value_network='copy', **defaults)
set_global_seeds(defaults['seed'])
obs_space = observation_placeholder(env.observation_space)
pi = policy(observ_placeholder=obs_space)
if defaults['trained_path'] is not None:
pi.load_var(defaults['trained_path'])
obs = env.reset()
loop = True
while loop:
actions = pi.step_deterministic(obs)[0]
obs, reward, done, info = env.step_runtime(actions)
print("Action: ", actions)
print("Reward: ", reward)
print("ee_translation[x, y, z]: ", obs[0][6:9])
print("ee_orientation[w, x, y, z]: ", obs[0][9:13])
def learn(
*,
network,
env,
total_timesteps,
timesteps_per_batch=1024, # what to train on
max_kl=0.002,
cg_iters=10,
gamma=0.99,
lam=1.0, # advantage estimation
seed=None,
ent_coef=0.00,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=3,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
load_path=None,
num_reward=1,
**network_kwargs):
'''
learn a policy function with TRPO algorithm
Parameters:
----------
network neural network to learn. Can be either string ('mlp', 'cnn', 'lstm', 'lnlstm' for basic types)
or function that takes input placeholder and returns tuple (output, None) for feedforward nets
or (output, (state_placeholder, state_output, mask_placeholder)) for recurrent nets
env environment (one of the gym environments or wrapped via baselines.common.vec_env.VecEnv-type class
timesteps_per_batch timesteps per gradient estimation batch
max_kl max KL divergence between old policy and new policy ( KL(pi_old || pi) )
ent_coef coefficient of policy entropy term in the optimization objective
cg_iters number of iterations of conjugate gradient algorithm
cg_damping conjugate gradient damping
vf_stepsize learning rate for adam optimizer used to optimie value function loss
vf_iters number of iterations of value function optimization iterations per each policy optimization step
total_timesteps max number of timesteps
max_episodes max number of episodes
max_iters maximum number of policy optimization iterations
callback function to be called with (locals(), globals()) each policy optimization step
load_path str, path to load the model from (default: None, i.e. no model is loaded)
**network_kwargs keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
Returns:
-------
learnt model
'''
if MPI is not None:
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
else:
nworkers = 1
rank = 0
cpus_per_worker = 1
U.get_session(
config=tf.ConfigProto(allow_soft_placement=True,
inter_op_parallelism_threads=cpus_per_worker,
intra_op_parallelism_threads=cpus_per_worker))
set_global_seeds(seed)
# 创建policy
policy = build_policy(env,
network,
value_network='copy',
num_reward=num_reward,
**network_kwargs)
process_dir = logger.get_dir()
save_dir = process_dir.split(
'Data')[-2] + 'log/l1/seed' + process_dir[-1] + '/'
os.makedirs(save_dir, exist_ok=True)
coe_save = []
impro_save = []
grad_save = []
adj_save = []
coe = np.ones((num_reward)) / num_reward
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
#################################################################
# ob ac ret atarg 都是 placeholder
# ret atarg 此处应该是向量形式
ob = observation_placeholder(ob_space)
# 创建pi和oldpi
with tf.variable_scope("pi"):
pi = policy(observ_placeholder=ob)
with tf.variable_scope("oldpi"):
oldpi = policy(observ_placeholder=ob)
# 每个reward都可以算一个atarg
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32,
shape=[None, num_reward]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
#此处的KL div和entropy与reward无关
##################################
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
# entbonus 是entropy loss
entbonus = ent_coef * meanent
#################################
###########################################################
# vferr 用来更新 v 网络
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac))
# advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
# optimgain 用来更新 policy 网络, 应该每个reward有一个
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
###########################################################
dist = meankl
# 定义要优化的变量和 V 网络 adam 优化器
all_var_list = get_trainable_variables("pi")
# var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
# vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
vfadam = MpiAdam(vf_var_list)
# 把变量展开成一个向量的类
get_flat = U.GetFlat(var_list)
# 这个类可以把一个向量分片赋值给var_list里的变量
set_from_flat = U.SetFromFlat(var_list)
# kl散度的梯度
klgrads = tf.gradients(dist, var_list)
####################################################################
# 拉直的向量
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
# 把拉直的向量重新分成很多向量
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
start += sz
####################################################################
####################################################################
# 把kl散度梯度与变量乘积相加
gvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(klgrads, tangents)
]) #pylint: disable=E1111
# 把gvp的梯度展成向量
fvp = U.flatgrad(gvp, var_list)
####################################################################
# 用学习后的策略更新old策略
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(get_variables("oldpi"), get_variables("pi"))
])
# 计算loss
compute_losses = U.function([ob, ac, atarg], losses)
# 计算loss和梯度
compute_lossandgrad = U.function([ob, ac, atarg], losses +
[U.flatgrad(optimgain, var_list)])
# 计算fvp
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
# 计算值网络的梯度
compute_vflossandgrad = U.function([ob, ret],
U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - tstart),
color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
if MPI is not None:
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
else:
out = np.copy(x)
return out
# 初始化variable
U.initialize()
if load_path is not None:
pi.load(load_path)
# 得到初始化的参数向量
th_init = get_flat()
if MPI is not None:
MPI.COMM_WORLD.Bcast(th_init, root=0)
# 把向量the_init的值分片赋值给var_list
set_from_flat(th_init)
#同步
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
# 这是一个生成数据的迭代器
seg_gen = traj_segment_generator(pi,
env,
timesteps_per_batch,
stochastic=True,
num_reward=num_reward)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
# 双端队列
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
if sum([max_iters > 0, total_timesteps > 0, max_episodes > 0]) == 0:
# noththing to be done
return pi
assert sum([max_iters>0, total_timesteps>0, max_episodes>0]) < 2, \
'out of max_iters, total_timesteps, and max_episodes only one should be specified'
while True:
if callback: callback(locals(), globals())
if total_timesteps and timesteps_so_far >= total_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
logger.log("********** Iteration %i ************" % iters_so_far)
with timed("sampling"):
seg = seg_gen.__next__()
# 计算累积回报
add_vtarg_and_adv(seg, gamma, lam, num_reward=num_reward)
###########$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$$ToDo
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
# ob, ac, atarg, tdlamret 的类型都是ndarray
#ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg["tdlamret"]
_, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
#print(seg['ob'].shape,type(seg['ob']))
#print(seg['ac'],type(seg['ac']))
#print(seg['adv'],type(seg['adv']))
#print(seg["tdlamret"].shape,type(seg['tdlamret']))
vpredbefore = seg["vpred"] # predicted value function before udpate
# 标准化
#print("============================== atarg =========================================================")
#print(atarg)
atarg = (atarg - np.mean(atarg, axis=0)) / np.std(
atarg, axis=0) # standardized advantage function estimate
#atarg = (atarg) / np.max(np.abs(atarg),axis=0)
#print('======================================= standardized atarg ====================================')
#print(atarg)
if hasattr(pi, "ret_rms"): pi.ret_rms.update(tdlamret)
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
## set old parameter values to new parameter values
assign_old_eq_new()
G = None
S = None
mr_lossbefore = np.zeros((num_reward, len(loss_names)))
grad_norm = np.zeros((num_reward + 1))
for i in range(num_reward):
args = seg["ob"], seg["ac"], atarg[:, i]
#print(atarg[:,i])
# 算是args的一个sample，每隔5个取出一个
fvpargs = [arr[::5] for arr in args]
# 这个函数计算fisher matrix 与向量 p 的 乘积
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
with timed("computegrad of " + str(i + 1) + ".th reward"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
mr_lossbefore[i] = lossbefore
g = allmean(g)
#print("***************************************************************")
#print(g)
#print('==================='+str(i+1)+"=====================",np.linalg.norm(g))
#print(atarg[:,i])
if isinstance(G, np.ndarray):
G = np.vstack((G, g))
else:
G = g
# g是目标函数的梯度
# 利用共轭梯度获得更新方向
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg of " + str(i + 1) + ".th reward"):
# stepdir 是更新方向
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=rank == 0)
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
#print(np.linalg.norm(fullstep))
grad_norm[i] = np.linalg.norm(fullstep)
assert np.isfinite(stepdir).all()
if isinstance(S, np.ndarray):
S = np.vstack((S, stepdir))
else:
S = stepdir
#print('======================================= G ====================================')
#print(G)
#print('======================================= S ====================================')
#print(S)
new_coe = get_coefficient(G, S)
#coe = 0.99 * coe + 0.01 * new_coe
coe = new_coe
coe_save.append(coe)
#根据梯度的夹角调整参数
try:
GG = np.dot(S, S.T)
D = np.sqrt(np.diag(1 / np.diag(GG)))
GG = np.dot(np.dot(D, GG), D)
#print('======================================= inner product ====================================')
#print(GG)
adj = np.sum(GG) / (num_reward**2)
except:
adj = 1
#print('======================================= adj ====================================')
#print(adj)
try:
adj = 1
adj_save.append(adj)
adj_max_kl = adj * max_kl
#################################################################
grad_norm = grad_norm * np.sqrt(adj)
stepdir = np.dot(coe, S)
g = np.dot(coe, G)
lossbefore = np.dot(coe, mr_lossbefore)
#################################################################
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / adj_max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
grad_norm[num_reward] = np.linalg.norm(fullstep)
grad_save.append(grad_norm)
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
def compute_mr_losses():
mr_losses = np.zeros((num_reward, len(loss_names)))
for i in range(num_reward):
args = seg["ob"], seg["ac"], atarg[:, i]
one_reward_loss = allmean(np.array(compute_losses(*args)))
mr_losses[i] = one_reward_loss
mr_loss = np.dot(coe, mr_losses)
return mr_loss, mr_losses
# 做10次搜索
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
mr_loss_new, mr_losses_new = compute_mr_losses()
mr_impro = mr_losses_new - mr_lossbefore
meanlosses = surr, kl, *_ = allmean(np.array(mr_loss_new))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > adj_max_kl * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
impro_save.append(np.hstack((mr_impro[:, 0], improve)))
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(), vfadam.getflat().sum())) # list of tuples
assert all(
np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
with timed("vf"):
#print('======================================= tdlamret ====================================')
#print(seg["tdlamret"])
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=64):
#with tf.Session() as sess:
# sess.run(tf.global_variables_initializer())
# aaa = sess.run(pi.vf,feed_dict={ob:mbob,ret:mbret})
# print("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa")
# print(aaa.shape)
# print(mbret.shape)
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
#print(mbob,mbret)
except:
print('error')
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
if MPI is not None:
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
else:
listoflrpairs = [lrlocal]
lens, rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if rank == 0:
logger.dump_tabular()
#pdb.set_trace()
np.save(save_dir + 'coe.npy', coe_save)
np.save(save_dir + 'grad.npy', grad_save)
np.save(save_dir + 'improve.npy', impro_save)
np.save(save_dir + 'adj.npy', adj_save)
return pi
def learn(
*,
network,
env,
seed=None,
beta,
total_timesteps,
sil_update,
sil_loss,
timesteps_per_batch=2048, # what to train on
epsilon=0.01,
cg_iters=10,
gamma=0.99,
lam=0.98, # advantage estimation
entcoeff=0.0,
lr=3e-4,
cg_damping=0.1,
vf_stepsize=1e-3,
vf_iters=5,
sil_value=0.01,
sil_alpha=0.6,
sil_beta=0.1,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
save_interval=0,
load_path=None,
model_fn=None,
update_fn=None,
init_fn=None,
mpi_rank_weight=1,
comm=None,
vf_coef=0.5,
max_grad_norm=0.5,
log_interval=1,
nminibatches=4,
noptepochs=4,
cliprange=0.2,
TRPO=False,
**network_kwargs):
set_global_seeds(seed)
if isinstance(lr, float): lr = constfn(lr)
else: assert callable(lr)
if isinstance(cliprange, float): cliprange = constfn(cliprange)
else: assert callable(cliprange)
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
policy = build_policy(env,
network,
value_network='copy',
copos=True,
**network_kwargs)
nenvs = env.num_envs
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
nbatch = nenvs * timesteps_per_batch
nbatch_train = nbatch // nminibatches
is_mpi_root = (MPI is None or MPI.COMM_WORLD.Get_rank() == 0)
if model_fn is None:
model_fn = Model
discrete_ac_space = isinstance(ac_space, gym.spaces.Discrete)
ob = observation_placeholder(ob_space)
with tf.variable_scope("pi", reuse=tf.AUTO_REUSE):
pi = policy(observ_placeholder=ob)
#sil_model=policy(None, None, sess=get_session)
make_model = lambda: Model(
policy=policy,
ob_space=ob_space,
ac_space=ac_space,
nbatch_act=nenvs,
nbatch_train=nbatch_train,
nsteps=timesteps_per_batch,
ent_coef=entcoeff,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm,
sil_update=sil_update,
sil_value=sil_value,
sil_alpha=sil_alpha,
sil_beta=sil_beta,
sil_loss=sil_loss,
# fn_reward=env.process_reward,
fn_reward=None,
# fn_obs=env.process_obs,
fn_obs=None,
ppo=False,
prev_pi='pi',
silm=pi)
model = make_model()
if load_path is not None:
model.load(load_path)
with tf.variable_scope("oldpi", reuse=tf.AUTO_REUSE):
oldpi = policy(observ_placeholder=ob)
make_old_model = lambda: Model(
policy=policy,
ob_space=ob_space,
ac_space=ac_space,
nbatch_act=nenvs,
nbatch_train=nbatch_train,
nsteps=timesteps_per_batch,
ent_coef=entcoeff,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm,
sil_update=sil_update,
sil_value=sil_value,
sil_alpha=sil_alpha,
sil_beta=sil_beta,
sil_loss=sil_loss,
# fn_reward=env.process_reward,
fn_reward=None,
# fn_obs=env.process_obs,
fn_obs=None,
ppo=False,
prev_pi='oldpi',
silm=oldpi)
old_model = make_old_model()
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
old_entropy = oldpi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "Entropy"]
dist = meankl
#all_var_list = pi.get_trainable_variables()
#all_var_list = [v for v in all_var_list if v.name.split("/")[0].startswith("pi")]
#var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
#vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
all_var_list = get_trainable_variables("pi")
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
vfadam = MpiAdam(vf_var_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + sz], shape))
start += sz
#????gvp and fvp???
gvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(klgrads, tangents)
]) #pylint: disable=E1111
fvp = U.flatgrad(gvp, var_list)
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(get_variables("oldpi"), get_variables("pi"))
])
compute_losses = U.function([ob, ac, atarg], losses)
compute_lossandgrad = U.function([ob, ac, atarg], losses +
[U.flatgrad(optimgain, var_list)])
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob, ret],
U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - tstart),
color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
return out
U.initialize()
if load_path is not None:
pi.load(load_path)
th_init = get_flat()
if MPI is not None:
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Initialize eta, omega optimizer
if discrete_ac_space:
init_eta = 1
init_omega = 0.5
eta_omega_optimizer = EtaOmegaOptimizerDiscrete(
beta, epsilon, init_eta, init_omega)
else:
init_eta = 0.5
init_omega = 2.0
#????eta_omega_optimizer details?????
eta_omega_optimizer = EtaOmegaOptimizer(beta, epsilon, init_eta,
init_omega)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(env,
timesteps_per_batch,
model,
stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
assert sum([max_iters > 0, total_timesteps > 0, max_episodes > 0]) == 1
while True:
if callback: callback(locals(), globals())
if total_timesteps and timesteps_so_far >= total_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
logger.log("********** Iteration %i ************" % iters_so_far)
with timed("sampling"):
seg = seg_gen.__next__()
add_vtarg_and_adv(seg, gamma, lam)
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
model = seg["model"]
atarg = (atarg - atarg.mean()
) / atarg.std() # standardized advantage function estimate
#print(ob[:20])
#print(ac[:20])
if hasattr(pi, "ret_rms"): pi.ret_rms.update(tdlamret)
if hasattr(pi, "rms"):
pi.rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
assign_old_eq_new() # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(stepdir).all()
if TRPO:
#
# TRPO specific code.
# Find correct step size using line search
#
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / epsilon)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > epsilon * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
else:
#
# COPOS specific implementation.
#
copos_update_dir = stepdir
# Split direction into log-linear 'w_theta' and non-linear 'w_beta' parts
w_theta, w_beta = pi.split_w(copos_update_dir)
tmp_ob = np.zeros(
(1, ) + env.observation_space.shape
) # We assume that entropy does not depend on the NN
# Optimize eta and omega
if discrete_ac_space:
entropy = lossbefore[4]
#entropy = - 1/timesteps_per_batch * np.sum(np.sum(pi.get_action_prob(ob) * pi.get_log_action_prob(ob), axis=1))
eta, omega = eta_omega_optimizer.optimize(
pi.compute_F_w(ob, copos_update_dir),
pi.get_log_action_prob(ob), timesteps_per_batch,
entropy)
else:
Waa, Wsa = pi.w2W(w_theta)
wa = pi.get_wa(ob, w_beta)
varphis = pi.get_varphis(ob)
#old_ent = old_entropy.eval({oldpi.ob: tmp_ob})[0]
old_ent = lossbefore[4]
eta, omega = eta_omega_optimizer.optimize(
w_theta, Waa, Wsa, wa, varphis, pi.get_kt(),
pi.get_prec_matrix(), pi.is_new_policy_valid, old_ent)
logger.log("Initial eta: " + str(eta) + " and omega: " +
str(omega))
current_theta_beta = get_flat()
prev_theta, prev_beta = pi.all_to_theta_beta(
current_theta_beta)
if discrete_ac_space:
# Do a line search for both theta and beta parameters by adjusting only eta
eta = eta_search(w_theta, w_beta, eta, omega, allmean,
compute_losses, get_flat, set_from_flat,
pi, epsilon, args, discrete_ac_space)
logger.log("Updated eta, eta: " + str(eta))
set_from_flat(pi.theta_beta_to_all(prev_theta, prev_beta))
# Find proper omega for new eta. Use old policy parameters first.
eta, omega = eta_omega_optimizer.optimize(
pi.compute_F_w(ob, copos_update_dir),
pi.get_log_action_prob(ob), timesteps_per_batch,
entropy, eta)
logger.log("Updated omega, eta: " + str(eta) +
" and omega: " + str(omega))
# do line search for ratio for non-linear "beta" parameter values
#ratio = beta_ratio_line_search(w_theta, w_beta, eta, omega, allmean, compute_losses, get_flat, set_from_flat, pi,
# epsilon, beta, args)
# set ratio to 1 if we do not use beta ratio line search
ratio = 1
#print("ratio from line search: " + str(ratio))
cur_theta = (eta * prev_theta +
w_theta.reshape(-1, )) / (eta + omega)
cur_beta = prev_beta + ratio * w_beta.reshape(-1, ) / eta
else:
for i in range(2):
# Do a line search for both theta and beta parameters by adjusting only eta
eta = eta_search(w_theta, w_beta, eta, omega, allmean,
compute_losses, get_flat,
set_from_flat, pi, epsilon, args)
logger.log("Updated eta, eta: " + str(eta) +
" and omega: " + str(omega))
# Find proper omega for new eta. Use old policy parameters first.
set_from_flat(
pi.theta_beta_to_all(prev_theta, prev_beta))
eta, omega = \
eta_omega_optimizer.optimize(w_theta, Waa, Wsa, wa, varphis, pi.get_kt(),
pi.get_prec_matrix(), pi.is_new_policy_valid, old_ent, eta)
logger.log("Updated omega, eta: " + str(eta) +
" and omega: " + str(omega))
# Use final policy
logger.log("Final eta: " + str(eta) + " and omega: " +
str(omega))
cur_theta = (eta * prev_theta +
w_theta.reshape(-1, )) / (eta + omega)
cur_beta = prev_beta + w_beta.reshape(-1, ) / eta
set_from_flat(pi.theta_beta_to_all(cur_theta, cur_beta))
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
##copos specific over
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(), vfadam.getflat().sum())) # list of tuples
assert all(
np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
#cg over
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
#policy update over
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
with timed('SIL'):
lrnow = lr(1.0 - timesteps_so_far / total_timesteps)
l_loss, sil_adv, sil_samples, sil_nlogp = model.sil_train(lrnow)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
print("Reward max: " + str(max(rewbuffer)))
print("Reward min: " + str(min(rewbuffer)))
logger.record_tabular(
"EpLenMean",
np.mean(lenbuffer) if np.sum(lenbuffer) != 0.0 else 0.0)
logger.record_tabular(
"EpRewMean",
np.mean(rewbuffer) if np.sum(rewbuffer) != 0.0 else 0.0)
logger.record_tabular(
"AverageReturn",
np.mean(rewbuffer) if np.sum(rewbuffer) != 0.0 else 0.0)
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if sil_update > 0:
logger.record_tabular("SilSamples", sil_samples)
if rank == 0:
logger.dump_tabular()
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None,
encoded_x=None):
ob_space = env.observation_space
extra_tensors = {}
if observ_placeholder is None:
X = observation_placeholder(ob_space, batch_size=nbatch)
if normalize_observations and X.dtype == tf.float32:
new_encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
new_encoded_x = X
new_encoded_x = encode_observation(ob_space, new_encoded_x)
new_encoded_x = get_network_builder("cnn")(
**policy_kwargs)(new_encoded_x)
else:
X = observ_placeholder
new_encoded_x = encoded_x
with tf.variable_scope('pi' + str(head), reuse=tf.AUTO_REUSE):
policy_latent = policy_network(new_encoded_x)
if isinstance(policy_latent, tuple):
policy_latent, recurrent_tensors = policy_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
new_encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf' + str(head), reuse=tf.AUTO_REUSE):
vf_latent, _ = _v_net(new_encoded_x)
policy = PolicyWithValue(
env=env,
observations=X,
latent=policy_latent,
head=head,
vf_latent=vf_latent, #this is the same as policy_latent...
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
#print(policy.vf)
return policy, X, new_encoded_x
def learn(*,
network,
env,
total_timesteps,
timesteps_per_batch=1024, # what to train on
max_kl=0.001,
cg_iters=10,
gamma=0.99,
lam=1.0, # advantage estimation
seed=None,
ent_coef=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters =3,
max_episodes=0, max_iters=0, # time constraint
callback=None,
load_path=None,
**network_kwargs
):
'''
learn a policy function with TRPO algorithm
Parameters:
----------
network neural network to learn. Can be either string ('mlp', 'cnn', 'lstm', 'lnlstm' for basic types)
or function that takes input placeholder and returns tuple (output, None) for feedforward nets
or (output, (state_placeholder, state_output, mask_placeholder)) for recurrent nets
env environment (one of the gym environments or wrapped via baselines.common.vec_env.VecEnv-type class
timesteps_per_batch timesteps per gradient estimation batch
max_kl max KL divergence between old policy and new policy ( KL(pi_old || pi) )
ent_coef coefficient of policy entropy term in the optimization objective
cg_iters number of iterations of conjugate gradient algorithm
cg_damping conjugate gradient damping
vf_stepsize learning rate for adam optimizer used to optimie value function loss
vf_iters number of iterations of value function optimization iterations per each policy optimization step
total_timesteps max number of timesteps
max_episodes max number of episodes
max_iters maximum number of policy optimization iterations
callback function to be called with (locals(), globals()) each policy optimization step
load_path str, path to load the model from (default: None, i.e. no model is loaded)
**network_kwargs keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
Returns:
-------
learnt model
'''
if MPI is not None:
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
else:
nworkers = 1
rank = 0
cpus_per_worker = 1
U.get_session(config=tf.ConfigProto(
allow_soft_placement=True,
inter_op_parallelism_threads=cpus_per_worker,
intra_op_parallelism_threads=cpus_per_worker
))
policy = build_policy(env, network, value_network='copy', **network_kwargs)
set_global_seeds(seed)
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
ob = observation_placeholder(ob_space)
with tf.variable_scope("pi"):
pi = policy(observ_placeholder=ob)
with tf.variable_scope("oldpi"):
oldpi = policy(observ_placeholder=ob)
atarg = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = ent_coef * meanent
vferr = tf.reduce_mean(tf.square(pi.vf - ret))
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = tf.reduce_mean(ratio * atarg)
optimgain = surrgain + entbonus
losses = [optimgain, meankl, entbonus, surrgain, meanent]
loss_names = ["optimgain", "meankl", "entloss", "surrgain", "entropy"]
dist = meankl
all_var_list = get_trainable_variables("pi")
# var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("pol")]
# vf_var_list = [v for v in all_var_list if v.name.split("/")[1].startswith("vf")]
var_list = get_pi_trainable_variables("pi")
vf_var_list = get_vf_trainable_variables("pi")
vfadam = MpiAdam(vf_var_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
flat_tangent = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start+sz], shape))
start += sz
gvp = tf.add_n([tf.reduce_sum(g*tangent) for (g, tangent) in zipsame(klgrads, tangents)]) #pylint: disable=E1111
fvp = U.flatgrad(gvp, var_list)
assign_old_eq_new = U.function([],[], updates=[tf.assign(oldv, newv)
for (oldv, newv) in zipsame(get_variables("oldpi"), get_variables("pi"))])
compute_losses = U.function([ob, ac, atarg], losses)
compute_lossandgrad = U.function([ob, ac, atarg], losses + [U.flatgrad(optimgain, var_list)])
compute_fvp = U.function([flat_tangent, ob, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob, ret], U.flatgrad(vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
tstart = time.time()
yield
print(colorize("done in %.3f seconds"%(time.time() - tstart), color='magenta'))
else:
yield
def allmean(x):
assert isinstance(x, np.ndarray)
if MPI is not None:
out = np.empty_like(x)
MPI.COMM_WORLD.Allreduce(x, out, op=MPI.SUM)
out /= nworkers
else:
out = np.copy(x)
return out
U.initialize()
if load_path is not None:
pi.load(load_path)
th_init = get_flat()
if MPI is not None:
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi, env, timesteps_per_batch, stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
if sum([max_iters>0, total_timesteps>0, max_episodes>0])==0:
# noththing to be done
return pi
assert sum([max_iters>0, total_timesteps>0, max_episodes>0]) < 2, \
'out of max_iters, total_timesteps, and max_episodes only one should be specified'
while True:
if callback: callback(locals(), globals())
if total_timesteps and timesteps_so_far >= total_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
logger.log("********** Iteration %i ************"%iters_so_far)
with timed("sampling"):
seg = seg_gen.__next__()
add_vtarg_and_adv(seg, gamma, lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
ob, ac, atarg, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg["tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()) / atarg.std() # standardized advantage function estimate
if hasattr(pi, "ret_rms"): pi.ret_rms.update(tdlamret)
if hasattr(pi, "ob_rms"): pi.ob_rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
assign_old_eq_new() # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = cg(fisher_vector_product, g, cg_iters=cg_iters, verbose=rank==0)
assert np.isfinite(stepdir).all()
shs = .5*stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lm
expectedimprove = g.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
meanlosses = surr, kl, *_ = allmean(np.array(compute_losses(*args)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f"%(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl > max_kl * 1.5:
logger.log("violated KL constraint. shrinking step.")
elif improve < 0:
logger.log("surrogate didn't improve. shrinking step.")
else:
logger.log("Stepsize OK!")
break
stepsize *= .5
else:
logger.log("couldn't compute a good step")
set_from_flat(thbefore)
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather((thnew.sum(), vfadam.getflat().sum())) # list of tuples
assert all(np.allclose(ps, paramsums[0]) for ps in paramsums[1:])
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches((seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False, batch_size=64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
logger.record_tabular("ev_tdlam_before", explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"]) # local values
if MPI is not None:
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
else:
listoflrpairs = [lrlocal]
lens, rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if rank==0:
logger.dump_tabular()
return pi
def dynamics_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None,
index=None):
ob_space = env.observation_space
# ac_space = env.action_space
# print("shape", (64,) + (ob_space.shape[0] + ac_space.shape[0], ))
# Assume we have the same type for state and action space (Continuous - Continuous, Discrete - Discrete)
# assert isinstance(ob_space, Discrete) or isinstance(ob_space, Box) or isinstance(ob_space, MultiDiscrete), \
# 'Can only deal with Discrete and Box observation spaces for now'
#
# dtype = ob_space.dtype
# if dtype == np.int8:
# dtype = np.uint8
#X = tf.placeholder(shape=(nbatch,) + (ob_space.shape[0] + ac_space.shape[0], ), dtype=dtype, name='dyn_input')
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
encoded_x = encode_observation(
ob_space, encoded_x
) # Encode input in the way that is appropriate to the observation space(float)
with tf.variable_scope('dyn%s' % index, reuse=tf.AUTO_REUSE):
dynamics_latent = dynamics_network(encoded_x)
if isinstance(dynamics_latent, tuple):
dynamics_latent, recurrent_tensors = dynamics_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent dynamics: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
dynamics_latent, recurrent_tensors = dynamics_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
# print('dynamics%s'%character, train_dynamics_model[i])
### original
# with tf.variable_scope('dyn', reuse=tf.AUTO_REUSE):
# dynamics_latent = dynamics_network(encoded_x)
# if isinstance(dynamics_latent, tuple):
# dynamics_latent, recurrent_tensors = dynamics_latent
#
# if recurrent_tensors is not None:
# # recurrent architecture, need a few more steps
# nenv = nbatch // nsteps
# assert nenv > 0, 'Bad input for recurrent dynamics: batch size {} smaller than nsteps {}'.format(nbatch, nsteps)
# dynamics_latent, recurrent_tensors = dynamics_network(encoded_x, nenv)
# extra_tensors.update(recurrent_tensors)
### original delete tf.variable_scope (first line)
# dynamics_latent = dynamics_network(encoded_x)
# if isinstance(dynamics_latent, tuple):
# dynamics_latent, recurrent_tensors = dynamics_latent
#
# if recurrent_tensors is not None:
# # recurrent architecture, need a few more steps
# nenv = nbatch // nsteps
# assert nenv > 0, 'Bad input for recurrent dynamics: batch size {} smaller than nsteps {}'.format(nbatch, nsteps)
# dynamics_latent, recurrent_tensors = dynamics_network(encoded_x, nenv)
# extra_tensors.update(recurrent_tensors)
# _v_net = value_network
#
# if _v_net is None or _v_net == 'shared':
# vf_latent = dynamics_latent
# else:
# if _v_net == 'copy':
# _v_net = dynamics_network
# else:
# assert callable(_v_net)
#
# with tf.variable_scope('dyn_vf', reuse=tf.AUTO_REUSE):
# vf_latent = _v_net(encoded_x)
dynamics = DynamicsWithValue(
env=env,
observations=X,
latent=dynamics_latent,
sess=sess,
index=index, ### added
**extra_tensors)
return dynamics
def policy_fn(
nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None,
goal_placeholder=None,
concat_on_latent=False,
goal_encoded=None,
):
ob_space = env.observation_space
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=nbatch)
extra_tensors = {}
if normalize_observations and X.dtype == tf.float32:
encoded_x, rms = _normalize_clip_observation(X)
extra_tensors['rms'] = rms
else:
encoded_x = X
if goal_encoded is not None:
assert goal_placeholder is not None
encoded_goal = goal_encoded
else:
encoded_goal = goal_placeholder
encoded_x = tf.to_float(encoded_x)
encoded_goal = tf.to_float(
encoded_goal) if goal_placeholder is not None else goal_placeholder
if goal_placeholder is not None and not concat_on_latent:
assert encoded_x.get_shape().as_list(
)[:-1] == encoded_goal.get_shape().as_list()[:-1]
encoded_x = tf.concat([encoded_x, encoded_goal],
axis=-1,
name="concat_obs")
logger.info("concat obs and goals on inputs")
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent = policy_network(encoded_x)
if goal_placeholder is not None and concat_on_latent:
assert policy_latent.get_shape().as_list(
)[:-1] == policy_latent.get_shape().as_list()[:-1]
policy_latent = tf.concat([policy_latent, encoded_goal],
axis=-1,
name="concat_latent")
logger.info("concat obs and goals on latent")
if isinstance(policy_latent, tuple):
policy_latent, recurrent_tensors = policy_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
# TODO recurrent architectures are not supported with value_network=copy yet
vf_latent = _v_net(encoded_x)
policy = PolicyWithValue(env=env,
observations=observ_placeholder,
goals=goal_placeholder,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def animation(args):
# sess = tf.Session()
# configure logger, disable logging in child MPI processes (with rank > 0)
arg_parser = common_arg_parser()
args, unknown_args = arg_parser.parse_known_args(args)
extra_args = parse_cmdline_kwargs(unknown_args)
env_type, env_id = get_env_type(args)
print('env_type: {}'.format(env_type))
alg_kwargs = get_learn_function_defaults(args.alg, env_type)
load_path = extra_args['load_path']
extra_args.pop('load_path', None)
alg_kwargs.update(extra_args)
env = build_env(args)
if args.save_video_interval != 0:
env = VecVideoRecorder(
env,
osp.join(logger.get_dir(), "videos"),
record_video_trigger=lambda x: x % args.save_video_interval == 0,
video_length=args.save_video_length)
if args.network:
network = args.network
else:
if alg_kwargs.get('network') is None:
network = get_default_network(env_type)
model_ = policies.build_policy(env,
network,
value_network='copy',
**alg_kwargs)
ob_space = env.observation_space
ob = input.observation_placeholder(ob_space)
with tf.variable_scope("pi"):
model = model_(observ_placeholder=ob)
U.initialize()
# model.load_variables(load_path)
model.load(load_path)
if args.play:
logger.log("Running trained model")
obs = env.reset()
state = model.initial_state if hasattr(model,
'initial_state') else None
dones = np.zeros((1, ))
episode_rew = 0
while True:
if state is not None:
actions, _, state, _ = model.step(obs, S=state, M=dones)
else:
actions, _, _, _ = model.step(obs)
obs, rew, done, _ = env.step(actions)
# print(rew)
episode_rew += rew[0] if isinstance(env, VecEnv) else rew
env.render()
done = done.any() if isinstance(done, np.ndarray) else done
if done:
print('episode_rew={}'.format(episode_rew))
episode_rew = 0
obs = env.reset()
env.close()
def policy_fn(nbatch=None,
nsteps=None,
sess=None,
observ_placeholder=None,
randomization=True):
ob_space = env.observation_space
extra_tensors = {}
X = observ_placeholder if observ_placeholder is not None else observation_placeholder(
ob_space, batch_size=None)
encoded_x = encode_observation(ob_space, X)
# Randomization
if randomization:
encoded_x = tf.layers.conv2d(
encoded_x / 255.,
3,
3,
padding='same',
kernel_initializer=tf.initializers.glorot_normal(),
trainable=False,
name='randcnn') * 255.
randcnn_param = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES,
scope="ppo2_model/randcnn")
extra_tensors['randcnn_param'] = randcnn_param
with tf.variable_scope('pi', reuse=tf.AUTO_REUSE):
policy_latent = policy_network(encoded_x)
extra_tensors['latent_fts'] = policy_latent
if isinstance(policy_latent, tuple):
policy_latent, recurrent_tensors = policy_latent
if recurrent_tensors is not None:
# recurrent architecture, need a few more steps
nenv = nbatch // nsteps
assert nenv > 0, 'Bad input for recurrent policy: batch size {} smaller than nsteps {}'.format(
nbatch, nsteps)
policy_latent, recurrent_tensors = policy_network(
encoded_x, nenv)
extra_tensors.update(recurrent_tensors)
_v_net = value_network
if _v_net is None or _v_net == 'shared':
vf_latent = policy_latent
else:
if _v_net == 'copy':
_v_net = policy_network
else:
assert callable(_v_net)
with tf.variable_scope('vf', reuse=tf.AUTO_REUSE):
# TODO recurrent architectures are not supported with value_network=copy yet
vf_latent = _v_net(encoded_x)
policy = PolicyWithValue(env=env,
observations=X,
latent=policy_latent,
vf_latent=vf_latent,
sess=sess,
estimate_q=estimate_q,
**extra_tensors)
return policy
def train(args, extra_args):
env_type, env_id = get_env_type(args)
print('env_type: {}'.format(env_type))
total_timesteps = int(args.num_timesteps)
seed = args.seed
set_global_seeds(seed)
learn = get_learn_function(args.alg)
alg_kwargs = get_learn_function_defaults(args.alg, env_type)
alg_kwargs.update(extra_args)
env = build_env(args, normalize_ob=False)
eval_env = build_env(args, normalize_ob=False, is_eval=True)
if args.save_video_interval != 0:
env = VecVideoRecorder(
env,
osp.join(logger.get_dir(), "videos"),
record_video_trigger=lambda x: x % args.save_video_interval == 0,
video_length=args.save_video_length)
if args.network:
alg_kwargs['network'] = args.network
else:
if alg_kwargs.get('network') is None:
alg_kwargs['network'] = get_default_network(env_type)
beta = -1
if beta < 0:
#print(alg_kwargs)
nr_episodes = total_timesteps // alg_kwargs['timesteps_per_batch']
# Automatically compute beta based on initial entropy and number of iterations
policy = build_policy(
env,
alg_kwargs['network'],
value_network='copy',
normalize_observations=alg_kwargs['normalize_observations'],
copos=True)
ob = observation_placeholder(env.observation_space)
sess = U.single_threaded_session()
sess.__enter__()
with tf.variable_scope("tmp_pi"):
tmp_pi = policy(observ_placeholder=ob)
sess.run(tf.global_variables_initializer())
tmp_ob = np.zeros((1, ) + env.observation_space.shape)
entropy = sess.run(tmp_pi.pd.entropy(), feed_dict={tmp_pi.X: tmp_ob})
#beta = 2 * entropy / nr_episodes
beta = 0
print("Initial entropy: " + str(entropy) + ", episodes: " +
str(nr_episodes))
print("Constantly set beta: " + str(beta))
print('Training {} on {}:{} with arguments \n{}'.format(
args.alg, env_type, env_id, alg_kwargs))
iters = 0
for model in learn(env=env,
env_id=env_id,
eval_env=eval_env,
make_eval_env=lambda: build_env(
args, normalize_ob=False, is_eval=True),
seed=seed,
beta=beta,
total_timesteps=total_timesteps,
**alg_kwargs):
if args.store_ckpt:
save_path = osp.join(logger.get_dir(), 'model-{}'.format(iters))
model.save(save_path)
if isinstance(env, VecNormalize):
rms_path = osp.join(logger.get_dir(), 'rms-{}'.format(iters))
with open(rms_path, 'wb') as f:
rms = (env.ob_rms, env.ret_rms)
pickle.dump(rms, f)
logger.log('Save {} model'.format(iters + 1))
iters += 1
return model, env