def vec_to_tri(vectors):
"""
Takes a DxM tensor vectors and maps it to a D x matrix_size x matrix_size
tensor where the lower triangle of each matrix_size x matrix_size matrix
is constructed by unpacking each M-vector.
"""
return tfd.fill_triangular(vectors)
def iaf_scale_from_matmul(shifted_codes, name=None):
"""Returns IAF scale matrix generated by lower-triangular mat-mul.
Args:
shifted_codes: Tensor with shape [num_samples, batch_size, latent_size,
num_codes], with the first latent_size dimension a tensor of zeros and the
others shifted down one (with the original last latent dimension missing).
name: String used for name scope.
Returns:
unconstrained_scale: Tensor with shape [latent_size, latent_size],
generated by multiplying padded codes by lower-triangular matrix. It can
take on any real value as it will later be passed through a softplus.
"""
with tf.name_scope(name, default_name="iaf_scale_from_transformer") as name:
latent_size = int(shifted_codes.shape[2])
scale_matrix = tf.get_variable(name + "scale_matrix",
[latent_size * (latent_size + 1) / 2],
dtype=tf.float32,
initializer=tf.zeros_initializer())
scale_bijector = tfb.Affine(
scale_tril=tfd.fill_triangular(scale_matrix))
unconstrained_scale = scale_bijector.forward(
tf.transpose(shifted_codes, [0, 1, 3, 2]))
# Transpose the bijector output since it performs a batch matmul.
unconstrained_scale = tf.transpose(unconstrained_scale, [0, 1, 3, 2])
return unconstrained_scale
def affine_flow_actor_critic(a, k):
d = r = act_dim = a.shape.as_list()[-1]
DTYPE = tf.float32
bijectors = []
initializer = tf.initializers.truncated_normal(0, 0.1)
for i in range(k):
with tf.variable_scope('bijector_%d' % i):
V = tf.get_variable('V', [d, r],
dtype=DTYPE,
initializer=initializer)
shift = tf.get_variable('shift', [d],
dtype=DTYPE,
initializer=initializer)
L = tf.get_variable('L', [d * (d + 1) / 2],
dtype=DTYPE,
initializer=initializer)
bijectors.append(
tfpb.Affine(
scale_tril=tfpd.fill_triangular(L),
scale_perturb_factor=V,
shift=shift,
))
alpha = tf.abs(tf.get_variable('alpha', [], dtype=DTYPE)) + .01
bijectors.append(PReLU(alpha=alpha))
mlp_bijector = tfpb.Chain(list(reversed(bijectors[:-1])),
name='mlp_bijector')
dist = tfpd.TransformedDistribution(
distribution=tfpd.MultivariateNormalDiag(loc=tf.zeros(act_dim),
scale_diag=0.1 *
tf.ones(act_dim)),
bijector=mlp_bijector)
pi = dist.sample(1)
logp_pi = tf.squeeze(dist.log_prob(pi))
logp = dist.log_prob(a)
return pi, logp, logp_pi
def get_coefs(self, output):
prior, mu, sigma = tf.split(output,
[self.n_mix, self.n_mix * self.n_pred, -1],
axis=1)
with tf.control_dependencies(
self._debug_nan([prior, mu, sigma],
names=['prior', 'mu', 'sigma'])):
prior = tf.nn.softmax(prior, axis=-1) + 1e-9
# Reshape tensors so that elements remain in the correct locations
mu = tf.stack(tf.split(mu, self.n_mix, 1), 1)
sigma = tf.stack(tf.split(sigma, self.n_mix, 1), 1)
sigma = tfd.fill_triangular(sigma, upper=False)
# Explicitly set the shapes
prior.set_shape((None, self.n_mix))
mu.set_shape((None, self.n_mix, self.n_pred))
sigma.set_shape((None, self.n_mix, self.n_pred, self.n_pred))
# Independent outputs
if self.distribution == 'MultivariateNormalDiag':
sigma = tf.exp(tf.compat.v1.matrix_diag_part(sigma))
norm = tf.ones((1, 1, self.n_pred))
# Full covariance estimation
else:
sigma = tf.einsum('abij,abjk->abik',
tf.transpose(sigma, perm=[0, 1, 3, 2]),
sigma)
norm = tf.linalg.diag(tf.ones((1, 1, self.n_pred)))
# Minimum uncertainty on covariance diagonal - prevents
# matrix inversion errors, and regularizes the model
sigma += self.epsilon * norm
# _,var = tf.nn.moments(tf.stop_gradient(mu), [0])
# var = tf.abs(tf.tile(tf.expand_dims(tf.linalg.diag(var), 0), [tf.shape(sigma)[0], 1, 1, 1]))
# eps = sigma * tf.reshape(tf.eye(self.n_pred), (1, 1, self.n_pred, self.n_pred))
# sigma-= eps
# sigma+= tf.clip_by_value(eps, var * 1e-3, np.inf) + 1e-8
# sigma+= tf.clip_by_value(eps, tf.abs(tf.linalg.diag(mu)) * 1e-2, np.inf) + 1e-8
# sigma += tf.abs(tf.linalg.diag(tf.stop_gradient(mu))) * 0.5
# Store for model loading
prior = tf.identity(prior, name='prior')
mu = tf.identity(mu, name='mu')
sigma = tf.identity(sigma, name='sigma')
self.most_likely = tf.identity(self.get_top(prior, mu),
name='most_likely')
self.avg_estimate = tf.identity(tf.reduce_sum(
mu * tf.expand_dims(prior, -1), 1),
name='avg_estimate')
self.thresholded = tf.identity(tf.compat.v2.where(
tf.expand_dims(
tf.math.greater(
tf.reduce_max(prior, 1) / self.T,
tf.math.sign(self.T)), -1), self.most_likely,
self.avg_estimate),
name='thresholded')
# For a given confidence level probability p (0<p<1), and number of dimensions d, rho is the error bar coefficient: rho=sqrt(2)*erfinv(p ** (1/d))
# https://faculty.ucmerced.edu/mcarreira-perpinan/papers/cs-99-03.pdf
top_sigma = self.get_top(prior, sigma)
avg_sigma = tf.reduce_sum(
tf.expand_dims(tf.expand_dims(prior, -1), -1) *
(sigma + tf.matmul(
tf.transpose(mu - tf.expand_dims(self.avg_estimate, 1),
(0, 2, 1)),
mu - tf.expand_dims(self.avg_estimate, 1))),
axis=1)
s_top, u_top, v_top = tf.linalg.svd(top_sigma)
s_avg, u_avg, v_avg = tf.linalg.svd(avg_sigma)
rho = 2**0.5 * erfinv(self.C**(1. / self.n_pred))
self.top_confidence = tf.identity(
rho * 2 * s_top**0.5, name='top_confidence'
) # confidence interval centered on top mu
self.avg_confidence = tf.identity(
rho * 2 * s_avg**0.5, name='avg_confidence'
) # confidence interval centered on avg mu
return prior, mu, sigma