def _get_action(self, s, visual_s, cell_state):
with tf.device(self.device):
feat, cell_state = self.get_feature(s, visual_s, cell_state=cell_state, record_cs=True)
if self.is_continuous:
mu, log_std = self.actor_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
pi, _ = tsallis_squash_rsample(mu, log_std, self.entropic_index)
mu = tf.tanh(mu) # squash mu
else:
logits = self.actor_net(feat)
mu = tf.argmax(logits, axis=1)
cate_dist = tfp.distributions.Categorical(logits)
pi = cate_dist.sample()
return mu, pi, cell_state
def _get_action(self, s, visual_s):
with tf.device(self.device):
feat = tf.concat([self.encoder(visual_s), s], axis=-1)
if self.is_continuous:
mu, log_std = self.actor_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min,
self.log_std_max)
pi, _ = squash_rsample(mu, log_std)
mu = tf.tanh(mu) # squash mu
else:
logits = self.actor_net(feat)
mu = tf.argmax(logits, axis=1)
cate_dist = tfp.distributions.Categorical(logits)
pi = cate_dist.sample()
return mu, pi
def _get_action(self, adj, x, visual_s, evaluation):
adj,x , visual_s = self.cast(adj, x, visual_s)
with tf.device(self.device):
if self.is_continuous:
mu, log_std = self.actor_net(adj, x, visual_s)
log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
pi, _ = squash_rsample(mu, log_std)
mu = tf.tanh(mu) # squash mu
else:
logits = self.actor_net(adj, x, visual_s)
mu = tf.argmax(logits, axis=1)
cate_dist = tfp.distributions.Categorical(logits)
pi = cate_dist.sample()
if evaluation == True:
return mu
else:
return pi
def train_persistent(self, s, visual_s, a, r, s_, visual_s_, done):
s, visual_s, a, r, s_, visual_s_, done = self.cast(s, visual_s, a, r, s_, visual_s_, done)
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
if self.is_continuous:
mu, log_std = self.actor_net(s, visual_s)
log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
pi, log_pi = squash_rsample(mu, log_std)
entropy = gaussian_entropy(log_std)
target_mu, target_log_std = self.actor_net(s_, visual_s_)
target_log_std = clip_nn_log_std(target_log_std)
target_pi, target_log_pi = squash_rsample(target_mu, target_log_std)
else:
logits = self.actor_net(s, visual_s)
logp_all = tf.nn.log_softmax(logits)
gumbel_noise = tf.cast(self.gumbel_dist.sample([a.shape[0], self.a_counts]), dtype=tf.float32)
_pi = tf.nn.softmax((logp_all + gumbel_noise) / self.discrete_tau)
_pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1), self.a_counts)
_pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
pi = _pi_diff + _pi
log_pi = tf.reduce_sum(tf.multiply(logp_all, pi), axis=1, keepdims=True)
entropy = -tf.reduce_mean(tf.reduce_sum(tf.exp(logp_all) * logp_all, axis=1, keepdims=True))
target_logits = self.actor_net(s_, visual_s_)
target_cate_dist = tfp.distributions.Categorical(target_logits)
target_pi = target_cate_dist.sample()
target_pi = tf.one_hot(target_pi, self.a_counts, dtype=tf.float32)
target_log_pi = target_cate_dist.log_prob(target_pi)
q1 = self.q1_net(s, visual_s, a)
q1_target = self.q1_target_net(s_, visual_s_, target_pi)
q2 = self.q2_net(s, visual_s, a)
q2_target = self.q2_target_net(s_, visual_s_, target_pi)
q1_s_pi = self.q1_net(s, visual_s, pi)
q2_s_pi = self.q2_net(s, visual_s, pi)
dc_r_q1 = tf.stop_gradient(r + self.gamma * (1 - done) * (q1_target - tf.exp(self.log_alpha) * target_log_pi))
dc_r_q2 = tf.stop_gradient(r + self.gamma * (1 - done) * (q2_target - tf.exp(self.log_alpha) * target_log_pi))
td_error1 = q1 - dc_r_q1
td_error2 = q2 - dc_r_q2
q1_loss = tf.reduce_mean(tf.square(td_error1) * self.IS_w)
q2_loss = tf.reduce_mean(tf.square(td_error2) * self.IS_w)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss
actor_loss = -tf.reduce_mean(tf.minimum(q1_s_pi, q2_s_pi) - tf.exp(self.log_alpha) * log_pi)
if self.auto_adaption:
alpha_loss = -tf.reduce_mean(self.log_alpha * tf.stop_gradient(log_pi - self.a_counts))
critic_grads = tape.gradient(critic_loss, self.q1_net.tv + self.q2_net.tv)
self.optimizer_critic.apply_gradients(
zip(critic_grads, self.q1_net.tv + self.q2_net.tv)
)
actor_grads = tape.gradient(actor_loss, self.actor_net.tv)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_net.tv)
)
if self.auto_adaption:
alpha_grads = tape.gradient(alpha_loss, [self.log_alpha])
self.optimizer_alpha.apply_gradients(
zip(alpha_grads, [self.log_alpha])
)
self.global_step.assign_add(1)
summaries = dict([
['LOSS/actor_loss', actor_loss],
['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss],
['LOSS/critic_loss', critic_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', tf.exp(self.log_alpha)],
['Statistics/entropy', entropy],
['Statistics/q_min', tf.reduce_min(tf.minimum(q1, q2))],
['Statistics/q_mean', tf.reduce_mean(tf.minimum(q1, q2))],
['Statistics/q_max', tf.reduce_max(tf.maximum(q1, q2))]
])
if self.auto_adaption:
summaries.update({
'LOSS/alpha_loss': alpha_loss
})
return td_error1 + td_error2 / 2, summaries
def train_persistent(self, memories, isw, crsty_loss, cell_state):
ss, vvss, a, r, done = memories
batch_size = tf.shape(a)[0]
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
feat, feat_ = self.get_feature(ss, vvss, cell_state=cell_state, s_and_s_=True)
if self.is_continuous:
mu, log_std = self.actor_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min, self.log_std_max)
pi, log_pi = tsallis_squash_rsample(mu, log_std, self.entropic_index)
entropy = gaussian_entropy(log_std)
target_mu, target_log_std = self.actor_net(feat_)
target_log_std = clip_nn_log_std(target_log_std)
target_pi, target_log_pi = tsallis_squash_rsample(target_mu, target_log_std, self.entropic_index)
else:
logits = self.actor_net(feat)
logp_all = tf.nn.log_softmax(logits)
gumbel_noise = tf.cast(self.gumbel_dist.sample([batch_size, self.a_dim]), dtype=tf.float32)
_pi = tf.nn.softmax((logp_all + gumbel_noise) / self.discrete_tau)
_pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1), self.a_dim)
_pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
pi = _pi_diff + _pi
log_pi = tf.reduce_sum(tf.multiply(logp_all, pi), axis=1, keepdims=True)
entropy = -tf.reduce_mean(tf.reduce_sum(tf.exp(logp_all) * logp_all, axis=1, keepdims=True))
target_logits = self.actor_net(feat_)
target_cate_dist = tfp.distributions.Categorical(target_logits)
target_pi = target_cate_dist.sample()
target_pi = tf.one_hot(target_pi, self.a_dim, dtype=tf.float32)
target_log_pi = target_cate_dist.log_prob(target_pi)
q1, q2 = self.critic_net(feat, a)
q1_target, q2_target = self.critic_target_net(feat_, target_pi)
q_s_pi = self.critic_net.get_min(feat, pi)
dc_r_q1 = tf.stop_gradient(r + self.gamma * (1 - done) * (q1_target - self.alpha * target_log_pi))
dc_r_q2 = tf.stop_gradient(r + self.gamma * (1 - done) * (q2_target - self.alpha * target_log_pi))
td_error1 = q1 - dc_r_q1
td_error2 = q2 - dc_r_q2
q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss + crsty_loss
actor_loss = -tf.reduce_mean(q_s_pi - self.alpha * log_pi)
if self.auto_adaption:
alpha_loss = -tf.reduce_mean(self.alpha * tf.stop_gradient(log_pi - self.target_entropy))
critic_grads = tape.gradient(critic_loss, self.critic_tv)
self.optimizer_critic.apply_gradients(
zip(critic_grads, self.critic_tv)
)
actor_grads = tape.gradient(actor_loss, self.actor_tv)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_tv)
)
if self.auto_adaption:
alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
self.optimizer_alpha.apply_gradients(
[(alpha_grad, self.log_alpha)]
)
self.global_step.assign_add(1)
summaries = dict([
['LOSS/actor_loss', actor_loss],
['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss],
['LOSS/critic_loss', critic_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', self.alpha],
['Statistics/entropy', entropy],
['Statistics/q_min', tf.reduce_min(tf.minimum(q1, q2))],
['Statistics/q_mean', tf.reduce_mean(tf.minimum(q1, q2))],
['Statistics/q_max', tf.reduce_max(tf.maximum(q1, q2))]
])
if self.auto_adaption:
summaries.update({
'LOSS/alpha_loss': alpha_loss
})
return (td_error1 + td_error2) / 2, summaries
def train(self, memories, isw, crsty_loss, cell_state):
s, visual_s, a, r, s_, visual_s_, done, pos = memories
batch_size = tf.shape(a)[0]
with tf.device(self.device):
with tf.GradientTape(persistent=True) as tape:
vis_feat = self.encoder(visual_s)
vis_feat_ = self.encoder(visual_s_)
target_vis_feat_ = self.encoder_target(visual_s_)
feat = tf.concat([vis_feat, s], axis=-1)
feat_ = tf.concat([vis_feat_, s_], axis=-1)
target_feat_ = tf.concat([target_vis_feat_, s_], axis=-1)
if self.is_continuous:
target_mu, target_log_std = self.actor_net(feat_)
target_log_std = clip_nn_log_std(target_log_std)
target_pi, target_log_pi = squash_rsample(
target_mu, target_log_std)
else:
target_logits = self.actor_net(feat_)
target_cate_dist = tfp.distributions.Categorical(
target_logits)
target_pi = target_cate_dist.sample()
target_log_pi = target_cate_dist.log_prob(target_pi)
target_pi = tf.one_hot(target_pi,
self.a_dim,
dtype=tf.float32)
q1, q2 = self.critic_net(feat, a)
q1_target, q2_target = self.critic_target_net(feat_, target_pi)
dc_r_q1 = tf.stop_gradient(
r + self.gamma * (1 - done) *
(q1_target - self.alpha * target_log_pi))
dc_r_q2 = tf.stop_gradient(
r + self.gamma * (1 - done) *
(q2_target - self.alpha * target_log_pi))
td_error1 = q1 - dc_r_q1
td_error2 = q2 - dc_r_q2
q1_loss = tf.reduce_mean(tf.square(td_error1) * isw)
q2_loss = tf.reduce_mean(tf.square(td_error2) * isw)
critic_loss = 0.5 * q1_loss + 0.5 * q2_loss + crsty_loss
z_a = vis_feat # [B, N]
z_out = self.encoder_target(pos)
logits = tf.matmul(
z_a, tf.matmul(self.curl_w, tf.transpose(z_out, [1, 0])))
logits -= tf.reduce_max(logits, axis=-1, keepdims=True)
curl_loss = tf.reduce_mean(
tf.keras.losses.sparse_categorical_crossentropy(
tf.range(self.batch_size), logits))
critic_grads = tape.gradient(critic_loss, self.critic_tv)
self.optimizer_critic.apply_gradients(
zip(critic_grads, self.critic_tv))
curl_grads = tape.gradient(curl_loss, [self.curl_w] +
self.encoder.trainable_variables)
self.optimizer_curl.apply_gradients(
zip(curl_grads,
[self.curl_w] + self.encoder.trainable_variables))
with tf.GradientTape() as tape:
if self.is_continuous:
mu, log_std = self.actor_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min,
self.log_std_max)
pi, log_pi = squash_rsample(mu, log_std)
entropy = gaussian_entropy(log_std)
else:
logits = self.actor_net(feat)
logp_all = tf.nn.log_softmax(logits)
gumbel_noise = tf.cast(self.gumbel_dist.sample(
[batch_size, self.a_dim]),
dtype=tf.float32)
_pi = tf.nn.softmax(
(logp_all + gumbel_noise) / self.discrete_tau)
_pi_true_one_hot = tf.one_hot(tf.argmax(_pi, axis=-1),
self.a_dim)
_pi_diff = tf.stop_gradient(_pi_true_one_hot - _pi)
pi = _pi_diff + _pi
log_pi = tf.reduce_sum(tf.multiply(logp_all, pi),
axis=1,
keepdims=True)
entropy = -tf.reduce_mean(
tf.reduce_sum(tf.exp(logp_all) * logp_all,
axis=1,
keepdims=True))
q_s_pi = self.critic_net.get_min(feat, pi)
actor_loss = -tf.reduce_mean(q_s_pi - self.alpha * log_pi)
actor_grads = tape.gradient(actor_loss, self.actor_tv)
self.optimizer_actor.apply_gradients(
zip(actor_grads, self.actor_tv))
if self.auto_adaption:
with tf.GradientTape() as tape:
if self.is_continuous:
mu, log_std = self.actor_net(feat)
log_std = clip_nn_log_std(log_std, self.log_std_min,
self.log_std_max)
norm_dist = tfp.distributions.Normal(
loc=mu, scale=tf.exp(log_std))
log_pi = tf.reduce_sum(norm_dist.log_prob(
norm_dist.sample()),
axis=-1)
else:
logits = self.actor_net(feat)
cate_dist = tfp.distributions.Categorical(logits)
log_pi = cate_dist.log_prob(cate_dist.sample())
alpha_loss = -tf.reduce_mean(
self.alpha *
tf.stop_gradient(log_pi + self.target_entropy))
alpha_grad = tape.gradient(alpha_loss, self.log_alpha)
self.optimizer_alpha.apply_gradients([(alpha_grad,
self.log_alpha)])
self.global_step.assign_add(1)
summaries = dict(
[['LOSS/actor_loss', actor_loss], ['LOSS/q1_loss', q1_loss],
['LOSS/q2_loss', q2_loss], ['LOSS/critic_loss', critic_loss],
['LOSS/curl_loss', curl_loss],
['Statistics/log_alpha', self.log_alpha],
['Statistics/alpha', self.alpha],
['Statistics/entropy', entropy],
['Statistics/q_min',
tf.reduce_min(tf.minimum(q1, q2))],
['Statistics/q_mean',
tf.reduce_mean(tf.minimum(q1, q2))],
['Statistics/q_max',
tf.reduce_max(tf.maximum(q1, q2))]])
if self.auto_adaption:
summaries.update({'LOSS/alpha_loss': alpha_loss})
return (td_error1 + td_error2) / 2., summaries
