def compute_vae_loss(z, mean, logvar, sample_logit, labels):
# mean, logvar = model.encode(x)
# z = model.reparameterize(mean, logvar)
# x_logit = model.decode(z)
cross_ent = tf.nn.sigmoid_cross_entropy_with_logits(logits=sample_logit,
labels=labels)
logpx_z = -tf.reduce_sum(cross_ent, axis=[1, 2, 3])
logpz = log_normal_pdf(z, 0., 0.)
logqz_x = log_normal_pdf(z, mean, logvar)
return -tf.reduce_mean(logpx_z + logpz - logqz_x)
def grad_step_towards_output(self, latent, known_occ, known_free, belief):
#TODO: Pass in acceptable prob. Perhaps just the original prob
acceptable_prob = belief.quantiles_log_pdf[25]
with tf.GradientTape() as tape:
# predicted_occ = self.decode(latent, apply_sigmoid=True)
# loss = tf.reduce_sum(known_output - known_output * predicted_occ)
# loss = tf.exp(loss)
# predicted_occ = self.decode(latent)
predicted_occ = self.decode(latent, apply_sigmoid=True)
#TODO: Remove hardcoded numbers and choose grad step size better
loss_known_occ = -tf.reduce_sum(known_occ * predicted_occ)
loss_known_free = tf.reduce_sum(
tf.clip_by_value(known_free * predicted_occ - 0.4, 0.0, 1.0))
log_pdf = log_normal_pdf(latent, belief.latent_prior_mean,
belief.latent_prior_logvar)
loss_latent_prob = -tf.clip_by_value(log_pdf, -10000,
acceptable_prob) / 10
# loss_latent_prob = -log_pdf/1000
# loss_latent_prob = tf.losses.mse(acceptable_prob, log_pdf)/100
loss = 1 + loss_known_occ + loss_known_free + loss_latent_prob
variables = [latent]
gradients = tape.gradient(loss, variables)
self.contact_optimizer.apply_gradients(zip(gradients, variables))
return loss
def complete_shape(req: CompleteShapeRequest):
print("Complete Shape Request Received")
pssnet = model_runner.model # type: PSSNet
elem = get_elem(test_records, 90)
# Sample random latent tables
# for i in range(100):
# # latent = tf.Variable(tf.random.normal(shape=[1, pssnet.hparams["num_latent_layers"]]))
# latent = tf.Variable(pssnet.sample_latent(elem))
# vg = pssnet.decode(latent, apply_sigmoid=True)
# VG_PUB.publish('predicted_occ', vg)
# # PT_PUB.publish('predicted_occ', vg)
# rospy.sleep(0.1)
# latent = tf.Variable(tf.random.normal(shape=[1, pssnet.hparams["num_latent_layers"]]))
latent = tf.Variable(pssnet.sample_latent(elem))
mean, logvar = pssnet.encode(stack_known(elem))
latent = tf.Variable(pssnet.apply_flow_to_latent_box(sample_gaussian(mean, logvar)))
log_normal_pdf(latent, mean, logvar)
p = lambda v: np.exp(log_normal_pdf(v, mean, logvar).numpy()[0])
known_contact = tf.Variable(tf.zeros((1, 64, 64, 64, 1)))
# known_free = tf.Variable(tf.zeros((1, 64, 64, 64, 1)))
transformed_free = PT_PUB.transformer.transform_to_frame(req.known_free, "object")
known_free = pointcloud2_msg_to_vox(transformed_free, scale=0.05)
# VG_PUB.publish('known_free', known_free)
# time.sleep(3)
for i in range(10):
vg = pssnet.decode(latent, apply_sigmoid=True)
VG_PUB.publish('predicted_occ', vg)
VG_PUB.publish('known_free', known_free)
if i == 0:
rospy.sleep(3)
# PT_PUB.publish('predicted_occ', vg)
pssnet.grad_step_towards_output(latent, known_contact, known_free)
rospy.sleep(0.1)
return CompleteShapeResponse()
def enforce_contact(latent: tf.Variable, known_free, chss, pssnet: PSSNet, belief: ParticleBelief,
vg_pub: VoxelgridPublisher, verbose=True):
success = True
vg_pub.publish("aux", 0 * known_free)
vg_pub.publish('predicted_occ', pssnet.decode(latent, apply_sigmoid=True))
pred_occ = pssnet.decode(latent, apply_sigmoid=True)
known_contact = get_assumed_occ(pred_occ, chss)
# any_chs = tf.reduce_max(chss, axis=-1, keepdims=True)
vg_pub.publish('known_free', known_free)
if chss is None:
vg_pub.publish('chs', known_free * 0)
else:
vg_pub.publish('chs', chss)
if verbose:
print()
log_pdf = log_normal_pdf(latent, belief.latent_prior_mean, belief.latent_prior_logvar)
quantile = belief.get_quantile(log_pdf)
if verbose:
print(f"Before optimization logprob: {log_pdf} ({quantile} quantile)")
if satisfies_constraints(pred_occ, known_contact, known_free):
if verbose:
print(f"{Fore.GREEN}Particle satisfies constraints without any gradient needed{Fore.RESET}")
return latent, True
prev_loss = 0.0
for i in range(GRADIENT_UPDATE_ITERATION_LIMIT):
# single_free = get_most_wrong_free(pred_occ, known_free)
loss = pssnet.grad_step_towards_output(latent, known_contact, known_free, belief)
if verbose:
print('\rloss: {}'.format(loss), end='')
pred_occ = pssnet.decode(latent, apply_sigmoid=True)
known_contact = get_assumed_occ(pred_occ, chss)
vg_pub.publish('predicted_occ', pred_occ)
vg_pub.publish('known_contact', known_contact)
if satisfies_constraints(pred_occ, known_contact, known_free):
if verbose:
print(
f"\t{Fore.GREEN}All known free have less that {KNOWN_FREE_LIMIT} prob occupancy, "
f"and chss have value > {KNOWN_OCC_LIMIT}{Fore.RESET}")
break
if loss == prev_loss:
if verbose:
print("\tNo progress made. Accepting shape as is")
success = False
break
prev_loss = loss
if tf.math.is_nan(loss):
if verbose:
print("\tLoss is nan. There is a problem I am not addressing")
success = False
break
else:
success = False
if verbose:
print()
if np.max(pred_occ * known_free) > KNOWN_FREE_LIMIT:
if verbose:
print("Optimization did not satisfy known freespace: ")
vg_pub.publish("aux", pred_occ * known_free)
if tf.reduce_min(tf.boolean_mask(pred_occ, known_contact)) < KNOWN_OCC_LIMIT:
if verbose:
print("Optimization did not satisfy assumed occ: ")
vg_pub.publish("aux", (1 - pred_occ) * known_contact)
if verbose:
print('\tWarning, enforcing contact terminated due to max iterations, not actually satisfying contact')
log_pdf = log_normal_pdf(latent, belief.latent_prior_mean, belief.latent_prior_logvar)
quantile = belief.get_quantile(log_pdf)
if verbose:
print(f"Latent logprob {log_pdf} ({quantile} quantile)")
return latent, success
def logp(sample, mean, var):
u = np.array([mean], dtype=np.float32)
sigma = np.array([var], dtype=np.float32)
return log_normal_pdf(sample, u, np.log(sigma)).numpy()[0]
def p(dist_from_mean, logvar):
u = np.array([[0, 0]], dtype=np.float32)
v = np.array([logvar], dtype=np.float32)
return log_normal_pdf([dist_from_mean], u, v).numpy()[0]
def p(dist_from_mean, logvar):
u = np.array([[0]], dtype=np.float32)
v = np.array([[logvar]], dtype=np.float32)
return np.exp(log_normal_pdf(dist_from_mean, u, v).numpy()[0])
def p(dist_from_mean, var):
u = np.array([[0]], dtype=np.float32)
root_sigma = np.array([[var]], dtype=np.float32)
return np.exp(log_normal_pdf(dist_from_mean, u, np.log(root_sigma)).numpy()[0])
def p(sample, mean, var):
u = np.array([[mean]], dtype=np.float32)
sigma = np.array([[var]], dtype=np.float32)
return np.exp(log_normal_pdf(sample, u, np.log(sigma)).numpy()[0])
def compute_angle_loss(true, mean, logvar):
return tf.reduce_mean(
-log_normal_pdf(tf.expand_dims(true, 1), mean, logvar))
def compute_box_loss(gt_latent_box, inferred_latent_box_mean,
inferred_latent_box_logvar):
losses = -log_normal_pdf(gt_latent_box, inferred_latent_box_mean,
inferred_latent_box_logvar)
return tf.reduce_mean(losses)
