def testSampleFromDiscretizedMixLogistic(self): batch = 2 height = 4 width = 4 num_mixtures = 5 seed = 42 logits = tf.concat( # assign all probability mass to first component [ tf.ones([batch, height, width, 1]) * 1e8, tf.zeros([batch, height, width, num_mixtures - 1]) ], axis=-1) locs = tf.random_uniform([batch, height, width, num_mixtures * 3], minval=-.9, maxval=.9) log_scales = tf.ones([batch, height, width, num_mixtures * 3]) * -1e8 coeffs = tf.atanh(tf.zeros([batch, height, width, num_mixtures * 3])) pred = tf.concat([logits, locs, log_scales, coeffs], axis=-1) locs_0 = locs[..., :3] expected_sample = tf.clip_by_value(locs_0, -1., 1.) actual_sample = common_layers.sample_from_discretized_mix_logistic( pred, seed=seed) actual_sample_val, expected_sample_val = self.evaluate( [actual_sample, expected_sample]) # Use a low tolerance: samples numerically differ, as the actual # implementation clips log-scales so they always contribute to sampling. self.assertAllClose(actual_sample_val, expected_sample_val, atol=1e-2)
def init_actor_loss(self): if self.policy_type == "GaussianPolicy": mu, std = self.policy.mu_and_std(self.states_ph) norm_dist = tfd.Normal(loc=mu, scale=std) entropy = norm_dist.entropy() log_actions = norm_dist.log_prob(tf_v1.atanh(self.actions_ph)) log_actions -= tf_v1.log(1.0 - self.actions_ph**2 + 1e-8) log_actions = tf_v1.reduce_sum(log_actions, axis=-1, keepdims=True) # log_actions = self.policy.log_prob(self.states_ph, self.actions_ph) elif self.policy_type == "DiscretePolicy": action_probs = self.policy.model(self.states_ph) entropy = -tf_v1.reduce_sum( tf_v1.multiply(tf_v1.log(action_probs), action_probs), axis=-1) hot_encoded = tf_v1.one_hot(self.actions_ph, self.action_size) log_actions = tf_v1.log( tf_v1.reduce_sum(hot_encoded * action_probs, axis=-1)) else: raise NotImplementedError( f"Received {self.policy_type}. This should have never happened!" ) log_loss = -log_actions * self.targets_ph entropy_loss = -self.alpha * entropy loss = log_loss + entropy_loss optimizer = tf_v1.train.AdamOptimizer(learning_rate=self.policy.lr_ph) train_op = get_clipped_train_op(loss, optimizer, var_list=self.policy.trainable_vars, clip_norm=self.policy.clip_norm) self.policy.setup_loss(loss, train_op)
def testDiscretizedMixLogisticLoss(self): batch = 2 height = 4 width = 4 channels = 3 num_mixtures = 5 logits = tf.concat( # assign all probability mass to first component [tf.ones([batch, height, width, 1]) * 1e8, tf.zeros([batch, height, width, num_mixtures - 1])], axis=-1) locs = tf.random_uniform([batch, height, width, num_mixtures * 3], minval=-.9, maxval=.9) log_scales = tf.random_uniform([batch, height, width, num_mixtures * 3], minval=-1., maxval=1.) coeffs = tf.atanh(tf.zeros([batch, height, width, num_mixtures * 3])) pred = tf.concat([logits, locs, log_scales, coeffs], axis=-1) # Test labels that don't satisfy edge cases where 8-bit value is 0 or 255. labels = tf.random_uniform([batch, height, width, channels], minval=-.9, maxval=.9) locs_0 = locs[..., :3] log_scales_0 = log_scales[..., :3] centered_labels = labels - locs_0 inv_stdv = tf.exp(-log_scales_0) plus_in = inv_stdv * (centered_labels + 1. / 255.) min_in = inv_stdv * (centered_labels - 1. / 255.) cdf_plus = tf.nn.sigmoid(plus_in) cdf_min = tf.nn.sigmoid(min_in) expected_loss = -tf.reduce_sum(tf.log(cdf_plus - cdf_min), axis=-1) actual_loss = common_layers.discretized_mix_logistic_loss( pred=pred, labels=labels) actual_loss_val, expected_loss_val = self.evaluate( [actual_loss, expected_loss]) self.assertAllClose(actual_loss_val, expected_loss_val, rtol=1e-5)
def log_prob(self, state, action): mu, std = self.mu_and_std(state) norm_dist = tfd.Normal(loc=mu, scale=std) log_probs = norm_dist.log_prob(tf_v1.atanh(action)) log_probs -= tf_v1.log(1.0 - action**2 + 1e-8) log_probs = tf_v1.reduce_sum(log_probs, axis=-1, keepdims=True) return log_probs