def learn(*,
network,
env,
total_timesteps,
starting_positions,
env_name,
win_percentage=0.5,
eval_env=None,
seed=None,
nsteps=2048,
ent_coef=0.0,
lr=3e-4,
vf_coef=0.5,
max_grad_norm=0.5,
gamma=0.99,
lam=0.95,
log_interval=10,
nminibatches=4,
noptepochs=4,
cliprange=0.2,
save_interval=0,
load_path=None,
model_fn=None,
**network_kwargs):
'''
Learn policy using PPO algorithm (https://arxiv.org/abs/1707.06347)
Parameters:
----------
network: policy network architecture. Either string (mlp, lstm, lnlstm, cnn_lstm, cnn, cnn_small, conv_only - see baselines.common/models.py for full list)
specifying the standard network architecture, or a function that takes tensorflow tensor as input and returns
tuple (output_tensor, extra_feed) where output tensor is the last network layer output, extra_feed is None for feed-forward
neural nets, and extra_feed is a dictionary describing how to feed state into the network for recurrent neural nets.
See common/models.py/lstm for more details on using recurrent nets in policies
env: baselines.common.vec_env.VecEnv environment. Needs to be vectorized for parallel environment simulation.
The environments produced by gym.make can be wrapped using baselines.common.vec_env.DummyVecEnv class.
nsteps: int number of steps of the vectorized environment per update (i.e. batch size is nsteps * nenv where
nenv is number of environment copies simulated in parallel)
total_timesteps: int number of timesteps (i.e. number of actions taken in the environment)
ent_coef: float policy entropy coefficient in the optimization objective
lr: float or function learning rate, constant or a schedule function [0,1] -> R+ where 1 is beginning of the
training and 0 is the end of the training.
vf_coef: float value function loss coefficient in the optimization objective
max_grad_norm: float or None gradient norm clipping coefficient
gamma: float discounting factor
lam: float advantage estimation discounting factor (lambda in the paper)
log_interval: int number of timesteps between logging events
nminibatches: int number of training minibatches per update. For recurrent policies,
should be smaller or equal than number of environments run in parallel.
noptepochs: int number of training epochs per update
cliprange: float or function clipping range, constant or schedule function [0,1] -> R+ where 1 is beginning of the training
and 0 is the end of the training
save_interval: int number of timesteps between saving events
load_path: str path to load the model from
**network_kwargs: keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
For instance, 'mlp' network architecture has arguments num_hidden and num_layers.
'''
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)
total_timesteps = int(total_timesteps)
policy = build_policy(env, network, **network_kwargs)
# Get the nb of env
nenvs = env.num_envs
# Get state_space and action_space
ob_space = env.observation_space
ac_space = env.action_space
# Calculate the batch_size
nbatch = nenvs * nsteps
nbatch_train = nbatch // nminibatches
# Instantiate the model object (that creates act_model and train_model)
if model_fn is None:
from ppo2.model import Model
model_fn = Model
model = model_fn(policy=policy,
ob_space=ob_space,
ac_space=ac_space,
nbatch_act=nenvs,
nbatch_train=nbatch_train,
nsteps=nsteps,
ent_coef=ent_coef,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm)
if load_path is not None:
model.load(load_path)
current_starting_position = starting_positions.pop()
# Instantiate the runner object
runner = Runner(env=env,
model=model,
nsteps=nsteps,
gamma=gamma,
lam=lam,
starting_position=current_starting_position)
if eval_env is not None:
eval_runner = Runner(env=eval_env,
model=model,
nsteps=nsteps,
gamma=gamma,
lam=lam,
starting_position=current_starting_position)
epinfobuf = deque(maxlen=100)
if eval_env is not None:
eval_epinfobuf = deque(maxlen=100)
# Start total timer
tfirststart = time.time()
start_changes = []
reached_goal = []
nupdates = total_timesteps // nbatch
for update in range(1, nupdates + 1):
assert nbatch % nminibatches == 0
# Start timer
tstart = time.time()
frac = 1.0 - (update - 1.0) / nupdates
# Calculate the learning rate
lrnow = lr(frac)
# Calculate the cliprange
cliprangenow = cliprange(frac)
# Get minibatch
obs, returns, masks, actions, values, neglogpacs, states, epinfos = runner.run(
) #pylint: disable=E0632
if eval_env is not None:
eval_obs, eval_returns, eval_masks, eval_actions, eval_values, eval_neglogpacs, eval_states, eval_epinfos = eval_runner.run(
) #pylint: disable=E0632
if env_name == "MountainCar-v0":
done_obs = obs[masks]
# Number of episodes past
n_eps = done_obs.shape[0]
# Reached goal if pos is > 0.5
n_goal_reached = (done_obs[:, 0] >= 0.5).sum()
reached_goal.extend([
done + update * nsteps - nsteps
for done in np.where(done_obs[:, 0] >= 0.5)[0]
])
if (n_goal_reached /
n_eps) > win_percentage and len(starting_positions) > 0:
start_changes.append(update * nsteps)
current_starting_position = starting_positions.pop()
runner.env.starting_position = current_starting_position
if eval_env is not None:
eval_runner.env.starting_position = current_starting_position
epinfobuf.extend(epinfos)
if eval_env is not None:
eval_epinfobuf.extend(eval_epinfos)
# Here what we're going to do is for each minibatch calculate the loss and append it.
mblossvals = []
if states is None: # nonrecurrent version
# Index of each element of batch_size
# Create the indices array
inds = np.arange(nbatch)
for _ in range(noptepochs):
# Randomize the indexes
np.random.shuffle(inds)
# 0 to batch_size with batch_train_size step
for start in range(0, nbatch, nbatch_train):
end = start + nbatch_train
mbinds = inds[start:end]
slices = (arr[mbinds]
for arr in (obs, returns, masks, actions, values,
neglogpacs))
mblossvals.append(model.train(lrnow, cliprangenow,
*slices))
else: # recurrent version
assert nenvs % nminibatches == 0
envsperbatch = nenvs // nminibatches
envinds = np.arange(nenvs)
flatinds = np.arange(nenvs * nsteps).reshape(nenvs, nsteps)
envsperbatch = nbatch_train // nsteps
for _ in range(noptepochs):
np.random.shuffle(envinds)
for start in range(0, nenvs, envsperbatch):
end = start + envsperbatch
mbenvinds = envinds[start:end]
mbflatinds = flatinds[mbenvinds].ravel()
slices = (arr[mbflatinds]
for arr in (obs, returns, masks, actions, values,
neglogpacs))
mbstates = states[mbenvinds]
mblossvals.append(
model.train(lrnow, cliprangenow, *slices, mbstates))
# Feedforward --> get losses --> update
lossvals = np.mean(mblossvals, axis=0)
# End timer
tnow = time.time()
# Calculate the fps (frame per second)
fps = int(nbatch / (tnow - tstart))
if update % log_interval == 0 or update == 1:
# Calculates if value function is a good predicator of the returns (ev > 1)
# or if it's just worse than predicting nothing (ev =< 0)
ev = explained_variance(values, returns)
logger.logkv("serial_timesteps", update * nsteps)
logger.logkv("nupdates", update)
logger.logkv("total_timesteps", update * nbatch)
logger.logkv("fps", fps)
logger.logkv("explained_variance", float(ev))
logger.logkv('eprewmean',
safemean([epinfo['r'] for epinfo in epinfobuf]))
logger.logkv('eplenmean',
safemean([epinfo['l'] for epinfo in epinfobuf]))
logger.logkv('start_changes',
"_".join([str(s) for s in start_changes]))
logger.logkv('reached_goal',
"_".join([str(goal) for goal in reached_goal]))
if eval_env is not None:
logger.logkv(
'eval_eprewmean',
safemean([epinfo['r'] for epinfo in eval_epinfobuf]))
logger.logkv(
'eval_eplenmean',
safemean([epinfo['l'] for epinfo in eval_epinfobuf]))
logger.logkv('time_elapsed', tnow - tfirststart)
for (lossval, lossname) in zip(lossvals, model.loss_names):
logger.logkv(lossname, lossval)
if MPI is None or MPI.COMM_WORLD.Get_rank() == 0:
logger.dumpkvs()
if save_interval and (update % save_interval == 0
or update == 1) and logger.get_dir() and (
MPI is None
or MPI.COMM_WORLD.Get_rank() == 0):
checkdir = osp.join(logger.get_dir(), 'checkpoints')
os.makedirs(checkdir, exist_ok=True)
savepath = osp.join(checkdir, '%.5i' % update)
print('Saving to', savepath)
model.save(savepath)
return model
def learn(
env,
policy_func,
discriminator,
expert_dataset,
embedding_z,
pretrained,
pretrained_weight,
*,
g_step,
d_step,
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,
d_stepsize=3e-4,
vf_iters=3,
max_timesteps=0,
max_episodes=0,
max_iters=0, # time constraint
callback=None,
save_per_iter=100,
ckpt_dir=None,
log_dir=None,
load_model_path=None,
task_name=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")
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")
]
d_adam = MpiAdam(discriminator.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(
[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
writer = U.FileWriter(log_dir)
U.initialize()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
d_adam.sync()
vfadam.sync()
print("Init param sum", th_init.sum(), flush=True)
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
discriminator,
embedding=embedding_z,
timesteps_per_batch=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(discriminator.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())
# if provieded model path
if load_model_path is not None:
U.load_state(load_model_path)
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 iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
U.save_state(os.path.join(ckpt_dir, task_name),
counter=iters_so_far)
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__()
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)))
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, discriminator.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 discriminator
if hasattr(discriminator, "obs_rms"):
discriminator.obs_rms.update(
np.concatenate((ob_batch, ob_expert), 0))
*newlosses, g = discriminator.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()
g_loss_stats.add_all_summary(writer, g_losses, iters_so_far)
d_loss_stats.add_all_summary(writer, np.mean(d_losses, axis=0),
iters_so_far)
ep_stats.add_all_summary(writer, [
np.mean(true_rewbuffer),
np.mean(rewbuffer),
np.mean(lenbuffer)
], iters_so_far)
def learn(self,
total_timesteps,
callback=None,
log_interval=1,
tb_log_name="PPO2",
reset_num_timesteps=True,
rho=0.0):
self.rho = rho
values_s0 = []
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
self.cliprange = get_schedule_fn(self.cliprange)
cliprange_vf = get_schedule_fn(self.cliprange_vf)
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
callback = self._init_callback(callback)
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
self._setup_learn()
t_first_start = time.time()
n_updates = total_timesteps // self.n_batch
callback.on_training_start(locals(), globals())
for update in range(1, n_updates + 1):
assert self.n_batch % self.nminibatches == 0, (
"The number of minibatches (`nminibatches`) "
"is not a factor of the total number of samples "
"collected per rollout (`n_batch`), "
"some samples won't be used.")
batch_size = self.n_batch // self.nminibatches
t_start = time.time()
frac = 1.0 - (update - 1.0) / n_updates
lr_now = self.learning_rate(frac)
cliprange_now = self.cliprange(frac)
cliprange_vf_now = cliprange_vf(frac)
callback.on_rollout_start()
# true_reward is the reward without discount
rollout = self.runner.run(callback)
# Unpack
obs, returns, masks, actions, values, neglogpacs, states, ep_infos, true_reward = rollout
_, val_s0, _ = self.predict(observation=state_s0,
deterministic=True)
values_s0.append(val_s0[0])
callback.on_rollout_end()
# Early stopping due to the callback
if not self.runner.continue_training:
break
self.ep_info_buf.extend(ep_infos)
mb_loss_vals = []
if states is None: # nonrecurrent version
update_fac = max(
self.n_batch // self.nminibatches // self.noptepochs,
1)
inds = np.arange(self.n_batch)
for epoch_num in range(self.noptepochs):
np.random.shuffle(inds)
for start in range(0, self.n_batch, batch_size):
timestep = self.num_timesteps // update_fac + (
(epoch_num * self.n_batch + start) //
batch_size)
end = start + batch_size
mbinds = inds[start:end]
slices = (arr[mbinds]
for arr in (obs, returns, masks, actions,
values, neglogpacs))
mb_loss_vals.append(
self._train_step(
lr_now,
cliprange_now,
*slices,
writer=writer,
update=timestep,
cliprange_vf=cliprange_vf_now))
else: # recurrent version
update_fac = max(
self.n_batch // self.nminibatches // self.noptepochs //
self.n_steps, 1)
assert self.n_envs % self.nminibatches == 0
env_indices = np.arange(self.n_envs)
flat_indices = np.arange(self.n_envs *
self.n_steps).reshape(
self.n_envs, self.n_steps)
envs_per_batch = batch_size // self.n_steps
for epoch_num in range(self.noptepochs):
np.random.shuffle(env_indices)
for start in range(0, self.n_envs, envs_per_batch):
timestep = self.num_timesteps // update_fac + (
(epoch_num * self.n_envs + start) //
envs_per_batch)
end = start + envs_per_batch
mb_env_inds = env_indices[start:end]
mb_flat_inds = flat_indices[mb_env_inds].ravel()
slices = (arr[mb_flat_inds]
for arr in (obs, returns, masks, actions,
values, neglogpacs))
mb_states = states[mb_env_inds]
mb_loss_vals.append(
self._train_step(
lr_now,
cliprange_now,
*slices,
update=timestep,
writer=writer,
states=mb_states,
cliprange_vf=cliprange_vf_now))
loss_vals = np.mean(mb_loss_vals, axis=0)
t_now = time.time()
fps = int(self.n_batch / (t_now - t_start))
if writer is not None:
total_episode_reward_logger(
self.episode_reward,
true_reward.reshape((self.n_envs, self.n_steps)),
masks.reshape((self.n_envs, self.n_steps)), writer,
self.num_timesteps)
if self.verbose >= 1 and (update % log_interval == 0
or update == 1):
explained_var = explained_variance(values, returns)
logger.logkv("serial_timesteps", update * self.n_steps)
logger.logkv("n_updates", update)
logger.logkv("total_timesteps", self.num_timesteps)
logger.logkv("fps", fps)
logger.logkv("explained_variance", float(explained_var))
if len(self.ep_info_buf) > 0 and len(
self.ep_info_buf[0]) > 0:
logger.logkv(
'ep_reward_mean',
safe_mean([
ep_info['r'] for ep_info in self.ep_info_buf
]))
logger.logkv(
'ep_len_mean',
safe_mean([
ep_info['l'] for ep_info in self.ep_info_buf
]))
logger.logkv('time_elapsed', t_start - t_first_start)
for (loss_val, loss_name) in zip(loss_vals,
self.loss_names):
logger.logkv(loss_name, loss_val)
logger.dumpkvs()
filename = directory + 'values_s0_iter' + str(iteration) + '.npy'
np.save(filename, values_s0)
callback.on_training_end()
return self
def learn(network,
env,
seed=None,
nsteps=5,
total_timesteps=int(80e6),
vf_coef=0.5,
ent_coef=0.01,
max_grad_norm=0.5,
lr=7e-4,
lrschedule='linear',
epsilon=1e-5,
alpha=0.99,
gamma=0.99,
log_interval=100,
load_path=None,
**network_kwargs):
'''
Main entrypoint for A2C algorithm. Train a policy with given network architecture on a given environment using a2c algorithm.
Parameters:
-----------
network: policy network architecture. Either string (mlp, lstm, lnlstm, cnn_lstm, cnn, cnn_small, conv_only - see baselines.common/models.py for full list)
specifying the standard network architecture, or a function that takes tensorflow tensor as input and returns
tuple (output_tensor, extra_feed) where output tensor is the last network layer output, extra_feed is None for feed-forward
neural nets, and extra_feed is a dictionary describing how to feed state into the network for recurrent neural nets.
See baselines.common/policies.py/lstm for more details on using recurrent nets in policies
env: RL environment. Should implement interface similar to VecEnv (baselines.common/vec_env) or be wrapped with DummyVecEnv (baselines.common/vec_env/dummy_vec_env.py)
seed: seed to make random number sequence in the alorightm reproducible. By default is None which means seed from system noise generator (not reproducible)
nsteps: int, number of steps of the vectorized environment per update (i.e. batch size is nsteps * nenv where
nenv is number of environment copies simulated in parallel)
total_timesteps: int, total number of timesteps to train on (default: 80M)
vf_coef: float, coefficient in front of value function loss in the total loss function (default: 0.5)
ent_coef: float, coeffictiant in front of the policy entropy in the total loss function (default: 0.01)
max_gradient_norm: float, gradient is clipped to have global L2 norm no more than this value (default: 0.5)
lr: float, learning rate for RMSProp (current implementation has RMSProp hardcoded in) (default: 7e-4)
lrschedule: schedule of learning rate. Can be 'linear', 'constant', or a function [0..1] -> [0..1] that takes fraction of the training progress as input and
returns fraction of the learning rate (specified as lr) as output
epsilon: float, RMSProp epsilon (stabilizes square root computation in denominator of RMSProp update) (default: 1e-5)
alpha: float, RMSProp decay parameter (default: 0.99)
gamma: float, reward discounting parameter (default: 0.99)
log_interval: int, specifies how frequently the logs are printed out (default: 100)
**network_kwargs: keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
For instance, 'mlp' network architecture has arguments num_hidden and num_layers.
'''
set_global_seeds(seed)
# Get the nb of env
nenvs = env.num_envs
policy = build_policy(env, network, **network_kwargs)
# Instantiate the model object (that creates step_model and train_model)
model = Model(policy=policy,
env=env,
nsteps=nsteps,
ent_coef=ent_coef,
vf_coef=vf_coef,
max_grad_norm=max_grad_norm,
lr=lr,
alpha=alpha,
epsilon=epsilon,
total_timesteps=total_timesteps,
lrschedule=lrschedule)
if load_path is not None:
model.load(load_path)
# Instantiate the runner object
runner = Runner(env, model, nsteps=nsteps, gamma=gamma)
# Calculate the batch_size
nbatch = nenvs * nsteps
# Start total timer
tstart = time.time()
for update in range(1, total_timesteps // nbatch + 1):
# Get mini batch of experiences
obs, states, rewards, masks, actions, values = runner.run()
policy_loss, value_loss, policy_entropy = model.train(
obs, states, rewards, masks, actions, values)
nseconds = time.time() - tstart
# Calculate the fps (frame per second)
fps = int((update * nbatch) / nseconds)
if update % log_interval == 0 or update == 1:
# Calculates if value function is a good predicator of the returns (ev > 1)
# or if it's just worse than predicting nothing (ev =< 0)
ev = explained_variance(values, rewards)
logger.record_tabular("nupdates", update)
logger.record_tabular("total_timesteps", update * nbatch)
logger.record_tabular("fps", fps)
logger.record_tabular("policy_entropy", float(policy_entropy))
logger.record_tabular("value_loss", float(value_loss))
logger.record_tabular("explained_variance", float(ev))
logger.dump_tabular()
return model
def learn(*, network, env, total_timesteps,n_of_paths=10, eval_env = None, seed=None, nsteps=2048, ent_coef=0.0, lr=3e-4,
vf_coef=0.5,lf_coef=0.5, max_grad_norm=0.5, gamma=0.99, lam=0.95, alpha3 = .005, init_labda = 1.,labda_clip_range = 0.1,
log_interval=10, nminibatches=4, noptepochs=4, cliprange=0.2, use_adaptive_alpha3=True,approximate_value_function=False,
save_interval=0, load_path=None, model_fn=None, **network_kwargs):
'''
Learn policy using PPO algorithm (https://arxiv.org/abs/1707.06347)
Parameters:
----------
network: policy network architecture. Either string (mlp, lstm, lnlstm, cnn_lstm, cnn, cnn_small, conv_only - see baselines.common/models.py for full list)
specifying the standard network architecture, or a function that takes tensorflow tensor as input and returns
tuple (output_tensor, extra_feed) where output tensor is the last network layer output, extra_feed is None for feed-forward
neural nets, and extra_feed is a dictionary describing how to feed state into the network for recurrent neural nets.
See common/models.py/lstm for more details on using recurrent nets in policies
env: baselines.common.vec_env.VecEnv environment. Needs to be vectorized for parallel environment simulation.
The environments produced by gym.make can be wrapped using baselines.common.vec_env.DummyVecEnv class.
nsteps: int number of steps of the vectorized environment per update (i.e. batch size is nsteps * nenv where
nenv is number of environment copies simulated in parallel)
total_timesteps: int number of timesteps (i.e. number of actions taken in the environment)
ent_coef: float policy entropy coefficient in the optimization objective
lr: float or function learning rate, constant or a schedule function [0,1] -> R+ where 1 is beginning of the
training and 0 is the end of the training.
vf_coef: float value function loss coefficient in the optimization objective
max_grad_norm: float or None gradient norm clipping coefficient
gamma: float discounting factor
lam: float advantage estimation discounting factor (lambda in the paper)
log_interval: int number of timesteps between logging events
nminibatches: int number of training minibatches per update. For recurrent policies,
should be smaller or equal than number of environments run in parallel.
noptepochs: int number of training epochs per update
cliprange: float or function clipping range, constant or schedule function [0,1] -> R+ where 1 is beginning of the training
and 0 is the end of the training
save_interval: int number of timesteps between saving events
load_path: str path to load the model from
**network_kwargs: keyword arguments to the policy / network builder. See baselines.common/policies.py/build_policy and arguments to a particular type of network
For instance, 'mlp' network architecture has arguments num_hidden and num_layers.
'''
set_global_seeds(seed)
logger.logkv('lr_a', lr)
if isinstance(lr, float): lr = constfn(lr)
else: assert callable(lr)
if isinstance(cliprange, float): cliprange = constfn(cliprange)
else: assert callable(cliprange)
total_timesteps = int(total_timesteps)
# 构建网络
policy = build_policy(env, network, **network_kwargs)
print("network build success")
# Get the nb of env
nenvs = env.num_envs
# Get state_space and action_space
ob_space = env.observation_space
ob_space_ = env.observation_space
ac_space = env.action_space
# Calculate the batch_size
nbatch = nenvs * nsteps
nbatch_train = nbatch // nminibatches
# Instantiate the model object (that creates act_model and train_model)
if model_fn is None:
from ppo2_lyapunov.model import Model
model_fn = Model
model = model_fn(policy=policy, ob_space=ob_space, ac_space=ac_space, nbatch_act=nenvs, nbatch_train=nbatch_train,
nsteps=nsteps, ent_coef=ent_coef, vf_coef=vf_coef,lf_coef=lf_coef, ALPHA3= alpha3, init_labda=init_labda,
max_grad_norm=max_grad_norm, use_adaptive_alpha3=use_adaptive_alpha3,approximate_value_function=approximate_value_function)
print("model build success")
if load_path is not None:
model.load(load_path)
# Instantiate the runner object
runner = Runner(env=env, model=model, nsteps=nsteps, gamma=gamma, lam=lam)
print("runner build success")
if eval_env is not None:
eval_runner = Runner(env = eval_env, model = model, nsteps = nsteps, gamma = gamma, lam= lam, n_of_paths=n_of_paths)
epinfobuf = deque(maxlen=100)
if eval_env is not None:
eval_epinfobuf = deque(maxlen=100)
logger.logkv('ent_coef', ent_coef)
logger.logkv('vf_coef', vf_coef)
logger.logkv('lf_coef', lf_coef)
logger.logkv('max_grad_norm', max_grad_norm)
logger.logkv('gamma', gamma)
logger.logkv('advantage_lam', 0.95)
logger.logkv('cliprange', cliprange)
# Start total timer
tfirststart = time.time()
global_step = 0
nupdates = total_timesteps//nbatch
ep_L_R_threshold = 0
ALPHA_MIN = 1000
for update in range(1, nupdates+1):
assert nbatch % nminibatches == 0
# Start timer
tstart = time.time()
frac = 1.0 - (update - 1.0) / nupdates
# Calculate the learning rate
lrnow = lr(frac)
# Calculate the cliprange
cliprangenow = cliprange(frac)
# Get minibatch
# 得到一连串的s,r,l_r,是否死了,a,v,l_v,mb_neglogpacs,s_,info
obs,obs_, returns, l_returns,masks, actions, values, l_values,mb_l_rewards,neglogpacs, states, epinfos = runner.run() #pylint: disable=E0632
if eval_env is not None:
eval_epinfos = eval_runner.eval_run() #pylint: disable=E0632
epinfobuf.extend(epinfos)
if eval_env is not None:
eval_epinfobuf.extend(eval_epinfos)
# Here what we're going to do is for each minibatch calculate the loss and append it.
mblossvals = []
if states is None: # nonrecurrent version
# Index of each element of batch_size
# Create the indices array
inds = np.arange(nbatch)
for _ in range(noptepochs):
# Randomize the indexes
np.random.shuffle(inds)
# 0 to batch_size with batch_train_size step
for start in range(0, nbatch, nbatch_train):
end = start + nbatch_train
mbinds = inds[start:end]
slices = (arr[mbinds] for arr in (obs, obs_, returns, l_returns, masks, actions, values, l_values, mb_l_rewards, neglogpacs))
# print(**slices)
mblossvals.append(model.train(lrnow, cliprangenow, *slices))
else: # recurrent version
assert nenvs % nminibatches == 0
envsperbatch = nenvs // nminibatches
envinds = np.arange(nenvs)
flatinds = np.arange(nenvs * nsteps).reshape(nenvs, nsteps)
for _ in range(noptepochs):
np.random.shuffle(envinds)
for start in range(0, nenvs, envsperbatch):
end = start + envsperbatch
mbenvinds = envinds[start:end]
mbflatinds = flatinds[mbenvinds].ravel()
slices = (arr[mbflatinds] for arr in (obs, obs_,returns,l_returns, masks, actions, values,l_values,mb_l_rewards, neglogpacs))
mbstates = states[mbenvinds]
mblossvals.append(model.train(lrnow, cliprangenow, *slices, mbstates))
# Feedforward --> get losses --> update
lossvals = np.mean(mblossvals, axis=0)
# End timer
tnow = time.time()
# Calculate the fps (frame per second)
fps = int(nbatch / (tnow - tstart))
ep_L_R = safemean([epinfo['lr'] for epinfo in epinfobuf])
if model.use_adaptive_alpha3:
if ep_L_R > ep_L_R_threshold:
ep_L_R_threshold = ep_L_R
# model.ALPHA=min(model.ALPHA*1.05,0.1) good with out nothing and ini 10-5
model.alpha3 = min(model.alpha3 * 1.5, labda_clip_range)
# model.alpha3 = min(model.alpha3 * 1.1, labda_clip_range)
# model.ALPHA = min(model.ALPHA * 1.001, 0.1)
model.alpha3 = min(model.alpha3 * 1.01, labda_clip_range)
if update % log_interval == 0 or update == 1:
# Calculates if value function is a good predicator of the returns (ev > 1)
# or if it's just worse than predicting nothing (ev =< 0)
ev_v = explained_variance(values, returns)
ev_l = explained_variance(l_values, l_returns)
logger.logkv("serial_timesteps", update*nsteps)
logger.logkv("nupdates", update)
logger.logkv("total_timesteps", update*nbatch)
logger.logkv("fps", fps)
logger.logkv("explained_variance_v", float(ev_v))
logger.logkv("explained_variance_l", float(ev_l))
logger.logkv('eprewmean', safemean([epinfo['r'] for epinfo in epinfobuf]))
logger.logkv('violation_times', safemean([epinfo['violation'] for epinfo in epinfobuf]))
logger.logkv('eplrewmean', safemean([epinfo['lr'] for epinfo in epinfobuf]))
logger.logkv('eplenmean', safemean([epinfo['l'] for epinfo in epinfobuf]))
logger.logkv('alpha3', model.alpha3)
if eval_env is not None:
logger.logkv('eval_eprewmean', eval_epinfos['mean_rewards'] )
logger.logkv('eval_eprewmin', eval_epinfos['min_rewards'])
logger.logkv('eval_eprewmax', eval_epinfos['max_rewards'])
logger.logkv('eval_eplrewmean', eval_epinfos['mean_lrewards'])
logger.logkv('eval_eplrewmin', eval_epinfos['min_lrewards'])
logger.logkv('eval_eplrewmax', eval_epinfos['max_lrewards'])
logger.logkv('eval_eplenmean', eval_epinfos['mean_length'])
logger.logkv('time_elapsed', tnow - tfirststart)
for (lossval, lossname) in zip(lossvals, model.loss_names):
logger.logkv(lossname, lossval)
if MPI is None or MPI.COMM_WORLD.Get_rank() == 0:
logger.dumpkvs()
if save_interval and (update % save_interval == 0 or update == 1) and logger.get_dir() and (MPI is None or MPI.COMM_WORLD.Get_rank() == 0):
checkdir = osp.join(logger.get_dir(), 'checkpoints')
os.makedirs(checkdir, exist_ok=True)
savepath = osp.join(checkdir, '%.5i'%update)
# savepath = 'Model/1.ckpt'
print('Saving to', savepath)
model.save(savepath)
return model
def learn(
env,
policy_func,
discriminator,
expert_dataset,
timesteps_per_batch,
*,
g_step,
d_step, # timesteps per actor per update
clip_param,
entcoeff, # clipping parameter epsilon, entropy coeff
optim_epochs,
optim_stepsize,
optim_batchsize, # optimization hypers
gamma,
lam, # advantage estimation
max_timesteps=0,
max_episodes=0,
max_iters=0,
max_seconds=0, # time constraint
callback=None, # you can do anything in the callback, since it takes locals(), globals()
adam_epsilon=1e-5,
d_stepsize=3e-4,
schedule='constant', # annealing for stepsize parameters (epsilon and adam)
save_per_iter=100,
ckpt_dir=None,
task="train",
sample_stochastic=True,
load_model_path=None,
task_name=None,
max_sample_traj=1500):
nworkers = MPI.COMM_WORLD.Get_size()
rank = MPI.COMM_WORLD.Get_rank()
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", ob_space,
ac_space) # Construct network for new policy
oldpi = policy_func("oldpi", ob_space, ac_space) # Network for old policy
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
lrmult = tf.placeholder(
name='lrmult', dtype=tf.float32,
shape=[]) # learning rate multiplier, updated with schedule
clip_param = clip_param * lrmult # Annealed cliping parameter epislon
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)
pol_entpen = (-entcoeff) * meanent
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac)) # pnew / pold
surr1 = ratio * atarg # surrogate from conservative policy iteration
surr2 = U.clip(ratio, 1.0 - clip_param, 1.0 + clip_param) * atarg #
pol_surr = -U.mean(tf.minimum(
surr1, surr2)) # PPO's pessimistic surrogate (L^CLIP)
vf_loss = U.mean(tf.square(pi.vpred - ret))
total_loss = pol_surr + pol_entpen + vf_loss
losses = [pol_surr, pol_entpen, vf_loss, meankl, meanent]
loss_names = ["pol_surr", "pol_entpen", "vf_loss", "kl", "ent"]
var_list = pi.get_trainable_variables()
lossandgrad = U.function([ob, ac, atarg, ret, lrmult],
losses + [U.flatgrad(total_loss, var_list)])
d_adam = MpiAdam(discriminator.get_trainable_variables())
adam = MpiAdam(var_list, epsilon=adam_epsilon)
get_flat = U.GetFlat(var_list)
set_from_flat = U.SetFromFlat(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, ret, lrmult], losses)
U.initialize()
th_init = get_flat()
MPI.COMM_WORLD.Bcast(th_init, root=0)
set_from_flat(th_init)
d_adam.sync()
adam.sync()
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
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
env,
discriminator,
timesteps_per_batch,
stochastic=True)
traj_gen = traj_episode_generator(pi,
env,
timesteps_per_batch,
stochastic=sample_stochastic)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=100) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=100) # rolling buffer for episode rewards
true_rewbuffer = deque(maxlen=100)
assert sum(
[max_iters > 0, max_timesteps > 0, max_episodes > 0,
max_seconds > 0]) == 1, "Only one time constraint permitted"
if task == 'sample_trajectory':
# not elegant, i know :(
sample_trajectory(load_model_path, max_sample_traj, traj_gen,
task_name, sample_stochastic)
sys.exit()
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
elif max_seconds and time.time() - tstart >= max_seconds:
break
if schedule == 'constant':
cur_lrmult = 1.0
elif schedule == 'linear':
cur_lrmult = max(1.0 - float(timesteps_so_far) / max_timesteps, 0)
else:
raise NotImplementedError
# Save model
if iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
U.save_state(os.path.join(ckpt_dir, task_name),
counter=iters_so_far)
logger.log("********** Iteration %i ************" % iters_so_far)
for _ in range(g_step):
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
d = Dataset(dict(ob=ob, ac=ac, atarg=atarg, vtarg=tdlamret),
shuffle=not pi.recurrent)
optim_batchsize = optim_batchsize or ob.shape[0]
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
assign_old_eq_new(
) # set old parameter values to new parameter values
logger.log("Optimizing...")
logger.log(fmt_row(13, loss_names))
# Here we do a bunch of optimization epochs over the data
for _ in range(optim_epochs):
losses = [
] # list of tuples, each of which gives the loss for a minibatch
for batch in d.iterate_once(optim_batchsize):
*newlosses, g = lossandgrad(batch["ob"], batch["ac"],
batch["atarg"], batch["vtarg"],
cur_lrmult)
adam.update(g, optim_stepsize * cur_lrmult)
losses.append(newlosses)
logger.log(fmt_row(13, np.mean(losses, axis=0)))
logger.log("Evaluating losses...")
losses = []
for batch in d.iterate_once(optim_batchsize):
newlosses = compute_losses(batch["ob"], batch["ac"],
batch["atarg"], batch["vtarg"],
cur_lrmult)
losses.append(newlosses)
meanlosses, _, _ = mpi_moments(losses, axis=0)
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, discriminator.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
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 discriminator
if hasattr(discriminator, "obs_rms"):
discriminator.obs_rms.update(
np.concatenate((ob_batch, ob_expert), 0))
*newlosses, g = discriminator.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)))
# ----------------- logger --------------------
logger.log(fmt_row(13, meanlosses))
for (lossval, name) in zipsame(meanlosses, loss_names):
logger.record_tabular("loss_" + name, 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_rews = map(flatten_lists, zip(*listoflrpairs))
lenbuffer.extend(lens)
rewbuffer.extend(rews)
true_rewbuffer.extend(true_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 MPI.COMM_WORLD.Get_rank() == 0:
logger.dump_tabular()
def run(self):
# switch to train mode
self.train()
# Prepare for rollouts
seg_generator = self.traj_segment_generator(self.pi, self.env, self.timesteps_per_batch)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=100) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=100) # rolling buffer for episode rewards
self.check_time_constraints()
while True:
if self.callback: self.callback(locals(), globals())
if self.max_timesteps and timesteps_so_far >= self.max_timesteps:
break
elif self.max_episodes and episodes_so_far >= self.max_episodes:
break
elif self.max_iters and iters_so_far >= self.max_iters:
break
elif self.max_seconds and time.time() - tstart >= self.max_seconds:
break
cur_lrmult = self.get_lr_multiplier(timesteps_so_far)
logger.log("********** Iteration %i ************"%iters_so_far)
segment = seg_generator.__next__()
self.add_vtarg_and_adv(segment, self.gamma, self.lam)
ob, ac, atarg, tdlamret = segment["ob"], segment["ac"], segment["adv"], segment["tdlamret"]
vpredbefore = segment["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()) / atarg.std() # standardized advantage function estimate
d = Dataset(dict(ob=ob, ac=ac, atarg=atarg, vtarg=tdlamret), shuffle=not self.pi.recurrent)
optim_batchsize = self.optim_batchsize or ob.shape[0]
# update running mean/std for policy
# if hasattr(self.pi, "ob_rms"): self.pi.ob_rms.update(ob)
# set old parameter values to new parameter values
self.oldpi.load_state_dict(self.pi.state_dict())
logger.log("Optimizing...")
logger.log(fmt_row(13, self.loss_names))
# Here we do a bunch of optimization epochs over the data
for _ in range(self.optim_epochs):
losses = [] # list of tuples, each of which gives the loss for a minibatch
for batch in d.iterate_once(self.optim_batchsize):
self.optimizer.zero_grad()
batch['ob'] = rearrange_batch_image(batch['ob'])
batch = self.convert_batch_tensor(batch)
total_loss, *newlosses = self.forward(batch["ob"], batch["ac"], batch["atarg"], batch["vtarg"], cur_lrmult)
total_loss.backward()
self.optimizer.step(_step_size=self.optim_stepsize * cur_lrmult)
losses.append(torch.stack(newlosses[0], dim=0).view(-1))
mean_losses = torch.mean(torch.stack(losses, dim=0), dim=0).data.cpu().numpy()
logger.log(fmt_row(13, mean_losses))
logger.log("Evaluating losses...")
losses = []
for batch in d.iterate_once(self.optim_batchsize):
batch['ob'] = rearrange_batch_image(batch['ob'])
batch = self.convert_batch_tensor(batch)
_, *newlosses = self.forward(batch["ob"], batch["ac"], batch["atarg"], batch["vtarg"], cur_lrmult)
losses.append(torch.stack(newlosses[0], dim=0).view(-1))
mean_losses = torch.mean(torch.stack(losses, dim=0), dim=0).data.cpu().numpy()
logger.log(fmt_row(13, mean_losses))
for (lossval, name) in zipsame(mean_losses, self.loss_names):
logger.record_tabular("loss_"+name, lossval)
logger.record_tabular("ev_tdlam_before", explained_variance(vpredbefore, tdlamret))
lrlocal = (segment["ep_lens"], segment["ep_rets"]) # local values
lens, rews = map(flatten_lists, zip(*[lrlocal]))
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.dump_tabular()
def learn(
env,
policy_func,
*,
timesteps=4,
timesteps_per_batch, # timesteps per actor per update
clip_param,
entcoeff, # clipping parameter epsilon, entropy coeff
optim_epochs,
optim_stepsize,
optim_batchsize, # optimization hypers
gamma,
lam, # advantage estimation
max_timesteps=0,
max_episodes=0,
max_iters=0,
max_seconds=0, # time constraint
callback=None, # you can do anything in the callback, since it takes locals(), globals()
adam_epsilon=1e-5,
schedule='constant', # annealing for stepsize parameters (epsilon and adam)
save_per_iter=100,
ckpt_dir=None,
task="train",
sample_stochastic=True,
load_model_path=None,
task_name=None,
max_sample_traj=1500):
# Setup losses and stuff
# ----------------------------------------
ob_space = env.observation_space
ac_space = env.action_space
pi = policy_func("pi", timesteps, ob_space,
ac_space) # Construct network for new policy
oldpi = policy_func("oldpi", timesteps, ob_space,
ac_space) # Network for old policy
atarg = tf.placeholder(
dtype=tf.float32,
shape=[None]) # Target advantage function (if applicable)
ret = tf.placeholder(dtype=tf.float32, shape=[None]) # Empirical return
pi_vpred = tf.placeholder(dtype=tf.float32, shape=[None])
lrmult = tf.placeholder(
name='lrmult', dtype=tf.float32,
shape=[]) # learning rate multiplier, updated with schedule
clip_param = clip_param * lrmult # Annealed cliping parameter epislon
ob = U.get_placeholder_cached(name="ob")
# ob_now = tf.placeholder(dtype=tf.float32, shape=[optim_batchsize, list(ob_space.shape)[0]])
ac = pi.pdtype.sample_placeholder([None])
kloldnew = oldpi.pd.kl(pi.pd)
ent = pi.pd.entropy()
meankl = U.mean(kloldnew)
meanent = U.mean(ent)
pol_entpen = (-entcoeff) * meanent
ratio = tf.exp(pi.pd.logp(ac) - oldpi.pd.logp(ac)) # pnew / pold
surr1 = ratio * atarg # surrogate from conservative policy iteration
surr2 = U.clip(ratio, 1.0 - clip_param, 1.0 + clip_param) * atarg #
pol_surr = -U.mean(tf.minimum(
surr1, surr2)) # PPO's pessimistic surrogate (L^CLIP)
vf_loss = U.mean(tf.square(pi.vpred - ret))
# total_loss = pol_surr + pol_entpen + vf_loss
total_loss = pol_surr + pol_entpen
losses = [pol_surr, pol_entpen, meankl, meanent]
loss_names = ["pol_surr", "pol_entpen", "kl", "ent"]
var_list = pi.get_trainable_variables()
vf_var_list = [
v for v in var_list if v.name.split("/")[1].startswith("vf")
]
pol_var_list = [
v for v in var_list if not v.name.split("/")[1].startswith("vf")
]
# lossandgrad = U.function([ob, ac, atarg ,ret, lrmult], losses + [U.flatgrad(total_loss, var_list)])
lossandgrad = U.function([ob, ac, atarg, ret, lrmult],
losses + [U.flatgrad(total_loss, pol_var_list)])
vf_grad = U.function([ob, ac, atarg, ret, lrmult],
U.flatgrad(vf_loss, vf_var_list))
# adam = MpiAdam(var_list, epsilon=adam_epsilon)
pol_adam = MpiAdam(pol_var_list, epsilon=adam_epsilon)
vf_adam = MpiAdam(vf_var_list, epsilon=adam_epsilon)
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, ret, lrmult], losses)
U.initialize()
#adam.sync()
pol_adam.sync()
vf_adam.sync()
# Prepare for rollouts
# ----------------------------------------
seg_gen = traj_segment_generator(pi,
timesteps,
env,
timesteps_per_batch,
stochastic=True)
traj_gen = traj_episode_generator(pi,
env,
timesteps_per_batch,
stochastic=sample_stochastic)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
tstart = time.time()
lenbuffer = deque(maxlen=100) # rolling buffer for episode lengths
rewbuffer = deque(maxlen=100) # rolling buffer for episode rewards
EpRewMean_MAX = 2.5e3
assert sum(
[max_iters > 0, max_timesteps > 0, max_episodes > 0,
max_seconds > 0]) == 1, "Only one time constraint permitted"
if task == 'sample_trajectory':
# not elegant, i know :(
sample_trajectory(load_model_path, max_sample_traj, traj_gen,
task_name, sample_stochastic)
sys.exit()
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
elif max_seconds and time.time() - tstart >= max_seconds:
break
if schedule == 'constant':
cur_lrmult = 1.0
elif schedule == 'linear':
cur_lrmult = max(1.0 - float(timesteps_so_far) / max_timesteps, 0)
else:
raise NotImplementedError
# Save model
if iters_so_far % save_per_iter == 0 and ckpt_dir is not None:
U.save_state(os.path.join(ckpt_dir, task_name),
counter=iters_so_far)
logger.log("********** Iteration %i ************" % iters_so_far)
# if(iters_so_far == 1):
# a = 1
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, vpred, tdlamret = seg["ob"], seg["ac"], seg["adv"], seg[
"vpred"], seg["tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before udpate
atarg = (atarg - atarg.mean()
) / atarg.std() # standardized advantage function estimate
d = Dataset(
dict(ob=ob, ac=ac, atarg=atarg, vpred=vpred, vtarg=tdlamret),
shuffle=False
) #d = Dataset(dict(ob=ob, ac=ac, atarg=atarg, vpred = vpred, vtarg=tdlamret), shuffle=not pi.recurrent)
optim_batchsize = optim_batchsize or ob.shape[0]
if hasattr(pi, "ob_rms"):
pi.ob_rms.update(ob) # update running mean/std for policy
assign_old_eq_new() # set old parameter values to new parameter values
logger.log("Optimizing...")
logger.log(fmt_row(13, loss_names))
# Here we do a bunch of optimization epochs over the data
for _ in range(optim_epochs):
losses = [
] # list of tuples, each of which gives the loss for a minibatch
pre_obs = [seg["ob_reset"] for jmj in range(timesteps - 1)]
for batch in d.iterate_once(optim_batchsize):
##feed ob, 重新处理一下ob,在batch["ob"]的最前面插入timesteps-1个env.reset的ob,然后滑动串口划分一下batch['ob]
ob_now = np.append(pre_obs, batch['ob']).reshape(
optim_batchsize + timesteps - 1,
list(ob_space.shape)[0])
pre_obs = ob_now[-(timesteps - 1):]
ob_fin = []
for jmj in range(optim_batchsize):
ob_fin.append(ob_now[jmj:jmj + timesteps])
*newlosses, g = lossandgrad(ob_fin, batch["ac"],
batch["atarg"], batch["vtarg"],
cur_lrmult) ###这里的g好像都是0
#adam.update(g, optim_stepsize * cur_lrmult)
pol_adam.update(g, optim_stepsize * cur_lrmult)
vf_g = vf_grad(ob_fin, batch["ac"], batch["atarg"],
batch["vtarg"], cur_lrmult)
vf_adam.update(vf_g, optim_stepsize * cur_lrmult)
losses.append(newlosses)
logger.log(fmt_row(13, np.mean(losses, axis=0)))
pre_obs = [seg["ob_reset"] for jmj in range(timesteps - 1)]
for batch in d.iterate_once(optim_batchsize):
##feed ob, 重新处理一下ob,在batch["ob"]的最前面插入timesteps-1个env.reset的ob,然后滑动串口划分一下batch['ob]
ob_now = np.append(pre_obs, batch['ob']).reshape(
optim_batchsize + timesteps - 1,
list(ob_space.shape)[0])
pre_obs = ob_now[-(timesteps - 1):]
ob_fin = []
for jmj in range(optim_batchsize):
ob_fin.append(ob_now[jmj:jmj + timesteps])
*newlosses, g = lossandgrad(ob_fin, batch["ac"],
batch["atarg"], batch["vtarg"],
cur_lrmult) ###这里的g好像都是0
#adam.update(g, optim_stepsize * cur_lrmult)
pol_adam.update(g, optim_stepsize * cur_lrmult)
vf_g = vf_grad(ob_fin, batch["ac"], batch["atarg"],
batch["vtarg"], cur_lrmult)
vf_adam.update(vf_g, optim_stepsize * cur_lrmult)
logger.log("Evaluating losses...")
losses = []
loss_pre_obs = [seg["ob_reset"] for jmj in range(timesteps - 1)]
for batch in d.iterate_once(optim_batchsize):
### feed ob
ob_now = np.append(loss_pre_obs, batch['ob']).reshape(
optim_batchsize + timesteps - 1,
list(ob_space.shape)[0])
loss_pre_obs = ob_now[-(timesteps - 1):]
ob_fin = []
for jmj in range(optim_batchsize):
ob_fin.append(ob_now[jmj:jmj + timesteps])
newlosses = compute_losses(ob_fin, batch["ac"], batch["atarg"],
batch["vtarg"], cur_lrmult)
losses.append(newlosses)
meanlosses, _, _ = mpi_moments(losses, axis=0)
logger.log(fmt_row(13, meanlosses))
for (lossval, name) in zipsame(meanlosses, loss_names):
logger.record_tabular("loss_" + name, lossval)
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))
if (np.mean(rewbuffer) > EpRewMean_MAX):
EpRewMean_MAX = np.mean(rewbuffer)
print(iters_so_far)
print(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 MPI.COMM_WORLD.Get_rank() == 0:
logger.dump_tabular()