def build_learner(pre, post, ws, act_space, num_frames):
global_step = tf.train.get_or_create_global_step()
init_lr = FLAGS.init_lr
decay = FLAGS.lr_decay
warmup_steps = FLAGS.warmup_steps
gamma = FLAGS.gamma
n_step = FLAGS.n_step
use_soft = FLAGS.use_soft
time_scale = FLAGS.time_scale
use_hrnn = FLAGS.use_hrnn
use_reward_prediction = FLAGS.use_reward_prediction
after_rnn = FLAGS.after_rnn
use_pixel_control = FLAGS.use_pixel_control
pq_kl_coef = FLAGS.pq_kl_coef
p_kl_coef = FLAGS.p_kl_coef
global_step_float = tf.cast(global_step, tf.float32)
lr = tf.train.polynomial_decay(
init_lr, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
init_lr / 10.)
is_warmup = tf.cast(global_step_float < warmup_steps, tf.float32)
lr = is_warmup * global_step_float / warmup_steps * init_lr + (
1.0 - is_warmup) * (init_lr * (1.0 - decay) + lr * decay)
optimizer = tf.train.AdamOptimizer(lr)
if FLAGS.zero_init:
pre["state_in"] = tf.zeros_like(pre["state_in"])
if use_hrnn:
rnn = TmpHierRNN(time_scale,
64,
4,
2,
8,
'lstm',
'rmc',
return_sequences=True,
return_state=True,
name="hrnn")
else:
rnn = tf.compat.v1.keras.layers.LSTM(256,
return_sequences=True,
return_state=True,
name="lstm")
pre_model = Model(act_space, gamma, n_step, use_soft, rnn, use_hrnn,
use_reward_prediction, after_rnn, use_pixel_control,
False, **pre)
post["state_in"] = tf.stop_gradient(pre_model.state_out)
post_model = Model(act_space, gamma, n_step, use_soft, rnn, use_hrnn,
use_reward_prediction, after_rnn, use_pixel_control,
True, **post)
v_loss = mse(
post_model.qa,
tf.stop_gradient(
rescaleTarget(post_model.n_step_rewards, gamma**n_step,
post_model.qa1)))
v_loss = FLAGS.vf_coef * tf.reduce_mean(
v_loss * post_model.mask[:, :-n_step] * ws[:, None])
add_loss = 0.0
if use_hrnn:
pq_kl_loss = KL_from_gaussians(post_model.q_mus, post_model.q_sigmas,
post_model.p_mus, post_model.p_sigmas)
pq_kl_loss = tf.reduce_mean(pq_kl_loss * post_model.mask)
p_kl_loss = KL_from_gaussians(post_model.p_mus, post_model.p_sigmas,
tf.zeros_like(post_model.p_mus),
0.01 * tf.ones_like(post_model.p_sigmas))
p_kl_loss = tf.reduce_mean(p_kl_loss * post_model.mask)
with tf.name_scope("hierarchy_loss"):
tf.summary.scalar("kl_div_pq", pq_kl_loss)
tf.summary.scalar("kl_div_prior", p_kl_loss)
add_loss += pq_kl_coef * pq_kl_loss
add_loss += p_kl_coef * p_kl_loss
if use_reward_prediction:
r_loss = tf.reduce_mean(
mse(post_model.reward_prediction, post_model.r[:, 1:1 - n_step]) *
post_model.mask[:, :-n_step])
tf.summary.scalar("r_loss", r_loss)
add_loss += r_loss
if use_pixel_control:
s = tf.cast(post_model.s[:, :1 - n_step, :, :, :], tf.float32) / 255.0
target = s[:, 1:, :, :, :] - s[:, :-1, :, :, :]
shape = get_shape(target)
target = tf.reshape(
target,
(shape[0], shape[1], 4, shape[2] // 4, 4, shape[3] // 4, shape[4]))
target = tf.reduce_mean(target, axis=(2, 4))
pixel_loss = tf.reduce_mean(
mse(post_model.pixel_control, target) *
post_model.mask[:, :-n_step, None, None, None])
with tf.name_scope("control_loss"):
tf.summary.scalar("pixel_control_loss", pixel_loss)
add_loss += pixel_loss
loss = FLAGS.vf_coef * v_loss + add_loss
abs_td = post_model.mask[:, :-n_step] * tf.abs(
post_model.qa - rescaleTarget(post_model.n_step_rewards, gamma**n_step,
post_model.qa1))
avg_p = tf.reduce_mean(abs_td, axis=-1)
max_p = tf.reduce_max(abs_td, axis=-1)
priority = 0.9 * max_p + 0.1 * avg_p
beta = tf.train.polynomial_decay(
0.4, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
1.0)
train_op = miniOp(optimizer, loss, FLAGS.grad_clip)
target_op = assignOp(1.0, {"q": "q_target"})
dependency = [train_op]
if use_soft:
qf_entropy = entropy_from_logits(post_model.qf_logits)
target_entropy = tf.train.polynomial_decay(
0.9 * np.log(act_space), global_step,
FLAGS.total_environment_frames //
(FLAGS.batch_size * FLAGS.seqlen), 0.5 * np.log(act_space))
ent_loss = tf.reduce_mean(
mse(qf_entropy,
tf.cast(target_entropy, tf.float32)[None, None]))
with tf.name_scope("ent_loss"):
tf.summary.scalar("ent_loss", ent_loss)
ent_op = miniOp(optimizer,
ent_loss,
grad_clip=FLAGS.grad_clip,
var_scope="temperature")
dependency.append(ent_op)
new_frames = tf.reduce_sum(post["mask"])
with tf.control_dependencies(dependency):
num_frames_and_train = tf.assign_add(num_frames, new_frames)
global_step_and_train = tf.assign_add(global_step, 1)
tf.summary.scalar("learning_rate", lr)
tf.summary.scalar("v_loss", v_loss)
tf.summary.scalar("all_loss", loss)
return num_frames_and_train, global_step_and_train, target_op, priority, beta
def build_learner(pre, post, act_space, num_frames, batch_weights):
global_step = tf.train.get_or_create_global_step()
init_lr = FLAGS.init_lr
decay = FLAGS.lr_decay
warmup_steps = FLAGS.warmup_steps
gamma = FLAGS.gamma
n_step = FLAGS.n_step
time_scale = FLAGS.time_scale
use_hrnn = FLAGS.use_hrnn
use_rmc = FLAGS.use_rmc
use_amc = FLAGS.use_amc
use_beta = FLAGS.use_beta
use_retrace = FLAGS.use_retrace
use_reward_prediction = FLAGS.use_reward_prediction
after_rnn = FLAGS.after_rnn
use_pixel_control = FLAGS.use_pixel_control
pq_kl_coef = FLAGS.pq_kl_coef
p_kl_coef = FLAGS.p_kl_coef
pi_coef = FLAGS.pi_coef
vf_coef = FLAGS.vf_coef
ent_coef = FLAGS.ent_coef
qf_coef = FLAGS.qf_coef
ppo_clip = FLAGS.ppo_clip
vf_clip = FLAGS.vf_clip
global_step_float = tf.cast(global_step, tf.float32)
lr = tf.train.polynomial_decay(
init_lr, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
init_lr / 10.)
is_warmup = tf.cast(global_step_float < warmup_steps, tf.float32)
lr = is_warmup * global_step_float / warmup_steps * init_lr + (
1.0 - is_warmup) * (init_lr * (1.0 - decay) + lr * decay)
ent_coef = tf.train.polynomial_decay(
ent_coef, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
ent_coef / 10.)
optimizer = tf.train.AdamOptimizer(lr)
if FLAGS.zero_init:
pre["state_in"] = tf.zeros_like(pre["state_in"])
if use_hrnn:
rnn = TmpHierRNN(time_scale,
64,
4,
2,
8,
'lstm',
'rmc',
return_sequences=True,
return_state=True,
name="hrnn")
elif use_rmc:
rnn = RMCRNN(64,
4,
64,
return_sequences=True,
return_state=True,
name="rmc")
elif use_amc:
rnn = AMCRNN(64,
4,
64,
return_sequences=True,
return_state=True,
name="amc")
else:
rnn = tf.compat.v1.keras.layers.CuDNNLSTM(256,
return_sequences=True,
return_state=True,
name="lstm")
pre_model = Model(act_space, gamma, n_step, rnn, use_hrnn, use_rmc,
use_amc, use_beta, use_reward_prediction, after_rnn,
use_pixel_control, False, **pre)
post["state_in"] = tf.stop_gradient(pre_model.state_out)
post_model = Model(act_space, gamma, n_step, rnn, use_hrnn, use_rmc,
use_amc, use_beta, use_reward_prediction, after_rnn,
use_pixel_control, True, **post)
tf.summary.scalar("adv_mean", post_model.adv_mean)
tf.summary.scalar("adv_std", post_model.adv_std)
if use_retrace:
q_loss = mse(post_model.qa, post_model.retrace_qs)
else:
q_loss = mse(post_model.qa, post_model.n_step_qs)
# q_loss = mse(
# post_model.qa,
# tf.stop_gradient(
# post_model.current_value[:, :-n_step] + post_model.adv))
q_loss = tf.reduce_mean(q_loss * post_model.mask[:, :-n_step] *
batch_weights[:, None]) + 3.0 * tf.reduce_mean(
q_loss * post_model.mask[:, :-n_step] *
(1.0 - batch_weights[:, None]))
ent_loss = tf.reduce_mean(
entropy_from_logits(post_model.current_act_logits) * post_model.mask *
batch_weights[:, None])
losses = dPPOcC(
act=post_model.a[:, 1:1 - n_step],
policy_logits=post_model.current_act_logits[:, :-n_step, :],
behavior_logits=post_model.behavior_logits[:, :-n_step, :],
advantage=post_model.adv,
policy_clip=ppo_clip,
vf=post_model.current_value[:, :-n_step],
vf_target=post_model.vs,
value_clip=vf_clip,
old_vf=post_model.old_vf[:, :-n_step])
p_loss = tf.reduce_mean(losses.p_loss * post_model.mask[:, :-n_step] *
batch_weights[:, None])
v_loss = tf.reduce_mean(losses.v_loss * post_model.mask[:, :-n_step] *
batch_weights[:, None])
add_loss = 0.0
if use_hrnn:
pq_kl_loss = KL_from_gaussians(post_model.q_mus, post_model.q_sigmas,
post_model.p_mus, post_model.p_sigmas)
pq_kl_loss = tf.reduce_mean(pq_kl_loss * post_model.mask)
p_kl_loss = KL_from_gaussians(post_model.p_mus, post_model.p_sigmas,
tf.zeros_like(post_model.p_mus),
0.01 * tf.ones_like(post_model.p_sigmas))
p_kl_loss = tf.reduce_mean(p_kl_loss * post_model.mask)
with tf.name_scope("hierarchy_loss"):
tf.summary.scalar("kl_div_pq", pq_kl_loss)
tf.summary.scalar("kl_div_prior", p_kl_loss)
add_loss += pq_kl_coef * pq_kl_loss
add_loss += p_kl_coef * p_kl_loss
if use_reward_prediction:
r_loss = tf.reduce_mean(
mse(post_model.reward_prediction, post_model.r[:, 1:1 - n_step]) *
post_model.mask[:, :-n_step])
tf.summary.scalar("r_loss", r_loss)
add_loss += r_loss
if use_pixel_control:
s = tf.cast(post_model.s[:, :1 - n_step, :, :, :], tf.float32) / 255.0
target = s[:, 1:, :, :, :] - s[:, :-1, :, :, :]
shape = get_shape(target)
target = tf.reshape(
target,
(shape[0], shape[1], 4, shape[2] // 4, 4, shape[3] // 4, shape[4]))
target = tf.reduce_mean(target, axis=(2, 4))
pixel_loss = tf.reduce_mean(
mse(post_model.pixel_control, target) *
post_model.mask[:, :-n_step, None, None, None])
with tf.name_scope("control_loss"):
tf.summary.scalar("pixel_control_loss", pixel_loss)
add_loss += pixel_loss
loss = (qf_coef * q_loss + vf_coef * v_loss + pi_coef * p_loss -
ent_coef * ent_loss + add_loss)
abs_td = post_model.mask[:, :-n_step] * tf.abs(
post_model.qa - post_model.n_step_rewards +
gamma**n_step * post_model.qa1)
avg_p = tf.reduce_mean(abs_td, axis=-1)
max_p = tf.reduce_max(abs_td, axis=-1)
priority = 0.9 * max_p + 0.1 * avg_p
beta = tf.train.polynomial_decay(
0.4, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
1.0)
train_op = miniOp(optimizer, loss, FLAGS.grad_clip)
if FLAGS.smooth_update:
init_target_op = assignOp(1.0, {"q": "q_target"})
target_op = assignOp(1.0 / FLAGS.target_update, {"q": "q_target"})
else:
init_target_op = assignOp(1.0, {"q": "q_target"})
target_op = tf.no_op()
dependency = [train_op, target_op]
new_frames = tf.reduce_sum(post["mask"])
with tf.control_dependencies(dependency):
num_frames_and_train = tf.assign_add(num_frames, new_frames)
global_step_and_train = tf.assign_add(global_step, 1)
tf.summary.scalar("learning_rate", lr)
tf.summary.scalar("pi_loss", p_loss)
tf.summary.scalar("q_loss", q_loss)
tf.summary.scalar("v_loss", v_loss)
tf.summary.scalar("ent_loss", ent_loss)
tf.summary.scalar("all_loss", loss)
return num_frames_and_train, global_step_and_train, init_target_op, priority, beta