def set_layer_flat(old_var_list, var_list):
dtype = tf.float32
shapes = list(map(U.var_shape, old_var_list))
total_size = np.sum([U.intprod(shape) for shape in shapes])
theta = theta = tf.placeholder(dtype, [total_size])
start = 0
assigns = []
for (shape, v) in zip(shapes, old_var_list):
size = U.intprod(shape)
assigns.append(tf.assign(v, tf.reshape(theta[start:start + size], shape)))
start += size
op = tf.group(*assigns)
return tf.get_default_session().run(op, feed_dict={theta: var_list})
def compute_fvp(self, flat_tangent, ob, ac, atarg):
shapes = [var.get_shape().as_list() for var in self.pi_var_list]
with tf.GradientTape() as outter_tape:
with tf.GradientTape() as inner_tape:
old_policy_latent = self.oldpi.policy_network(ob)
old_pd, _ = self.oldpi.pdtype.pdfromlatent(old_policy_latent)
policy_latent = self.pi.policy_network(ob)
pd, _ = self.pi.pdtype.pdfromlatent(policy_latent)
kloldnew = old_pd.kl(pd)
meankl = tf.reduce_mean(kloldnew)
klgrads = inner_tape.gradient(meankl, self.pi_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)
])
hessians_products = outter_tape.gradient(gvp, self.pi_var_list)
fvp = U.flatgrad(hessians_products, self.pi_var_list)
return fvp
def reshape_from_flat(self, flat_tangents):
shapes = [var.get_shape().as_list() for var in self.pi_var_list]
start = 0
tangents = []
for shape in shapes:
sz = U.intprod(shape)
tangents.append(tf.reshape(flat_tangents[start:start + sz], shape))
start += sz
return tangents
def learn(env, policy_func, *,
timesteps_per_batch, # what to train on
max_kl, cg_iters,
gamma, lam, # advantage estimation
entcoeff=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters =3,
max_timesteps=0, max_episodes=0, max_iters=0, # time constraint
callback=None
):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", ob_space, ac_space)
oldpi = policy_func("oldpi", ob_space, ac_space)
atarg = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = U.mean(kloldnew)
meanent = U.mean(ent)
entbonus = entcoeff * meanent
vferr = U.mean(tf.square(pi.vpred - ret))
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = U.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()
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")]
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([U.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(oldpi.get_variables(), pi.get_variables())])
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()
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
# ----------------------------------------
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
assert sum([max_iters>0, max_timesteps>0, max_episodes>0])==1
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_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
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
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()
def learn(env, make_policy, *,
n_episodes,
horizon,
delta,
gamma,
max_iters,
use_natural_gradient=False, #can be 'exact', 'approximate'
fisher_reg=1e-2,
iw_method='is',
iw_norm='none',
bound='J',
line_search_type='parabola',
save_weights=False,
improvement_tol=0.,
center_return=False,
render_after=None,
max_offline_iters=100,
callback=None):
np.set_printoptions(precision=3)
max_samples = horizon * n_episodes
if line_search_type == 'binary':
line_search = line_search_binary
elif line_search_type == 'parabola':
line_search = line_search_parabola
else:
raise ValueError()
# Building the environment
ob_space = env.observation_space
ac_space = env.action_space
# Building the policy
pi = make_policy('pi', ob_space, ac_space)
oldpi = make_policy('oldpi', ob_space, ac_space)
all_var_list = pi.get_trainable_variables()
var_list = [v for v in all_var_list if v.name.split('/')[1].startswith('pol')]
shapes = [U.intprod(var.get_shape().as_list()) for var in var_list]
n_parameters = sum(shapes)
# Placeholders
ob_ = ob = U.get_placeholder_cached(name='ob')
ac_ = pi.pdtype.sample_placeholder([max_samples], name='ac')
mask_ = tf.placeholder(dtype=tf.float32, shape=(max_samples), name='mask')
disc_rew_ = tf.placeholder(dtype=tf.float32, shape=(max_samples), name='disc_rew')
gradient_ = tf.placeholder(dtype=tf.float32, shape=(n_parameters, 1), name='gradient')
# Policy densities
target_log_pdf = pi.pd.logp(ac_)
behavioral_log_pdf = oldpi.pd.logp(ac_)
log_ratio = target_log_pdf - behavioral_log_pdf
# Split operations
disc_rew_split = tf.stack(tf.split(disc_rew_ * mask_, n_episodes))
log_ratio_split = tf.stack(tf.split(log_ratio * mask_, n_episodes))
target_log_pdf_split = tf.stack(tf.split(target_log_pdf * mask_, n_episodes))
mask_split = tf.stack(tf.split(mask_, n_episodes))
# Renyi divergence
emp_d2_split = tf.stack(tf.split(pi.pd.renyi(oldpi.pd, 2) * mask_, n_episodes))
emp_d2_cum_split = tf.reduce_sum(emp_d2_split, axis=1)
empirical_d2 = tf.reduce_mean(tf.exp(emp_d2_cum_split))
# Return
ep_return = tf.reduce_sum(mask_split * disc_rew_split, axis=1)
if center_return:
ep_return = ep_return - tf.reduce_mean(ep_return)
return_mean = tf.reduce_mean(ep_return)
return_std = U.reduce_std(ep_return)
return_max = tf.reduce_max(ep_return)
return_min = tf.reduce_min(ep_return)
return_abs_max = tf.reduce_max(tf.abs(ep_return))
if iw_method == 'pdis':
raise NotImplementedError()
elif iw_method == 'is':
iw = tf.exp(tf.reduce_sum(log_ratio_split, axis=1))
if iw_norm == 'none':
iwn = iw / n_episodes
w_return_mean = tf.reduce_sum(iwn * ep_return)
elif iw_norm == 'sn':
iwn = iw / tf.reduce_sum(iw)
w_return_mean = tf.reduce_sum(iwn * ep_return)
elif iw_norm == 'regression':
iwn = iw / n_episodes
mean_iw = tf.reduce_mean(iw)
beta = tf.reduce_sum((iw - mean_iw) * ep_return * iw) / (tf.reduce_sum((iw - mean_iw) ** 2) + 1e-24)
w_return_mean = tf.reduce_mean(iw * ep_return - beta * (iw - 1))
else:
raise NotImplementedError()
ess_classic = tf.linalg.norm(iw, 1) ** 2 / tf.linalg.norm(iw, 2) ** 2
sqrt_ess_classic = tf.linalg.norm(iw, 1) / tf.linalg.norm(iw, 2)
ess_renyi = n_episodes / empirical_d2
else:
raise NotImplementedError()
if bound == 'J':
bound_ = w_return_mean
elif bound == 'std-d2':
bound_ = w_return_mean - tf.sqrt((1 - delta) / (delta * ess_renyi)) * return_std
elif bound == 'max-d2':
bound_ = w_return_mean - tf.sqrt((1 - delta) / (delta * ess_renyi)) * return_abs_max
elif bound == 'max-ess':
bound_ = w_return_mean - tf.sqrt((1 - delta) / delta) / sqrt_ess_classic * return_abs_max
elif bound == 'std-ess':
bound_ = w_return_mean - tf.sqrt((1 - delta) / delta) / sqrt_ess_classic * return_std
else:
raise NotImplementedError()
losses = [bound_, return_mean, return_max, return_min, return_std, empirical_d2, w_return_mean,
tf.reduce_max(iwn), tf.reduce_min(iwn), tf.reduce_mean(iwn), U.reduce_std(iwn), tf.reduce_max(iw),
tf.reduce_min(iw), tf.reduce_mean(iw), U.reduce_std(iw), ess_classic, ess_renyi]
loss_names = ['Bound', 'InitialReturnMean', 'InitialReturnMax', 'InitialReturnMin', 'InitialReturnStd',
'EmpiricalD2', 'ReturnMeanIW', 'MaxIWNorm', 'MinIWNorm', 'MeanIWNorm', 'StdIWNorm',
'MaxIW', 'MinIW', 'MeanIW', 'StdIW', 'ESSClassic', 'ESSRenyi']
if use_natural_gradient:
p = tf.placeholder(dtype=tf.float32, shape=[None])
target_logpdf_episode = tf.reduce_sum(target_log_pdf_split * mask_split, axis=1)
grad_logprob = U.flatgrad(tf.stop_gradient(iwn) * target_logpdf_episode, var_list)
dot_product = tf.reduce_sum(grad_logprob * p)
hess_logprob = U.flatgrad(dot_product, var_list)
compute_linear_operator = U.function([p, ob_, ac_, disc_rew_, mask_], [-hess_logprob])
assign_old_eq_new = U.function([], [], updates=[tf.assign(oldv, newv)
for (oldv, newv) in zipsame(oldpi.get_variables(), pi.get_variables())])
compute_lossandgrad = U.function([ob_, ac_, disc_rew_, mask_], losses + [U.flatgrad(bound_, var_list)])
compute_grad = U.function([ob_, ac_, disc_rew_, mask_], [U.flatgrad(bound_, var_list)])
compute_bound = U.function([ob_, ac_, disc_rew_, mask_], [bound_])
compute_losses = U.function([ob_, ac_, disc_rew_, mask_], losses)
set_parameter = U.SetFromFlat(var_list)
get_parameter = U.GetFlat(var_list)
seg_gen = traj_segment_generator(pi, env, n_episodes, horizon, stochastic=True, gamma=gamma)
sampler = type("SequentialSampler", (object,), {"collect": lambda self, _: seg_gen.__next__()})()
U.initialize()
# Starting optimizing
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=n_episodes)
rewbuffer = deque(maxlen=n_episodes)
while True:
iters_so_far += 1
if render_after is not None and iters_so_far % render_after == 0:
if hasattr(env, 'render'):
render(env, pi, horizon)
if callback:
callback(locals(), globals())
if iters_so_far >= max_iters:
print('Finised...')
break
logger.log('********** Iteration %i ************' % iters_so_far)
theta = get_parameter()
with timed('sampling'):
seg = sampler.collect(theta)
lens, rets = seg['ep_lens'], seg['ep_rets']
lenbuffer.extend(lens)
rewbuffer.extend(rets)
episodes_so_far += len(lens)
timesteps_so_far += sum(lens)
args = ob, ac, disc_rew, mask = seg['ob'], seg['ac'], seg['disc_rew'], seg['mask']
assign_old_eq_new()
def evaluate_loss():
loss = compute_bound(*args)
return loss[0]
def evaluate_gradient():
gradient = compute_grad(*args)
return gradient[0]
if use_natural_gradient:
def evaluate_fisher_vector_prod(x):
return compute_linear_operator(x, *args)[0] + fisher_reg * x
def evaluate_natural_gradient(g):
return cg(evaluate_fisher_vector_prod, g, cg_iters=10, verbose=0)
else:
evaluate_natural_gradient = None
with timed('summaries before'):
logger.record_tabular("Itaration", iters_so_far)
logger.record_tabular("InitialBound", evaluate_loss())
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
if save_weights:
logger.record_tabular('Weights', str(get_parameter()))
with timed("offline optimization"):
theta, improvement = optimize_offline(theta,
set_parameter,
line_search,
evaluate_loss,
evaluate_gradient,
evaluate_natural_gradient,
max_offline_ite=max_offline_iters)
set_parameter(theta)
with timed('summaries after'):
meanlosses = np.array(compute_losses(*args))
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
logger.dump_tabular()
env.close()
def learn(env,
policy_func,
rank,
pretrained,
pretrained_weight,
*,
g_step,
d_step,
entcoeff,
save_per_iter,
ckpt_dir,
log_dir,
timesteps_per_batch,
task_name,
gamma,
lam,
max_kl,
cg_iters,
cg_damping=1e-2,
vf_stepsize=1e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0,
callback=None):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func(
"pi",
ob_space,
ac_space,
) # reuse=(pretrained_weight != None)
oldpi = policy_func("oldpi", ob_space, ac_space)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
#ob = U.get_placeholder_cached(name="ob")
ob_config = U.get_placeholder_cached(name="ob")
ob_target = U.get_placeholder_cached(name="goal")
obs_pos = U.get_placeholder_cached(name="obs_pos")
#obs_pos2 = U.get_placeholder_cached(name="obs_pos2")
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 = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(pi.vpred - 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()
var_list = [
v for v in all_var_list if v.name.startswith("pi/pol")
or v.name.startswith("pi/logstd") or v.name.startswith("pi/obs")
]
vf_var_list = [
v for v in all_var_list
if v.name.startswith("pi/vf") or v.name.startswith("pi/obs")
]
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(oldpi.get_variables(), pi.get_variables())
])
compute_losses = U.function([ob_config, ob_target, obs_pos, ac, atarg],
losses)
compute_lossandgrad = U.function(
[ob_config, ob_target, obs_pos, ac, atarg],
losses + [U.flatgrad(optimgain, var_list)])
compute_fvp = U.function(
[flat_tangent, ob_config, ob_target, obs_pos, ac, atarg], fvp)
compute_vflossandgrad = U.function([ob_config, ob_target, obs_pos, 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()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
if rank == 0:
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
true_rewbuffer = deque(maxlen=40)
max_trm = -5
true_reward_mean = 0
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
g_loss_stats = stats(loss_names)
ep_stats = stats(["True_rewards", "Rewards", "Episode_length"])
# if provide pretrained weight
if pretrained_weight is not None:
#U.load_variables(pretrained_weight, variables=pi.get_variables())
saver = tf.train.Saver()
saver.restore(tf.get_default_session(), pretrained_weight)
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
# Save model
if rank == 0 and ckpt_dir is not None and true_reward_mean > max_trm:
fname = os.path.join(ckpt_dir, task_name)
os.makedirs(os.path.dirname(fname), exist_ok=True)
saver = tf.train.Saver()
saver.save(tf.get_default_session(), fname)
max_trm = true_reward_mean
logger.log("********** Iteration %i ************" % iters_so_far)
def fisher_vector_product(p):
v1 = allmean(compute_fvp(p, *fvpargs))
# print("norm(v1):%.2e, norm(p):%.2e, cg_damping:%.2e"%(np.linalg.norm(v1), np.linalg.norm(p), cg_damping))
return v1 + cg_damping * p
# ------------------ Update G ------------------
logger.log("Optimizing Policy...")
for _ in range(g_step):
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, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
config, goal, obstacle_pos = [], [], []
for o in seg["ob"]:
config.append(o["joint"])
goal.append(o["target"])
obstacle_pos.append(o["obstacle_pos1"])
#obstacle_pos2.append(o["obstacle_pos2"])
config, goal, obstacle_pos = map(np.array,
[config, goal, obstacle_pos])
args = config, goal, obstacle_pos, seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
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)
logger.log(
'iter:{:d}, norm of g: {:.4f}, error of cg: {:.4f}'.
format(
cg_iters, np.linalg.norm(g),
np.linalg.norm(g -
compute_fvp(stepdir, *fvpargs))))
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:])
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbg, mbop, mbret) in dataset.iterbatches(
(config, goal, obstacle_pos, seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=128):
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(
mbob) # update running mean/std for policy
g = allmean(
compute_vflossandgrad(mbob, mbg, mbop, mbret))
vfadam.update(g, vf_stepsize)
g_losses = meanlosses
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
lrlocal = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]
) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*listoflrpairs))
true_rewbuffer.extend(true_rets)
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpTrueRewMean", np.mean(true_rewbuffer))
true_reward_mean = np.mean(true_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()
def learn(
env,
policy_fn,
*,
timesteps_per_batch, # what to train on
epsilon,
beta,
cg_iters,
gamma,
lam, # advantage estimation
entcoeff=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
TRPO=False,
n_policy=1,
policy_type=0,
filepath='',
session,
retrace=False):
'''
:param TRPO: True: TRPO, False: COPOS
:param n_policy: Number of periodic policy parts
:param policy_type: 0: Optimize 'n_policy' policies that are executed periodically. All the policies are updated.
1: The last 'n_policy' policies are executed periodically but only the last one is optimized.
2: The policy is spread over 'n_policy' time steps.
'''
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pis = [
policy_fn("pi_" + str(i), ob_space, ac_space) for i in range(n_policy)
]
oldpis = [
policy_fn("oldpi_" + str(i), ob_space, ac_space)
for i in range(n_policy)
]
pi_vf = policy_fn("pi_vf", ob_space, ac_space)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
if policy_type == 0:
print(
"Policy type: " + str(policy_type) +
". Optimize 'n_policy' policies that are executed periodically. All the policies are updated."
)
elif policy_type == 1:
print(
"Policy type: " + str(policy_type) +
". The last 'n_policy' policies are executed periodically but only the last one is optimized."
)
elif policy_type == 2:
print("Policy type: " + str(policy_type) +
". The policy is spread over 'n_policy' time steps.")
else:
print("Policy type: " + str(policy_type) + " is not supported.")
# Compute variables for each policy separately
old_entropy = []
get_flat = []
set_from_flat = []
assign_old_eq_new = []
copy_policy_back = []
compute_losses = []
compute_lossandgrad = []
compute_fvp = []
for i in range(n_policy):
pi = pis[i]
oldpi = oldpis[i]
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
old_entropy.append(oldpi.pd.entropy())
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = entcoeff * meanent
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
if retrace:
surrgain = tf.reduce_mean(
atarg) # atarg incorporates pnew / pold already
else:
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")
]
#
# fvp: Fisher Information Matrix / vector product based on Hessian of KL-divergence
# fvp = F * v, where F = - E \partial_1 \partial_2 KL_div(p1 || p2)
#
get_flat.append(U.GetFlat(var_list))
set_from_flat.append(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)
#
# fvpll: Fisher Information Matrix / vector product based on exact FIM
# fvpll = F * v, where F = E[\partial_1 \log p * \partial_2 \log p]
#
# Mean: (\partial \mu^T / \partial param1) * Precision * (\partial \mu / \partial param1)
# Covariance: 0.5 * Trace[Precision * (\partial Cov / \partial param1) *
# Precision * (\partial Cov / \partial param2)]
if i > 0:
# Only needed for policy_type == 1 for copying policy 'i' to policy 'i-1'
copy_policy_back.append(
U.function(
[], [],
updates=[
tf.assign(oldv, newv) for (oldv, newv) in zipsame(
pis[i - 1].get_variables(), pi.get_variables())
]))
assign_old_eq_new.append(
U.function([], [],
updates=[
tf.assign(oldv, newv) for (oldv, newv) in zipsame(
oldpi.get_variables(), pi.get_variables())
]))
compute_losses.append(U.function([ob, ac, atarg], losses))
compute_lossandgrad.append(
U.function([ob, ac, atarg],
losses + [U.flatgrad(optimgain, var_list)]))
compute_fvp.append(U.function([flat_tangent, ob, ac, atarg], fvp))
# Value function is global to all policies
vferr = tf.reduce_mean(tf.square(pi_vf.vpred - ret))
all_var_list = pi_vf.get_trainable_variables()
vf_var_list = [
v for v in all_var_list if v.name.split("/")[1].startswith("vf")
]
vfadam = MpiAdam(vf_var_list)
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 policy_type == 1:
# Initialize policies to identical values
th_init = get_flat[0]()
for i in range(n_policy):
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat[i](th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
else:
for i in range(n_policy):
th_init = get_flat[i]()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat[i](th_init)
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Initialize eta, omega optimizer
init_eta = 0.5
init_omega = 2.0
eta_omega_optimizer = EtaOmegaOptimizer(beta, epsilon, init_eta,
init_omega)
# Prepare for rollouts
# ----------------------------------------
seg_gen = []
for i in range(len(pis)):
seg_gen.append(
traj_segment_generator(pis,
i + 1,
pi_vf,
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
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
n_saves = 0
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_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)
if max_timesteps > 0 and (timesteps_so_far >=
(n_saves * max_timesteps // 5)):
# Save policy
saver = tf.train.Saver()
saver.save(session, filepath + "_" + str(iters_so_far))
n_saves += 1
with timed("sampling"):
if policy_type == 1 and iters_so_far < len(pis):
all_seg = seg_gen[iters_so_far].__next__(
) # For four time steps use the four policies
else:
all_seg = seg_gen[-1].__next__()
if policy_type == 1 and retrace:
act_pi_ids = np.empty_like(all_seg["vpred"], dtype=int)
act_pi_ids[:] = n_policy - 1 # Always update the last policy
add_vtarg_and_adv_retrace(all_seg, gamma, lam, act_pi_ids)
else:
add_vtarg_and_adv(all_seg, gamma, lam)
# Split the advantage functions etc. among the policies
segs = split_traj_segment(pis, all_seg)
# Update all policies
for pi_id in range(n_policy):
if policy_type == 1:
# Update only last policy
pi_id = n_policy - 1
# Using all the samples
seg = all_seg
else:
seg = segs[pi_id]
pi = pis[pi_id]
oldpi = oldpis[pi_id]
# 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 update
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[pi_id](p, *
fvpargs)) + cg_damping * p
assign_old_eq_new[pi_id](
) # set old parameter values to new parameter values
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad[pi_id](*args)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
logger.log("Got zero gradient. not updating")
if policy_type == 1:
# Update only the last policy
break
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[pi_id]()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat[pi_id](thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses[pi_id](*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[pi_id](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)
# q_beta(s,a) = \grad_beta \log \pi(a|s) * w_beta
# = features_beta(s) * K^T * Prec * a
# q_beta = self.target.get_q_beta(features_beta, actions)
Waa, Wsa = pi.w2W(w_theta)
wa = pi.get_wa(ob, w_beta)
varphis = pi.get_varphis(ob)
# Optimize eta and omega
tmp_ob = np.zeros(
(1, ) + env.observation_space.shape
) # We assume that entropy does not depend on the NN
old_ent = old_entropy[pi_id].eval({oldpi.ob: tmp_ob})[0]
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[pi_id]()
prev_theta, prev_beta = pi.all_to_theta_beta(
current_theta_beta)
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[pi_id],
get_flat[pi_id], set_from_flat[pi_id],
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_id](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
thnew = pi.theta_beta_to_all(cur_theta, cur_beta)
set_from_flat[pi_id](thnew)
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses[pi_id](*args)))
if nworkers > 1 and iters_so_far % 20 == 0:
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(),
vfadam[pi_id].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)
logger.record_tabular(
"ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
if policy_type == 1:
# Update only the last policy
break
if policy_type == 1:
# Copy policies 1, ..., i to 0, ..., i-1
for j in range(n_policy - 1):
copy_policy_back[j]()
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(all_seg["ob"], all_seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
lrlocal = (all_seg["ep_lens"], all_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)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("AverageReturn", 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()
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(
env,
policy_fn,
*,
timesteps_per_batch, # what to train on
epsilon,
beta,
cg_iters,
gamma,
lam, # advantage estimation
trial,
method,
entcoeff=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
TRPO=False):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_fn("pi", ob_space, ac_space)
oldpi = policy_fn("oldpi", ob_space, ac_space)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
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.vpred - ret))
# ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac)) # advantage * pnew / pold
# surrgain = tf.reduce_mean(ratio * atarg)
surrgain = tf.reduce_mean(pi.pd.logp(ac) * 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")
]
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(oldpi.get_variables(), pi.get_variables())
])
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()
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)
# Initialize eta, omega optimizer
init_eta = 0.5
init_omega = 2.0
eta_omega_optimizer = EtaOmegaOptimizer(beta, epsilon, init_eta,
init_omega)
# 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
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_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()
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)
# q_beta(s,a) = \grad_beta \log \pi(a|s) * w_beta
# = features_beta(s) * K^T * Prec * a
# q_beta = self.target.get_q_beta(features_beta, actions)
Waa, Wsa = pi.w2W(w_theta)
wa = pi.get_wa(ob, w_beta)
varphis = pi.get_varphis(ob)
# Optimize eta and omega
tmp_ob = np.zeros(
(1, ) + env.observation_space.shape
) # We assume that entropy does not depend on the NN
old_ent = old_entropy.eval({oldpi.ob: tmp_ob})[0]
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)
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)))
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
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
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)
logger.record_tabular("Name", method)
logger.record_tabular("Iteration", iters_so_far)
logger.record_tabular("trial", trial)
if rank == 0:
logger.dump_tabular()
def learn(env, last_ob, last_jpos, run_reach, policy_func, reward_giver, expert_dataset, rank,
pretrained, pretrained_weight, *,
g_step, d_step, entcoeff, save_per_iter,
ckpt_dir, log_dir, timesteps_per_batch, task_name,
gamma, lam,
max_kl, cg_iters, cg_damping=1e-2,
vf_stepsize=3e-4, d_stepsize=3e-4, vf_iters=3,
max_timesteps=0, max_episodes=0, max_iters=0,
callback=None
):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi_grasp", ob_space, ac_space, reuse=(pretrained_weight != None))
oldpi = policy_func("oldpi", ob_space, ac_space)
atarg = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.compat.v1.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
ac = pi.pdtype.sample_placeholder([None])
# Changes are made in order to use tensorboard
# -------------------------------------------
#train_writer = tf.compat.v1.summary.FileWriter('../../logs/trpo_mpi') # sets log dir to GailPart folder
#sess = tf.compat.v1.Session() # create a session??
# -------------------------------------------
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(pi.vpred - 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()
var_list = [v for v in all_var_list if v.name.startswith("pi_grasp/pol") or v.name.startswith("pi_grasp/logstd")]
vf_var_list = [v for v in all_var_list if v.name.startswith("pi_grasp/vff")]
assert len(var_list) == len(vf_var_list) + 1
d_adam = MpiAdam(reward_giver.get_trainable_variables())
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.compat.v1.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.compat.v1.assign(oldv, newv)
for (oldv, newv) in zipsame(oldpi.get_variables(), pi.get_variables())])
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()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
d_adam.sync()
vfadam.sync()
if rank == 0:
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi, last_ob, last_jpos, run_reach, policy_func, env, reward_giver, 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
true_rewbuffer = deque(maxlen=40)
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
g_loss_stats = stats(loss_names)
d_loss_stats = stats(reward_giver.loss_name)
ep_stats = stats(["True_rewards", "Rewards", "Episode_length"])
# if provide pretrained weight
if pretrained_weight is not None:
U.load_state(pretrained_weight, var_list=pi.get_variables())
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
# Save model
if rank == 0 and iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
t_name = task_name + "_" + str(iters_so_far)
fname = os.path.join(ckpt_dir, t_name) # changed from task_name
os.makedirs(os.path.dirname(fname), exist_ok=True)
saver = tf.compat.v1.train.Saver()
saver.save(tf.compat.v1.get_default_session(), fname)
logger.log("********** Iteration %i ************" % iters_so_far)
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
# ------------------ Update G ------------------
logger.log("Optimizing Policy...")
for _ in range(g_step):
with timed("sampling"):
seg = seg_gen.__next__()
#print("trpo_mpi, seg = seg_gen.__next__() call output: ", 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"]
vpredbefore = seg["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()) / atarg.std() # standardized advantage function estimate
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]
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)))
#logger.log("trpo_mpi.py, what should be logged with loss names ie. meanlosses:_", meanlosses)
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:])
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=128):
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(mbob) # update running mean/std for policy
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
g_losses = meanlosses
#logger.log("trpo_mpi.py, mean losses before logging wiht loss names: \n")
#logger.log(meanlosses)
# This is where the nan values are tabulated for some of the entries
#logger.log("trpo_mpi.py, view whats being printed with (loss_names, lossvalues)")
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
logger.record_tabular("ev_tdlam_before", explained_variance(vpredbefore, tdlamret))
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, reward_giver.loss_name))
ob_expert, ac_expert = expert_dataset.get_next_batch(len(ob))
batch_size = len(ob) // d_step
d_losses = [] # list of tuples, each of which gives the loss for a minibatch
for ob_batch, ac_batch in dataset.iterbatches((ob, ac),
include_final_partial_batch=False,
batch_size=batch_size):
ob_expert, ac_expert = expert_dataset.get_next_batch(len(ob_batch))
# update running mean/std for reward_giver
if hasattr(reward_giver, "obs_rms"): reward_giver.obs_rms.update(np.concatenate((ob_batch, ob_expert), 0))
*newlosses, g = reward_giver.lossandgrad(ob_batch, ac_batch, ob_expert, ac_expert)
d_adam.update(allmean(g), d_stepsize)
d_losses.append(newlosses)
# This is to see what the d_losses are
#logger.log("trpo_mpi.py, see what is being logged in d_losses")
#logger.log("trpo_mpi.py, d_losses")
#logger.log(d_losses)
logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
# For Tensorboard Logging
# ---------------------------
#tf.compat.v1.summary.scalar("Generator Accuracy", tf.convert_to_tensor( np.mean(d_losses, axis=0)[4] ) ) # 5 position
#tf.compat.v1.summary.scalar("Expert Accuracy", tf.convert_to_tensor( np.mean(d_losses, axis=0)[5] ) ) # 6 position
#tf.compat.v1.summary.scalar("Entropy Loss", tf.convert_to_tensor( np.mean(d_losses, axis=0)[3] ) ) # 4 position
#merge = tf.compat.v1.summary.merge_all() # merge summaries
#summary = sess.run([merge])
#train_writer.add_summary(summary, iters_so_far)
# Is there a need to reset metric after every epoch? I dont think so?
# ---------------------------
#logger.log("trpo_mpi.py, after logging, but before recordeing timesteps so far")
lrlocal = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]) # local values, truly confirmed is empty after call
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*listoflrpairs))
true_rewbuffer.extend(true_rets)
lenbuffer.extend(lens)
rewbuffer.extend(rews)
# Could it be that the seg locals for lens and rets are ommitted since has no use in gail algorithm?
# Probably dont have to worry about it, check the scalar part
logger.record_tabular("EpLenMean", np.mean(lenbuffer)) # This has nan values
logger.record_tabular("EpRewMean", np.mean(rewbuffer)) # This has nan values
logger.record_tabular("EpTrueRewMean", np.mean(true_rewbuffer)) # This has nan values
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
#timesteps_so_far += sum(lens)
timesteps_so_far += seg["steps"] # changed to match setup with no finishing condition
iters_so_far += 1
#env.reset() #reset the environment after a new iteration, therefore in traj generator check ob
logger.record_tabular("EpisodesSoFar", episodes_so_far) # This is 0 ? if lens which is the number of entries for episode length doesnt exist, doesnt make sense for it to have a return.
logger.record_tabular("TimestepsSoFar", timesteps_so_far)
logger.record_tabular("TimeElapsed", time.time() - tstart)
# I think the entloss, entrpoy, ev_.... and the useful ones arent from the environment called using the trpo
if rank == 0:
logger.dump_tabular()
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 learn(env,
policy_func,
reward_giver,
expert_dataset,
rank,
g_step,
d_step,
entcoeff,
save_per_iter,
timesteps_per_batch,
ckpt_dir,
log_dir,
task_name,
gamma,
lam,
max_kl,
cg_iters,
cg_damping=1e-2,
vf_stepsize=3e-4,
d_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0,
callback=None):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", ob_space, ac_space)
saver = tf.train.Saver(
var_list=tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, scope='pi'))
saver.restore(tf.get_default_session(), U_.getPath() + '/model/bc.ckpt')
oldpi = policy_func("oldpi", ob_space, ac_space)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
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 = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(pi.vpred - 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()
var_list = [
v for v in all_var_list
if v.name.startswith("pi/pol") or v.name.startswith("pi/logstd")
]
vf_var_list = [v for v in all_var_list if v.name.startswith("pi/vff")]
assert len(var_list) == len(vf_var_list) + 1
d_adam = MpiAdam(reward_giver.get_trainable_variables())
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(oldpi.get_variables(), pi.get_variables())
])
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()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
d_adam.sync()
vfadam.sync()
if rank == 0:
print("Init param sum", th_init.sum())
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
reward_giver,
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
true_rewbuffer = deque(maxlen=40)
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
g_loss_stats = stats(loss_names)
d_loss_stats = stats(reward_giver.loss_name)
ep_stats = stats(["True_rewards", "Rewards", "Episode_length"])
# if provide pretrained weight
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
# Save model
if rank == 0 and iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
fname = os.path.join(ckpt_dir, task_name)
print('save model as ', fname)
try:
os.makedirs(os.path.dirname(fname))
except OSError:
# folder already exists
pass
saver = tf.train.Saver()
saver.save(tf.get_default_session(), fname)
print("********** Iteration %i ************" % iters_so_far)
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
# ------------------ Update G ------------------
print("Optimizing Policy...")
for _ in range(g_step):
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, "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]
assign_old_eq_new(
) # set old parameter values to new parameter values
with timed("computegrad"):
tmp_result = compute_lossandgrad(seg["ob"], seg["ac"], atarg)
lossbefore = tmp_result[:-1]
g = tmp_result[-1]
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
if np.allclose(g, 0):
print("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)
# print("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 = allmean(
np.array(compute_losses(seg["ob"], seg["ac"], atarg)))
surr = meanlosses[0]
kl = meanlosses[1]
improve = surr - surrbefore
print("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(meanlosses).all():
print("Got non-finite value of losses -- bad!")
elif kl > max_kl * 1.5:
print("violated KL constraint. shrinking step.")
elif improve < 0:
print("surrogate didn't improve. shrinking step.")
else:
print("Stepsize OK!")
break
stepsize *= .5
else:
print("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:])
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=128):
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(
mbob) # update running mean/std for policy
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
print("ev_tdlam_before", explained_variance(vpredbefore, tdlamret))
# ------------------ Update D ------------------
print("Optimizing Discriminator...")
print(fmt_row(13, reward_giver.loss_name))
ob_expert, ac_expert = expert_dataset.get_next_batch(len(ob))
batch_size = len(ob) // d_step
d_losses = [
] # list of tuples, each of which gives the loss for a minibatch
for ob_batch, ac_batch in tqdm(
dataset.iterbatches((ob, ac),
include_final_partial_batch=False,
batch_size=batch_size)):
ob_expert, ac_expert = expert_dataset.get_next_batch(len(ob_batch))
# update running mean/std for reward_giver
if hasattr(reward_giver, "obs_rms"):
reward_giver.obs_rms.update(
np.concatenate((ob_batch, ob_expert), 0))
tmp_result = reward_giver.lossandgrad(ob_batch, ac_batch,
ob_expert, ac_expert)
newlosses = tmp_result[:-1]
g = tmp_result[-1]
d_adam.update(allmean(g), d_stepsize)
d_losses.append(newlosses)
print(fmt_row(13, np.mean(d_losses, axis=0)))
timesteps_so_far += len(seg['ob'])
iters_so_far += 1
print("EpisodesSoFar", episodes_so_far)
print("TimestepsSoFar", timesteps_so_far)
print("TimeElapsed", time.time() - tstart)
def learn(env, policy_fn, *,
batch_size, # what to train on
task_horizon,
max_kl, cg_iters,
gamma, lam, # advantage estimation
entcoeff=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters =3,
max_timesteps=0, max_episodes=0, max_iters=0, # time constraint
callback=None,
weights_dir='.',
per_decision = True,
normalize = False,
truncate_at = np.infty
):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
timesteps_per_batch = batch_size * task_horizon
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_fn("pi", ob_space, ac_space)
oldpi = policy_fn("oldpi", ob_space, ac_space)
atarg = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
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 = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(pi.vpred - 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()
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")]
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(oldpi.get_variables(), pi.get_variables())])
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()
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
# ----------------------------------------
seg_gen = traj_segment_generator(pi, env, task_horizon, 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
assert sum([max_iters>0, max_timesteps>0, max_episodes>0])==1
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_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)
#Params
#"""
params = pi.eval_param()
#print(params)
np.save(weights_dir+'/weights_'+str(iters_so_far), params)
#"""
# 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)
print('DOT: %s' % np.dot(stepdir, g))
# 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"], seg["ob"],
seg["ac"],seg["rew"]) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, states, actions, rewards = map(flatten_lists, zip(*listoflrpairs))
disc_rews = []
start = 0
for ep_len in lens:
end = start + ep_len
disc = gamma + np.zeros(ep_len)
disc[0] = 1
disc = np.cumprod(disc)
disc_rewards = np.array(rewards[start:end]) * disc
disc_rews.append(np.sum(disc_rewards))
start = end
#Save importance weights
simple_iw = pi.eval_simple_iw(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi)
np.save(weights_dir+'/iws_'+str(iters_so_far), simple_iw)
#print(len(simple_iw), simple_iw)
#Save returns
ep_rets = np.array(disc_rews)
np.save(weights_dir+'/rets_'+str(iters_so_far), ep_rets)
#print(len(ep_rets), ep_rets)
#lenbuffer.extend(lens)
#rewbuffer.extend(rews)
#Renyi
"""
renyi_4 = np.mean(pi.eval_renyi(states, oldpi, 4))
#print('Renyi:', renyi)
#"""
#Importance weights stats
"""
avg_iw, var_iw, max_iw, ess = pi.eval_iw_stats(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi,
per_decision=per_decision,
normalize=normalize,
truncate_at=truncate_at)
#"""
#Returns stats
"""
avg_ret, var_ret, max_ret = pi.eval_ret_stats(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi,
per_decision=per_decision,
normalize=normalize,
truncate_at=truncate_at)
#"""
#Performance
#"""
bound_delta = .2
batch_size = len(lens)
J = pi.eval_J(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi,
per_decision=per_decision,
normalize=normalize,
truncate_at=truncate_at)
var_J = pi.eval_var_J(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi,
per_decision=per_decision,
normalize=normalize,
truncate_at=truncate_at)
"""
bound = pi.eval_bound(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi,
per_decision=per_decision,
normalize=normalize,
truncate_at=truncate_at,
delta=bound_delta,
use_ess=True)
#"""
#Sample Renyi
d2s = pi.eval_renyi(states, oldpi, 2)
d2s_by_episode = []
start = 0
for ep_len in lens:
end = start + ep_len
d2s_by_episode = np.sum(d2s[start:end])
start = end
sample_d2 = np.mean(np.exp(d2s_by_episode))
"""
grad_bound = pi.eval_grad_bound(states,
actions,
rewards,
lens,
gamma=gamma,
behavioral=oldpi,
per_decision=per_decision,
normalize=normalize,
truncate_at=truncate_at,
delta=bound_delta,
use_ess=True)
print(grad_bound)
#print('Target performance', J, '+-', np.sqrt(var_J/len(lens)))
#"""
#Gradients
"""
grad_J = pi.eval_grad_J(states,
actions,
rewards,
lens,
behavioral=oldpi,
per_decision=True)
grad_var_J = pi.eval_grad_var_J(states,
actions,
rewards,
lens,
behavioral=oldpi,
per_decision=True)
print('Target performance grads', grad_J, grad_var_J)
#"""
#Student-t bound
"""
bound = pi.eval_bound(states,
actions,
rewards,
lens,
behavioral=oldpi,
per_decision=True)
#print('Bound comp. time', time.time() - checkpoint)
print("StudentTBound", bound)
#"""
#Student-t bound grad
"""
bound_grad = pi.eval_bound_grad(states,
actions,
rewards,
lens,
behavioral=oldpi,
per_decision=True)
print("StudentTBound grad", bound_grad)
#"""
#Fisher
"""
checkpoint = time.time()
fisher = oldpi.eval_fisher(states, actions, lens, behavioral=None)
#print(fisher)
assert np.array_equal(fisher, fisher.T)
print('Fisher comp. time', time.time() - checkpoint)
checkpoint = time.time()
natural = np.linalg.solve(fisher + 1e-12*np.eye(fisher.shape[0]), grad_J)
print(natural)
#print('Fisher vector product time:', time.time() - checkpoint)
#"""
#Logging
logger.record_tabular("Step_size", stepsize)
#logger.record_tabular("Our_bound", bound)
#logger.record_tabular("Reny_4", renyi_4)
logger.record_tabular("SampleRenyi2", sample_d2)
#logger.record_tabular("Max_iw", max_iw)
#logger.record_tabular("Ess", ess)
#logger.record_tabular("Avg_iw", avg_iw)
#logger.record_tabular("Var_iw", var_iw)
#logger.record_tabular("Max_ret", max_ret)
#logger.record_tabular("Avg_ret", avg_ret)
#logger.record_tabular("Var_ret", var_ret)
logger.record_tabular("EpLenMean", np.mean(lens))
logger.record_tabular("DiscEpRewMean", np.mean(disc_rews))
logger.record_tabular("EpRewMean", np.mean(rews))
logger.record_tabular("EpThisIter", len(lens))
logger.record_tabular("J_hat", J)
logger.record_tabular("Var_J", var_J)
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()
def learn(
env,
policy_func,
*,
timesteps_per_batch,
max_kl,
cg_iters,
gamma,
lam,
entcoeff=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0,
callback=None,
# GAIL Params
pretrained_weight=None,
reward_giver=None,
expert_dataset=None,
rank=0,
save_per_iter=1,
ckpt_dir="/tmp/gail/ckpt/",
g_step=1,
d_step=1,
task_name="task_name",
d_stepsize=3e-4,
using_gail=True):
"""
learns a GAIL policy using the given environment
:param env: (Gym Environment) the environment
:param policy_func: (function (str, Gym Space, Gym Space, bool): MLPPolicy) policy generator
:param timesteps_per_batch: (int) the number of timesteps to run per batch (horizon)
:param max_kl: (float) the kullback leiber loss threashold
:param cg_iters: (int) the number of iterations for the conjugate gradient calculation
:param gamma: (float) the discount value
:param lam: (float) GAE factor
:param entcoeff: (float) the weight for the entropy loss
:param cg_damping: (float) the compute gradient dampening factor
:param vf_stepsize: (float) the value function stepsize
:param vf_iters: (int) the value function's number iterations for learning
:param max_timesteps: (int) the maximum number of timesteps before halting
:param max_episodes: (int) the maximum number of episodes before halting
:param max_iters: (int) the maximum number of training iterations before halting
:param callback: (function (dict, dict)) the call back function, takes the local and global attribute dictionary
:param pretrained_weight: (str) the save location for the pretrained weights
:param reward_giver: (TransitionClassifier) the reward predicter from obsevation and action
:param expert_dataset: (MujocoDset) the dataset manager
:param rank: (int) the rank of the mpi thread
:param save_per_iter: (int) the number of iterations before saving
:param ckpt_dir: (str) the location for saving checkpoints
:param g_step: (int) number of steps to train policy in each epoch
:param d_step: (int) number of steps to train discriminator in each epoch
:param task_name: (str) the name of the task (can be None)
:param d_stepsize: (float) the reward giver stepsize
:param using_gail: (bool) using the GAIL model
"""
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
sess = tf_util.single_threaded_session()
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
policy = policy_func("pi", ob_space, ac_space, sess=sess)
old_policy = policy_func("oldpi",
ob_space,
ac_space,
sess=sess,
placeholders={
"obs": policy.obs_ph,
"stochastic": policy.stochastic_ph
})
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
observation = policy.obs_ph
action = policy.pdtype.sample_placeholder([None])
kloldnew = old_policy.proba_distribution.kl(policy.proba_distribution)
ent = policy.proba_distribution.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
entbonus = entcoeff * meanent
vferr = tf.reduce_mean(tf.square(policy.vpred - ret))
# advantage * pnew / pold
ratio = tf.exp(
policy.proba_distribution.logp(action) -
old_policy.proba_distribution.logp(action))
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 = policy.get_trainable_variables()
if using_gail:
var_list = [
v for v in all_var_list
if v.name.startswith("pi/pol") or v.name.startswith("pi/logstd")
]
vf_var_list = [v for v in all_var_list if v.name.startswith("pi/vff")]
assert len(var_list) == len(vf_var_list) + 1
d_adam = MpiAdam(reward_giver.get_trainable_variables())
else:
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")
]
vfadam = MpiAdam(vf_var_list, sess=sess)
get_flat = tf_util.GetFlat(var_list, sess=sess)
set_from_flat = tf_util.SetFromFlat(var_list, sess=sess)
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:
var_size = tf_util.intprod(shape)
tangents.append(tf.reshape(flat_tangent[start:start + var_size],
shape))
start += var_size
gvp = tf.add_n([
tf.reduce_sum(grad * tangent)
for (grad, tangent) in zipsame(klgrads, tangents)
]) # pylint: disable=E1111
fvp = tf_util.flatgrad(gvp, var_list)
assign_old_eq_new = tf_util.function(
[], [],
updates=[
tf.assign(oldv, newv) for (oldv, newv) in zipsame(
old_policy.get_variables(), policy.get_variables())
])
compute_losses = tf_util.function([observation, action, atarg], losses)
compute_lossandgrad = tf_util.function(
[observation, action, atarg],
losses + [tf_util.flatgrad(optimgain, var_list)])
compute_fvp = tf_util.function([flat_tangent, observation, action, atarg],
fvp)
compute_vflossandgrad = tf_util.function([observation, ret],
tf_util.flatgrad(
vferr, vf_var_list))
@contextmanager
def timed(msg):
if rank == 0:
print(colorize(msg, color='magenta'))
start_time = time.time()
yield
print(
colorize("done in %.3f seconds" % (time.time() - start_time),
color='magenta'))
else:
yield
def allmean(arr):
assert isinstance(arr, np.ndarray)
out = np.empty_like(arr)
MPI.COMM_WORLD.Allreduce(arr, out, op=MPI.SUM)
out /= nworkers
return out
tf_util.initialize(sess=sess)
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
if using_gail:
d_adam.sync()
vfadam.sync()
if rank == 0:
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
if using_gail:
seg_gen = traj_segment_generator(policy,
env,
timesteps_per_batch,
stochastic=True,
reward_giver=reward_giver,
gail=True)
else:
seg_gen = traj_segment_generator(policy,
env,
timesteps_per_batch,
stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
t_start = 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, max_timesteps > 0, max_episodes > 0]) == 1
if using_gail:
true_rewbuffer = deque(maxlen=40)
# Stats not used for now
# g_loss_stats = Stats(loss_names)
# d_loss_stats = Stats(reward_giver.loss_name)
# ep_stats = Stats(["True_rewards", "Rewards", "Episode_length"])
# if provide pretrained weight
if pretrained_weight is not None:
tf_util.load_state(pretrained_weight,
var_list=policy.get_variables())
while True:
if callback:
callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
# Save model
if using_gail and rank == 0 and iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
fname = os.path.join(ckpt_dir, task_name)
os.makedirs(os.path.dirname(fname), exist_ok=True)
saver = tf.train.Saver()
saver.save(sess, fname)
logger.log("********** Iteration %i ************" % iters_so_far)
def fisher_vector_product(vec):
return allmean(compute_fvp(vec, *fvpargs,
sess=sess)) + cg_damping * vec
# ------------------ Update G ------------------
logger.log("Optimizing Policy...")
# g_step = 1 when not using GAIL
for _ in range(g_step):
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))
observation, action, 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(policy, "ret_rms"):
policy.ret_rms.update(tdlamret)
if hasattr(policy, "ob_rms"):
policy.ob_rms.update(
observation) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
assign_old_eq_new(sess=sess)
with timed("computegrad"):
*lossbefore, grad = compute_lossandgrad(*args, sess=sess)
lossbefore = allmean(np.array(lossbefore))
grad = allmean(grad)
if np.allclose(grad, 0):
logger.log("Got zero gradient. not updating")
else:
with timed("cg"):
stepdir = conjugate_gradient(fisher_vector_product,
grad,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
# abs(shs) to avoid taking square root of negative values
lagrange_multiplier = np.sqrt(abs(shs) / max_kl)
# logger.log("lagrange multiplier:", lm, "gnorm:", np.linalg.norm(g))
fullstep = stepdir / lagrange_multiplier
expectedimprove = grad.dot(fullstep)
surrbefore = lossbefore[0]
stepsize = 1.0
thbefore = get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
set_from_flat(thnew)
mean_losses = surr, kl_loss, *_ = allmean(
np.array(compute_losses(*args, sess=sess)))
improve = surr - surrbefore
logger.log("Expected: %.3f Actual: %.3f" %
(expectedimprove, improve))
if not np.isfinite(mean_losses).all():
logger.log("Got non-finite value of losses -- bad!")
elif kl_loss > 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:])
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=128):
if hasattr(policy, "ob_rms"):
policy.ob_rms.update(
mbob) # update running mean/std for policy
grad = allmean(
compute_vflossandgrad(mbob, mbret, sess=sess))
vfadam.update(grad, vf_stepsize)
for (loss_name, loss_val) in zip(loss_names, mean_losses):
logger.record_tabular(loss_name, loss_val)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
if using_gail:
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, reward_giver.loss_name))
ob_expert, ac_expert = expert_dataset.get_next_batch(
len(observation))
batch_size = len(observation) // d_step
d_losses = [
] # list of tuples, each of which gives the loss for a minibatch
for ob_batch, ac_batch in dataset.iterbatches(
(observation, action),
include_final_partial_batch=False,
batch_size=batch_size):
ob_expert, ac_expert = expert_dataset.get_next_batch(
len(ob_batch))
# update running mean/std for reward_giver
if hasattr(reward_giver, "obs_rms"):
reward_giver.obs_rms.update(
np.concatenate((ob_batch, ob_expert), 0))
*newlosses, grad = reward_giver.lossandgrad(
ob_batch, ac_batch, ob_expert, ac_expert)
d_adam.update(allmean(grad), d_stepsize)
d_losses.append(newlosses)
logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
lrlocal = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]
) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*listoflrpairs))
true_rewbuffer.extend(true_rets)
else:
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)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
if using_gail:
logger.record_tabular("EpTrueRewMean", np.mean(true_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() - t_start)
if rank == 0:
logger.dump_tabular()
def learn(
env,
policy_fn,
*,
timesteps_per_batch, # what to train on
epsilon,
beta,
cg_iters,
gamma,
lam, # advantage estimation
trial,
sess,
method,
entcoeff=0.0,
cg_damping=1e-2,
kl_target=0.01,
crosskl_coeff=0.01,
vf_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
TRPO=False):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
total_space = env.total_space
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_fn("pi", ob_space, ac_space, ob_name="ob")
oldpi = policy_fn("oldpi", ob_space, ac_space, ob_name="ob")
gpi = policy_fn("gpi", total_space, ac_space, ob_name="gob")
goldpi = policy_fn("goldpi", total_space, ac_space, ob_name="gob")
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
gatarg = tf.placeholder(dtype=tf.float32, shape=[None])
gret = tf.placeholder(dtype=tf.float32, shape=[None])
ob = U.get_placeholder_cached(name="ob")
gob = U.get_placeholder_cached(name='gob')
ac = pi.pdtype.sample_placeholder([None])
crosskl_c = tf.placeholder(dtype=tf.float32, shape=[])
# crosskl_c = 0.01
kloldnew = oldpi.pd.kl(pi.pd)
gkloldnew = goldpi.pd.kl(gpi.pd)
#TODO: check if it can work in this way
# crosskl_ob = pi.pd.kl(goldpi.pd)
# crosskl_gob = gpi.pd.kl(oldpi.pd)
crosskl_gob = pi.pd.kl(gpi.pd)
crosskl_ob = gpi.pd.kl(pi.pd)
# crosskl
pdmean = pi.pd.mean
pdstd = pi.pd.std
gpdmean = gpi.pd.mean
gpdstd = gpi.pd.std
ent = pi.pd.entropy()
gent = gpi.pd.entropy()
old_entropy = oldpi.pd.entropy()
gold_entropy = goldpi.pd.entropy()
meankl = tf.reduce_mean(kloldnew)
meanent = tf.reduce_mean(ent)
meancrosskl = tf.reduce_mean(crosskl_ob)
# meancrosskl = tf.maximum(tf.reduce_mean(crosskl_ob - 100), 0)
gmeankl = tf.reduce_mean(gkloldnew)
gmeanent = tf.reduce_mean(gent)
gmeancrosskl = tf.reduce_mean(crosskl_gob)
vferr = tf.reduce_mean(tf.square(pi.vpred - ret))
gvferr = tf.reduce_mean(tf.square(gpi.vpred - gret))
ratio = tf.exp(pi.pd.logp(ac) -
oldpi.pd.logp(ac)) # advantage * pnew / pold
gratio = tf.exp(gpi.pd.logp(ac) - goldpi.pd.logp(ac))
# Ratio objective
# surrgain = tf.reduce_mean(ratio * atarg)
# gsurrgain = tf.reduce_mean(gratio * gatarg)
# Log objective
surrgain = tf.reduce_mean(pi.pd.logp(ac) * atarg)
gsurrgain = tf.reduce_mean(gpi.pd.logp(ac) * gatarg)
# optimgain = surrgain + crosskl_c * meancrosskl
optimgain = surrgain
losses = [
optimgain, meankl, meancrosskl, surrgain, meanent,
tf.reduce_mean(ratio)
]
loss_names = [
"optimgain", "meankl", "meancrosskl", "surrgain", "entropy", "ratio"
]
# goptimgain = gsurrgain + crosskl_c * gmeancrosskl
goptimgain = gsurrgain
glosses = [
goptimgain, gmeankl, gmeancrosskl, gsurrgain, gmeanent,
tf.reduce_mean(gratio)
]
gloss_names = [
"goptimgain", "gmeankl", "gmeancrosskl", "gsurrgain", "gentropy",
"gratio"
]
dist = meankl
gdist = gmeankl
all_pi_var_list = pi.get_trainable_variables()
all_var_list = [
v for v in all_pi_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")
]
vfadam = MpiAdam(vf_var_list)
poladam = MpiAdam(var_list)
gall_gpi_var_list = gpi.get_trainable_variables()
gall_var_list = [
v for v in gall_gpi_var_list if v.name.split("/")[0].startswith("gpi")
]
gvar_list = [
v for v in gall_var_list if v.name.split("/")[1].startswith("pol")
]
gvf_var_list = [
v for v in gall_var_list if v.name.split("/")[1].startswith("vf")
]
gvfadam = MpiAdam(gvf_var_list)
# gpoladpam = MpiAdam(gvar_list)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(var_list)
klgrads = tf.gradients(dist, var_list)
# crossklgrads = tf.gradients(meancrosskl, 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)
gget_flat = U.GetFlat(gvar_list)
gset_from_flat = U.SetFromFlat(gvar_list)
gklgrads = tf.gradients(gdist, gvar_list)
# gcrossklgrads = tf.gradients(gmeancrosskl, gvar_list)
gflat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="gflat_tan")
gshapes = [var.get_shape().as_list() for var in gvar_list]
gstart = 0
gtangents = []
for shape in gshapes:
sz = U.intprod(shape)
gtangents.append(tf.reshape(gflat_tangent[gstart:gstart + sz], shape))
gstart += sz
ggvp = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(gklgrads, gtangents)
]) #pylint: disable=E1111
gfvp = U.flatgrad(ggvp, gvar_list)
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(oldpi.get_variables(), pi.get_variables())
])
gassign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv,
newv) in zipsame(goldpi.get_variables(), gpi.get_variables())
])
compute_losses = U.function([crosskl_c, gob, ob, ac, atarg], losses)
compute_lossandgrad = U.function([crosskl_c, gob, 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))
compute_crossklandgrad = U.function([ob, gob],
U.flatgrad(meancrosskl, var_list))
gcompute_losses = U.function([crosskl_c, ob, gob, ac, gatarg], glosses)
gcompute_lossandgrad = U.function([crosskl_c, ob, gob, ac, gatarg],
glosses +
[U.flatgrad(goptimgain, gvar_list)])
gcompute_fvp = U.function([gflat_tangent, gob, ac, gatarg], gfvp)
gcompute_vflossandgrad = U.function([gob, gret],
U.flatgrad(gvferr, gvf_var_list))
# compute_gcrossklandgrad = U.function([gob, ob], U.flatgrad(gmeancrosskl, gvar_list))
saver = tf.train.Saver()
@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()
guided_initilizer(gpol=gvar_list,
gvf=gvf_var_list,
fpol=var_list,
fvf=vf_var_list)
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
poladam.sync()
print("Init final policy param sum", th_init.sum(), flush=True)
gth_init = gget_flat()
MPI.COMM_WORLD.Bcast(gth_init, root=0)
gset_from_flat(gth_init)
gvfadam.sync()
# gpoladpam.sync()
print("Init guided policy param sum", gth_init.sum(), flush=True)
# Initialize eta, omega optimizer
init_eta = 0.5
init_omega = 2.0
eta_omega_optimizer = EtaOmegaOptimizer(beta, epsilon, init_eta,
init_omega)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
gpi,
env,
timesteps_per_batch,
stochastic=True)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
num_iters = max_timesteps // timesteps_per_batch
lenbuffer = deque(maxlen=40) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=40) # rolling buffer for episode rewards
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_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"]
gob, gatarg, gtdlamret = seg["gob"], seg["gadv"], seg["gtdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
gvpredbefore = seg["gvpred"]
atarg = (atarg - atarg.mean()
) / atarg.std() # standardized advantage function estimate
gatarg = (gatarg - gatarg.mean()) / gatarg.std()
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
if hasattr(gpi, "ret_rms"): gpi.ret_rms.update(gtdlamret)
if hasattr(gpi, "ob_rms"): gpi.ob_rms.update(gob)
args = crosskl_coeff, seg["gob"], seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args[2:]]
gargs = crosskl_coeff, seg["ob"], seg["gob"], seg["ac"], gatarg
gfvpargs = [arr[::5] for arr in gargs[2:]]
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
def gfisher_vector_product(p):
return allmean(gcompute_fvp(p, *gfvpargs)) + cg_damping * p
assign_old_eq_new() # set old parameter values to new parameter values
gassign_old_eq_new()
with timed("computegrad"):
*lossbefore, g = compute_lossandgrad(*args)
*glossbefore, gg = gcompute_lossandgrad(*gargs)
lossbefore = allmean(np.array(lossbefore))
g = allmean(g)
glossbefore = allmean(np.array(glossbefore))
gg = allmean(gg)
if np.allclose(g, 0) or np.allclose(gg, 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)
gstepdir = cg(gfisher_vector_product,
gg,
cg_iters=cg_iters,
verbose=rank == 0)
assert np.isfinite(gstepdir).all()
assert np.isfinite(stepdir).all()
if TRPO:
#
# TRPO specific code.
# Find correct step size using line search
#
#TODO: also enable guided learning for TRPO
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
gcopos_update_dir = gstepdir
# Split direction into log-linear 'w_theta' and non-linear 'w_beta' parts
w_theta, w_beta = pi.split_w(copos_update_dir)
gw_theta, gw_beta = gpi.split_w(gcopos_update_dir)
# q_beta(s,a) = \grad_beta \log \pi(a|s) * w_beta
# = features_beta(s) * K^T * Prec * a
# q_beta = self.target.get_q_beta(features_beta, actions)
Waa, Wsa = pi.w2W(w_theta)
wa = pi.get_wa(ob, w_beta)
gWaa, gWsa = gpi.w2W(gw_theta)
gwa = gpi.get_wa(gob, gw_beta)
varphis = pi.get_varphis(ob)
gvarphis = gpi.get_varphis(gob)
# Optimize eta and omega
tmp_ob = np.zeros(
(1, ) + ob_space.shape
) # We assume that entropy does not depend on the NN
old_ent = old_entropy.eval({oldpi.ob: tmp_ob})[0]
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 of final policy: " + str(eta) +
" and omega: " + str(omega))
gtmp_ob = np.zeros((1, ) + total_space.shape)
gold_ent = gold_entropy.eval({goldpi.ob: gtmp_ob})[0]
geta, gomega = eta_omega_optimizer.optimize(
gw_theta, gWaa, gWsa, gwa, gvarphis, gpi.get_kt(),
gpi.get_prec_matrix(), gpi.is_new_policy_valid, gold_ent)
logger.log("Initial eta of guided policy: " + str(geta) +
" and omega: " + str(gomega))
current_theta_beta = get_flat()
prev_theta, prev_beta = pi.all_to_theta_beta(
current_theta_beta)
gcurrent_theta_beta = gget_flat()
gprev_theta, gprev_beta = gpi.all_to_theta_beta(
gcurrent_theta_beta)
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 of final policy, 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 of final policy, eta: " +
str(eta) + " and omega: " + str(omega))
geta = eta_search(gw_theta, gw_beta, geta, gomega, allmean,
gcompute_losses, gget_flat,
gset_from_flat, gpi, epsilon, gargs)
logger.log("updated eta of guided policy, eta:" +
str(geta) + "and omega:" + str(gomega))
gset_from_flat(
gpi.theta_beta_to_all(gprev_theta, gprev_beta))
geta, gomega = eta_omega_optimizer.optimize(
gw_theta, gWaa, gWsa, gwa, gvarphis, gpi.get_kt(),
gpi.get_prec_matrix(), gpi.is_new_policy_valid,
gold_ent, geta)
logger.log("Updated omega of guided policy, eta:" +
str(geta) + "and omega:" + str(gomega))
# Use final policy
logger.log("Final eta of final policy: " + str(eta) +
" and omega: " + str(omega))
logger.log("Final eta of guided policy: " + str(geta) +
"and omega:" + str(gomega))
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))
gcur_theta = (geta * gprev_theta +
gw_theta.reshape(-1, )) / (geta + gomega)
gcur_beta = gprev_beta + gw_beta.reshape(-1, ) / geta
gset_from_flat(gpi.theta_beta_to_all(gcur_theta, gcur_beta))
meanlosses = surr, kl, crosskl, *_ = allmean(
np.array(compute_losses(*args)))
gmeanlosses = gsurr, gkl, gcrosskl, *_ = allmean(
np.array(gcompute_losses(*gargs)))
# poladam.update(allmean(compute_crossklandgrad(ob, gob)), vf_stepsize)
# gpoladpam.update(allmean(compute_gcrossklandgrad(gob, ob)), vf_stepsize)
for _ in range(vf_iters):
for (mbob, mbgob) in dataset.iterbatches(
(seg["ob"], seg["gob"]),
include_final_partial_batch=False,
batch_size=64):
g = allmean(compute_crossklandgrad(mbob, mbgob))
poladam.update(g, vf_stepsize)
# pd_crosskl = np.mean((crosskl, gcrosskl))
# pd_crosskl = crosskl
# if pd_crosskl < kl_target / 2:
# print("KL divergence between guided policy and final control policy is small, reduce the coefficient")
# crosskl_coeff /= 1.5
# elif pd_crosskl > kl_target * 2:
# print("KL divergence between guided policy and final control policy is large, increse the coefficient")
# crosskl_coeff *= 1.5
# crosskl_coeff = np.clip(crosskl_coeff, 1e-4, 30)
# 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)
for (lossname, lossval) in zip(gloss_names, gmeanlosses):
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)
for (mbob, mbret) in dataset.iterbatches(
(seg["gob"], seg["gtdlamret"]),
include_final_partial_batch=False,
batch_size=64):
gg = allmean(gcompute_vflossandgrad(mbob, mbret))
gvfadam.update(gg, vf_stepsize)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
logger.record_tabular("gev_tdlam_before",
explained_variance(gvpredbefore, gtdlamret))
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)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpThisIter", len(lens))
logger.record_tabular("CrossKLCoeff :", crosskl_coeff)
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)
logger.record_tabular("Name", method)
logger.record_tabular("Iteration", iters_so_far)
logger.record_tabular("trial", trial)
if rank == 0:
logger.dump_tabular()
if iters_so_far % 100 == 0 or iters_so_far == 1 or iters_so_far == num_iters:
# sess = tf.get_default_session()
checkdir = get_dir(osp.join(logger.get_dir(), 'checkpoints'))
savepath = osp.join(checkdir, '%.5i.ckpt' % iters_so_far)
saver.save(sess, save_path=savepath)
print("save model to path:", savepath)
def hybrid_learn(env, policy_func, reward_giver, rank,
*,
policy_step, boundary_step, entcoeff, save_per_iter,
ckpt_dir, log_dir, timesteps_per_batch, task_name_1, task_name_2,
gamma, lam,
max_kl, cg_iters, cg_damping=1e-2,
vf_stepsize=3e-4, d_stepsize=3e-4, vf_iters=3,
max_timesteps=0, max_episodes=0, max_iters=0,
callback=None
):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi_task1 = policy_func("pi_task1", ob_space, ac_space)
oldpi_task1 = policy_func("oldpi_task1", ob_space, ac_space)
pi_task2 = policy_func("pi_task2", ob_space, ac_space)
oldpi_task2 = policy_func("oldpi_task2", ob_space, ac_space)
atarg_task1 = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret_task1 = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
atarg_task2 = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret_task2 = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob_task1 = U.get_placeholder_cached(name="ob_task1")
ac_task1 = pi_task1.pdtype.sample_placeholder([None])
ob_task2 = U.get_placeholder_cached(name="ob_task2")
ac_task2 = pi_task2.pdtype.sample_placeholder([None])
kloldnew_task1 = oldpi_task1.pd.kl(pi_task1.pd)
ent_task1 = pi_task1.pd.entropy()
meankl_task1 = tf.reduce_mean(kloldnew_task1)
meanent_task1 = tf.reduce_mean(ent_task1)
entbonus_task1 = entcoeff_task1 * meanent_task1
kloldnew_task2 = oldpi_task2.pd.kl(pi_task2.pd)
ent_task2 = pi_task2.pd.entropy()
meankl_task2 = tf.reduce_mean(kloldnew_task2)
meanent_task2 = tf.reduce_mean(ent_task2)
entbonus_task2 = entcoeff_task2 * meanent_task2
vferr_task1 = tf.reduce_mean(tf.square(pi_task1.vpred - ret_task1))
vferr_task2 = tf.reduce_mean(tf.square(pi_task2.vpred - ret_task2))
ratio_task1 = tf.exp(pi_task1.pd.logp(ac) - oldpi_task1.pd.logp(ac)) # advantage * pnew / pold
ratio_task2 = tf.exp(pi_task2.pd.logp(ac) - oldpi_task2.pd.logp(ac)) # advantage * pnew / pold
surrgain_task1 = tf.reduce_mean(ratio_task1 * atarg_task1)
surrgain_task2 = tf.reduce_mean(ratio_task2 * atarg_task2)
optimgain_task1 = surrgain_task1 + entbonus_task1
optimgain_task2 = surrgain_task2 + entbonus_task2
optimgain = optimgain_task1 + optimgain_task2
meankl = meankl_task1 + meankl_task2
entbonus = entbonus_task1 + entbonus_task2
surrgain = surrgain_task1 + surrgain_task2
meanent = meanent_task1 + meanent_task2
losses_task1 = [optimgain_task1, meankl_task1,
entbonus_task1, surrgain_task1, meanent_task1]
losses_task2 = [optimgain_task2, meankl_task2,
entbonus_task2, surrgain_task2, meanent_task2]
# losses = [optimgain_task1, optimgain_task2, meankl_task1, meankl_task2,
# entbonus_task1, entbonus_task2, surrgain_task1, surrgain_task2,
# meanent_task1, meanent_task2]
loss_names = ['optimgain_task1', 'optimgain_task2', 'meankl_task1', 'meankl_task2',
'entbonus_task1', 'entbonus_task2', 'surrgain_task1', 'surrgain_task2',
'meanent_task1', 'meanent_task2']
dist_task1 = meankl_task1
dist_task2 = meankl_task2
all_var_list_task1 = pi_task1.get_trainable_variables()
var_list_task1 = [v for v in all_var_list_task1 if v.name.startswith("pi/pol") or v.name.startswith("pi/logstd")]
vf_var_list_task1 = [v for v in all_var_list_task1 if v.name.startswith("pi/vff")]
assert len(var_list_task1) == len(vf_var_list_task1) + 1
# d_adam = MpiAdam(reward_giver.get_trainable_variables())
vfadam_task1 = MpiAdam(vf_var_list_task1)
all_var_list_task2 = pi_task2.get_trainable_variables()
var_list_task2 = [v for v in all_var_list_task2 if v.name.startswith("pi/pol") or v.name.startswith("pi/logstd")]
vf_var_list_task2 = [v for v in all_var_list_task2 if v.name.startswith("pi/vff")]
assert len(var_list_task2) == len(vf_var_list_task2) + 1
# d_adam = MpiAdam(reward_giver.get_trainable_variables())
vfadam_task2 = MpiAdam(vf_var_list_task2)
get_flat_task1 = U.GetFlat(var_list_task1)
set_from_flat_task1 = U.SetFromFlat(var_list_task1)
klgrads_task1 = tf.gradients(dist_task1, var_list_task1)
flat_tangent_task1 = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan_task1")
shapes_task1 = [var.get_shape().as_list() for var in var_list_task1]
start = 0
tangents_task1 = []
for shape in shapes_task1:
sz = U.intprod(shape)
tangents_task1.append(tf.reshape(flat_tangent_task1[start:start+sz], shape))
start += sz
gvp_task1 = tf.add_n([tf.reduce_sum(g*tangent) for (g, tangent) in zipsame(klgrads_task1, tangents_task1)]) # pylint: disable=E1111
fvp_task1 = U.flatgrad(gvp_task1, var_list_task1)
get_flat_task2 = U.GetFlat(var_list_task2)
set_from_flat_task2 = U.SetFromFlat(var_list_task2)
klgrads_task2 = tf.gradients(dist_task2, var_list_task2)
flat_tangent_task2 = tf.placeholder(dtype=tf.float32, shape=[None], name="flat_tan_task2")
shapes_task2 = [var.get_shape().as_list() for var in var_list_task2]
start = 0
tangents_task2 = []
for shape in shapes_task2:
sz = U.intprod(shape)
tangents_task2.append(tf.reshape(flat_tangent_task2[start:start+sz], shape))
start += sz
gvp_task2 = tf.add_n([tf.reduce_sum(g*tangent) for (g, tangent) in zipsame(klgrads_task2, tangents_task2)]) # pylint: disable=E1111
fvp_task2 = U.flatgrad(gvp_task2, var_list)_task2
assign_old_eq_new_task1 = U.function([], [], updates=[tf.assign(oldv, newv)
for (oldv, newv) in zipsame(oldpi_task1.get_variables(), pi_task1.get_variables())])
assign_old_eq_new_task2 = U.function([], [], updates=[tf.assign(oldv, newv)
for (oldv, newv) in zipsame(oldpi_task2.get_variables(), pi_task2.get_variables())])
compute_losses_task1 = U.function([ob, ac, atarg_task1], losses_task1)
compute_lossandgrad_task1 = U.function([ob, ac, atarg_task1], losses_task1 + [U.flatgrad(optimgain_task1, var_list_task1)])
compute_fvp_task1 = U.function([flat_tangent_task1, ob, ac, atarg_task1], fvp_task1)
compute_vflossandgrad_task1 = U.function([ob, ret_task1], U.flatgrad(vferr_task1, vf_var_list_task1))
compute_losses_task2 = U.function([ob, ac, atarg_task2], losses_task2)
compute_lossandgrad_task2 = U.function([ob, ac, atarg_task2], losses_task2 + [U.flatgrad(optimgain_task2, var_list_task2)])
compute_fvp_task2 = U.function([flat_tangent_task2, ob, ac, atarg_task2], fvp_task2)
compute_vflossandgrad_task2 = U.function([ob, ret_task2], U.flatgrad(vferr_task2, vf_var_list_task2))
@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()
th_init_task1 = get_flat_task1()
MPI.COMM_WORLD.Bcast(th_init_task1, root=0)
set_from_flat_task1(th_init_task1)
# d_adam.sync()
vfadam_task1.sync()
th_init_task2 = get_flat_task2()
MPI.COMM_WORLD.Bcast(th_init_task2, root=0)
set_from_flat_task2(th_init_task2)
# d_adam.sync()
vfadam_task2.sync()
if rank == 0:
print("Init param sum", th_init_task1.sum(), flush=True)
print("Init param sum", th_init_task2.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator_composed(pi_task1, pi_task2, env, boundary_condition, 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
true_rewbuffer = deque(maxlen=40)
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
g_loss_stats = stats(loss_names)
# d_loss_stats = stats(reward_giver.loss_name)
ep_stats = stats(["True_rewards", "Rewards", "Episode_length"])
fname = os.path.join(ckpt_dir, task_name)
weight_file = tf.train.latest_checkpoint(ckpt_dir)
print("fname: {} weight_file: {}".format(fname, weight_file))
if weight_file is not None:
U.load_state(weight_file)#, var_list=pi.get_variables())
tf.logging.info('%s loaded' % weight_file)
else:
print("from scratch")
tf.logging.info('Training from the scratch (no pre-trained weight_filets)..')
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_timesteps:
break
elif max_episodes and episodes_so_far >= max_episodes:
break
elif max_iters and iters_so_far >= max_iters:
break
# Save model
if rank == 0 and iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
fname = os.path.join(ckpt_dir, task_name)
os.makedirs(os.path.dirname(fname), exist_ok=True)
saver = tf.train.Saver()
saver.save(tf.get_default_session(), fname)
print("============= SAVE ===============")
logger.log("********** Iteration %i ************" % iters_so_far)
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
# ------------------ Update G ------------------
logger.log("Optimizing Policy...")
# First, we optimize each policy. and then next step we find optimized boundary for given policy.
for _ in range(policy_step):
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_task1, ac_task1, atarg_task1, tdlamret_task1 = seg["ob_task1"], seg["ac_task1"], seg["adv_task1"], seg["tdlamret_task1"]
ob_task2, ac_task2, atarg_task2, tdlamret_task2 = seg["ob_task2"], seg["ac_task2"], seg["adv_task2"], seg["tdlamret_task2"]
vpredbefore_task1 = seg["vpred_task1"] # predicted value function before udpate
vpredbefore_task2 = seg["vpred_task2"] # predicted value function before udpate
atarg_task1 = (atarg_task1 - atarg_task1.mean()) / atarg_task1.std() # standardized advantage function estimate
atarg_task2 = (atarg_task2 - atarg_task2.mean()) / atarg_task2.std() # standardized advantage function estimate
if hasattr(pi_task1, "ob_rms_task1"): pi_task1.ob_rms.update(ob_task1) # update running mean/std for policy
if hasattr(pi_task2, "ob_rms_task2"): pi_task2.ob_rms.update(ob_task2) # update running mean/std for policy
args_task1 = seg["ob_task1"], seg["ac_task1"], atarg_task1
fvpargs_task1 = [arr[::5] for arr in args_task1]
args_task2 = seg["ob_task2"], seg["ac_task2"], atarg_task2
fvpargs_task2 = [arr[::5] for arr in args_task2]
assign_old_eq_new_task1() # set old parameter values to new parameter values
assign_old_eq_new_task2() # 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:])
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=128):
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(mbob) # update running mean/std for policy
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
g_losses = meanlosses
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
logger.record_tabular("ev_tdlam_before", explained_variance(vpredbefore, tdlamret))
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, reward_giver.loss_name))
ob_expert, ac_expert = expert_dataset.get_next_batch(len(ob))
batch_size = len(ob) // d_step
d_losses = [] # list of tuples, each of which gives the loss for a minibatch
for ob_batch, ac_batch in dataset.iterbatches((ob, ac),
include_final_partial_batch=False,
batch_size=batch_size):
ob_expert, ac_expert = expert_dataset.get_next_batch(len(ob_batch))
# update running mean/std for reward_giver
if hasattr(reward_giver, "obs_rms"): reward_giver.obs_rms.update(np.concatenate((ob_batch, ob_expert), 0))
*newlosses, g = reward_giver.lossandgrad(ob_batch, ac_batch, ob_expert, ac_expert)
d_adam.update(allmean(g), d_stepsize)
d_losses.append(newlosses)
logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
lrlocal = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*listoflrpairs))
true_rewbuffer.extend(true_rets)
lenbuffer.extend(lens)
rewbuffer.extend(rews)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpTrueRewMean", np.mean(true_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()
def render_evaluate(env, policy_func, *,
timesteps_per_batch, # what to train on
max_kl, cg_iters,
gamma, lam, # advantage estimation
entcoeff=0.0,
cg_damping=1e-2,
vf_stepsize=3e-4,
vf_iters =3,
max_timesteps=0, max_episodes=0, max_iters=0, # time constraint
callback=None
):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", ob_space, ac_space)
oldpi = policy_func("oldpi", ob_space, ac_space)
atarg = tf.placeholder(dtype=tf.float32, shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
ob = U.get_placeholder_cached(name="ob")
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.compute_kl(pi.pd)
ent = pi.pd.entropy()
meankl = U.mean(kloldnew)
meanent = U.mean(ent)
entbonus = entcoeff * meanent
vferr = U.mean(tf.square(pi.vpred - ret))
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac)) # advantage * pnew / pold
surrgain = U.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()
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")]
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([U.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(oldpi.get_variables(), pi.get_variables())])
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
# set up saver
sess = tf.get_default_session()
saver = tf.train.Saver()
U.initialize()
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)
print("loading pretrained model")
saver.restore(sess, callback.model_dir)
# 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
assert sum([max_iters>0, max_timesteps>0, max_episodes>0])==1
import gym
env = gym.make('Ant-v1')
if True:
obsall = []
for _ in range(50):
obs = []
done = False
ob = env.reset()
#env.render()
stochastic = 1
obs.append(env.unwrapped.get_body_com('torso')[:2].copy())
while not done:
ac, vpred = pi.act(stochastic, ob)
ob, rew, done, _ = env.step(ac)
#env.render()
obs.append(env.unwrapped.get_body_com('torso')[:2].copy())
obsall.append(obs)
if rank==0:
logger.dump_tabular()
if callback is not None:
callback(locals(), globals())
"""
def learn(env,
policy_func,
reward_giver,
expert_dataset,
rank,
pretrained,
pretrained_weight,
*,
g_step,
d_step,
entcoeff,
ckpt_dir,
timesteps_per_batch,
task_name,
gamma,
lam,
max_kl,
cg_iters,
cg_damping=1e-2,
vf_stepsize=3e-4,
d_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0,
rnd_iter=200,
callback=None,
dyn_norm=False,
mmd=False):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
np.set_printoptions(precision=3)
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", ob_space, ac_space)
oldpi = policy_func("oldpi", ob_space, ac_space)
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ob = U.get_placeholder_cached(name="ob")
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 = entcoeff * meanent
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()
var_list = [
v for v in all_var_list
if v.name.startswith("pi/pol") or v.name.startswith("pi/logstd")
]
vf_var_list = [v for v in all_var_list if v.name.startswith("pi/vff")]
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(oldpi.get_variables(), pi.get_variables())
])
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 = pi.vlossandgrad
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()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
vfadam.sync()
if rank == 0:
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
reward_giver,
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
true_rewbuffer = deque(maxlen=40)
assert sum([max_iters > 0, max_timesteps > 0, max_episodes > 0]) == 1
ep_stats = stats(["True_rewards", "Rewards", "Episode_length"])
# if provide pretrained weight
if pretrained_weight is not None:
U.load_variables(pretrained_weight, variables=pi.get_variables())
else:
if not dyn_norm:
pi.ob_rms.update(expert_dataset[0])
if not mmd:
reward_giver.train(*expert_dataset, iter=rnd_iter)
#inspect the reward learned
# for batch in iterbatches(expert_dataset, batch_size=32):
# print(reward_giver.get_reward(*batch))
best = -2000
save_ind = 0
max_save = 3
while True:
if callback: callback(locals(), globals())
if max_timesteps and timesteps_so_far >= max_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)
def fisher_vector_product(p):
return allmean(compute_fvp(p, *fvpargs)) + cg_damping * p
# ------------------ Update G ------------------
# logger.log("Optimizing Policy...")
for _ in range(g_step):
seg = seg_gen.__next__()
#mmd reward
if mmd:
reward_giver.set_b2(seg["ob"], seg["ac"])
seg["rew"] = reward_giver.get_reward(seg["ob"], seg["ac"])
#report stats and save policy if any good
lrlocal = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]
) # local values
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*listoflrpairs))
true_rewbuffer.extend(true_rets)
lenbuffer.extend(lens)
rewbuffer.extend(rews)
true_rew_avg = np.mean(true_rewbuffer)
logger.record_tabular("EpLenMean", np.mean(lenbuffer))
logger.record_tabular("EpRewMean", np.mean(rewbuffer))
logger.record_tabular("EpTrueRewMean", true_rew_avg)
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)
logger.record_tabular("Best so far", best)
# Save model
if ckpt_dir is not None and true_rew_avg >= best:
best = true_rew_avg
fname = os.path.join(ckpt_dir, task_name)
os.makedirs(os.path.dirname(fname), exist_ok=True)
pi.save_policy(fname + "_" + str(save_ind))
save_ind = (save_ind + 1) % max_save
#compute gradient towards next policy
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
atarg = (atarg - atarg.mean()) / atarg.std(
) # standardized advantage function estimate
if hasattr(pi, "ob_rms") and dyn_norm:
pi.ob_rms.update(ob) # update running mean/std for policy
args = seg["ob"], seg["ac"], atarg
fvpargs = [arr[::5] for arr in args]
assign_old_eq_new(
) # set old parameter values to new parameter values
*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:
stepdir = cg(fisher_vector_product,
g,
cg_iters=cg_iters,
verbose=False)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
lm = np.sqrt(shs / max_kl)
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:])
if pi.use_popart:
pi.update_popart(tdlamret)
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(seg["ob"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=128):
if hasattr(pi, "ob_rms") and dyn_norm:
pi.ob_rms.update(
mbob) # update running mean/std for policy
vfadam.update(allmean(compute_vflossandgrad(mbob, mbret)),
vf_stepsize)
g_losses = meanlosses
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
if rank == 0:
logger.dump_tabular()
def learn(env_name, make_env, seed, make_policy, *,
max_iters,
horizon,
drho,
delta,
gamma,
multiple_init=None,
sampler=None,
feature_fun=None,
iw_norm='none',
bound_type='max-ess',
max_offline_iters=10,
save_weights=False,
render_after=None,
grid_size_1d=None,
mu_min=None,
mu_max=None,
truncated_mise=True,
delta_t=None,
k=2,
filename=None,
find_optimal_arm=False,
plot_bound=False,
plot_ess_profile=False,
trainable_std=False,
rescale_ep_return=False):
"""
Learns a policy from scratch
make_env: environment maker
make_policy: policy maker
horizon: max episode length
delta: probability of failure
gamma: discount factor
max_iters: total number of learning iteration
"""
# Print options
np.set_printoptions(precision=3)
losses_with_name = []
# Build the environment
env = make_env()
ob_space = env.observation_space
ac_space = env.action_space
env.seed(seed)
# Build the higher level target and behavioral policies
pi = make_policy('pi', ob_space, ac_space)
oldpi = make_policy('oldpi', ob_space, ac_space)
logger.record_tabular('NumTrainableParams', int(pi._n_higher_params))
# Get all pi's learnable parameters
all_var_list = pi.get_trainable_variables()
var_list = \
[v for v in all_var_list if v.name.split('/')[1].startswith('higher')]
shapes = [U.intprod(var.get_shape().as_list()) for var in var_list]
d = sum(shapes)
# Get all oldpi's learnable parameters
all_var_list_old = oldpi.get_trainable_variables()
var_list_old = \
[v for v in all_var_list_old
if v.name.split('/')[1].startswith('higher')]
# Get hyperpolicy's logstd
higher_logstd_list = [pi.get_higher_logstd()]
# My Placeholders
actor_params_ = tf.placeholder(shape=[max_iters, pi._n_actor_weights],
name='actor_params', dtype=tf.float32)
last_actor_param_ = tf.placeholder(shape=(pi._n_actor_weights),
name='last_actor_params',
dtype=tf.float32)
den_mise_log_ = tf.placeholder(shape=[max_iters], dtype=tf.float32,
name='den_mise')
renyi_bound_ = tf.placeholder(dtype=tf.float32, name='renyi_bound')
ret_ = tf.placeholder(dtype=tf.float32, shape=(max_iters), name='ret')
disc_ret_ = tf.placeholder(dtype=tf.float32, shape=(max_iters),
name='disc_ret')
n_ = tf.placeholder(dtype=tf.float32, name='iter_number')
n_int = tf.cast(n_, dtype=tf.int32)
mask_iters_ = tf.placeholder(dtype=tf.float32, shape=(max_iters),
name='mask_iters')
# grad_ = tf.placeholder(dtype=tf,.float32,
# shape=(d, 1), name='grad')
# Multiple importance weights (with balance heuristic)
target_log_pdf = tf.reduce_sum(
pi.pd.independent_logps(actor_params_), axis=1)
behavioral_log_pdf = tf.reduce_sum(
oldpi.pd.independent_logps(actor_params_), axis=1)
behavioral_log_pdf_last_sample = tf.reduce_sum(
oldpi.pd.independent_logps(last_actor_param_))
log_ratio = target_log_pdf - den_mise_log_
miw = tf.exp(log_ratio) * mask_iters_
den_mise_log_mean = tf.reduce_sum(den_mise_log_) / n_
den_mise_log_last = den_mise_log_[n_int-1]
losses_with_name.extend([(den_mise_log_mean, 'DenMISEMeanLog'),
(den_mise_log_[0], 'DenMISELogFirst'),
(den_mise_log_last, 'DenMISELogLast'),
(miw[0], 'IWFirstEpisode'),
(miw[n_int-1], 'IWLastEpisode'),
(tf.reduce_sum(miw)/n_, 'IWMean'),
(tf.reduce_max(miw), 'IWMax'),
(tf.reduce_min(miw), 'IWMin')])
# Return
ep_return = disc_ret_
return_mean = tf.reduce_sum(ep_return) / n_
return_last = ep_return[n_int - 1]
return_max = tf.reduce_max(ep_return[:n_int])
return_min = tf.reduce_min(ep_return[:n_int])
return_abs_max = tf.reduce_max(tf.abs(ep_return[:n_int]))
regret = n_ * 5 - tf.reduce_sum(ep_return)
regret_over_t = 5 - return_mean
losses_with_name.extend([(return_mean, 'ReturnMean'),
(return_max, 'ReturnMax'),
(return_min, 'ReturnMin'),
(return_last, 'ReturnLastEpisode'),
(return_abs_max, 'ReturnAbsMax'),
(regret, 'Regret'),
(regret_over_t, 'Regret/t')])
# Regret
# Exponentiated Renyi divergence between the target and one behavioral
renyi_component = pi.pd.renyi(oldpi.pd)
renyi_component = tf.cond(tf.is_nan(renyi_component),
lambda: tf.constant(np.inf),
lambda: renyi_component)
renyi_component = tf.cond(renyi_component < 0.,
lambda: tf.constant(np.inf),
lambda: renyi_component)
renyi_component = tf.exp(renyi_component)
if truncated_mise:
# Bound to d2(target || mixture of behaviorals)/n
mn = tf.sqrt((n_**2 * renyi_bound_) / tf.log(1 / delta))
mn_broadcasted = \
tf.ones(shape=miw.get_shape().as_list(), dtype=np.float32) * mn
min = tf.where(tf.less(miw, mn_broadcasted), miw, mn_broadcasted)
mise = tf.reduce_sum(min * ep_return * mask_iters_)/n_
else:
# MISE
mise = tf.reduce_sum(miw * ep_return * mask_iters_)/n_
losses_with_name.append((mise, 'MISE'))
# Bounds
if delta_t == 'continuous':
tau = tf.ceil(n_**(1 / k))
delta_cst = delta
delta = 6 * delta / ((np.pi * n_)**2 * (1 + tau**d))
elif delta_t == 'discrete':
delta_cst = delta
delta = 3 * delta / ((np.pi * n_)**2 * grid_size_1d)
elif delta_t is None:
grid_size_1d = 100 # ToDo correggiiiiiiiiiiiiiiii
delta_cst = delta
delta = tf.constant(delta)
else:
raise NotImplementedError
losses_with_name.append((delta, 'Delta'))
if bound_type == 'J':
bound = mise
elif bound_type == 'max-renyi':
if truncated_mise:
const = return_abs_max * (np.sqrt(2) + 1 / 3) \
* tf.sqrt(tf.log(1 / delta))
exploration_bonus = const * tf.sqrt(renyi_bound_)
bound = mise + exploration_bonus
else:
const = return_abs_max * tf.sqrt(1 / delta - 1)
exploration_bonus = const * tf.sqrt(renyi_bound_)
bound = mise + exploration_bonus
else:
raise NotImplementedError
losses_with_name.append((mise, 'BoundMISE'))
losses_with_name.append((exploration_bonus, 'BoundBonus'))
losses_with_name.append((bound, 'Bound'))
# ESS estimation by d2
ess_d2 = n_ / renyi_bound_
# ESS estimation by miw norms
eps = 1e-18 # for eps<1e-18 miw_2=0 if weights are zero
miw_ess = (tf.exp(log_ratio) + eps) * mask_iters_
miw_1 = tf.linalg.norm(miw_ess, ord=1)
miw_2 = tf.linalg.norm(miw_ess, ord=2)
ess_miw = miw_1**2 / miw_2**2
# Infos
losses, loss_names = map(list, zip(*losses_with_name))
# TF functions
set_parameters = U.SetFromFlat(var_list)
get_parameters = U.GetFlat(var_list)
set_parameters_old = U.SetFromFlat(var_list_old)
# set_higher_logstd = U.SetFromFlat(higher_logstd_list)
# set_higher_logstd(np.log([0.15, 0.2]))
compute_behav = U.function(
[actor_params_], behavioral_log_pdf)
compute_behav_last_sample = U.function(
[last_actor_param_], behavioral_log_pdf_last_sample)
compute_renyi = U.function(
[], renyi_component)
compute_bound = U.function(
[actor_params_, disc_ret_, ret_, n_,
mask_iters_, den_mise_log_, renyi_bound_], bound)
compute_grad = U.function(
[actor_params_, disc_ret_, ret_, n_,
mask_iters_, den_mise_log_, renyi_bound_],
U.flatgrad(bound, var_list))
compute_return_mean = U.function(
[actor_params_, disc_ret_, ret_, n_,
mask_iters_], return_mean)
compute_losses = U.function(
[actor_params_, disc_ret_, ret_, n_,
mask_iters_, den_mise_log_, renyi_bound_], losses)
compute_roba = U.function(
[actor_params_, disc_ret_, ret_, n_,
mask_iters_, den_mise_log_, renyi_bound_],
[mise, exploration_bonus, ess_d2, ess_miw])
# Tf initialization
U.initialize()
# Store behaviorals' params and their trajectories
old_rhos_list = []
all_eps = {}
all_eps['actor_params'] = np.zeros(shape=[max_iters, pi._n_actor_weights])
all_eps['disc_ret'] = np.zeros(max_iters)
all_eps['ret'] = np.zeros(max_iters)
mask_iters = np.zeros(max_iters)
# Set learning loop variables
den_mise = np.zeros(mask_iters.shape).astype(np.float32)
if delta_t == 'continuous':
renyi_components_sum = None
else:
renyi_components_sum = np.zeros(grid_size_1d**d)
new_grid = True
grid_size_1d_old = 0
iters_so_far = 0
lens = []
tstart = time.time()
# Sample actor's params before entering the learning loop
rho = get_parameters()
theta = pi.resample()
all_eps['actor_params'][iters_so_far, :] = theta
# Establish grid dimension if needed
grid_dimension = ob_space.shape[0]
# Learning loop ###########################################################
while True:
iters_so_far += 1
mask_iters[:iters_so_far] = 1
# Render one episode
if render_after is not None and iters_so_far % render_after == 0:
if hasattr(env, 'render'):
render(env, pi, horizon)
# Exit loop in the end
if iters_so_far - 1 >= max_iters:
print('Finished...')
break
# Learning iteration
logger.log('********** Iteration %i ************' % iters_so_far)
# Generate one trajectory
with timed('sampling'):
# Sample a trajectory with the newly parametrized actor
ret, disc_ret, ep_len = eval_trajectory(
env, pi, gamma, horizon, feature_fun, rescale_ep_return)
all_eps['ret'][iters_so_far-1] = ret
all_eps['disc_ret'][iters_so_far-1] = disc_ret
lens.append(ep_len)
# Store the parameters of the behavioral hyperpolicy
old_rhos_list.append(rho)
with timed('summaries before'):
logger.record_tabular("Iteration", iters_so_far)
logger.record_tabular("NumTrajectories", iters_so_far)
logger.record_tabular("TimestepsSoFar", np.sum(lens))
logger.record_tabular('AvgEpLen', np.mean(lens))
logger.record_tabular('MinEpLen', np.min(lens))
logger.record_tabular("TimeElapsed", time.time() - tstart)
# Save policy parameters to disk
if save_weights:
logger.record_tabular('Weights', str(get_parameters()))
import pickle
file = open('checkpoint.pkl', 'wb')
pickle.dump(rho, file)
# Tensor evaluations
def evaluate_behav():
return compute_behav(all_eps['actor_params'])
def evaluate_behav_last_sample():
args_behav_last = [all_eps['actor_params'][iters_so_far - 1]]
return compute_behav_last_sample(*args_behav_last)
def evaluate_renyi_component():
return compute_renyi()
args = all_eps['actor_params'], all_eps['disc_ret'], \
all_eps['ret'], iters_so_far, mask_iters
def evaluate_bound(den_mise_log, renyi_bound):
args_bound = args + (den_mise_log, renyi_bound, )
return compute_bound(*args_bound)
def evaluate_grad(den_mise_log, renyi_bound):
args_grad = args + (den_mise_log, renyi_bound, )
return compute_grad(*args_grad)
def evaluate_roba(den_mise_log, renyi_bound):
args_roba = args + (den_mise_log, renyi_bound, )
return compute_roba(*args_roba)
if bound_type == 'J':
evaluate_renyi = None
elif bound_type == 'max-renyi':
evaluate_renyi = evaluate_renyi_component
else:
raise NotImplementedError
with timed("Optimization"):
if find_optimal_arm:
pass
elif multiple_init:
bound = 0
improvement = 0
check = False
for i in range(multiple_init):
rho_init = [np.arctanh(np.random.uniform(
pi.min_mean, pi.max_mean))]
rho_i, improvement_i, den_mise_log_i, bound_i = \
optimize_offline(evaluate_roba, pi,
rho_init, drho,
old_rhos_list,
iters_so_far,
mask_iters, set_parameters,
set_parameters_old,
evaluate_behav, evaluate_renyi,
evaluate_bound,
evaluate_grad,
max_offline_ite=max_offline_iters)
if bound_i > bound:
check = True
rho = rho_i
improvement = improvement_i
den_mise_log = den_mise_log_i
if not check:
den_mise_log = den_mise_log_i
else:
if delta_t == 'continuous':
grid_size_1d = int(np.ceil(iters_so_far**(1 / k)))
if grid_size_1d > grid_size_1d_old:
new_grid = True
renyi_components_sum = np.zeros(grid_size_1d**d)
grid_size_1d_old = grid_size_1d
rho, improvement, den_mise_log, den_mise, \
renyi_components_sum, renyi_bound = \
best_of_grid(pi, grid_size_1d, mu_min, mu_max,
grid_dimension, trainable_std,
rho, old_rhos_list,
iters_so_far, mask_iters,
set_parameters, set_parameters_old,
delta_cst, renyi_components_sum,
evaluate_behav, den_mise,
evaluate_behav_last_sample,
evaluate_bound, evaluate_renyi, evaluate_roba,
filename, plot_bound,
plot_ess_profile, delta_t, new_grid)
new_grid = False
set_parameters(rho)
with timed('summaries after'):
# Sample actor's parameters from hyperpolicy and assign to actor
if iters_so_far < max_iters:
theta = pi.resample()
all_eps['actor_params'][iters_so_far, :] = theta
if env.spec is not None:
if env.spec.id == 'LQG1D-v0':
mu1_actor = pi.eval_actor_mean([[1]])[0][0]
mu1_higher = pi.eval_higher_mean()[0]
sigma = pi.eval_higher_std()[0]
logger.record_tabular("LQGmu1_actor", mu1_actor)
logger.record_tabular("LQGmu1_higher", mu1_higher)
logger.record_tabular("LQGsigma_higher", sigma)
elif env.spec.id == 'MountainCarContinuous-v0':
ac1 = pi.eval_actor_mean([[1, 1]])[0][0]
mu1_higher = pi.eval_higher_mean()
sigma = pi.eval_higher_std()
logger.record_tabular("ActionIn1", ac1)
logger.record_tabular("MountainCar_mu0_higher", mu1_higher[0])
logger.record_tabular("MountainCar_mu1_higher", mu1_higher[1])
logger.record_tabular("MountainCar_std0_higher", sigma[0])
logger.record_tabular("MountainCar_std1_higher", sigma[1])
elif env.id is not None:
if env.id == 'inverted_pendulum':
ac1 = pi.eval_actor_mean([[1, 1, 1, 1]])[0][0]
mu1_higher = pi.eval_higher_mean()
sigma = pi.eval_higher_std()
logger.record_tabular("ActionIn1", ac1)
logger.record_tabular("InvPendulum_mu0_higher", mu1_higher[0])
logger.record_tabular("InvPendulum_mu1_higher", mu1_higher[1])
logger.record_tabular("InvPendulum_mu2_higher", mu1_higher[2])
logger.record_tabular("InvPendulum_mu3_higher", mu1_higher[3])
logger.record_tabular("InvPendulum_std0_higher", sigma[0])
logger.record_tabular("InvPendulum_std1_higher", sigma[1])
logger.record_tabular("InvPendulum_std2_higher", sigma[2])
logger.record_tabular("InvPendulum_std3_higher", sigma[3])
if find_optimal_arm:
ret_mean = compute_return_mean(*args)
logger.record_tabular('ReturnMean', ret_mean)
else:
args_losses = args + (den_mise_log, renyi_bound, )
meanlosses = np.array(compute_losses(*args_losses))
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
# Print all info in a table
logger.dump_tabular()
# Close environment in the end
env.close()
def run_hoof_no_lamgam(
network,
env,
total_timesteps,
timesteps_per_batch, # what to train on
kl_range,
gamma_range,
lam_range, # advantage estimation
num_kl,
num_gamma_lam,
cg_iters=10,
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
'''
MPI = None
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
# +2 for gamma, lambda
ob = tf.placeholder(shape=(None, env.observation_space.shape[0] + 2),
dtype=env.observation_space.dtype,
name='Ob')
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 = 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_ratio = U.function(
[ob, ac, atarg], ratio) # IS ratio - used for computing IS weights
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_with_gl(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'
kl_range = np.atleast_1d(kl_range)
gamma_range = np.atleast_1d(gamma_range)
lam_range = np.atleast_1d(lam_range)
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__()
thbefore = get_flat()
rand_gamma = gamma_range[0] + (
gamma_range[-1] - gamma_range[0]) * np.random.rand(num_gamma_lam)
rand_lam = lam_range[0] + (
lam_range[-1] - lam_range[0]) * np.random.rand(num_gamma_lam)
rand_kl = kl_range[0] + (kl_range[-1] -
kl_range[0]) * np.random.rand(num_kl)
opt_polval = -10**8
est_polval = np.zeros((num_gamma_lam, num_kl))
ob_lam_gam = []
tdlamret = []
vpred = []
for gl in range(num_gamma_lam):
oblg, vpredbefore, atarg, tdlr = add_vtarg_and_adv_without_gl(
pi, seg, rand_gamma[gl], rand_lam[gl])
ob_lam_gam += [oblg]
tdlamret += [tdlr]
vpred += [vpredbefore]
atarg = (atarg - atarg.mean()) / atarg.std(
) # standardized advantage function estimate
pol_ob = np.concatenate(
(seg['ob'], np.zeros(seg['ob'].shape[:-1] + (2, ))), axis=-1)
args = pol_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=False)
assert np.isfinite(stepdir).all()
shs = .5 * stepdir.dot(fisher_vector_product(stepdir))
surrbefore = lossbefore[0]
for m, kl in enumerate(rand_kl):
lm = np.sqrt(shs / kl)
fullstep = stepdir / lm
thnew = thbefore + fullstep
set_from_flat(thnew)
# compute the IS estimates
lik_ratio = compute_ratio(*args)
est_polval[gl, m] = wis_estimate(seg, lik_ratio)
# update best policy found so far
if est_polval[gl, m] > opt_polval:
opt_polval = est_polval[gl, m]
opt_th = thnew
opt_kl = kl
opt_gamma = rand_gamma[gl]
opt_lam = rand_lam[gl]
opt_vpredbefore = vpredbefore
opt_tdlr = tdlr
meanlosses = surr, kl, *_ = allmean(
np.array(compute_losses(*args)))
improve = surr - surrbefore
expectedimprove = g.dot(fullstep)
set_from_flat(thbefore)
logger.log("Expected: %.3f Actual: %.3f" % (expectedimprove, improve))
set_from_flat(opt_th)
for (lossname, lossval) in zip(loss_names, meanlosses):
logger.record_tabular(lossname, lossval)
ob_lam_gam = np.concatenate(ob_lam_gam, axis=0)
tdlamret = np.concatenate(tdlamret, axis=0)
vpred = np.concatenate(vpred, axis=0)
with timed("vf"):
for _ in range(vf_iters):
for (mbob, mbret) in dataset.iterbatches(
(ob_lam_gam, tdlamret),
include_final_partial_batch=False,
batch_size=num_gamma_lam * 64):
g = allmean(compute_vflossandgrad(mbob, mbret))
vfadam.update(g, vf_stepsize)
logger.record_tabular("ev_tdlam_before",
explained_variance(vpred, 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)
logger.record_tabular("Opt_KL", opt_kl)
logger.record_tabular("gamma", opt_gamma)
logger.record_tabular("lam", opt_lam)
if rank == 0:
logger.dump_tabular()
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 learn(env,
policy,
vf,
gamma,
lam,
timesteps_per_batch,
num_timesteps,
animate=False,
callback=None,
desired_kl=0.002,
lr=0.03,
momentum=0.9):
ob_dim, ac_dim = policy.ob_dim, policy.ac_dim
dbpi = GaussianMlpPolicy(ob_dim, ac_dim, 'dbp')
oldpi = GaussianMlpPolicy(ob_dim, ac_dim, 'oe')
dboldpi = GaussianMlpPolicy(ob_dim, ac_dim, 'doi')
# with tf.variable_scope('dbp'):
# with tf.variable_scope('oe'):
# with tf.variable_scope('doi'):
pi = policy
do_std = U.function([], [pi.std_1a, pi.logstd_1a])
kloldnew = oldpi.pd.kl(pi.pd)
dbkloldnew = dboldpi.pd.kl(dbpi.pd)
dist = meankl = tf.reduce_mean(kloldnew)
dbkl = tf.reduce_mean(dbkloldnew)
obfilter = ZFilter(env.observation_space.shape)
max_pathlength = env.spec.timestep_limit
stepsize = tf.Variable(initial_value=np.float32(np.array(lr)),
name='stepsize')
inputs, loss, loss_sampled = policy.update_info
var_list = [v for v in tf.global_variables() if "pi" in v.name]
db_var_list = [v for v in tf.global_variables() if "dbp" in v.name]
old_var_list = [v for v in tf.global_variables() if "oe" in v.name]
db_old_var_list = [v for v in tf.global_variables() if "doi" in v.name]
print(len(var_list), len(db_var_list), len(old_var_list),
len(db_old_var_list))
assign_old_eq_new = U.function(
[], [],
updates=[
tf.assign(oldv, newv)
for (oldv, newv) in zipsame(old_var_list, var_list)
])
assign_db = U.function(
[], [],
updates=[
tf.assign(db, o) for (db, o) in zipsame(db_var_list, var_list)
] + [
tf.assign(dbold, dbnew)
for (dbold, dbnew) in zipsame(db_old_var_list, old_var_list)
])
assign_old_eq_newr = U.function(
[], [],
updates=[
tf.assign(newv, oldv)
for (oldv, newv) in zipsame(old_var_list, var_list)
])
# assign_dbr = U.function([], [], updates=
# [tf.assign(o, db) for (db, o) in zipsame(db_var_list, var_list)] +
# [tf.assign(dbnew, dbold) for (dbold, dbnew) in zipsame(db_old_var_list, old_var_list)])
klgrads = tf.gradients(dist, var_list)
dbklgrads = tf.gradients(dbkl, db_var_list)
p_grads = [tf.ones_like(v) for v in dbklgrads]
get_flat = U.GetFlat(var_list)
get_old_flat = U.GetFlat(old_var_list)
set_from_flat = U.SetFromFlat(var_list)
flat_tangent2 = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan2")
shapes = [var.get_shape().as_list() for var in var_list]
start = 0
tangents2 = []
for shape in shapes:
sz = U.intprod(shape)
tangents2.append(tf.reshape(flat_tangent2[start:start + sz], shape))
start += sz
gvp2 = tf.add_n([
tf.reduce_sum(g * tangent2)
for (g, tangent2) in zipsame(dbklgrads, tangents2)
])
gvp2_grads = tf.gradients(gvp2, db_var_list)
neg_term = tf.add_n([
tf.reduce_sum(g * tangent2)
for (g, tangent2) in zipsame(gvp2_grads, tangents2)
]) / 2.
ng1 = tf.gradients(neg_term, db_var_list)
ng2 = tf.gradients(neg_term, db_old_var_list)
neg_term_grads = [
a + b for (a, b) in zip(tf.gradients(neg_term, db_var_list),
tf.gradients(neg_term, db_old_var_list))
]
neg_term = neg_term_grads
# neg_term = tf.concat(axis=0, values=[tf.reshape(v, [U.numel(v)]) for v in neg_term_grads])
pos_term = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(gvp2_grads, p_grads)
])
pos_term_grads = [
a + b for (a, b) in zip(tf.gradients(pos_term, db_var_list),
tf.gradients(pos_term, db_old_var_list))
]
pos_term_sum = tf.add_n([
tf.reduce_sum(g * tangent)
for (g, tangent) in zipsame(pos_term_grads, tangents2)
])
pos_term_grads = tf.gradients(pos_term_sum, p_grads)
pos_term = pos_term_grads
# pos_term = tf.concat(axis=0, values=[tf.reshape(v, [U.numel(v)]) for v in pos_term_grads])
geo_term = [(p - n) * 0.5 for p, n in zip(pos_term, neg_term)]
optim = kfac.KfacOptimizer(learning_rate=stepsize, cold_lr=stepsize*(1-0.9), momentum=momentum, kfac_update=2,\
epsilon=1e-2, stats_decay=0.99, async=1, cold_iter=1,
weight_decay_dict=policy.wd_dict, max_grad_norm=None)
pi_var_list = []
for var in tf.trainable_variables():
if "pi" in var.name:
pi_var_list.append(var)
grads = optim.compute_gradients(loss, var_list=pi_var_list)
update_op, q_runner = optim.minimize(loss,
loss_sampled,
var_list=pi_var_list)
geo_term = [g1 + g2[0] for g1, g2 in zip(geo_term, grads)]
geo_grads = list(zip(geo_term, var_list))
update_geo_op, q_runner_geo = optim.apply_gradients(geo_grads)
do_update = U.function(inputs, update_op)
inputs_tangent = list(inputs) + [flat_tangent2]
do_update_geo = U.function(inputs_tangent, update_geo_op)
do_get_geo_term = U.function(inputs_tangent, [ng1, ng2])
U.initialize()
# start queue runners
enqueue_threads = []
coord = tf.train.Coordinator()
for qr in [q_runner, vf.q_runner, q_runner_geo]:
assert (qr != None)
enqueue_threads.extend(
qr.create_threads(tf.get_default_session(),
coord=coord,
start=True))
i = 0
timesteps_so_far = 0
while True:
if timesteps_so_far > num_timesteps:
break
logger.log("********** Iteration %i ************" % i)
# Collect paths until we have enough timesteps
timesteps_this_batch = 0
paths = []
while True:
path = rollout(env,
policy,
max_pathlength,
animate=(len(paths) == 0 and (i % 10 == 0)
and animate),
obfilter=obfilter)
paths.append(path)
n = pathlength(path)
timesteps_this_batch += n
timesteps_so_far += n
if timesteps_this_batch > timesteps_per_batch:
break
# Estimate advantage function
vtargs = []
advs = []
for path in paths:
rew_t = path["reward"]
return_t = common.discount(rew_t, gamma)
vtargs.append(return_t)
vpred_t = vf.predict(path)
vpred_t = np.append(vpred_t,
0.0 if path["terminated"] else vpred_t[-1])
delta_t = rew_t + gamma * vpred_t[1:] - vpred_t[:-1]
adv_t = common.discount(delta_t, gamma * lam)
advs.append(adv_t)
# Update value function
vf.fit(paths, vtargs)
# Build arrays for policy update
ob_no = np.concatenate([path["observation"] for path in paths])
action_na = np.concatenate([path["action"] for path in paths])
oldac_dist = np.concatenate([path["action_dist"] for path in paths])
adv_n = np.concatenate(advs)
standardized_adv_n = (adv_n - adv_n.mean()) / (adv_n.std() + 1e-8)
assign_old_eq_new() # set old parameter values to new parameter values
assign_db()
# Policy update
do_update(ob_no, action_na, standardized_adv_n)
# ft2 = get_flat() - get_old_flat()
# assign_old_eq_newr() # assign back
# gnp = do_get_geo_term(ob_no, action_na, standardized_adv_n, ft2)
# def check_nan(bs):
# return [~np.isnan(b).all() for b in bs]
# print(gnp[0])
# print('.....asdfasdfadslfkadsjfaksdfalsdkfjaldskf')
# print(gnp[1])
# do_update_geo(ob_no, action_na, standardized_adv_n, ft2)
min_stepsize = np.float32(1e-8)
max_stepsize = np.float32(1e0)
# Adjust stepsize
kl = policy.compute_kl(ob_no, oldac_dist)
# if kl > desired_kl * 2:
# logger.log("kl too high")
# tf.assign(stepsize, tf.maximum(min_stepsize, stepsize / 1.5)).eval()
# elif kl < desired_kl / 2:
# logger.log("kl too low")
# tf.assign(stepsize, tf.minimum(max_stepsize, stepsize * 1.5)).eval()
# else:
# logger.log("kl just right!")
logger.record_tabular(
"EpRewMean", np.mean([path["reward"].sum() for path in paths]))
logger.record_tabular(
"EpRewSEM",
np.std([
path["reward"].sum() / np.sqrt(len(paths)) for path in paths
]))
logger.record_tabular("EpLenMean",
np.mean([pathlength(path) for path in paths]))
logger.record_tabular("KL", kl)
print(do_std())
if callback:
callback()
logger.dump_tabular()
i += 1
coord.request_stop()
coord.join(enqueue_threads)
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)))
def _build_trpo(self):
pi = self._policy
oldpi = self._old_policy
other_pis = self._pis
# input placeholders
ob = pi.ob
ac = pi.pdtype.sample_placeholder([None], name='action')
atarg = tf.placeholder(
dtype=tf.float32, shape=[None],
name='advantage') # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None],
name='return') # Empirical return
# policy
all_var_list = pi.get_trainable_variables()
pol_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")
]
self._vf_adam = MpiAdam(vf_var_list)
kl_oldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
mean_kl = tf.reduce_mean(kl_oldnew)
mean_ent = tf.reduce_mean(ent)
pol_entpen = -self._config.ent_coeff * mean_ent
vf_loss = tf.reduce_mean(tf.square(pi.vpred - ret))
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac))
pol_surr = tf.reduce_mean(ratio * atarg)
# divergence
other_obs = [] # put id-th data
for other_pi in other_pis:
other_obs.extend(other_pi.obs[self.id])
my_obs_for_other = flatten_lists(pi.obs) # put i-th data
other_obs_for_other = [] # put i-th data
for i, other_pi in enumerate(other_pis):
other_obs_for_other.extend(other_pi.obs[i])
pairwise_divergence = [tf.constant(0.)]
for i, other_pi in enumerate(other_pis):
if i != self.id:
pairwise_divergence.append(
tf.reduce_mean(pi.pds[self.id].kl(other_pi.pds[self.id])))
pairwise_divergence.append(
tf.reduce_mean(other_pi.pds[i].kl(pi.pds[i])))
pol_divergence = self._config.divergence_coeff * tf.reduce_mean(
pairwise_divergence)
pol_loss = pol_surr + pol_entpen + pol_divergence
pol_losses = {
'pol_loss': pol_loss,
'pol_surr': pol_surr,
'pol_entpen': pol_entpen,
'pol_divergence': pol_divergence,
'kl': mean_kl,
'entropy': mean_ent
}
self.summary_name += ['vf_loss']
self.summary_name += pol_losses.keys()
self._get_flat = U.GetFlat(pol_var_list)
self._set_from_flat = U.SetFromFlat(pol_var_list)
klgrads = tf.gradients(mean_kl, pol_var_list)
flat_tangent = tf.placeholder(dtype=tf.float32,
shape=[None],
name="flat_tan")
shapes = [var.get_shape().as_list() for var in pol_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, pol_var_list)
self._update_oldpi = U.function(
[], [],
updates=[
tf.assign(oldv, newv) for (
oldv,
newv) in zipsame(oldpi.get_variables(), pi.get_variables())
])
obs_pairwise = other_obs + my_obs_for_other + other_obs_for_other + ob
self._compute_losses = U.function(obs_pairwise + [ac, atarg],
pol_losses)
pol_losses = dict(pol_losses)
pol_losses.update({'g': U.flatgrad(pol_loss, pol_var_list)})
self._compute_lossandgrad = U.function(obs_pairwise + [ac, atarg],
pol_losses)
self._compute_fvp = U.function([flat_tangent] + obs_pairwise +
[ac, atarg], fvp)
self._compute_vflossandgrad = U.function(
ob + [ret], U.flatgrad(vf_loss, vf_var_list))
self._compute_vfloss = U.function(ob + [ret], vf_loss)
# initialize and sync
U.initialize()
th_init = self._get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
self._set_from_flat(th_init)
self._vf_adam.sync()
rank = MPI.COMM_WORLD.Get_rank()
if self._config.debug:
logger.log(
"[worker: {} local net: {}] Init pol param sum: {}".format(
rank, self.id, th_init.sum()))
logger.log(
"[worker: {} local net: {}] Init vf param sum: {}".format(
rank, self.id,
self._vf_adam.getflat().sum()))