def build_evaluator_model(kwargs):
Model = warp_Model()
import tensorflow as tf
from module import TmpHierRNN
from module import AMCRNN
from module import RMCRNN
from module import NewLSTM
frames = kwargs["frames"]
image_size = kwargs["image_size"]
act_space = kwargs["act_space"]
gamma = kwargs["gamma"]
n_step = kwargs["n_step"]
time_scale = kwargs["time_scale"]
state_size = kwargs["state_size"]
use_hrnn = kwargs["use_hrnn"]
use_rmc = kwargs["use_rmc"]
use_amc = kwargs["use_amc"]
use_beta = kwargs["use_beta"]
use_reward_prediction = kwargs["use_reward_prediction"]
after_rnn = kwargs["after_rnn"]
use_pixel_control = kwargs["use_pixel_control"]
phs = dict()
phs["s"] = tf.placeholder(
dtype=tf.uint8, shape=[None, None, image_size, image_size, frames])
phs["a"] = tf.placeholder(dtype=tf.int32, shape=[None, None])
phs["r"] = tf.placeholder(dtype=tf.float32, shape=[None, None])
phs["state_in"] = tf.placeholder(dtype=tf.float32, shape=[None, state_size])
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 = NewLSTM(
256, return_sequences=True, return_state=True, name="lstm")
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, **phs)
return model
def build_evaluator_model(kwargs):
Model = warp_Model()
import tensorflow as tf
from module import RMCRNN
from module import TmpHierRMCRNN
from module import TmpHierRNN
frames = kwargs["frames"]
act_space = kwargs["act_space"]
state_size = kwargs["state_size"]
use_rmc = kwargs["use_rmc"]
use_hrmc = kwargs["use_hrmc"]
use_hrnn = kwargs["use_hrnn"]
use_reward_prediction = kwargs["use_reward_prediction"]
after_rnn = kwargs["after_rnn"]
use_pixel_control = kwargs["use_pixel_control"]
use_pixel_reconstruction = kwargs["use_pixel_reconstruction"]
phs = dict()
phs["s"] = tf.placeholder(dtype=tf.float32,
shape=[None, None, 84, 84, frames])
phs["prev_a"] = tf.placeholder(dtype=tf.int32, shape=[None, None])
phs["prev_r"] = tf.placeholder(dtype=tf.float32, shape=[None, None])
phs["state_in"] = tf.placeholder(dtype=tf.float32,
shape=[None, state_size])
if use_hrnn:
rnn = TmpHierRNN(4,
64,
4,
2,
8,
'lstm',
'rmc',
return_sequences=True,
return_state=True,
name="hrnn")
elif use_hrmc:
rnn = TmpHierRMCRNN(4,
64,
4,
4,
return_sequences=True,
return_state=True,
name="hrmcrnn")
elif use_rmc:
rnn = RMCRNN(64,
4,
4,
return_sequences=True,
return_state=True,
name="rmcrnn")
else:
rnn = tf.compat.v1.keras.layers.LSTM(256,
return_sequences=True,
return_state=True,
name="lstm")
model = Model(act_space, rnn, use_rmc, use_hrmc or use_hrnn,
use_reward_prediction, after_rnn, use_pixel_control,
use_pixel_reconstruction, "agent", **phs)
return model
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):
global_step = tf.train.get_or_create_global_step()
init_lr = FLAGS.init_lr
decay = FLAGS.lr_decay
warmup_steps = FLAGS.warmup_steps
use_rmc = FLAGS.use_rmc
use_hrmc = FLAGS.use_hrmc
use_hrnn = FLAGS.use_hrnn
use_icm = FLAGS.use_icm
use_coex = FLAGS.use_coex
use_reward_prediction = FLAGS.use_reward_prediction
after_rnn = FLAGS.after_rnn
use_pixel_control = FLAGS.use_pixel_control
use_pixel_reconstruction = FLAGS.use_pixel_reconstruction
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)
ent_coef = tf.train.polynomial_decay(
FLAGS.ent_coef, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
FLAGS.ent_coef / 10.)
if FLAGS.zero_init:
pre["state_in"] = tf.zeros_like(pre["state_in"])
if use_hrnn:
rnn = TmpHierRNN(4,
64,
4,
2,
8,
'lstm',
'rmc',
return_sequences=True,
return_state=True,
name="hrnn")
elif use_hrmc:
rnn = TmpHierRMCRNN(4,
64,
4,
4,
return_sequences=True,
return_state=True,
name="hrmcrnn")
elif use_rmc:
rnn = RMCRNN(64,
4,
4,
return_sequences=True,
return_state=True,
name="rmcrnn")
else:
rnn = tf.compat.v1.keras.layers.LSTM(256,
return_sequences=True,
return_state=True,
name="lstm")
pre_model = Model(act_space, rnn, use_rmc, use_hrmc or use_hrnn,
use_reward_prediction, after_rnn,
use_pixel_reconstruction, "agent", **pre)
post["state_in"] = tf.stop_gradient(pre_model.state_out)
post_model = Model(act_space, rnn, use_rmc, use_hrmc or use_hrnn,
use_reward_prediction, after_rnn,
use_pixel_reconstruction, "agent", **post)
tf.summary.scalar("adv_mean", post_model.adv_mean)
tf.summary.scalar("adv_std", post_model.adv_std)
losses = dPPOcC(act=post_model.a_t,
policy_logits=post_model.current_act_logits,
old_policy_logits=post_model.old_act_logits,
advantage=post_model.advantage,
policy_clip=FLAGS.ppo_clip,
vf=post_model.current_value,
vf_target=post_model.ret,
value_clip=FLAGS.vf_clip,
old_vf=post_model.old_current_value)
entropy_loss = tf.reduce_mean(
entropy(post_model.current_act_logits) * post_model.slots)
p_loss = tf.reduce_mean(losses.p_loss * post_model.slots)
v_loss = tf.reduce_mean(losses.v_loss * post_model.slots)
add_loss = 0.0
if use_icm:
icmloss = icm(post_model.cnn_feature[:, :-1, :],
post_model.cnn_feature[:, 1:, :], post_model.a_t[:, :-1],
act_space)
add_loss += 0.2 * tf.reduce_mean(
icmloss.f_loss * post_model.slots[:, :-1]) + 0.8 * tf.reduce_mean(
icmloss.i_loss * post_model.slots[:, :-1])
if use_coex:
coexloss = coex(post_model.image_feature[:, :-1, :, :, :],
post_model.image_feature[:, 1:, :, :, :],
post_model.a_t[:, :-1], act_space)
add_loss += tf.reduce_mean(coexloss * post_model.slots[:, :-1])
if use_hrmc or 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.slots)
tf.summary.scalar("kl_div", pq_kl_loss)
add_loss += pq_kl_coef * pq_kl_loss
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.slots)
tf.summary.scalar("kl_div_prior", p_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_t) *
post_model.slots)
tf.summary.scalar("r_loss", r_loss)
add_loss += r_loss
if use_pixel_control:
change_of_cells = tf.reduce_mean(post_model.s_t[:, 1:, :, :, :] -
post_model.s_t[:, :-1, :, :, :],
axis=-1)
s_shape = get_shape(change_of_cells)
s_H, s_W = s_shape[2:]
ctr_H, ctr_W = get_shape(post_model.pixel_control)[2:4]
change_of_cells = tf.reduce_mean(tf.reshape(
change_of_cells,
shape=s_shape[:2] + [ctr_H, s_H // ctr_H, ctr_W, s_W // ctr_W]),
axis=(3, 5))
ctr = tf.reduce_sum(
tf.transpose(post_model.pixel_control, perm=(0, 1, 4, 2, 3)) *
tf.one_hot(post_model.a_t,
depth=post_model.act_space,
dtype=tf.float32)[:, :, :, None, None],
axis=2)[:, :-1, :, :]
ctr_loss = tf.reduce_mean(mse(ctr, change_of_cells))
tf.summary.scalar("pixel_control_loss", ctr_loss)
add_loss += ctr_loss
if use_pixel_reconstruction:
rec_loss = tf.reduce_mean(
mse(post_model.pixel_reconstruction, post_model.s_t) *
post_model.slots[:, :, None, None, None])
tf.summary.scalar("rec_loss", rec_loss)
add_loss += rec_loss
loss = (FLAGS.pi_coef * p_loss + FLAGS.vf_coef * v_loss -
ent_coef * entropy_loss + add_loss)
train_op = miniOp(optimizer, loss, FLAGS.grad_clip)
new_frames = tf.reduce_sum(post["slots"])
with tf.control_dependencies([train_op]):
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("ent_coef", ent_coef)
tf.summary.scalar("ent_loss", entropy_loss)
tf.summary.scalar("p_loss", p_loss)
tf.summary.scalar("v_loss", v_loss)
tf.summary.scalar("all_loss", loss)
return num_frames_and_train, global_step_and_train
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
def build_learner(pre, post, 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
use_hrnn = FLAGS.use_hrnn
use_reward_prediction = FLAGS.use_reward_prediction
after_rnn = FLAGS.after_rnn
use_pixel_control = FLAGS.use_pixel_control
use_pixel_reconstruction = FLAGS.use_pixel_reconstruction
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)
ent_coef = tf.train.polynomial_decay(
FLAGS.ent_coef, global_step,
FLAGS.total_environment_frames // (FLAGS.batch_size * FLAGS.seqlen),
FLAGS.ent_coef / 10.)
if FLAGS.zero_init:
pre["state_in"] = tf.zeros_like(pre["state_in"])
if use_hrnn:
rnn = TmpHierRNN(4,
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, rnn, use_hrnn, use_reward_prediction,
after_rnn, use_pixel_control, use_pixel_reconstruction,
False, **pre)
post["state_in"] = tf.stop_gradient(pre_model.state_out)
post_model = Model(act_space, gamma, rnn, use_hrnn, use_reward_prediction,
after_rnn, use_pixel_control, use_pixel_reconstruction,
True, **post)
tf.summary.scalar("adv_mean", post_model.adv_mean)
tf.summary.scalar("adv_std", post_model.adv_std)
losses = dPPOcC(act=post_model.a,
policy_logits=post_model.current_act_logits,
behavior_logits=post_model.behavior_logits,
advantage=post_model.adv,
policy_clip=FLAGS.ppo_clip,
vf=post_model.current_value,
vf_target=post_model.vs,
value_clip=FLAGS.vf_clip,
old_vf=post_model.old_vf)
entropy_loss = tf.reduce_mean(
entropy(post_model.current_act_logits) * post_model.mask[:, :-1])
p_loss = tf.reduce_mean(losses.p_loss * post_model.mask[:, :-1])
v_loss = tf.reduce_mean(losses.v_loss * post_model.mask[:, :-1])
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) *
post_model.mask[:, :-1])
tf.summary.scalar("r_loss", r_loss)
add_loss += r_loss
if use_pixel_control:
s = tf.cast(post_model.s, 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[:, :-1, None, None, None])
with tf.name_scope("control_loss"):
tf.summary.scalar("pixel_control_loss", pixel_loss)
add_loss += pixel_loss
if use_pixel_reconstruction:
s = tf.cast(post_model.s, tf.float32) / 255.0
rec_loss = tf.reduce_mean(
mse(post_model.pixel_reconstruction, s[:, :-1, :, :, :]) *
post_model.mask[:, :-1, None, None, None])
tf.summary.scalar("rec_loss", rec_loss)
add_loss += rec_loss
loss = (FLAGS.pi_coef * p_loss + FLAGS.vf_coef * v_loss -
ent_coef * entropy_loss + add_loss)
train_op = miniOp(optimizer, loss, FLAGS.grad_clip)
new_frames = tf.reduce_sum(post["mask"])
with tf.control_dependencies([train_op]):
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("ent_coef", ent_coef)
tf.summary.scalar("ent_loss", entropy_loss)
tf.summary.scalar("p_loss", p_loss)
tf.summary.scalar("v_loss", v_loss)
tf.summary.scalar("all_loss", loss)
return num_frames_and_train, global_step_and_train