def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="A2C",
reset_num_timesteps=True):
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()
self.learning_rate_schedule = Scheduler(initial_value=self.learning_rate, n_values=total_timesteps,
schedule=self.lr_schedule)
t_start = time.time()
callback.on_training_start(locals(), globals())
for update in range(1, total_timesteps // self.n_batch + 1):
callback.on_rollout_start()
# true_reward is the reward without discount
rollout = self.runner.run(callback)
# unpack
obs, states, rewards, masks, actions, values, ep_infos, true_reward = rollout
callback.on_rollout_end()
# Early stopping due to the callback
if not self.runner.continue_training:
break
self.ep_info_buf.extend(ep_infos)
_, value_loss, policy_entropy = self._train_step(obs, states, rewards, masks, actions, values,
self.num_timesteps // self.n_batch, writer)
n_seconds = time.time() - t_start
fps = int((update * self.n_batch) / n_seconds)
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, rewards)
logger.record_tabular("nupdates", update)
logger.record_tabular("total_timesteps", self.num_timesteps)
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(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.dump_tabular()
callback.on_training_end()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=1,
tb_log_name="PPO2",
reset_num_timesteps=True):
# 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
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 = 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 + (
(self.noptepochs * self.n_batch + 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 = 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 + (
(self.noptepochs * self.n_envs + 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()
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="ACKTR",
reset_num_timesteps=True):
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()
self.n_batch = self.n_envs * self.n_steps
self.learning_rate_schedule = Scheduler(initial_value=self.learning_rate, n_values=total_timesteps,
schedule=self.lr_schedule)
# FIFO queue of the q_runner thread is closed at the end of the learn function.
# As a result, it needs to be redefinied at every call
with self.graph.as_default():
with tf.variable_scope("kfac_apply", reuse=self.trained,
custom_getter=tf_util.outer_scope_getter("kfac_apply")):
# Some of the variables are not in a scope when they are create
# so we make a note of any previously uninitialized variables
tf_vars = tf.global_variables()
is_uninitialized = self.sess.run([tf.is_variable_initialized(var) for var in tf_vars])
old_uninitialized_vars = [v for (v, f) in zip(tf_vars, is_uninitialized) if not f]
self.train_op, self.q_runner = self.optim.apply_gradients(list(zip(self.grads_check, self.params)))
# then we check for new uninitialized variables and initialize them
tf_vars = tf.global_variables()
is_uninitialized = self.sess.run([tf.is_variable_initialized(var) for var in tf_vars])
new_uninitialized_vars = [v for (v, f) in zip(tf_vars, is_uninitialized)
if not f and v not in old_uninitialized_vars]
if len(new_uninitialized_vars) != 0:
self.sess.run(tf.variables_initializer(new_uninitialized_vars))
self.trained = True
t_start = time.time()
coord = tf.train.Coordinator()
if self.q_runner is not None:
enqueue_threads = self.q_runner.create_threads(self.sess, coord=coord, start=True)
else:
enqueue_threads = []
callback.on_training_start(locals(), globals())
for update in range(1, total_timesteps // self.n_batch + 1):
callback.on_rollout_start()
# pytype:disable=bad-unpacking
# true_reward is the reward without discount
if isinstance(self.runner, PPO2Runner):
# We are using GAE
rollout = self.runner.run(callback)
obs, returns, masks, actions, values, _, states, ep_infos, true_reward = rollout
else:
rollout = self.runner.run(callback)
obs, states, returns, masks, actions, values, ep_infos, true_reward = rollout
# pytype:enable=bad-unpacking
callback.on_rollout_end()
# Early stopping due to the callback
if not self.runner.continue_training:
break
self.ep_info_buf.extend(ep_infos)
policy_loss, value_loss, policy_entropy = self._train_step(obs, states, returns, masks, actions, values,
self.num_timesteps // (self.n_batch + 1),
writer)
n_seconds = time.time() - t_start
fps = int((update * self.n_batch) / n_seconds)
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.record_tabular("nupdates", update)
logger.record_tabular("total_timesteps", self.num_timesteps)
logger.record_tabular("fps", fps)
logger.record_tabular("policy_entropy", float(policy_entropy))
logger.record_tabular("policy_loss", float(policy_loss))
logger.record_tabular("value_loss", float(value_loss))
logger.record_tabular("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.dump_tabular()
coord.request_stop()
coord.join(enqueue_threads)
callback.on_training_end()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=100,
tb_log_name="TRPO",
reset_num_timesteps=True):
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()
with self.sess.as_default():
callback.on_training_start(locals(), globals())
seg_gen = trajectories1st(self.env, self.policy_pi,
self.timesteps_per_batch, self.gamma,
self.lam)
# seg_gen = traj_segment_generator(self.policy_pi, self.env, self.timesteps_per_batch,
# reward_giver=self.reward_giver,
# gail=self.using_gail, callback=callback)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
t_start = time.time()
len_buffer = deque(
maxlen=40) # rolling buffer for episode lengths
reward_buffer = deque(
maxlen=40) # rolling buffer for episode rewards
true_reward_buffer = None
if self.using_gail:
true_reward_buffer = deque(maxlen=40)
self._initialize_dataloader()
# Stats not used for now
# TODO: replace with normal tb logging
# g_loss_stats = Stats(loss_names)
# d_loss_stats = Stats(reward_giver.loss_name)
# ep_stats = Stats(["True_rewards", "Rewards", "Episode_length"])
while True:
if timesteps_so_far >= total_timesteps:
break
logger.log("********** Iteration %i ************" %
iters_so_far)
def fisher_vector_product(vec):
return self.allmean(
self.compute_fvp(
vec, *fvpargs,
sess=self.sess)) + self.cg_damping * vec
# ------------------ Update G ------------------
logger.log("Optimizing Policy...")
# g_step = 1 when not using GAIL
mean_losses = None
vpredbefore = None
tdlamret = None
observation = None
action = None
seg = None
for k in range(self.g_step):
with self.timed("sampling"):
seg = seg_gen.__next__()
# Stop training early (triggered by the callback)
if not seg.get('continue_training', True): # pytype: disable=attribute-error
break
#add_vtarg_and_adv(seg, self.gamma, self.lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
observation, action = seg["observations"], seg[
"actions"]
atarg, tdlamret = seg["adv"], seg["tdlamret"]
vpredbefore = seg[
"vpred"] # predicted value function before update
atarg = (atarg - atarg.mean()) / (
atarg.std() + 1e-8
) # standardized advantage function estimate
# true_rew is the reward without discount
if writer is not None:
total_episode_reward_logger(
self.episode_reward,
seg["true_rewards"].reshape(
(self.n_envs, -1)), seg["dones"].reshape(
(self.n_envs, -1)), writer,
self.num_timesteps)
args = seg["observations"], seg["observations"], seg[
"actions"], atarg
# Subsampling: see p40-42 of John Schulman thesis
# http://joschu.net/docs/thesis.pdf
fvpargs = [arr[::5] for arr in args]
self.assign_old_eq_new(sess=self.sess)
with self.timed("computegrad"):
steps = self.num_timesteps + (k + 1) * (
seg["total_timestep"] / self.g_step)
run_options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata(
) if self.full_tensorboard_log else None
# run loss backprop with summary, and save the metadata (memory, compute time, ...)
if writer is not None:
summary, grad, *lossbefore = self.compute_lossandgrad(
*args,
tdlamret,
sess=self.sess,
options=run_options,
run_metadata=run_metadata)
if self.full_tensorboard_log:
writer.add_run_metadata(
run_metadata, 'step%d' % steps)
writer.add_summary(summary, steps)
else:
_, grad, *lossbefore = self.compute_lossandgrad(
*args,
tdlamret,
sess=self.sess,
options=run_options,
run_metadata=run_metadata)
lossbefore = self.allmean(np.array(lossbefore))
grad = self.allmean(grad)
if np.allclose(grad, 0):
logger.log("Got zero gradient. not updating")
else:
with self.timed("conjugate_gradient"):
stepdir = cg_via_np(fisher_vector_product,
grad,
self.cg_iters,
final_not_prime=False)
# stepdir = conjugate_gradient(fisher_vector_product, grad, cg_iters=self.cg_iters,
# verbose=self.rank == 0 and self.verbose >= 1)
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) / self.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 = self.get_flat()
for _ in range(10):
thnew = thbefore + fullstep * stepsize
self.set_from_flat(thnew)
mean_losses = surr, kl_loss, *_ = self.allmean(
np.array(
self.compute_losses(*args,
sess=self.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 > self.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")
self.set_from_flat(thbefore)
if self.nworkers > 1 and iters_so_far % 20 == 0:
# list of tuples
paramsums = MPI.COMM_WORLD.allgather(
(thnew.sum(), self.vfadam.getflat().sum()))
assert all(
np.allclose(ps, paramsums[0])
for ps in paramsums[1:])
for (loss_name,
loss_val) in zip(self.loss_names,
mean_losses):
logger.record_tabular(loss_name, loss_val)
with self.timed("vf"):
for _ in range(self.vf_iters):
# NOTE: for recurrent policies, use shuffle=False?
for (mbob, mbret) in dataset.iterbatches(
(seg["observations"], seg["tdlamret"]),
include_final_partial_batch=False,
batch_size=128,
shuffle=True):
grad = self.allmean(
self.compute_vflossandgrad(
mbob, mbob, mbret, sess=self.sess))
self.vfadam.update(grad, self.vf_stepsize)
# Stop training early (triggered by the callback)
if not seg.get('continue_training', True): # pytype: disable=attribute-error
break
logger.record_tabular(
"explained_variance_tdlam_before",
explained_variance(vpredbefore, tdlamret))
if self.using_gail:
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, self.reward_giver.loss_name))
assert len(observation) == self.timesteps_per_batch
batch_size = self.timesteps_per_batch // self.d_step
# NOTE: uses only the last g step for observation
d_losses = [
] # list of tuples, each of which gives the loss for a minibatch
# NOTE: for recurrent policies, use shuffle=False?
for ob_batch, ac_batch in dataset.iterbatches(
(observation, action),
include_final_partial_batch=False,
batch_size=batch_size,
shuffle=True):
ob_expert, ac_expert = self.expert_dataset.get_next_batch(
)
# update running mean/std for reward_giver
if self.reward_giver.normalize:
self.reward_giver.obs_rms.update(
np.concatenate((ob_batch, ob_expert), 0))
# Reshape actions if needed when using discrete actions
if isinstance(self.action_space,
gym.spaces.Discrete):
if len(ac_batch.shape) == 2:
ac_batch = ac_batch[:, 0]
if len(ac_expert.shape) == 2:
ac_expert = ac_expert[:, 0]
*newlosses, grad = self.reward_giver.lossandgrad(
ob_batch, ac_batch, ob_expert, ac_expert)
self.d_adam.update(self.allmean(grad),
self.d_stepsize)
d_losses.append(newlosses)
logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
# lr: lengths and rewards
lr_local = (seg["ep_lens"], seg["ep_rets"],
seg["ep_true_rets"]) # local values
list_lr_pairs = MPI.COMM_WORLD.allgather(
lr_local) # list of tuples
lens, rews, true_rets = map(flatten_lists,
zip(*list_lr_pairs))
true_reward_buffer.extend(true_rets)
else:
# lr: lengths and rewards
lr_local = (seg["ep_lens"], seg["ep_rets"]
) # local values
list_lr_pairs = MPI.COMM_WORLD.allgather(
lr_local) # list of tuples
lens, rews = map(flatten_lists, zip(*list_lr_pairs))
len_buffer.extend(lens)
reward_buffer.extend(rews)
if len(len_buffer) > 0:
logger.record_tabular("EpLenMean", np.mean(len_buffer))
logger.record_tabular("EpRewMean",
np.mean(reward_buffer))
if self.using_gail:
logger.record_tabular("EpTrueRewMean",
np.mean(true_reward_buffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
current_it_timesteps = MPI.COMM_WORLD.allreduce(
seg["total_timestep"])
timesteps_so_far += current_it_timesteps
self.num_timesteps += current_it_timesteps
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", self.num_timesteps)
logger.record_tabular("TimeElapsed", time.time() - t_start)
if self.verbose >= 1 and self.rank == 0:
logger.dump_tabular()
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="ACER",
reset_num_timesteps=True):
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()
self.learning_rate_schedule = Scheduler(initial_value=self.learning_rate, n_values=total_timesteps,
schedule=self.lr_schedule)
episode_stats = EpisodeStats(self.n_steps, self.n_envs)
if self.replay_ratio > 0:
buffer = Buffer(env=self.env, n_steps=self.n_steps, size=self.buffer_size)
else:
buffer = None
t_start = time.time()
callback.on_training_start(locals(), globals())
# n_batch samples, 1 on_policy call and multiple off-policy calls
for steps in range(0, total_timesteps, self.n_batch):
callback.on_rollout_start()
enc_obs, obs, actions, rewards, mus, dones, masks = self.runner.run(callback)
callback.update_locals(locals())
callback.on_rollout_end()
# Early stopping due to the callback
if not self.runner.continue_training:
break
episode_stats.feed(rewards, dones)
if buffer is not None:
buffer.put(enc_obs, actions, rewards, mus, dones, masks)
if writer is not None:
total_episode_reward_logger(self.episode_reward,
rewards.reshape((self.n_envs, self.n_steps)),
dones.reshape((self.n_envs, self.n_steps)),
writer, self.num_timesteps)
# reshape stuff correctly
obs = obs.reshape(self.runner.batch_ob_shape)
actions = actions.reshape([self.n_batch])
rewards = rewards.reshape([self.n_batch])
mus = mus.reshape([self.n_batch, self.n_act])
dones = dones.reshape([self.n_batch])
masks = masks.reshape([self.runner.batch_ob_shape[0]])
names_ops, values_ops = self._train_step(obs, actions, rewards, dones, mus, self.initial_state, masks,
self.num_timesteps, writer)
if self.verbose >= 1 and (int(steps / self.n_batch) % log_interval == 0):
logger.record_tabular("total_timesteps", self.num_timesteps)
logger.record_tabular("fps", int(steps / (time.time() - t_start)))
# IMP: In EpisodicLife env, during training, we get done=True at each loss of life,
# not just at the terminal state. Thus, this is mean until end of life, not end of episode.
# For true episode rewards, see the monitor files in the log folder.
logger.record_tabular("mean_episode_length", episode_stats.mean_length())
logger.record_tabular("mean_episode_reward", episode_stats.mean_reward())
for name, val in zip(names_ops, values_ops):
logger.record_tabular(name, float(val))
logger.dump_tabular()
if (self.replay_ratio > 0 and
buffer is not None and
buffer.has_atleast(self.replay_start)):
samples_number = np.random.poisson(self.replay_ratio)
for _ in range(samples_number):
# get obs, actions, rewards, mus, dones from buffer.
obs, actions, rewards, mus, dones, masks = buffer.get()
# reshape stuff correctly
obs = obs.reshape(self.runner.batch_ob_shape)
actions = actions.reshape([self.n_batch])
rewards = rewards.reshape([self.n_batch])
mus = mus.reshape([self.n_batch, self.n_act])
dones = dones.reshape([self.n_batch])
masks = masks.reshape([self.runner.batch_ob_shape[0]])
self._train_step(obs, actions, rewards, dones, mus, self.initial_state, masks,
self.num_timesteps)
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None,
log_interval=4, tb_log_name="TD3", reset_num_timesteps=True, replay_wrapper=None):
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
callback = self._init_callback(callback)
last_replay_update = 0
if replay_wrapper is not None:
self.replay_buffer = replay_wrapper(self.replay_buffer)
if isinstance(self.train_freq, tuple): # TODO: bug with optuna please FIX
self.train_freq = self.train_freq[0]
self.gradient_steps = self.gradient_steps[0]
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
self._setup_learn()
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
# Initial learning rate
current_lr = self.learning_rate(1)
start_time = time.time()
episode_rewards = [0.0]
episode_successes = []
if self.action_noise is not None:
self.action_noise.reset()
obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs().squeeze()
n_updates = 0
infos_values = []
self.active_sampling = False
initial_step = self.num_timesteps
episode_data = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
if self.buffer_is_prioritized and \
((replay_wrapper is not None and self.replay_buffer.replay_buffer.__name__ == "ReplayBuffer")
or (replay_wrapper is None and self.replay_buffer.__name__ == "ReplayBuffer")) \
and self.num_timesteps >= self.prioritization_starts:
self._set_prioritized_buffer()
if self.recurrent_policy:
done = False
policy_state = self.policy_tf_act.initial_state
prev_policy_state = self.policy_tf_act.initial_state # Keep track of this so it doesnt have to be recalculated when saving it to replay buffer
for step in range(initial_step, total_timesteps):
# Before training starts, randomly sample actions
# from a uniform distribution for better exploration.
# Afterwards, use the learned policy
# if random_exploration is set to 0 (normal setting)
if self.num_timesteps < self.learning_starts or np.random.rand() < self.random_exploration:
# actions sampled from action space are from range specific to the environment
# but algorithm operates on tanh-squashed actions therefore simple scaling is used
unscaled_action = self.env.action_space.sample()
action = scale_action(self.action_space, unscaled_action)
else:
if self.recurrent_policy:
action, policy_state = self.policy_tf_act.step(obs[None], state=policy_state, mask=np.array(done)[None])
action = action.flatten()
else:
action = self.policy_tf.step(obs[None]).flatten()
# Add noise to the action, as the policy
# is deterministic, this is required for exploration
if self.action_noise is not None:
action = np.clip(action + self.action_noise(), -1, 1)
# Rescale from [-1, 1] to the correct bounds
unscaled_action = unscale_action(self.action_space, action)
assert action.shape == self.env.action_space.shape
new_obs, reward, done, info = self.env.step(unscaled_action)
self.num_timesteps += 1
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback.on_step() is False:
break
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs().squeeze()
reward_ = self._vec_normalize_env.get_original_reward().squeeze()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, reward
if self.reward_transformation is not None:
reward = self.reward_transformation(reward)
# Store transition in the replay buffer.
extra_data = {}
if self.time_aware:
bootstrap = True
if done:
info_time_limit = info.get("TimeLimit.truncated", None)
bootstrap = info.get("termination", None) == "steps" or \
(info_time_limit is not None and info_time_limit)
extra_data["bootstrap"] = bootstrap
if hasattr(self.policy, "collect_data"):
if self.recurrent_policy:
extra_data.update(self.policy_tf_act.collect_data(locals(), globals()))
if self.policy_tf.save_target_state:
extra_data.update({"target_" + state_name: self.target_policy_tf.initial_state[0, :]
for state_name in (["state"] if self.target_policy_tf.share_lstm
else ["pi_state", "qf1_state", "qf2_state"])})
else:
extra_data.update(self.policy_tf.collect_data(locals(), globals()))
self.replay_buffer.add(obs, action, reward, new_obs, done, **extra_data) # Extra data must be sent as kwargs to support separate bootstrap and done signals (needed for HER style algorithms)
episode_data.append({"obs": obs, "action": action, "reward": reward, "obs_tp1": new_obs, "done": done, **extra_data})
obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
if ((replay_wrapper is not None and self.replay_buffer.replay_buffer.__name__ == "RankPrioritizedReplayBuffer")\
or self.replay_buffer.__name__ == "RankPrioritizedReplayBuffer") and \
self.num_timesteps % self.buffer_size == 0:
self.replay_buffer.rebalance()
# Retrieve reward and episode length if using Monitor wrapper
maybe_ep_info = info.get('episode')
if maybe_ep_info is not None and self.num_timesteps >= self.learning_starts:
self.ep_info_buf.extend([maybe_ep_info])
if writer is not None:
# Write reward per episode to tensorboard
ep_reward = np.array([reward_]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(self.episode_reward, ep_reward,
ep_done, writer, self.num_timesteps)
if self.num_timesteps % self.train_freq == 0:
callback.on_rollout_end()
mb_infos_vals = []
# Update policy, critics and target networks
for grad_step in range(self.gradient_steps):
# Break if the warmup phase is not over
# or if there are not enough samples in the replay buffer
if not self.replay_buffer.can_sample(self.batch_size) \
or self.num_timesteps < self.learning_starts:
break
n_updates += 1
# Compute current learning_rate
frac = 1.0 - self.num_timesteps / total_timesteps
current_lr = self.learning_rate(frac)
# Update policy and critics (q functions)
# Note: the policy is updated less frequently than the Q functions
# this is controlled by the `policy_delay` parameter
step_writer = writer if grad_step % self.write_freq == 0 else None
mb_infos_vals.append(
self._train_step(step, step_writer, current_lr, (step + grad_step) % self.policy_delay == 0))
# Log losses and entropy, useful for monitor training
if len(mb_infos_vals) > 0:
infos_values = np.mean(mb_infos_vals, axis=0)
callback.on_rollout_start()
episode_rewards[-1] += reward
if self.recurrent_policy:
prev_policy_state = policy_state
if done:
if isinstance(self.replay_buffer, DiscrepancyReplayBuffer) and n_updates - last_replay_update >= 5000:
self.replay_buffer.update_priorities()
last_replay_update = n_updates
if self.action_noise is not None:
self.action_noise.reset()
if not isinstance(self.env, VecEnv):
if self.active_sampling:
sample_obs, sample_state = self.env.get_random_initial_states(25)
obs_discrepancies = self.policy_tf.get_q_discrepancy(sample_obs)
obs = self.env.reset(**sample_state[np.argmax(obs_discrepancies)])
else:
obs = self.env.reset()
episode_data = []
episode_rewards.append(0.0)
if self.recurrent_policy:
prev_policy_state = self.policy_tf_act.initial_state
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if len(episode_rewards[-101:-1]) == 0:
mean_reward = -np.inf
else:
mean_reward = round(float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
self.num_timesteps += 1
if self.buffer_is_prioritized and \
((replay_wrapper is not None and self.replay_buffer.replay_buffer.__name__ == "ReplayBuffer")
or (replay_wrapper is None and self.replay_buffer.__name__ == "ReplayBuffer"))\
and self.num_timesteps >= self.prioritization_starts:
self._set_prioritized_buffer()
# Display training infos
if self.verbose >= 1 and done and log_interval is not None and len(episode_rewards) % log_interval == 0:
fps = int(step / (time.time() - start_time))
logger.logkv("episodes", num_episodes)
logger.logkv("mean 100 episode reward", mean_reward)
if len(self.ep_info_buf) > 0 and len(self.ep_info_buf[0]) > 0:
logger.logkv('ep_rewmean', safe_mean([ep_info['r'] for ep_info in self.ep_info_buf]))
logger.logkv('eplenmean', safe_mean([ep_info['l'] for ep_info in self.ep_info_buf]))
logger.logkv("n_updates", n_updates)
logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - start_time))
if len(episode_successes) > 0:
logger.logkv("success rate", np.mean(episode_successes[-100:]))
if len(infos_values) > 0:
for (name, val) in zip(self.infos_names, infos_values):
logger.logkv(name, val)
logger.logkv("total timesteps", self.num_timesteps)
logger.dumpkvs()
# Reset infos:
infos_values = []
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None, seed=None, log_interval=100, tb_log_name="MDPO",
reset_num_timesteps=True):
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
callback = self._init_callback(callback)
print("got seed {}, sgd_steps {}".format(seed, self.sgd_steps))
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
with self.sess.as_default():
callback.on_training_start(locals(), globals())
seg_gen = traj_segment_generator(self.old_policy, self.env, self.timesteps_per_batch,
reward_giver=self.reward_giver,
gail=self.using_gail, mdal=self.using_mdal, neural=self.neural,
action_space=self.action_space, gamma=self.gamma, callback=callback)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
t_start = time.time()
len_buffer = deque(maxlen=40) # rolling buffer for episode lengths
reward_buffer = deque(maxlen=40) # rolling buffer for episode rewards
self.episode_reward = np.zeros((self.n_envs,))
self.outer_learning_rate = get_schedule_fn(3e-4)
self.cliprange_vf = get_schedule_fn(0.2)
true_reward_buffer = None
if self.using_gail or self.using_mdal:
true_reward_buffer = deque(maxlen=40)
# Initialize dataloader
batchsize = self.timesteps_per_batch // self.d_step
self.expert_dataset.init_dataloader(batchsize)
# Stats not used for now
# TODO: replace with normal tb logging
# g_loss_stats = Stats(loss_names)
# d_loss_stats = Stats(reward_giver.loss_name)
# ep_stats = Stats(["True_rewards", "Rewards", "Episode_length"])
while True:
# if callback is not None:
# # Only stop training if return value is False, not when it is None. This is for backwards
# # compatibility with callbacks that have no return statement.
# if callback(locals(), globals()) is False:
# break
if total_timesteps and timesteps_so_far >= total_timesteps:
break
logger.log("********** Iteration %i ************" % iters_so_far)
#def fisher_vector_product(vec):
# return self.allmean(self.compute_fvp(vec, *fvpargs, sess=self.sess)) + self.cg_damping * vec
# ------------------ Update G ------------------
# logger.log("Optimizing Policy...")
# g_step = 1 when not using GAIL
mean_losses = None
vpredbefore = None
tdlamret = None
observation = None
action = None
seg = None
for k in range(self.g_step):
with self.timed("sampling"):
seg = seg_gen.__next__()
if not seg.get('continue_training', True): # pytype: disable=attribute-error
break
add_vtarg_and_adv(seg, self.gamma, self.lam)
if self.using_mdal:
policy_successor_features = add_successor_features(seg, self.gamma,
is_action_features=self.is_action_features)
else:
policy_successor_features = add_successor_features(seg, self.gamma)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
observation, action = seg["observations"], seg["actions"]
atarg, tdlamret = seg["adv"], seg["tdlamret"]
vpredbefore = seg["vpred"] # predicted value function before update
atarg = (atarg - atarg.mean()) / atarg.std() # standardized advantage function estimate
# true_rew is the reward without discount
if writer is not None:
self.episode_reward = total_episode_reward_logger(self.episode_reward,
seg["true_rewards"].reshape(
(self.n_envs, -1)),
seg["dones"].reshape((self.n_envs, -1)),
writer, self.num_timesteps)
n_updates = int(total_timesteps / self.timesteps_per_batch)
lr_now = np.float32(1.0 - (iters_so_far - 1.0) / n_updates)
outer_lr_now = self.outer_learning_rate(1.0 - (iters_so_far - 1.0) / n_updates)
clip_now = self.cliprange_vf(1.0 - (iters_so_far - 1.0) / n_updates)
args = seg["observations"], seg["observations"], seg["actions"], atarg
# Subsampling: see p40-42 of John Schulman thesis
# http://joschu.net/docs/thesis.pdf
#fvpargs = [arr[::5] for arr in args]
with self.timed("computegrad"):
steps = self.num_timesteps + (k + 1) * (seg["total_timestep"] / self.g_step)
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata() if self.full_tensorboard_log else None
# run loss backprop with summary, and save the metadata (memory, compute time, ...)
if writer is not None:
summary, grad, *lossbefore = self.compute_lossandgrad(*args, tdlamret,
lr_now, seg["vpred"],
seg["observations"],
sess=self.sess,
options=run_options,
run_metadata=run_metadata)
if self.full_tensorboard_log:
writer.add_run_metadata(run_metadata, 'step%d' % steps)
writer.add_summary(summary, steps)
else:
_, grad, *lossbefore = self.compute_lossandgrad(*args, tdlamret,
lr_now, seg["vpred"],
seg["observations"],
sess=self.sess,
options=run_options,
run_metadata=run_metadata)
td_map = {self.policy_pi.obs_ph: seg["observations"],
self.old_policy.obs_ph: seg["observations"],
self.closed_policy.obs_ph: seg["observations"],
self.action: seg["actions"], self.atarg: atarg, self.ret: tdlamret,
self.learning_rate_ph: lr_now, self.outer_learning_rate_ph: outer_lr_now,
self.vtarg: seg["vpred"]}
for _ in range(int(self.sgd_steps)):
_ = self.sess.run(self._train, td_map)
#if self.method == "closed-KL":
# _ = self.sess.run(self._train_policy, td_map)
if np.allclose(grad, 0):
logger.log("Got zero gradient. not updating")
else:
for _ in range(1):
mean_losses = surr, kl_loss, *_ = self.allmean(
np.array(self.compute_losses(*args, lr_now, seg["vpred"], sess=self.sess)))
with self.timed("vf"):
for _ in range(self.vf_iters):
# NOTE: for recurrent policies, use shuffle=False?
for (mbob, mbret, mbval) in dataset.iterbatches((seg["observations"], seg["tdlamret"], seg["vpred"]),
include_final_partial_batch=False,
batch_size=128,
shuffle=True):
grad = self.allmean(self.compute_vflossandgrad(mbob, mbob, mbret, mbval, clip_now, sess=self.sess))
self.vfadam.update(grad, outer_lr_now) #self.vf_stepsize)
if iters_so_far % 1 == 0:
# print("updating theta now")
self.assign_old_eq_new(sess=self.sess)
for (loss_name, loss_val) in zip(self.loss_names, mean_losses):
logger.record_tabular(loss_name, loss_val)
logger.record_tabular("explained_variance_tdlam_before",
explained_variance(vpredbefore, tdlamret))
if self.using_gail:
# ------------------ Update D ------------------
logger.log("Optimizing Discriminator...")
logger.log(fmt_row(13, self.reward_giver.loss_name))
assert len(observation) == self.timesteps_per_batch
batch_size = self.timesteps_per_batch // self.d_step
# NOTE: uses only the last g step for observation
d_losses = [] # list of tuples, each of which gives the loss for a minibatch
# NOTE: for recurrent policies, use shuffle=False?
for ob_batch, ac_batch in dataset.iterbatches((observation, action),
include_final_partial_batch=False,
batch_size=batch_size,
shuffle=True):
ob_expert, ac_expert = self.expert_dataset.get_next_batch()
# update running mean/std for reward_giver
if self.reward_giver.normalize:
self.reward_giver.obs_rms.update(np.concatenate((ob_batch, ob_expert), 0))
# Reshape actions if needed when using discrete actions
if isinstance(self.action_space, gym.spaces.Discrete):
if len(ac_batch.shape) == 2:
ac_batch = ac_batch[:, 0]
if len(ac_expert.shape) == 2:
ac_expert = ac_expert[:, 0]
*newlosses, grad = self.reward_giver.lossandgrad(ob_batch, ac_batch, ob_expert, ac_expert)
self.d_adam.update(self.allmean(grad), self.d_stepsize)
d_losses.append(newlosses)
logger.log(fmt_row(13, np.mean(d_losses, axis=0)))
elif self.using_mdal:
batch_sampling = True
if self.neural:
if batch_sampling:
batch_size = self.timesteps_per_batch // self.d_step
# NOTE: uses only the last g step for observation
d_losses = [] # list of tuples, each of which gives the loss for a minibatch
# NOTE: for recurrent policies, use shuffle=False?
for ob_batch, ac_batch in dataset.iterbatches((observation, action),
include_final_partial_batch=False,
batch_size=batch_size,
shuffle=True):
# ob_batch, ac_batch, gamma_batch = np.array(batch_buffer['obs']), np.array(
# batch_buffer['acs']), np.array(batch_buffer['gammas'])
gamma_batch = np.ones((ob_batch.shape[0]))
ob_expert, ac_expert = self.expert_dataset.get_next_batch()
gamma_expert = np.ones((ob_expert.shape[0]))
# ob_expert, ac_expert, gamma_expert = np.concatenate(self.expert_dataset.ep_obs),\
# np.concatenate(self.expert_dataset.ep_acs),\
# np.concatenate(self.expert_dataset.ep_gammas)
# update running mean/std for reward_giver
if self.reward_giver.normalize:
self.reward_giver.obs_rms.update(np.concatenate((ob_batch, ob_expert), 0))
# Reshape actions if needed when using discrete actions
if isinstance(self.action_space, gym.spaces.Discrete):
if len(ac_batch.shape) == 2:
ac_batch = ac_batch[:, 0]
if len(ac_expert.shape) == 2:
ac_expert = ac_expert[:, 0]
ob_reg_expert, ac_reg_expert = np.array(ob_expert), np.array(ac_expert)
# while True:
# if ob_reg_expert.shape[0] == ob_batch.shape[0] and ac_reg_expert.shape[0] == \
# ac_batch.shape[0]:
# break
# ob_reg_expert, ac_reg_expert = self.expert_dataset.get_next_batch()
# ob_reg_expert, ac_reg_expert = np.array(ob_reg_expert), np.array(ac_reg_expert)
alpha = np.random.uniform(0.0, 1.0, size=(ob_reg_expert.shape[0], 1))
ob_mix_batch = alpha * ob_batch[:ob_reg_expert.shape[0]] + (1 - alpha) * ob_reg_expert
ac_mix_batch = alpha * ac_batch[:ac_reg_expert.shape[0]] + (1 - alpha) * ac_reg_expert
with self.sess.as_default():
# self.reward_giver.train(ob_batch, ac_batch, np.expand_dims(gamma_batch, axis=1),
# ob_expert, ac_expert, np.expand_dims(gamma_expert, axis=1))
*newlosses, grad = self.reward_giver.lossandgrad(
ob_batch, ac_batch, np.expand_dims(gamma_batch, axis=1),
ob_expert, ac_expert, np.expand_dims(gamma_expert, axis=1),
ob_mix_batch, ac_mix_batch)
self.d_adam.update(self.allmean(grad), self.d_stepsize)
else:
# assert len(observation) == self.timesteps_per_batch
# Comment out if you want only the latest rewards:
obs_batch, acs_batch, gammas_batch = seg['obs_batch'], seg['acs_batch'], seg['gammas_batch']
batch_successor_features = seg['successor_features_batch']
if self.reward_giver.normalize:
ob_reg_batch, ac_reg_batch = observation, action
ob_expert, _ = self.expert_dataset.get_next_batch()
self.reward_giver.obs_rms.update(np.concatenate((ob_reg_batch, ob_expert), 0))
# self.reward_giver.obs_rms.update(
# np.array(batch_successor_features)[:, :self.observation_space.shape[0]])
for idx, (ob_batch, ac_batch, gamma_batch) in enumerate(
zip(obs_batch, acs_batch, gammas_batch)):
rand_traj = np.random.randint(self.expert_dataset.num_traj)
ob_expert, ac_expert, gamma_expert = self.expert_dataset.ep_obs[rand_traj], \
self.expert_dataset.ep_acs[rand_traj], \
self.expert_dataset.ep_gammas[rand_traj]
ob_batch, ac_batch, gamma_batch = np.array(ob_batch), np.array(ac_batch), np.array(
gamma_batch)
while True:
ob_reg_expert, ac_reg_expert = self.expert_dataset.get_next_batch()
ob_reg_expert, ac_reg_expert = np.array(ob_reg_expert), np.array(ac_reg_expert)
if ob_reg_expert.shape[0] == ob_reg_batch.shape[0] and ac_reg_expert.shape[0] == \
ac_reg_batch.shape[0]:
break
alpha = np.random.uniform(0.0, 1.0, size=(ob_reg_batch.shape[0], 1))
ob_mix_batch = alpha * ob_reg_batch + (1 - alpha) * ob_reg_expert
ac_mix_batch = alpha * ac_reg_batch + (1 - alpha) * ac_reg_expert
with self.sess.as_default():
*newlosses, grad = self.reward_giver.lossandgrad(
ob_batch, ac_batch, np.expand_dims(gamma_batch, axis=1),
ob_expert, ac_expert, np.expand_dims(gamma_expert, axis=1),
ob_mix_batch, ac_mix_batch)
self.d_adam.update(self.allmean(grad), self.d_stepsize)
# self.reward_giver.train(ob_batch, ac_batch, np.expand_dims(gamma_batch, axis=1),
# ob_expert, ac_expert,
# np.expand_dims(gamma_expert, axis=1),
# ob_mix_batch, ac_mix_batch)
if self.using_gail or self.using_mdal:
# lr: lengths and rewards
lr_local = (seg["ep_lens"], seg["ep_rets"], seg["ep_true_rets"]) # local values
list_lr_pairs = MPI.COMM_WORLD.allgather(lr_local) # list of tuples
lens, rews, true_rets = map(flatten_lists, zip(*list_lr_pairs))
true_reward_buffer.extend(true_rets)
else:
lr_local = (seg["ep_lens"], seg["ep_rets"]) # local values
list_lr_pairs = MPI.COMM_WORLD.allgather(lr_local) # list of tuples
lens, rews = map(flatten_lists, zip(*list_lr_pairs))
len_buffer.extend(lens)
reward_buffer.extend(rews)
if len(len_buffer) > 0:
if self.using_gail or self.using_mdal:
logger.record_tabular("EpTrueRewMean", np.mean(true_reward_buffer))
logger.record_tabular("EpRewMean", np.mean(reward_buffer))
logger.record_tabular("EpLenMean", np.mean(len_buffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
current_it_timesteps = MPI.COMM_WORLD.allreduce(seg["total_timestep"])
timesteps_so_far += current_it_timesteps
self.num_timesteps += current_it_timesteps
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", self.num_timesteps)
logger.record_tabular("TimeElapsed", time.time() - t_start)
logger.record_tabular("Tsallis-q", self.tsallis_q)
logger.record_tabular("steps", self.num_timesteps)
logger.record_tabular("seed", self.seed)
if self.verbose >= 1 and self.rank == 0:
logger.dump_tabular()
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None, log_interval=1, tb_log_name="PPO2"):
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
self.cliprange = get_schedule_fn(self.cliprange)
with TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name) as writer:
self._setup_learn()
runner = Runner(env=self.env, model=self, n_steps=self.n_steps, gamma=self.gamma, lam=self.lam)
self.episode_reward = np.zeros((self.n_envs,))
ep_info_buf = deque(maxlen=100)
t_first_start = time.time()
n_timesteps = 0
# nupdates = total_timesteps // self.n_batch
for timestep in range(1, total_timesteps + 1):
assert self.n_batch % self.nminibatches == 0
batch_size = self.n_batch // self.nminibatches
t_start = time.time()
frac = 1.0 - timestep / total_timesteps
lr_now = self.learning_rate(frac)
cliprangenow = self.cliprange(frac)
# true_reward is the reward without discount
obs, returns, masks, actions, values, neglogpacs, states, ep_infos, true_reward = runner.run()
n_timesteps += len(obs)
ep_info_buf.extend(ep_infos)
mb_loss_vals = []
if states is None: # nonrecurrent version
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 = ((update * self.noptepochs * self.n_batch + 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, cliprangenow, *slices, writer=writer,
update=n_timesteps))
else: # recurrent version
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 stan_timestepsrt in range(0, self.n_envs, envs_per_batch):
# timestep = ((update * self.noptepochs * self.n_envs + 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, cliprangenow, *slices, update=n_timesteps,
writer=writer, states=mb_states))
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:
self.episode_reward = 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, n_timesteps)
if self.verbose >= 1 and (timestep % log_interval == 0 or timestep == 1):
explained_var = explained_variance(values, returns)
logger.logkv("total_timesteps", n_timesteps)
logger.logkv("fps", fps)
logger.logkv("explained_variance", float(explained_var))
logger.logkv('ep_rewmean', safe_mean([ep_info['r'] for ep_info in ep_info_buf]))
logger.logkv('eplenmean', safe_mean([ep_info['l'] for ep_info in 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()
if callback is not None:
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback(locals(), globals()) is False:
break
if n_timesteps > total_timesteps:
break
return self
def learn(self, total_timesteps, callback=None, log_interval=4, tb_log_name="BDQ",
reset_num_timesteps=True, replay_wrapper=None):
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()
# Create the replay buffer
if self.prioritized_replay:
self.replay_buffer = PrioritizedReplayBuffer(self.buffer_size, alpha=self.prioritized_replay_alpha)
if self.prioritized_replay_beta_iters is None:
prioritized_replay_beta_iters = total_timesteps
else:
prioritized_replay_beta_iters = self.prioritized_replay_beta_iters
self.beta_schedule = LinearSchedule(prioritized_replay_beta_iters,
initial_p=self.prioritized_replay_beta0,
final_p=1.0)
else:
self.replay_buffer = ReplayBuffer(self.buffer_size)
self.beta_schedule = None
if replay_wrapper is not None:
assert not self.prioritized_replay, "Prioritized replay buffer is not supported by HER"
self.replay_buffer = replay_wrapper(self.replay_buffer)
if self.epsilon_greedy:
approximate_num_iters = 2e6 / 4
# TODO Decide which schedule type to use
# self.exploration = PiecewiseSchedule([(0, 1.0),
# (approximate_num_iters / 50, 0.1),
# (approximate_num_iters / 5, 0.01)
# ], outside_value=0.01)
self.exploration = LinearSchedule(schedule_timesteps=int(self.exploration_fraction * total_timesteps),
initial_p=self.exploration_initial_eps,
final_p=self.exploration_final_eps)
else:
self.exploration = ConstantSchedule(value=0.0) # greedy policy
std_schedule = LinearSchedule(schedule_timesteps=self.timesteps_std,
initial_p=self.initial_std,
final_p=self.final_std)
episode_rewards = [0.0]
episode_successes = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
obs = self.env.reset()
reset = True
self.episode_reward = np.zeros((1,))
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs().squeeze()
for _ in range(total_timesteps):
# Take action and update exploration to the newest value
kwargs = {}
if not self.param_noise:
update_eps = self.exploration.value(self.num_timesteps)
update_param_noise_threshold = 0.
else:
update_eps = 0.
# Compute the threshold such that the KL divergence between perturbed and non-perturbed
# policy is comparable to eps-greedy exploration with eps = exploration.value(t).
# See Appendix C.1 in Parameter Space Noise for Exploration, Plappert et al., 2017
# for detailed explanation.
update_param_noise_threshold = \
-np.log(1. - self.exploration.value(self.num_timesteps) +
self.exploration.value(self.num_timesteps) / float(self.env.action_space.n))
kwargs['reset'] = reset
kwargs['update_param_noise_threshold'] = update_param_noise_threshold
kwargs['update_param_noise_scale'] = True
with self.sess.as_default():
# action = self.act(np.array(obs)[None], update_eps=update_eps, **kwargs)[0]
# print("time step {} and update eps {}".format(self.num_timesteps, update_eps))
action_idxes = np.array(self.act(np.array(obs)[None], update_eps=update_eps, **kwargs)) #update_eps=exploration.value(t)))
action = action_idxes / self.num_action_grains * self.actions_range + self.low
if not self.epsilon_greedy: # Gaussian noise
actions_greedy = action
action_idx_stoch = []
action = []
for index in range(len(actions_greedy)):
a_greedy = actions_greedy[index]
out_of_range_action = True
while out_of_range_action:
# Sample from a Gaussian with mean at the greedy action and a std following a schedule of choice
a_stoch = np.random.normal(loc=a_greedy, scale=std_schedule.value(self.num_timesteps))
# Convert sampled cont action to an action idx
a_idx_stoch = np.rint((a_stoch + self.high[index]) / self.actions_range[index] * self.num_action_grains)
# Check if action is in range
if a_idx_stoch >= 0 and a_idx_stoch < self.num_actions_pad:
action_idx_stoch.append(a_idx_stoch)
action.append(a_stoch)
out_of_range_action = False
action_idxes = action_idx_stoch
env_action = action
reset = False
new_obs, rew, done, info = self.env.step(env_action)
self.num_timesteps += 1
# Stop training if return value is False
if callback.on_step() is False:
break
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs().squeeze()
reward_ = self._vec_normalize_env.get_original_reward().squeeze()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, rew
# Store transition in the replay buffer.
self.replay_buffer.add(obs_, action_idxes, reward_, new_obs_, float(done))
obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
if writer is not None:
ep_rew = np.array([reward_]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(self.episode_reward, ep_rew, ep_done, writer,
self.num_timesteps)
# self.episode_reward = total_episode_reward_logger(self.episode_reward, ep_rew, ep_done, writer,
# self.num_timesteps)
# episode_rewards[-1] += rew
episode_rewards[-1] += reward_
if done:
# print("ep number", len(episode_rewards))
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
reset = True
# Do not train if the warmup phase is not over
# or if there are not enough samples in the replay buffer
can_sample = self.replay_buffer.can_sample(self.batch_size)
if can_sample and self.num_timesteps > self.learning_starts \
and self.num_timesteps % self.train_freq == 0:
callback.on_rollout_end()
# Minimize the error in Bellman's equation on a batch sampled from replay buffer.
# pytype:disable=bad-unpacking
if self.prioritized_replay:
assert self.beta_schedule is not None, \
"BUG: should be LinearSchedule when self.prioritized_replay True"
experience = self.replay_buffer.sample(self.batch_size,
beta=self.beta_schedule.value(self.num_timesteps))
(obses_t, actions, rewards, obses_tp1, dones, weights, batch_idxes) = experience
else:
obses_t, actions, rewards, obses_tp1, dones = self.replay_buffer.sample(self.batch_size)
weights, batch_idxes = np.ones_like(rewards), None
# pytype:enable=bad-unpacking
if writer is not None:
# run loss backprop with summary, but once every 100 steps save the metadata
# (memory, compute time, ...)
if (1 + self.num_timesteps) % 100 == 0:
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
summary, td_errors, mean_loss = self._train_step(obses_t, actions, rewards, obses_tp1, obses_tp1,
dones, weights, sess=self.sess, options=run_options,
run_metadata=run_metadata)
writer.add_run_metadata(run_metadata, 'step%d' % self.num_timesteps)
else:
summary, td_errors, mean_loss = self._train_step(obses_t, actions, rewards, obses_tp1, obses_tp1,
dones, weights, sess=self.sess)
writer.add_summary(summary, self.num_timesteps)
else:
_, td_errors, mean_loss = self._train_step(obses_t, actions, rewards, obses_tp1, obses_tp1, dones, weights,
sess=self.sess)
if self.prioritized_replay:
new_priorities = np.abs(td_errors) + self.prioritized_replay_eps
assert isinstance(self.replay_buffer, PrioritizedReplayBuffer)
self.replay_buffer.update_priorities(batch_idxes, new_priorities)
callback.on_rollout_start()
if can_sample and self.num_timesteps > self.learning_starts and \
self.num_timesteps % self.target_network_update_freq == 0:
# Update target network periodically.
self.update_target(sess=self.sess)
if len(episode_rewards[-101:-1]) == 0:
mean_100ep_reward = -np.inf
else:
mean_100ep_reward = round(float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
# Log training infos
kvs = {}
if self.verbose >= 1 and done and log_interval is not None \
and len(episode_rewards) % log_interval == 0 \
and self.num_timesteps > self.train_freq \
and self.num_timesteps > self.learning_starts:
if self.log_dir is not None:
kvs["episodes"] = num_episodes
kvs["mean_100rew"] = mean_100ep_reward
kvs["current_lr"] = self.learning_rate
kvs["success_rate"] = np.mean(episode_successes[-100:])
kvs["total_timesteps"] = self.num_timesteps
kvs["mean_loss"] = mean_loss
kvs["mean_td_errors"] = np.mean(td_errors)
kvs["time_spent_exploring"] = int(100 * self.exploration.value(self.num_timesteps))
self.log_csv.writekvs(kvs)
logger.record_tabular("steps", self.num_timesteps)
logger.record_tabular("episodes", num_episodes)
if len(episode_successes) > 0:
logger.logkv("success rate", np.mean(episode_successes[-100:]))
logger.record_tabular("mean 100 episode reward", mean_100ep_reward)
logger.record_tabular("% time spent exploring",
int(100 * self.exploration.value(self.num_timesteps)))
logger.dump_tabular()
callback.on_training_end()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=4,
tb_log_name="SAC",
reset_num_timesteps=True,
replay_wrapper=None):
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
if replay_wrapper is not None:
self.replay_buffer = replay_wrapper(self.replay_buffer)
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
self._setup_learn()
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
# Initial learning rate
current_lr = self.learning_rate(1)
start_time = time.time()
episode_rewards = [0.0]
episode_successes = []
if self.action_noise is not None:
self.action_noise.reset()
obs = self.env.reset()
#obs = tf.image.rgb_to_grayscale(obs)
#obs = obs.eval(session=tf.compat.v1.Session())
self.episode_reward = np.zeros((1, ))
ep_info_buf = deque(maxlen=100)
n_updates = 0
infos_values = []
prev_state = 0
for step in range(total_timesteps):
time.sleep(0.05)
if callback is not None:
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback(locals(), globals()) is False:
break
# Before training starts, randomly sample actions
# from a uniform distribution for better exploration.
# Afterwards, use the learned policy
# if random_exploration is set to 0 (normal setting)
print('num_timesteps:' + str(self.num_timesteps))
print('learning_starts' + str(self.learning_starts))
if self.num_timesteps < self.learning_starts or np.random.rand(
) < self.random_exploration:
# actions sampled from action space are from range specific to the environment
# but algorithm operates on tanh-squashed actions therefore simple scaling is used
unscaled_action = self.env.action_space.sample()
action = scale_action(self.action_space, unscaled_action)
else:
#time.sleep(1)
#continue
action = self.policy_tf.step(
obs[None], deterministic=False).flatten()
# Add noise to the action (improve exploration,
# not needed in general)
if self.action_noise is not None:
action = np.clip(action + self.action_noise(), -1, 1)
# inferred actions need to be transformed to environment action_space before stepping
unscaled_action = unscale_action(self.action_space, action)
assert action.shape == self.env.action_space.shape
new_obs, reward, done, info = self.env.step(unscaled_action)
print('got reward %s' % str(reward))
if done and reward >= 0:
print('ERROR')
reward = -10
if done:
self.env.reset()
# Store transition in the replay buffer.
self.replay_buffer.add(obs, action, reward, new_obs,
float(done))
obs = new_obs
# Retrieve reward and episode length if using Monitor wrapper
maybe_ep_info = info.get('episode')
if maybe_ep_info is not None:
ep_info_buf.extend([maybe_ep_info])
print('WRITER OUT %s' % str(writer))
if writer is not None:
print('WRITER')
# Write reward per episode to tensorboard
ep_reward = np.array([reward]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
self.episode_reward = total_episode_reward_logger(
self.episode_reward, ep_reward, ep_done, writer,
self.num_timesteps)
if True and step % self.train_freq == 0 and done:
mb_infos_vals = []
print('SAC: Update network...')
# Update policy, critics and target networks
for grad_step in range(self.gradient_steps):
# Break if the warmup phase is not over
# or if there are not enough samples in the replay buffer
if not self.replay_buffer.can_sample(self.batch_size) \
or self.num_timesteps < self.learning_starts:
break
n_updates += 1
# Compute current learning_rate
frac = 1.0 - step / total_timesteps
current_lr = self.learning_rate(frac)
# Update policy and critics (q functions)
mb_infos_vals.append(
self._train_step(step, writer, current_lr))
#self.target_update_interval = 4
# Update target network
if (step +
grad_step) % self.target_update_interval == 0:
#print('SAC: Update network...')
#print('SAC self.target_update_op%s'%str(self.target_update_op))
# Update target network
self.sess.run(self.target_update_op)
# Log losses and entropy, useful for monitor training
if len(mb_infos_vals) > 0:
infos_values = np.mean(mb_infos_vals, axis=0)
print('TEST4')
#episode_rewards[-1] += reward
episode_rewards[-1] = episode_rewards[-1] + reward
if done:
if self.action_noise is not None:
self.action_noise.reset()
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if len(episode_rewards[-101:-1]) == 0:
mean_reward = -np.inf
else:
mean_reward = round(
float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
self.num_timesteps += 1
# Display training infos
if self.verbose >= 1 and done and log_interval is not None and len(
episode_rewards) % log_interval == 0:
fps = int(step / (time.time() - start_time))
logger.logkv("episodes", num_episodes)
logger.logkv("mean 100 episode reward", mean_reward)
if len(ep_info_buf) > 0 and len(ep_info_buf[0]) > 0:
logger.logkv(
'ep_rewmean',
safe_mean(
[ep_info['r'] for ep_info in ep_info_buf]))
logger.logkv(
'eplenmean',
safe_mean(
[ep_info['l'] for ep_info in ep_info_buf]))
logger.logkv("n_updates", n_updates)
logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - start_time))
if len(episode_successes) > 0:
logger.logkv("success rate",
np.mean(episode_successes[-100:]))
if len(infos_values) > 0:
for (name, val) in zip(self.infos_names, infos_values):
logger.logkv(name, val)
logger.logkv("total timesteps", self.num_timesteps)
logger.dumpkvs()
# Reset infos:
infos_values = []
# check if game is over
if True and done:
while self.env.is_game_over():
print('waiting for control')
time.sleep(1)
return self
def learn(self, total_timesteps, callback=None,
log_interval=1, tb_log_name="SAC", print_freq=100,vae=None):
self.learning_rate = get_schedule_fn(self.learning_rate)
callback = self._init_callback(callback)
with TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name) as writer:
self._setup_learn()
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
start_time = time.time()
episode_rewards = [0.0]
obs = self.env.reset()
self.episode_reward = np.zeros((1,))
ep_info_buf = deque(maxlen=100)
ep_len = 0
self.n_updates = 0
infos_values = []
mb_infos_vals = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
for step in range(total_timesteps):
# Compute current learning_rate
frac = 1.0 - step / total_timesteps
current_lr = self.learning_rate(frac)
if callback is not None:
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback.on_step() is False:
break
# Before training starts, randomly sample actions
# from a uniform distribution for better exploration.
# Afterwards, use the learned policy.
if step < self.learning_starts:
action = self.env.action_space.sample()
# No need to rescale when sampling random action
rescaled_action = action
else:
action = self.policy_tf.step(obs[None], deterministic=False).flatten()
# Rescale from [-1, 1] to the correct bounds
rescaled_action = action * np.abs(self.action_space.low)
assert action.shape == self.env.action_space.shape
new_obs, reward, done, info = self.env.step(rescaled_action)
ep_len += 1
callback.update_locals(locals())
if callback.on_step() is False:
break
##################
arr = vae.decode(new_obs[:, :512].reshape(1, 512))
arr = np.round(arr).astype(np.uint8)
arr = arr.reshape(80, 160, 3)
#to visualize what car sees
#cv2.imwrite("decoded_img.png", arr)
###############3
if print_freq > 0 and ep_len % print_freq == 0 and ep_len > 0:
print("{} steps".format(ep_len))
# Store transition in the replay buffer.
self.replay_buffer.add(obs, action, reward, new_obs, float(done))
obs = new_obs
# Retrieve reward and episode length if using Monitor wrapper
maybe_ep_info = info.get('episode')
if maybe_ep_info is not None:
ep_info_buf.extend([maybe_ep_info])
if writer is not None:
# Write reward per episode to tensorboard
ep_reward = np.array([reward]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
self.episode_reward = total_episode_reward_logger(self.episode_reward, ep_reward,
ep_done, writer, step)
if ep_len > self.train_freq:
print("Additional training")
self.env.reset()
mb_infos_vals = self.optimize(step, writer, current_lr)
done = True
episode_rewards[-1] += reward
if done:
obs = self.env.reset()
print("Episode finished. Reward: {:.2f} {} Steps".format(episode_rewards[-1], ep_len))
episode_rewards.append(0.0)
ep_len = 0
mb_infos_vals = self.optimize(step, writer, current_lr)
# train VAE
train_start = time.time()
#training VAE with SAC
#vae.optimize()
print("VAE training duration:", time.time() - train_start)
obs = self.env.reset()
callback.on_rollout_end()
# Log losses and entropy, useful for monitor training
if len(mb_infos_vals) > 0:
infos_values = np.mean(mb_infos_vals, axis=0)
if len(episode_rewards[-101:-1]) == 0:
mean_reward = -np.inf
else:
mean_reward = round(float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
callback.on_rollout_start()
if self.verbose >= 1 and done and log_interval is not None and len(episode_rewards) % log_interval == 0:
fps = int(step / (time.time() - start_time))
logger.logkv("episodes", num_episodes)
logger.logkv("mean 100 episode reward", mean_reward)
logger.logkv('ep_rewmean', safe_mean([ep_info['r'] for ep_info in ep_info_buf]))
logger.logkv('eplenmean', safe_mean([ep_info['l'] for ep_info in ep_info_buf]))
logger.logkv("n_updates", self.n_updates)
logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', "{:.2f}".format(time.time() - start_time))
if len(infos_values) > 0:
for (name, val) in zip(self.infos_names, infos_values):
logger.logkv(name, val)
logger.logkv("total timesteps", step)
logger.dumpkvs()
# Reset infos:
infos_values = []
# Use last batch
print("Final optimization before saving")
self.env.reset()
mb_infos_vals = self.optimize(step, writer, current_lr)
callback.on_training_end()
return self
def learn(self, total_timesteps, log_dir, logger,
callback=None, log_interval=1, tb_log_name="PPO2",
reset_num_timesteps=True):
# 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)
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()
runner = Runner(env=self.env, model=self, n_steps=self.n_steps, gamma=self.gamma, lam=self.lam)
self.episode_reward = np.zeros((self.n_envs,))
ep_info_buf = deque(maxlen=100)
n_updates = total_timesteps // self.n_batch
for update in range(1, n_updates + 1):
# Do the following except keyboard interrupt the learning process.
try:
assert self.n_batch % self.nminibatches == 0
batch_size = self.n_batch // self.nminibatches
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)
t_start = time.time()
# Unpack
obs, returns, masks, actions, values, neglogpacs, states, ep_infos, true_reward = runner.run()
# # add by Yunlong
t_now = time.time()
fps = int(self.n_batch / (t_now - t_start))
self.ep_info_buf.extend(ep_infos)
mb_loss_vals = []
if states is None: # nonrecurrent version
update_fac = 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 + ((self.noptepochs * self.n_batch + 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 = 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 + ((self.noptepochs * self.n_envs + 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)
# # comment out by Yunlong
# 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)
logger.logkv('true_reward', np.mean(true_reward))
for (loss_val, loss_name) in zip(loss_vals, self.loss_names):
logger.logkv(loss_name, loss_val)
logger.dumpkvs()
if callback is not None:
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback(locals(), globals()) is False:
break
except KeyboardInterrupt:
print("You have stopped the learning process by keyboard interrupt. Model Parameter is saved. \n")
# You can actually save files using the instance of self. save the model parameters.
self.save(log_dir + "_Iteration_{}".format(update))
sys.exit()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=4,
tb_log_name="SAC",
reset_num_timesteps=True,
replay_wrapper=None,
planning_steps=0):
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
callback = self._init_callback(callback)
# TODO: use builtin log writer instead of this old lib
tb_configure(self.tensorboard_log)
action_log_csv = self.tensorboard_log + "_actions.csv"
action_log_df = pd.DataFrame(columns=np.concatenate((
["iteration"],
["p" + str(i) for i in range(24)],
["b" + str(i) for i in range(24)],
["e" + str(i) for i in range(24)],
)))
action_log_index = 0
steps_in_real_env = 0
person_data_dict = {}
if replay_wrapper is not None:
self.replay_buffer = replay_wrapper(self.replay_buffer)
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
self._setup_learn()
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
# Initial learning rate
current_lr = self.learning_rate(1)
start_time = time.time()
episode_rewards = [0.0]
episode_successes = []
if self.action_noise is not None:
self.action_noise.reset()
obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs().squeeze()
n_updates = 0
infos_values = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
for step in range(total_timesteps):
# Before training starts, randomly sample actions
# from a uniform distribution for better exploration.
# Afterwards, use the learned policy
# if random_exploration is set to 0 (normal setting)
if self.num_timesteps < self.learning_starts or np.random.rand(
) < self.random_exploration:
# actions sampled from action space are from range specific to the environment
# but algorithm operates on tanh-squashed actions therefore simple scaling is used
unscaled_action = self.env.action_space.sample()
action = scale_action(self.action_space, unscaled_action)
else:
action = self.policy_tf.step(
obs[None], deterministic=False).flatten()
# Add noise to the action (improve exploration,
# not needed in general)
if self.action_noise is not None:
action = np.clip(action + self.action_noise(), -1, 1)
# inferred actions need to be transformed to environment action_space before stepping
unscaled_action = unscale_action(self.action_space, action)
assert action.shape == self.env.action_space.shape
# if not planning:
# new_obs, reward, done, info = self.env.step(unscaled_action)
# else:
if not self.num_timesteps % (planning_steps + 1):
## TODO: work on this?
# if self.num_timesteps ==1:
# # form the control
# from sklearn.preprocessing import MinMaxScaler
# grid_price = self.non_vec_env.prices[self.non_vec_env.day - 1]
# scaler = MinMaxScaler(feature_range = (0, 10))
# scaled_grid_price = scaler.fit_transform(np.array(grid_price).reshape(-1, 1))
# scaled_grid_price = np.squeeze(scaled_grid_price)
# energy_consumptions = self.non_vec_env._simulate_humans(scaled_grid_price)
# person_data_dict["control"] = {
# "x" : list(range(8, 18)),
# "grid_price" : scaled_grid_price,
# "energy_consumption" : energy_consumptions["avg"],
# "reward" : self.non_vec_env._get_reward(price = grid_price, energy_consumptions = energy_consumptions),
# }
# # form the data_dict
# if self.num_timesteps in [100, 1000, 9500]:
# person_data_dict["Step " + str(self.num_timesteps)] = {
# "x" : list(range(8, 18)),
# "grid_price" : self.non_vec_env.prices[self.non_vec_env.day - 1],
# "action" : unscaled_action,
# "energy_consumption" : self.non_vec_env.prev_energy,
# "reward" : reward,
# }
# if self.num_timesteps == 9501 and self.people_reaction_log_dir and self.plotter_person_reaction:
# # call the plotting statement
# self.plotter_person_reaction(person_data_dict, self.people_reaction_log_dir)
new_obs, reward, done, info = self.env.step(
unscaled_action) #, step_num = self.num_timesteps)
steps_in_real_env += 1
else:
print("planning step")
new_obs, reward, done, info = self.non_vec_env.planning_step(
unscaled_action)
# write the action to a csv
# if ((not self.num_timesteps % 10) & (self.num_timesteps > 10000)) or self.num_timesteps>19500:
# ### get the battery charging
# battery_op = {}
# total_battery_consumption = np.zeros(24)
# total_energy_consumption = np.zeros(24)
# for prosumer_name in self.non_vec_env.prosumer_dict:
# #Get players response to agent's actions
# day = self.non_vec_env.day
# price = self.non_vec_env.price
# prosumer = self.non_vec_env.prosumer_dict[prosumer_name]
# prosumer_battery = prosumer.get_battery_operation(day, price)
# prosumer_demand = prosumer.get_response(day, price)
# total_battery_consumption += prosumer_battery
# total_energy_consumption += prosumer_demand
# action_log_df.loc[action_log_index] = np.concatenate(
# ([self.num_timesteps],
# price,
# total_battery_consumption,
# total_energy_consumption,))
# action_log_index += 1
# action_log_df.to_csv(action_log_csv)
# print("Iteration: " + str(self.num_timesteps))
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
callback.update_locals(locals())
if callback.on_step() is False:
break
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs(
).squeeze()
reward_ = self._vec_normalize_env.get_original_reward(
).squeeze()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, reward
if not self.num_timesteps % (planning_steps + 1):
tb_log_value("reward_in_environment", reward_,
steps_in_real_env)
# tb_log_value("reward_planning", reward_, self.num_timesteps)
self.num_timesteps += 1
# Store transition in the replay buffer.
self.replay_buffer_add(obs_, action, reward_, new_obs_, done,
info)
obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
# Retrieve reward and episode length if using Monitor wrapper
maybe_ep_info = info.get('episode')
if maybe_ep_info is not None:
self.ep_info_buf.extend([maybe_ep_info])
if writer is not None:
# Write reward per episode to tensorboard
ep_reward = np.array([reward_]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(
self.episode_reward, ep_reward, ep_done, writer,
self.num_timesteps)
if self.num_timesteps % 100 == 0 and not np.any(
unscaled_action == np.inf):
if self.action_to_prices_fn:
prices = self.action_to_prices_fn(unscaled_action)
# tf_util.log_histogram(writer, "action_vec_hist", unscaled_action, self.num_timesteps, bins=10, flush=False)
# tb_log_value("constant_load_price", np.sum(prices), self.num_timesteps)
# tf_util.log_vec_as_histogram(writer, "prices", prices, self.num_timesteps, flush=True)
if self.num_timesteps % self.train_freq == 0:
callback.on_rollout_end()
mb_infos_vals = []
# Update policy, critics and target networks
for grad_step in range(self.gradient_steps):
# Break if the warmup phase is not over
# or if there are not enough samples in the replay buffer
if not self.replay_buffer.can_sample(self.batch_size) \
or self.num_timesteps < self.learning_starts:
break
n_updates += 1
# Compute current learning_rate
frac = 1.0 - step / total_timesteps
current_lr = self.learning_rate(frac)
# Update policy and critics (q functions)
mb_infos_vals.append(
self._train_step(step, writer, current_lr))
# Update target network
if (step +
grad_step) % self.target_update_interval == 0:
# Update target network
self.sess.run(self.target_update_op)
# Log losses and entropy, useful for monitor training
if len(mb_infos_vals) > 0:
infos_values = np.mean(mb_infos_vals, axis=0)
callback.on_rollout_start()
episode_rewards[-1] += reward_
if done:
if self.action_noise is not None:
self.action_noise.reset()
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if len(episode_rewards[-101:-1]) == 0:
mean_reward = -np.inf
else:
mean_reward = round(
float(np.mean(episode_rewards[-101:-1])), 1)
# substract 1 as we appended a new term just now
num_episodes = len(episode_rewards) - 1
# Display training infos
if self.verbose >= 1 and done and log_interval is not None and num_episodes % log_interval == 0:
fps = int(step / (time.time() - start_time))
logger.logkv("episodes", num_episodes)
logger.logkv("mean 100 episode reward", mean_reward)
if len(self.ep_info_buf) > 0 and len(
self.ep_info_buf[0]) > 0:
logger.logkv(
'ep_rewmean',
safe_mean([
ep_info['r'] for ep_info in self.ep_info_buf
]))
logger.logkv(
'eplenmean',
safe_mean([
ep_info['l'] for ep_info in self.ep_info_buf
]))
logger.logkv("n_updates", n_updates)
logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - start_time))
if len(episode_successes) > 0:
logger.logkv("success rate",
np.mean(episode_successes[-100:]))
if len(infos_values) > 0:
for (name, val) in zip(self.infos_names, infos_values):
logger.logkv(name, val)
logger.logkv("total timesteps", self.num_timesteps)
logger.dumpkvs()
# Reset infos:
infos_values = []
callback.on_training_end()
return self #, ep_reward #, reward_
def learn(
self,
total_timesteps,
model_coworker,
role,
callback=None,
log_interval=100,
tb_log_name="DQN",
reset_num_timesteps=True,
replay_wrapper=None,
clipping_during_training=True,
):
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()
# Create the replay buffer
if self.prioritized_replay:
self.replay_buffer = PrioritizedReplayBuffer(
self.buffer_size, alpha=self.prioritized_replay_alpha)
if self.prioritized_replay_beta_iters is None:
prioritized_replay_beta_iters = total_timesteps
else:
prioritized_replay_beta_iters = self.prioritized_replay_beta_iters
self.beta_schedule = LinearSchedule(
prioritized_replay_beta_iters,
initial_p=self.prioritized_replay_beta0,
final_p=1.0,
)
else:
if self.replay_buffer is None:
self.replay_buffer = ReplayBuffer(self.buffer_size)
self.beta_schedule = None
if replay_wrapper is not None:
assert (not self.prioritized_replay
), "Prioritized replay buffer is not supported by HER"
self.replay_buffer = replay_wrapper(self.replay_buffer)
# Create the schedule for exploration starting from 1.
self.exploration = LinearSchedule(
schedule_timesteps=int(self.exploration_fraction *
total_timesteps),
initial_p=self.exploration_initial_eps,
final_p=self.exploration_final_eps,
)
episode_rewards = [0.0]
episode_successes = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
reset = True
obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs().squeeze()
for _ in range(total_timesteps):
# Take action and update exploration to the newest value
kwargs = {}
if not self.param_noise:
update_eps = self.exploration.value(self.num_timesteps)
update_param_noise_threshold = 0.0
else:
update_eps = 0.0
# Compute the threshold such that the KL divergence between perturbed and non-perturbed
# policy is comparable to eps-greedy exploration with eps = exploration.value(t).
# See Appendix C.1 in Parameter Space Noise for Exploration, Plappert et al., 2017
# for detailed explanation.
update_param_noise_threshold = -np.log(
1.0 - self.exploration.value(self.num_timesteps) +
self.exploration.value(self.num_timesteps) /
float(self.env.action_space.n))
kwargs["reset"] = reset
kwargs[
"update_param_noise_threshold"] = update_param_noise_threshold
kwargs["update_param_noise_scale"] = True
with self.sess.as_default():
action = self.act(np.array(obs)[None],
update_eps=update_eps,
**kwargs)[0]
turn, speed = None, None
if role == "turn":
turn = action
speed, nothing = model_coworker.predict(np.array(obs))
else:
turn, nothing = model_coworker.predict(np.array(obs))
speed = action
if clipping_during_training:
# check if next state (after action) would be outside of fish tank (CLIPPING)
env_state = self.env.get_state()
turn_speed = self.env.action([turn, speed])
global_turn = env_state[0][2] + turn_speed[0]
coords = np.array([
env_state[0][0] + turn_speed[1] * np.cos(global_turn),
env_state[0][1] + turn_speed[1] * np.sin(global_turn),
])
changed = False
if coords[0] < -0.49:
coords[0] = -0.47
changed = True
elif coords[0] > 0.49:
coords[0] = 0.47
changed = True
if coords[1] < -0.49:
coords[1] = -0.47
changed = True
elif coords[1] > 0.49:
coords[1] = 0.47
changed = True
if changed:
diff = coords - env_state[0, :2]
speed = np.linalg.norm(diff)
angles = np.arctan2(diff[1], diff[0])
turn = angles - env_state[0, 2]
turn = turn - 2 * np.pi if turn > np.pi else turn
turn = turn + 2 * np.pi if turn < -np.pi else turn
# convert to DQN output
dist_turn = np.abs(self.env.turn_rate_bins - turn)
dist_speed = np.abs(self.env.speed_bins - speed)
# convert to bins
turn = np.argmin(dist_turn, axis=0)
speed = np.argmin(dist_speed, axis=0)
if role == "turn":
action = turn
else:
action = speed
reset = False
new_obs, rew, done, info = self.env.step([turn, speed])
self.num_timesteps += 1
# Stop training if return value is False
if callback.on_step() is False:
break
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs(
).squeeze()
reward_ = self._vec_normalize_env.get_original_reward(
).squeeze()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, rew
# Store transition in the replay buffer, but change reward to 0 (use it for plot later though)
self.replay_buffer.add(obs_, action, 0, new_obs_, float(done))
# Also give transition to model coworker
if model_coworker.replay_buffer is None:
model_coworker.replay_buffer = ReplayBuffer(
self.buffer_size)
if role == "turn":
model_coworker.replay_buffer.add(obs_, speed, 0, new_obs_,
float(done))
else:
model_coworker.replay_buffer.add(obs_, turn, 0, new_obs_,
float(done))
obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
if writer is not None:
ep_rew = np.array([reward_]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(
self.episode_reward, ep_rew, ep_done, writer,
self.num_timesteps)
episode_rewards[-1] += reward_
if done:
maybe_is_success = info.get("is_success")
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
reset = True
# Do not train if the warmup phase is not over
# or if there are not enough samples in the replay buffer
can_sample = self.replay_buffer.can_sample(self.batch_size)
if (can_sample and self.num_timesteps > self.learning_starts
and self.num_timesteps % self.train_freq == 0):
callback.on_rollout_end()
# Minimize the error in Bellman's equation on a batch sampled from replay buffer.
# pytype:disable=bad-unpacking
if self.prioritized_replay:
assert (
self.beta_schedule is not None
), "BUG: should be LinearSchedule when self.prioritized_replay True"
experience = self.replay_buffer.sample(
self.batch_size,
beta=self.beta_schedule.value(self.num_timesteps),
env=self._vec_normalize_env,
)
(
obses_t,
actions,
rewards,
obses_tp1,
dones,
weights,
batch_idxes,
) = experience
else:
(
obses_t,
actions,
rewards,
obses_tp1,
dones,
) = self.replay_buffer.sample(
self.batch_size, env=self._vec_normalize_env)
# also sample from expert buffer
(
obses_t_exp,
actions_exp,
rewards_exp,
obses_tp1_exp,
dones_exp,
) = self.expert_buffer.sample(
self.batch_size, env=self._vec_normalize_env)
weights, batch_idxes = np.ones_like(rewards), None
weights_exp, batch_idxes_exp = np.ones_like(
rewards_exp), None
# pytype:enable=bad-unpacking
if writer is not None:
# run loss backprop with summary, but once every 100 steps save the metadata
# (memory, compute time, ...)
if (1 + self.num_timesteps) % 100 == 0:
run_options = tf.RunOptions(
trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
summary, td_errors = self._train_step(
np.append(obses_t, obses_t_exp, axis=0),
np.append(actions,
actions_exp.flatten(),
axis=0),
np.append(rewards,
rewards_exp.flatten(),
axis=0),
np.append(obses_tp1, obses_tp1_exp, axis=0),
np.append(obses_tp1, obses_tp1_exp, axis=0),
np.append(dones.flatten(),
dones_exp.flatten(),
axis=0),
np.append(weights, weights_exp),
sess=self.sess,
options=run_options,
run_metadata=run_metadata,
)
writer.add_run_metadata(
run_metadata, "step%d" % self.num_timesteps)
else:
summary, td_errors = self._train_step(
np.append(obses_t, obses_t_exp, axis=0),
np.append(actions,
actions_exp.flatten(),
axis=0),
np.append(rewards,
rewards_exp.flatten(),
axis=0),
np.append(obses_tp1, obses_tp1_exp, axis=0),
np.append(obses_tp1, obses_tp1_exp, axis=0),
np.append(dones.flatten(),
dones_exp.flatten(),
axis=0),
np.append(weights, weights_exp),
sess=self.sess,
options=run_options,
run_metadata=run_metadata,
)
writer.add_summary(summary, self.num_timesteps)
else:
_, td_errors = self._train_step(
np.append(obses_t, obses_t_exp, axis=0),
np.append(actions, actions_exp.flatten(), axis=0),
np.append(rewards, rewards_exp.flatten(), axis=0),
np.append(obses_tp1, obses_tp1_exp, axis=0),
np.append(obses_tp1, obses_tp1_exp, axis=0),
np.append(dones.flatten(),
dones_exp.flatten(),
axis=0),
np.append(weights, weights_exp),
sess=self.sess,
)
if self.prioritized_replay:
new_priorities = np.abs(
td_errors) + self.prioritized_replay_eps
assert isinstance(self.replay_buffer,
PrioritizedReplayBuffer)
self.replay_buffer.update_priorities(
batch_idxes, new_priorities)
callback.on_rollout_start()
if (can_sample and self.num_timesteps > self.learning_starts
and self.num_timesteps %
self.target_network_update_freq == 0):
# Update target network periodically.
self.update_target(sess=self.sess)
if len(episode_rewards[-101:-1]) == 0:
mean_100ep_reward = -np.inf
else:
mean_100ep_reward = round(
float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
if (self.verbose >= 1 and done and log_interval is not None
and len(episode_rewards) % log_interval == 0):
logger.record_tabular("steps", self.num_timesteps)
logger.record_tabular("episodes", num_episodes)
if len(episode_successes) > 0:
logger.logkv("success rate",
np.mean(episode_successes[-100:]))
logger.record_tabular("mean 100 episode reward",
mean_100ep_reward)
logger.record_tabular(
"% time spent exploring",
int(100 * self.exploration.value(self.num_timesteps)),
)
logger.dump_tabular()
callback.on_training_end()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=100,
tb_log_name="PPO1",
reset_num_timesteps=True):
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()
assert issubclass(self.policy, ActorCriticPolicy), "Error: the input policy for the PPO1 model must be " \
"an instance of common.policies.ActorCriticPolicy."
with self.sess.as_default():
self.adam.sync()
callback.on_training_start(locals(), globals())
# Prepare for rollouts
seg_gen = traj_segment_generator(self.policy_pi,
self.env,
self.timesteps_per_actorbatch,
callback=callback)
episodes_so_far = 0
timesteps_so_far = 0
iters_so_far = 0
t_start = time.time()
# rolling buffer for episode lengths
len_buffer = deque(maxlen=100)
# rolling buffer for episode rewards
reward_buffer = deque(maxlen=100)
while True:
if timesteps_so_far >= total_timesteps:
break
if self.schedule == 'constant':
cur_lrmult = 1.0
elif self.schedule == 'linear':
cur_lrmult = max(
1.0 - float(timesteps_so_far) / total_timesteps, 0)
else:
raise NotImplementedError
logger.log("********** Iteration %i ************" %
iters_so_far)
seg = seg_gen.__next__()
# Stop training early (triggered by the callback)
if not seg.get('continue_training', True): # pytype: disable=attribute-error
break
add_vtarg_and_adv(seg, self.gamma, self.lam)
# ob, ac, atarg, ret, td1ret = map(np.concatenate, (obs, acs, atargs, rets, td1rets))
observations, actions = seg["observations"], seg["actions"]
atarg, tdlamret = seg["adv"], seg["tdlamret"]
# true_rew is the reward without discount
if writer is not None:
total_episode_reward_logger(
self.episode_reward, seg["true_rewards"].reshape(
(self.n_envs, -1)), seg["dones"].reshape(
(self.n_envs, -1)), writer,
self.num_timesteps)
# predicted value function before udpate
vpredbefore = seg["vpred"]
# standardized advantage function estimate
atarg = (atarg - atarg.mean()) / atarg.std()
dataset = Dataset(dict(ob=observations,
ac=actions,
atarg=atarg,
vtarg=tdlamret),
shuffle=not self.policy.recurrent)
optim_batchsize = self.optim_batchsize or observations.shape[
0]
# set old parameter values to new parameter values
self.assign_old_eq_new(sess=self.sess)
logger.log("Optimizing...")
logger.log(fmt_row(13, self.loss_names))
# Here we do a bunch of optimization epochs over the data
for k in range(self.optim_epochs):
# list of tuples, each of which gives the loss for a minibatch
losses = []
for i, batch in enumerate(
dataset.iterate_once(optim_batchsize)):
steps = (
self.num_timesteps + k * optim_batchsize +
int(i *
(optim_batchsize / len(dataset.data_map))))
if writer is not None:
# run loss backprop with summary, but once every 10 runs save the metadata
# (memory, compute time, ...)
if self.full_tensorboard_log and (1 +
k) % 10 == 0:
run_options = tf.compat.v1.RunOptions(
trace_level=tf.compat.v1.RunOptions.
FULL_TRACE)
run_metadata = tf.compat.v1.RunMetadata()
summary, grad, *newlosses = self.lossandgrad(
batch["ob"],
batch["ob"],
batch["ac"],
batch["atarg"],
batch["vtarg"],
cur_lrmult,
sess=self.sess,
options=run_options,
run_metadata=run_metadata)
writer.add_run_metadata(
run_metadata, 'step%d' % steps)
else:
summary, grad, *newlosses = self.lossandgrad(
batch["ob"],
batch["ob"],
batch["ac"],
batch["atarg"],
batch["vtarg"],
cur_lrmult,
sess=self.sess)
writer.add_summary(summary, steps)
else:
_, grad, *newlosses = self.lossandgrad(
batch["ob"],
batch["ob"],
batch["ac"],
batch["atarg"],
batch["vtarg"],
cur_lrmult,
sess=self.sess)
self.adam.update(grad,
self.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 dataset.iterate_once(optim_batchsize):
newlosses = self.compute_losses(batch["ob"],
batch["ob"],
batch["ac"],
batch["atarg"],
batch["vtarg"],
cur_lrmult,
sess=self.sess)
losses.append(newlosses)
mean_losses, _, _ = mpi_moments(losses, axis=0)
logger.log(fmt_row(13, mean_losses))
for (loss_val, name) in zipsame(mean_losses,
self.loss_names):
logger.record_tabular("loss_" + name, loss_val)
logger.record_tabular(
"ev_tdlam_before",
explained_variance(vpredbefore, tdlamret))
# local values
lrlocal = (seg["ep_lens"], seg["ep_rets"])
# list of tuples
listoflrpairs = MPI.COMM_WORLD.allgather(lrlocal)
lens, rews = map(flatten_lists, zip(*listoflrpairs))
len_buffer.extend(lens)
reward_buffer.extend(rews)
if len(len_buffer) > 0:
logger.record_tabular("EpLenMean", np.mean(len_buffer))
logger.record_tabular("EpRewMean",
np.mean(reward_buffer))
logger.record_tabular("EpThisIter", len(lens))
episodes_so_far += len(lens)
current_it_timesteps = MPI.COMM_WORLD.allreduce(
seg["total_timestep"])
timesteps_so_far += current_it_timesteps
self.num_timesteps += current_it_timesteps
iters_so_far += 1
logger.record_tabular("EpisodesSoFar", episodes_so_far)
logger.record_tabular("TimestepsSoFar", self.num_timesteps)
logger.record_tabular("TimeElapsed", time.time() - t_start)
if self.verbose >= 1 and MPI.COMM_WORLD.Get_rank() == 0:
logger.dump_tabular()
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None, log_interval=100, tb_log_name="DQN",
reset_num_timesteps=True, replay_wrapper=None):
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()
# Create the replay buffer
if self.prioritized_replay:
self.replay_buffer = PrioritizedReplayBuffer(self.buffer_size, alpha=self.prioritized_replay_alpha)
if self.prioritized_replay_beta_iters is None:
prioritized_replay_beta_iters = total_timesteps
else:
prioritized_replay_beta_iters = self.prioritized_replay_beta_iters
self.beta_schedule = LinearSchedule(prioritized_replay_beta_iters,
initial_p=self.prioritized_replay_beta0,
final_p=1.0)
else:
self.replay_buffer = ReplayBuffer(self.buffer_size)
self.beta_schedule = None
if replay_wrapper is not None:
assert not self.prioritized_replay, "Prioritized replay buffer is not supported by HER"
self.replay_buffer = replay_wrapper(self.replay_buffer)
# Create the schedule for exploration starting from 1.
self.exploration = LinearSchedule(schedule_timesteps=int(self.exploration_fraction * total_timesteps),
initial_p=self.exploration_initial_eps,
final_p=self.exploration_final_eps)
episode_rewards = [0.0]
episode_successes = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
reset = True
obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs().squeeze()
for _ in range(total_timesteps):
# Take action and update exploration to the newest value
kwargs = {}
if not self.param_noise:
update_eps = self.exploration.value(self.num_timesteps)
update_param_noise_threshold = 0.
else:
update_eps = 0.
# Compute the threshold such that the KL divergence between perturbed and non-perturbed
# policy is comparable to eps-greedy exploration with eps = exploration.value(t).
# See Appendix C.1 in Parameter Space Noise for Exploration, Plappert et al., 2017
# for detailed explanation.
update_param_noise_threshold = \
-np.log(1. - self.exploration.value(self.num_timesteps) +
self.exploration.value(self.num_timesteps) / float(self.env.action_space.n))
kwargs['reset'] = reset
kwargs['update_param_noise_threshold'] = update_param_noise_threshold
kwargs['update_param_noise_scale'] = True
with self.sess.as_default():
action = self.act(np.array(obs)[None], update_eps=update_eps, **kwargs)[0]
env_action = action
reset = False
new_obs, rew, done, info = self.env.step(env_action)
self.num_timesteps += 1
# Stop training if return value is False
if callback.on_step() is False:
break
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs().squeeze()
reward_ = self._vec_normalize_env.get_original_reward().squeeze()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, rew
# Store transition in the replay buffer.
self.replay_buffer.add(obs_, action, reward_, new_obs_, float(done))
if self.expert_exp is not None:
self.add_expert_exp()
obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
if writer is not None:
ep_rew = np.array([reward_]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(self.episode_reward, ep_rew, ep_done, writer,
self.num_timesteps)
episode_rewards[-1] += reward_
if done:
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
reset = True
# Do not train if the warmup phase is not over
# or if there are not enough samples in the replay buffer
can_sample = self.replay_buffer.can_sample(self.batch_size)
if can_sample and self.num_timesteps > self.learning_starts \
and self.num_timesteps % self.train_freq == 0:
callback.on_rollout_end()
# Minimize the error in Bellman's equation on a batch sampled from replay buffer.
# pytype:disable=bad-unpacking
if self.prioritized_replay:
assert self.beta_schedule is not None, \
"BUG: should be LinearSchedule when self.prioritized_replay True"
experience = self.replay_buffer.sample(self.batch_size,
beta=self.beta_schedule.value(self.num_timesteps),
env=self._vec_normalize_env)
(obses_t, actions, rewards, obses_tp1, dones, weights, batch_idxes) = experience
else:
obses_t, actions, rewards, obses_tp1, dones = self.replay_buffer.sample(self.batch_size,
env=self._vec_normalize_env)
weights, batch_idxes = np.ones_like(rewards), None
# pytype:enable=bad-unpacking
if writer is not None:
# run loss backprop with summary, but once every 100 steps save the metadata
# (memory, compute time, ...)
if (1 + self.num_timesteps) % 100 == 0:
run_options = tf.RunOptions(trace_level=tf.RunOptions.FULL_TRACE)
run_metadata = tf.RunMetadata()
summary, td_errors = self._train_step(obses_t, actions, rewards, obses_tp1, obses_tp1,
dones, weights, sess=self.sess, options=run_options,
run_metadata=run_metadata)
writer.add_run_metadata(run_metadata, 'step%d' % self.num_timesteps)
else:
summary, td_errors = self._train_step(obses_t, actions, rewards, obses_tp1, obses_tp1,
dones, weights, sess=self.sess)
writer.add_summary(summary, self.num_timesteps)
else:
_, td_errors = self._train_step(obses_t, actions, rewards, obses_tp1, obses_tp1, dones, weights,
sess=self.sess)
if self.prioritized_replay:
new_priorities = np.abs(td_errors) + self.prioritized_replay_eps
assert isinstance(self.replay_buffer, PrioritizedReplayBuffer)
self.replay_buffer.update_priorities(batch_idxes, new_priorities)
callback.on_rollout_start()
if can_sample and self.num_timesteps > self.learning_starts and \
self.num_timesteps % self.target_network_update_freq == 0:
# Update target network periodically.
self.update_target(sess=self.sess)
if len(episode_rewards[-101:-1]) == 0:
mean_100ep_reward = -np.inf
else:
mean_100ep_reward = round(float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
if self.verbose >= 1 and done and log_interval is not None and len(episode_rewards) % log_interval == 0:
logger.record_tabular("steps", self.num_timesteps)
logger.record_tabular("episodes", num_episodes)
if len(episode_successes) > 0:
logger.logkv("success rate", np.mean(episode_successes[-100:]))
logger.record_tabular("mean 100 episode reward", mean_100ep_reward)
logger.record_tabular("% time spent exploring",
int(100 * self.exploration.value(self.num_timesteps)))
logger.dump_tabular()
callback.on_training_end()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=4,
tb_log_name="SAC_LR_CYCLED",
reset_num_timesteps=True,
replay_wrapper=None,
lr_cycler=False):
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
callback = self._init_callback(callback)
if replay_wrapper is not None:
self.replay_buffer = replay_wrapper(self.replay_buffer)
def cyclic_lr(step, num_cycle_steps=10000, base_lr=5e-4, max_lr=1e-2):
mod_step = step % num_cycle_steps
half = num_cycle_steps / 2
mod_step_half = mod_step % half
pct = mod_step_half / half
if mod_step < half:
diff = max_lr - base_lr
diff = pct * diff
return base_lr + diff
else:
diff = max_lr - base_lr
diff = pct * diff
return max_lr - diff
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
self._setup_learn()
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
# Initial learning rate
current_lr = self.learning_rate(1)
start_time = time.time()
episode_rewards = [0.0]
episode_successes = []
if self.action_noise is not None:
self.action_noise.reset()
obs = self.env.reset()
# Retrieve unnormalized observation for saving into the buffer
if self._vec_normalize_env is not None:
obs_ = self._vec_normalize_env.get_original_obs().squeeze()
n_updates = 0
infos_values = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
for step in range(total_timesteps):
# Before training starts, randomly sample actions
# from a uniform distribution for better exploration.
# Afterwards, use the learned policy
# if random_exploration is set to 0 (normal setting)
if self.num_timesteps < self.learning_starts or np.random.rand(
) < self.random_exploration:
# actions sampled from action space are from range specific to the environment
# but algorithm operates on tanh-squashed actions therefore simple scaling is used
unscaled_action = self.env.action_space.sample()
action = scale_action(self.action_space, unscaled_action)
else:
action = self.policy_tf.step(
obs[None], deterministic=False).flatten()
# Add noise to the action (improve exploration,
# not needed in general)
if self.action_noise is not None:
action = np.clip(action + self.action_noise(), -1, 1)
# inferred actions need to be transformed to environment action_space before stepping
unscaled_action = unscale_action(self.action_space, action)
assert action.shape == self.env.action_space.shape
new_obs, reward, done, info = self.env.step(unscaled_action)
self.num_timesteps += 1
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback.on_step() is False:
break
# Store only the unnormalized version
if self._vec_normalize_env is not None:
new_obs_ = self._vec_normalize_env.get_original_obs(
).squeeze()
reward_ = self._vec_normalize_env.get_original_reward(
).squeeze()
else:
# Avoid changing the original ones
obs_, new_obs_, reward_ = obs, new_obs, reward
# Store transition in the replay buffer.
self.replay_buffer.add(obs_, action, reward_, new_obs_,
float(done))
obs = new_obs
# Save the unnormalized observation
if self._vec_normalize_env is not None:
obs_ = new_obs_
# Retrieve reward and episode length if using Monitor wrapper
maybe_ep_info = info.get('episode')
if maybe_ep_info is not None:
self.ep_info_buf.extend([maybe_ep_info])
if writer is not None:
# Write reward per episode to tensorboard
ep_reward = np.array([reward_]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(
self.episode_reward, ep_reward, ep_done, writer,
self.num_timesteps)
if step % self.train_freq == 0:
callback.on_rollout_end()
mb_infos_vals = []
# Update policy, critics and target networks
for grad_step in range(self.gradient_steps):
# Break if the warmup phase is not over
# or if there are not enough samples in the replay buffer
if not self.replay_buffer.can_sample(self.batch_size) \
or self.num_timesteps < self.learning_starts:
break
n_updates += 1
# Compute current learning_rate
if lr_cycler:
current_lr = cyclic_lr(step)
else:
frac = 1.0 - step / total_timesteps
current_lr = self.learning_rate(frac)
# Update policy and critics (q functions)
mb_infos_vals.append(
self._train_step(step, writer, current_lr))
# Update target network
if (step +
grad_step) % self.target_update_interval == 0:
# Update target network
self.sess.run(self.target_update_op)
# Log losses and entropy, useful for monitor training
if len(mb_infos_vals) > 0:
infos_values = np.mean(mb_infos_vals, axis=0)
callback.on_rollout_start()
episode_rewards[-1] += reward_
if done:
if self.action_noise is not None:
self.action_noise.reset()
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if len(episode_rewards[-101:-1]) == 0:
mean_reward = -np.inf
else:
mean_reward = round(
float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
# Display training infos
if self.verbose >= 1 and done and log_interval is not None and len(
episode_rewards) % log_interval == 0:
fps = int(step / (time.time() - start_time))
logger.logkv("episodes", num_episodes)
logger.logkv("mean 100 episode reward", mean_reward)
if len(self.ep_info_buf) > 0 and len(
self.ep_info_buf[0]) > 0:
logger.logkv(
'ep_rewmean',
safe_mean([
ep_info['r'] for ep_info in self.ep_info_buf
]))
logger.logkv(
'eplenmean',
safe_mean([
ep_info['l'] for ep_info in self.ep_info_buf
]))
logger.logkv("n_updates", n_updates)
logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - start_time))
if len(episode_successes) > 0:
logger.logkv("success rate",
np.mean(episode_successes[-100:]))
if len(infos_values) > 0:
for (name, val) in zip(self.infos_names, infos_values):
logger.logkv(name, val)
logger.logkv("total timesteps", self.num_timesteps)
logger.dumpkvs()
# Reset infos:
infos_values = []
callback.on_training_end()
return self
def learn(self,
total_timesteps,
callback=None,
log_interval=4,
tb_log_name="TD3",
reset_num_timesteps=True,
replay_wrapper=None):
new_tb_log = self._init_num_timesteps(reset_num_timesteps)
callback = self._init_callback(callback)
if replay_wrapper is not None:
self.replay_buffer = replay_wrapper(self.replay_buffer)
with SetVerbosity(self.verbose), TensorboardWriter(self.graph, self.tensorboard_log, tb_log_name, new_tb_log) \
as writer:
self._setup_learn()
# Transform to callable if needed
self.learning_rate = get_schedule_fn(self.learning_rate)
# Initial learning rate
current_lr = self.learning_rate(1)
start_time = time.time()
episode_rewards = [0.0]
episode_successes = []
if self.action_noise is not None:
self.action_noise.reset()
obs = self.env.reset()
n_updates = 0
infos_values = []
callback.on_training_start(locals(), globals())
callback.on_rollout_start()
for step in range(total_timesteps):
# Before training starts, randomly sample actions
# from a uniform distribution for better exploration.
# Afterwards, use the learned policy
# if random_exploration is set to 0 (normal setting)
if self.num_timesteps < self.learning_starts or np.random.rand(
) < self.random_exploration:
# actions sampled from action space are from range specific to the environment
# but algorithm operates on tanh-squashed actions therefore simple scaling is used
unscaled_action = self.env.action_space.sample()
action = scale_action(self.action_space, unscaled_action)
else:
action = self.policy_tf.step(obs[None]).flatten()
# Add noise to the action, as the policy
# is deterministic, this is required for exploration
if self.action_noise is not None:
action = np.clip(action + self.action_noise(), -1, 1)
# Rescale from [-1, 1] to the correct bounds
unscaled_action = unscale_action(self.action_space, action)
assert action.shape == self.env.action_space.shape
new_obs, reward, done, info = self.env.step(unscaled_action)
self.num_timesteps += 1
# Only stop training if return value is False, not when it is None. This is for backwards
# compatibility with callbacks that have no return statement.
if callback.on_step() is False:
break
# Store transition in the replay buffer.
self.replay_buffer.add(obs, action, reward, new_obs,
float(done))
obs = new_obs
# Retrieve reward and episode length if using Monitor wrapper
maybe_ep_info = info.get('episode')
if maybe_ep_info is not None:
self.ep_info_buf.extend([maybe_ep_info])
if writer is not None:
# Write reward per episode to tensorboard
ep_reward = np.array([reward]).reshape((1, -1))
ep_done = np.array([done]).reshape((1, -1))
tf_util.total_episode_reward_logger(
self.episode_reward, ep_reward, ep_done, writer,
self.num_timesteps)
if step % self.train_freq == 0:
callback.on_rollout_end()
mb_infos_vals = []
# Update policy, critics and target networks
for grad_step in range(self.gradient_steps):
# Break if the warmup phase is not over
# or if there are not enough samples in the replay buffer
if not self.replay_buffer.can_sample(self.batch_size) \
or self.num_timesteps < self.learning_starts:
break
n_updates += 1
# Compute current learning_rate
frac = 1.0 - step / total_timesteps
current_lr = self.learning_rate(frac)
# Update policy and critics (q functions)
# Note: the policy is updated less frequently than the Q functions
# this is controlled by the `policy_delay` parameter
mb_infos_vals.append(
self._train_step(step, writer, current_lr,
(step + grad_step) %
self.policy_delay == 0))
# Log losses and entropy, useful for monitor training
if len(mb_infos_vals) > 0:
infos_values = np.mean(mb_infos_vals, axis=0)
callback.on_rollout_start()
episode_rewards[-1] += reward
if done:
if self.action_noise is not None:
self.action_noise.reset()
if not isinstance(self.env, VecEnv):
obs = self.env.reset()
episode_rewards.append(0.0)
maybe_is_success = info.get('is_success')
if maybe_is_success is not None:
episode_successes.append(float(maybe_is_success))
if len(episode_rewards[-101:-1]) == 0:
mean_reward = -np.inf
else:
mean_reward = round(
float(np.mean(episode_rewards[-101:-1])), 1)
num_episodes = len(episode_rewards)
# Display training infos
if self.verbose >= 1 and done and log_interval is not None and len(
episode_rewards) % log_interval == 0:
fps = int(step / (time.time() - start_time))
logger.logkv("episodes", num_episodes)
logger.logkv("mean 100 episode reward", mean_reward)
if len(self.ep_info_buf) > 0 and len(
self.ep_info_buf[0]) > 0:
logger.logkv(
'ep_rewmean',
safe_mean([
ep_info['r'] for ep_info in self.ep_info_buf
]))
logger.logkv(
'eplenmean',
safe_mean([
ep_info['l'] for ep_info in self.ep_info_buf
]))
logger.logkv("n_updates", n_updates)
logger.logkv("current_lr", current_lr)
logger.logkv("fps", fps)
logger.logkv('time_elapsed', int(time.time() - start_time))
if len(episode_successes) > 0:
logger.logkv("success rate",
np.mean(episode_successes[-100:]))
if len(infos_values) > 0:
for (name, val) in zip(self.infos_names, infos_values):
logger.logkv(name, val)
logger.logkv("total timesteps", self.num_timesteps)
logger.dumpkvs()
# Reset infos:
infos_values = []
callback.on_training_end()
return self
def learn(self, total_timesteps, callback=None, log_interval=1, tb_log_name="Dual",
reset_num_timesteps=True):
# 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)
top_callback = SaveOnTopRewardCallback(check_freq=self.n_steps, logdir=self.tensorboard_log, models_num=self.models_num)
callback.append(top_callback)
callback = self._init_callback(callback)
with SetVerbosity(self.verbose), TensorboardWriter(self.models[0].graph, self.tensorboard_log, tb_log_name, new_tb_log) as writer:
for model in self.models:
model._setup_learn(self)
t_first_start = time.time()
n_updates = total_timesteps // (self.n_envs * self.n_steps)
callback.on_training_start(locals(), globals())
for update in range(1, n_updates + 1):
assert (self.n_envs * self.n_steps) % 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_envs * self.n_steps) // 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()
rollouts = self.runner.run(callback) #execute episode
callback.on_rollout_end()
# Early stopping due to the callback
if not self.runner.continue_training:
break
# Unpack
i = 0
steps_used = rollouts[-1]
for rollout in rollouts[0]:
obs, returns, masks, actions, values, neglogpacs, states, ep_infos, true_reward, success_stages = rollout
model = self.models[i]
# calc = len(true_reward)
# model.n_batch = calc
if model.n_batch == 0:
b = 0
else:
self.ep_info_buf.extend(ep_infos)
mb_loss_vals = []
if states is None: # nonrecurrent version
update_fac = max(model.n_batch // self.nminibatches // self.noptepochs, 1)
inds = np.arange(len(obs))#np.arange(model.n_batch)
for epoch_num in range(self.noptepochs):
np.random.shuffle(inds)
for start in range(0, model.n_batch, batch_size):
timestep = self.num_timesteps // update_fac + ((epoch_num * model.n_batch + start) // batch_size)
end = start + batch_size
mbinds = inds[start:end]
if len(obs) > 1:
slices = (arr[mbinds] for arr in (obs, returns, masks, actions, values, neglogpacs))
mb_loss_vals.append(self._train_step(lr_now, cliprange_now, *slices, model=self.models[i], writer=writer, update=timestep, cliprange_vf=cliprange_vf_now))
else:
mb_loss_vals.append((0,0,0,0,0))
i+=1
else:
exit("does not support recurrent version")
loss_vals = np.mean(mb_loss_vals, axis=0)
t_now = time.time()
fps = int(model.n_batch / (t_now - t_start))
if writer is not None:
n_steps = model.n_batch
try:
total_episode_reward_logger(self.episode_reward, true_reward.reshape((self.n_envs, n_steps)), masks.reshape((self.n_envs, n_steps)), writer, self.num_timesteps)
except:
print("Failed to log episode reward of shape {}".format(true_reward.shape))
summary = tf.Summary(value=[tf.Summary.Value(tag='episode_reward/Successful stages',
simple_value=success_stages)])
writer.add_summary(summary, self.num_timesteps)
#@TODO plot in one graph:
for i, val in enumerate(steps_used):
summary = tf.Summary(value=[tf.Summary.Value(tag='episode_reward/Used steps net {}'.format(i),
simple_value=val)])
writer.add_summary(summary, 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("Steps", steps_used)
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, model.loss_names):
logger.logkv(loss_name, loss_val)
logger.dumpkvs()
callback.on_training_end()
return self