def _build(self, inputs):
locs = self._modules['locs'](inputs)
log_scales = self._modules['scales'](inputs)
logits = self._modules['logits'](inputs)
scales = tf.nn.softplus(log_scales + softplus_inverse(1.0))
locs = tf.reshape(locs, [-1, self.k, self.ndim])
scales = tf.reshape(scales, [-1, self.k, self.ndim])
logits = tf.reshape(logits, [-1, self.k])
# reshape so that the first dim is the mixture, because we are doing to unstack them
# also swap the batch size and the ones that come from the steps of this run
# (K x N x D)
mix_first_locs = tf.transpose(locs, [1, 0, 2])
mix_first_scales = tf.transpose(scales, [1, 0, 2])
outs = {'locs': locs, 'scales': scales, 'logits': logits}
outs['flattened'] = flatten_mdn(logits, locs, scales, self.FLAGS)
cat = tfd.Categorical(logits=logits)
components = []
eval_components = []
for loc, scale in zip(tf.unstack(mix_first_locs),
tf.unstack(mix_first_scales)):
normal = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)
components.append(normal)
eval_normal = tfd.MultivariateNormalDiag(loc=loc[..., :2],
scale_diag=scale[..., :2])
eval_components.append(eval_normal)
mixture = tfd.Mixture(cat=cat, components=components)
eval_cat = tfd.Categorical(logits=logits)
eval_mixture = tfd.Mixture(cat=eval_cat, components=eval_components)
outs['mixture'] = mixture
outs['eval_mixture'] = eval_mixture
return outs
def get_pdf(param_vec, vehicle_type):
# see https://ericmjl.github.io/blog/2019/5/29/reasoning-about-shapes-and-probability-distributions/
# for info on shapes
if vehicle_type == 'other_vehicle':
alpha, mus, sigmas = slice_pvector(param_vec, vehicle_type) # Unpack parameter vectors
mvn = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=alpha),
components_distribution=tfd.Normal(
loc=mus,
scale=sigmas))
if vehicle_type == 'merge_vehicle':
alphas, mus_long, sigmas_long, mus_lat, \
sigmas_lat, rhos = slice_pvector(param_vec, vehicle_type)
cov = get_CovMatrix(rhos, sigmas_long, sigmas_lat)
mus = tf.stack([mus_long, mus_lat], axis=3, name='mus')
mvn = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(
probs=alphas),
components_distribution=tfd.MultivariateNormalTriL(
loc=mus,
scale_tril=tf.linalg.cholesky(cov), name='MultivariateNormalTriL'))
# print('mus shape: ', mus.shape)
return mvn
def conditional_distribution(self, x, slice_in, slice_out):
marginal_in = ds.MixtureSameFamily(
mixture_distribution=ds.Categorical(probs=self._priors),
components_distribution=ds.MultivariateNormalFullCovariance(
loc=self._locs[:, slice_in],
covariance_matrix=self._covs[:, slice_in, slice_in]))
p_k_in = ds.Categorical(logits=tf_utils.log_normalize(
marginal_in.components_distribution.log_prob(x[:, None]) +
marginal_in.mixture_distribution.logits[None],
axis=1))
sigma_in_out = self._covs[:, slice_in, slice_out]
inv_sigma_in_in = tf.linalg.inv(self._covs[:, slice_in, slice_in])
inv_sigma_out_in = tf.matmul(sigma_in_out,
inv_sigma_in_in,
transpose_a=True)
A = inv_sigma_out_in
b = self._locs[:, slice_out] - tf.matmul(
inv_sigma_out_in, self._locs[:, slice_in, None])[:, :, 0]
cov_est = (self._covs[:, slice_out, slice_out] -
tf.matmul(inv_sigma_out_in, sigma_in_out))
ys = tf.einsum('aij,bj->abi', A, x) + b[:, None]
p_out_in_k = ds.MultivariateNormalFullCovariance(
tf.transpose(ys, perm=(1, 0, 2)), cov_est)
return ds.MixtureSameFamily(mixture_distribution=p_k_in,
components_distribution=p_out_in_k)
def call(self, x, evaluation=False, epsilon=0, temp=1):
self.logits = logits = self._layers(x)
if evaluation:
temp = tf.where(temp == 0, 1e-9, temp)
logits = logits / temp
dist = tfd.Categorical(logits)
action = dist.sample()
else:
if self._epsilon_scaled_logits and \
(isinstance(epsilon, tf.Tensor) or epsilon):
self.logits = epsilon_scaled_logits(logits, epsilon, temp)
else:
self.logits = logits / temp
dist = tfd.Categorical(self.logits)
action = dist.sample()
if not self._epsilon_scaled_logits and \
(isinstance(epsilon, tf.Tensor) or epsilon):
action = epsilon_greedy(action, epsilon, True, self.action_dim)
self._dist = dist
self._action = action
return action
def mix(gamma, eta, loc, scale, neg_inf, n):
return tfd.Mixture(
cat=tfd.Categorical(probs=tf.stack([gamma, 1 - gamma], axis=-1)),
components=[
tfd.Sample(tfd.Normal(np.float64(neg_inf), 1e-5), sample_shape=n),
tfd.Sample(tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=eta),
components_distribution=tfd.Normal(loc=loc, scale=scale)),
sample_shape=n)
])
def mix(gamma, eta, loc, scale, neg_inf):
_gamma = gamma[..., tf.newaxis]
# FIXME: Possible to use tfd.Blockwise?
return tfd.Mixture(
cat=tfd.Categorical(probs=tf.concat([_gamma, 1 - _gamma], axis=-1)),
components=[
tfd.Deterministic(np.float64(neg_inf)),
tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=eta),
components_distribution=tfd.Normal(loc=loc, scale=scale)),
])
def __call__(self):
"""Get the distribution object from the backend"""
if get_backend() == 'pytorch':
import torch.distributions as tod
raise NotImplementedError
else:
from tensorflow_probability import distributions as tfd
return tfd.HiddenMarkovModel(
initial_distribution=tfd.Categorical(self['initial']),
transition_distribution=tfd.Categorical(self['transition']),
observation_distribution=self['observation'],
num_steps=self['steps'])
def train_HMM(self):
# Define variable to represent the unknown log rates.
_trainable_log_rates = tf.Variable(
np.log(np.mean(self.observed_counts)) +
tf.random.normal([self.num_states]),
name='log_rates')
self.hmm = tfd.HiddenMarkovModel(
initial_distribution=tfd.Categorical(logits=initial_state_logits),
transition_distribution=tfd.Categorical(
probs=self._transition_probs),
observation_distribution=tfd.Poisson(
log_rate=_trainable_log_rates),
num_steps=len(observed_counts))
def call(self, inputs): prob, y = inputs uni_prob = K.ones_like(prob)/tf.cast(K.shape(prob)[-1], tf.float32) #variational distribution var_dist = tfp.Categorical(prob) #prior uniform distribution pri_dist = tfp.Categorical(uni_prob) ent_loss = K.mean(var_dist.entropy()) kl_loss = K.mean(tfp.kl_divergence(var_dist, pri_dist)) cent_loss = K.mean(K.categorical_crossentropy(y, prob)) self.add_loss( self.rho*cent_loss + \ (1-self.rho)*(self.lamb_kl*kl_loss+self.lamb_ent*ent_loss) ) return kl_loss
def mdn_loss_func(y_true, y_pred):
# Reshape inputs in case this is used in a TimeDistribued layer
y_pred = tf.reshape(y_pred,
[-1, (2 * num_mixes * output_dim) + num_mixes],
name='reshape_ypreds')
y_true = tf.reshape(y_true, [-1, output_dim], name='reshape_ytrue')
# Split the inputs into paramaters
out_mu, out_sigma, out_pi = tf.split(y_pred,
num_or_size_splits=[
num_mixes * output_dim,
num_mixes * output_dim,
num_mixes
],
axis=-1,
name='mdn_coef_split')
# Construct the mixture models
cat = tfd.Categorical(logits=out_pi)
component_splits = [output_dim] * num_mixes
mus = tf.split(out_mu, num_or_size_splits=component_splits, axis=1)
sigs = tf.split(out_sigma, num_or_size_splits=component_splits, axis=1)
coll = [
tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)
for loc, scale in zip(mus, sigs)
]
mixture = tfd.Mixture(cat=cat, components=coll)
loss = mixture.log_prob(y_true)
loss = tf.negative(loss)
loss = tf.reduce_mean(loss)
return loss
def generate_grm_data(n_sample, n_factor, n_item,
nu, ld, rho,
dtype = tf.float64):
if (n_item % n_factor) != 0:
n_item = n_factor * (n_item // n_factor)
item_per_factor = (n_item // n_factor)
n_category = len(nu) + 1
intercept = tf.tile(tf.constant([nu], dtype = dtype),
multiples = [n_item, 1])
loading = np.zeros((n_item, n_factor))
for i in range(n_factor):
for j in range(i * item_per_factor,
(i + 1) * item_per_factor):
loading[j, i] = ld
loading = tf.constant(loading, dtype = dtype)
if rho is None:
cor = tf.eye(n_factor, dtype = dtype)
else:
unit = tf.ones((n_factor, 1), dtype = dtype)
identity = tf.eye(n_factor, dtype = dtype)
cor = rho * (unit @ tf.transpose(unit)) + (1 - rho) * identity
dist_eta = tfd.MultivariateNormalTriL(
loc = tf.zeros(n_factor, dtype = dtype),
scale_tril = tf.linalg.cholesky(cor))
eta = dist_eta.sample(n_sample)
c, d = create_cd(n_category, dtype)
probs = grm_irf(eta, intercept, loading, c, d)
x = tfd.Categorical(probs=probs, dtype=dtype).sample()
return x
def mix(n, eta, loc, scale, name):
return tfd.Sample(
tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=eta),
components_distribution=tfd.Normal(loc=loc, scale=scale),
name=name),
sample_shape=n)
def call(self, inputs, **kwargs):
latent = self._base_model(inputs)
logits = self._policy_fn(latent)
deterministic = kwargs.get("deterministic", False)
if self._discrete:
pi = tfd.Categorical(logits)
if deterministic:
actions = pi.mode()
else:
actions = pi.sample()
logpacs = tf.nn.log_softmax(logits)
else:
means, logstds = tf.split(logits, 2, axis=-1)
logstds = tf.clip_by_value(logstds, MIN_LOG_NN_OUTPUT, MAX_LOG_NN_OUTPUT)
pi = tfd.MultivariateNormalDiag(means, tf.exp(logstds))
if deterministic:
actions = pi.mean()
else:
actions = pi.sample()
logpacs = pi.log_prob(actions)
# Adjust loglikelihoods for squashed actions
logpacs -= tf.reduce_sum(
2 * (np.log(2) - actions - tf.nn.softplus(-2 * actions)), axis=1
)
actions = tf.math.tanh(actions)
actions = tf.clip_by_value(actions, -1 + SMALL_NUMBER, 1 - SMALL_NUMBER)
actions *= self._action_limit
return actions, logpacs
def prepare_t_student_mixture(self):
cluster_weights_unnorm = self.get_cluster_unnormalized_weights()
cluster_weights_norm = cluster_weights_unnorm / np.sum(
cluster_weights_unnorm)
cat_distr = tpd.Categorical(
probs=tf.cast(cluster_weights_norm, dtype=tf.float32))
t_student_params = [
self.prepare_t_student_params(ind)
for ind in range(self.clusters_num)
]
dofs = tf.constant(
[t_student_params[ind]['df'] for ind in range(self.clusters_num)],
dtype=tf.float32)
means = tf.stack([
t_student_params[ind]['mean'] for ind in range(self.clusters_num)
])
cov_chols = tf.stack([
t_student_params[ind]['cov_chol']
for ind in range(self.clusters_num)
])
t_student_distr = tpd.MultivariateStudentTLinearOperator(
df=dofs,
loc=means,
scale=tf.linalg.LinearOperatorLowerTriangular(cov_chols))
t_student_mixture = tpd.MixtureSameFamily(
mixture_distribution=cat_distr,
components_distribution=t_student_distr)
return t_student_mixture
def output_function(self, state):
params = dense_layer(state.h3,
self.output_units,
scope='gmm',
reuse=tf.compat.v1.AUTO_REUSE)
pis, mus, sigmas, rhos, es = self._parse_parameters(params)
mu1, mu2 = tf.split(mus, 2, axis=1)
mus = tf.stack([mu1, mu2], axis=2)
sigma1, sigma2 = tf.split(sigmas, 2, axis=1)
covar_matrix = [
tf.square(sigma1), rhos * sigma1 * sigma2, rhos * sigma1 * sigma2,
tf.square(sigma2)
]
covar_matrix = tf.stack(covar_matrix, axis=2)
covar_matrix = tf.reshape(
covar_matrix,
(self.batch_size, self.num_output_mixture_components, 2, 2))
mvn = tfd.MultivariateNormalFullCovariance(
loc=mus, covariance_matrix=covar_matrix)
b = tfd.Bernoulli(probs=es)
c = tfd.Categorical(probs=pis)
sampled_e = b.sample()
sampled_coords = mvn.sample()
sampled_idx = c.sample()
idx = tf.stack([tf.range(self.batch_size), sampled_idx], axis=1)
coords = tf.gather_nd(sampled_coords, idx)
return tf.concat([coords, tf.cast(sampled_e, tf.float32)], axis=1)
def _train_auxiliary(self, obs, returns, old_pi_logits):
if self.policy.is_discrete:
old_pi = tfd.Categorical(old_pi_logits)
else:
mean, logstd = tf.split(old_pi_logits, 2, axis=-1)
old_pi = tfd.MultivariateNormalDiag(mean, tf.exp(logstd))
with tf.GradientTape() as tape:
pi, aux_value_preds = self.policy.auxiliary_heads(obs)
bc_loss = tf.reduce_mean(old_pi.kl_divergence(pi))
aux_value_loss = 0.5 * tf.reduce_mean((returns - aux_value_preds) ** 2)
joint_loss = aux_value_loss + self.bc_coef * bc_loss
grads = tape.gradient(joint_loss, self.policy.auxiliary_weights)
if self.max_grad_norm is not None:
grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm)
self.aux_optimizer.apply_gradients(zip(grads, self.policy.auxiliary_weights))
with tf.GradientTape() as tape:
value_preds = self.policy.value(obs)
value_loss = 0.5 * tf.reduce_mean((returns - value_preds) ** 2)
grads = tape.gradient(value_loss, self.policy.value_weights)
if self.max_grad_norm is not None:
grads, _ = tf.clip_by_global_norm(grads, self.max_grad_norm)
if self.value_optimizer is not None:
self.value_optimizer.apply_gradients(zip(grads, self.policy.value_weights))
else:
self.policy_optimizer.apply_gradients(zip(grads, self.policy.value_weights))
def get_resample_idx(self, log_W, sample_size=()):
"""
Get resample index a_t^k ~ Categorical(w_t^1, ..., w_t^K) with logits = log_W last axis
Input:
log_W.shape = (K, batch_size_0, ..., batch_size_last)
"""
nb_classes = log_W.shape.as_list()[0]
batch_shape = log_W.shape.as_list()[1:]
perm = list(range(1, len(batch_shape) + 1)) + [0]
log_W = tf.transpose(log_W, perm=perm)
categorical = tfd.Categorical(logits=log_W, validate_args=True, name="Categorical")
# sample multiple times to remove idx out of range
if sample_size == ():
idx_shape = batch_shape
else:
assert isinstance(sample_size, int), "sample_size should be int, {} is given".format(sample_size)
idx_shape = [sample_size] + batch_shape
idx = tf.ones(idx_shape, dtype=tf.int32) * nb_classes
for _ in range(1):
fixup_idx = categorical.sample(sample_size)
idx = tf.where(idx >= nb_classes, fixup_idx, idx)
# if still got idx out of range, replace them with idx from uniform distribution
final_fixup = tf.random.uniform(idx_shape, maxval=nb_classes, dtype=tf.int32)
idx = tf.where(idx >= nb_classes, final_fixup, idx)
batch_idx = np.meshgrid(*[range(i) for i in idx_shape], indexing='ij')
if sample_size != ():
batch_idx = batch_idx[1:]
resample_idx = tf.stack([idx] + batch_idx, axis=-1)
return resample_idx
def mdn_head(h, FLAGS):
with tf.variable_scope('mdn'):
locs = tf.reshape(tf.layers.dense(h, 2 * FLAGS['k'], activation=None),
[-1, FLAGS['k'], 2])
scales = tf.reshape(
tf.layers.dense(h, 2 * FLAGS['k'], activation=tf.exp),
[-1, FLAGS['k'], 2])
logits = tf.layers.dense(h, FLAGS['k'], activation=None)
cat = tfd.Categorical(logits=logits)
components = []
eval_components = []
for loc, scale in zip(tf.unstack(tf.transpose(locs, [1, 0, 2])),
tf.unstack(tf.transpose(scales, [1, 0, 2]))):
# TODO: does this need to be a more complex distribution?
normal = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)
components.append(normal)
eval_normal = tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)
eval_components.append(eval_normal)
mixture = tfd.Mixture(cat=cat, components=components)
eval_mixture = tfd.Mixture(cat=cat, components=eval_components)
return {
'locs': locs,
'scales': scales,
'logits': logits,
'mixture': mixture,
'eval_mixture': eval_mixture
}
def get_dist(self, y_pred):
"""turns an output into a distribution. Literally use this as you'd use the normal keras predict.
Args:
y_pred: nn output
Returns:
a probability distribution over outputs, for each input.
"""
num_mix = self.num_mix
output_dim = self.output_dim
y_pred = tf.reshape(y_pred, [-1, (2 * num_mix * output_dim) + num_mix],
name='reshape_ypreds')
out_mu, out_sigma, out_pi = tf.split(y_pred,
num_or_size_splits=[
num_mix * output_dim,
num_mix * output_dim, num_mix
],
axis=1,
name='mdn_coef_split')
cat = tfd.Categorical(logits=out_pi)
component_splits = [output_dim] * num_mix
mus = tf.split(out_mu, num_or_size_splits=component_splits, axis=1)
sigs = tf.split(out_sigma, num_or_size_splits=component_splits, axis=1)
coll = [
tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)
for loc, scale in zip(mus, sigs)
]
return tfd.Mixture(cat=cat, components=coll)
def learn(self):
actions = np.array(self.action_memory) # Get actions taken
G = self.calc_advantages() # Calc advantages of each action
with GradientTape() as tape: # Calculate gradients for all weights
loss = 0
for idx, (g, state) in enumerate(zip(G, self.state_memory)):
state = convert_to_tensor(
[state], dtype='float32') # Makes training faster
probs = self.policy(state, training=False)[
0] # Get probability predictions based on each step state
action_probs = distributions.Categorical(
probs=probs) # Create distribution based on probabilities
log_prob = action_probs.log_prob(
actions[idx]
) # Find logistic probability of action being taken
loss += -g * squeeze(
log_prob
) # Update loss based on reward and model's certainty
# Calculate gradients from observing tape, then apply them to the policy weights
gradient = tape.gradient(loss, self.policy.trainable_variables)
self.policy.optimizer.apply_gradients(
zip(gradient, self.policy.trainable_variables))
self.clear_steps() # Clear step memory
def _exploration(self, action, training):
if training:
amount = self._c.expl_amount
if self._c.expl_decay:
amount *= 0.5**(tf.cast(self._step, tf.float32) /
self._c.expl_decay)
if self._c.expl_min:
amount = tf.maximum(self._c.expl_min, amount)
self._metrics['expl_amount'].update_state(amount)
elif self._c.eval_noise:
amount = self._c.eval_noise
else:
return action
if self._c.expl == 'additive_gaussian':
return tf.clip_by_value(tfd.Normal(action, amount).sample(), -1, 1)
if self._c.expl == 'completely_random':
return tf.random.uniform(action.shape, -1, 1)
if self._c.expl == 'epsilon_greedy':
indices = tfd.Categorical(0 * action).sample()
rnd = tf.random.uniform(action.shape[:1], 0, 1) < amount
rnd = tf.expand_dims(rnd, axis=-1)
one_hot = tf.one_hot(indices, action.shape[-1], dtype=self._float)
return tf.where(tf.repeat(rnd, 4, axis=-1), one_hot, action)
# return tf.where(
# tf.random.uniform(action.shape[:1], 0, 1) < amount,
# tf.one_hot(indices, action.shape[-1], dtype=self._float),
# action)
raise NotImplementedError(self._c.expl)
def sample_is(self, x, n=1):
mixture_distribution, mixture_components = \
self._gate.conditional_mixture_distribution(x),\
self._experts.conditional_components_distribution(x)
y = mixture_components.sample(n)
# npdt = y.dtype.as_numpy_dtype
is_logits = self._is_function(mixture_distribution.logits)
is_mixture_distribution = ds.Categorical(logits=is_logits)
idx = is_mixture_distribution.sample(n)
# TODO check if we should not renormalize mixture.logits - tf.stop_...
weights = tf.batch_gather(
mixture_distribution.logits - tf.stop_gradient(is_logits),
tf.transpose(idx))
# TODO check axis
# weights = tf.batch_gather(
# log_normalize(mixture_distribution.logits - tf.stop_gradient(is_logits), axis=1),
# tf.transpose(idx))
if n == 1:
return tf.batch_gather(y, idx[:, :,
None])[0, :,
0], tf.transpose(weights)[0]
else:
return tf.batch_gather(y, idx[:, :,
None])[:, :,
0], tf.transpose(weights)
def get_distribution(self, x, **kwargs):
"""Build the mixture distribution implied by the set of oracles
that are trained in this module
Args:
x: tf.Tensor
a batch of training inputs shaped like [batch_size, channels]
Returns:
distribution: tfpd.Distribution
the mixture of gaussian distributions implied by the oracles
"""
# get the distribution parameters for all models
params = defaultdict(list)
for fm in self.forward_models:
for key, val in fm.get_params(x, **kwargs).items():
params[key].append(val)
# stack the parameters in a new component axis
for key, val in params.items():
params[key] = tf.stack(val, axis=-1)
# build the mixture distribution using the family of component one
weights = tf.fill([self.bootstraps], 1 / self.bootstraps)
return tfpd.MixtureSameFamily(
tfpd.Categorical(probs=weights),
self.forward_models[0].distribution(**params))
def _exploration(self, action, training):
if training:
amount = self._c.expl_amount
if self._c.expl_decay: # 0.3 is True
amount *= 0.5**(tf.cast(self._step, tf.float32) /
self._c.expl_decay)
if self._c.expl_min: # 0.0 is False
amount = tf.maximum(self._c.expl_min, amount)
self._metrics["expl_amount"].update_state(amount)
elif self._c.eval_noise:
amount = self._c.eval_noise
else:
return action
if self._c.expl == "additive_gaussian":
return tf.clip_by_value(tfd.Normal(action, amount).sample(), -1, 1)
if self._c.expl == "completely_random":
return tf.random.uniform(action.shape, -1, 1)
if self._c.expl == "epsilon_greedy":
# print(
# "0 * action:", 0 * action
# ) # 0 * action: Tensor("mul:0", shape=(1, 4), dtype=float16)
indices = tfd.Categorical(0 * action).sample()
# print("epsilon_greedy indices:", indices) # shape=(1,)
return tf.where(
tf.random.uniform(action.shape[:1], 0, 1) <
amount, # randomly decide do epsilon greedy or not
tf.one_hot(indices, action.shape[-1], dtype=tf.float32),
action,
)
raise NotImplementedError(self._c.expl)
def mse_func(y_true, y_pred):
# Reshape inputs in case this is used in a TimeDistribued layer
y_pred = tf.reshape(y_pred,
[-1, (2 * num_mixes * output_dim) + num_mixes],
name='reshape_ypreds')
y_true = tf.reshape(y_true, [-1, output_dim], name='reshape_ytrue')
out_mu, out_sigma, out_pi = tf.split(y_pred,
num_or_size_splits=[
num_mixes * output_dim,
num_mixes * output_dim,
num_mixes
],
axis=1,
name='mdn_coef_split')
cat = tfd.Categorical(logits=out_pi)
component_splits = [output_dim] * num_mixes
mus = tf.split(out_mu, num_or_size_splits=component_splits, axis=1)
sigs = tf.split(out_sigma, num_or_size_splits=component_splits, axis=1)
coll = [
tfd.MultivariateNormalDiag(loc=loc, scale_diag=scale)
for loc, scale in zip(mus, sigs)
]
mixture = tfd.Mixture(cat=cat, components=coll)
samp = mixture.sample()
mse = tf.reduce_mean(tf.square(samp - y_true), axis=-1)
# Todo: temperature adjustment for sampling functon.
return mse
def __call__(self):
"""Get the distribution object from the backend"""
if get_backend() == "pytorch":
# import torch.distributions as tod
raise NotImplementedError
else:
import tensorflow as tf
from tensorflow_probability import distributions as tfd
# Convert to tensorflow distributions if probflow distributions
if isinstance(self.distributions, BaseDistribution):
self.distributions = self.distributions()
# Broadcast probs/logits
shape = self.distributions.batch_shape
args = {"logits": None, "probs": None}
if self.logits is not None:
args["logits"] = tf.broadcast_to(self["logits"], shape)
else:
args["probs"] = tf.broadcast_to(self["probs"], shape)
# Return TFP distribution object
return tfd.MixtureSameFamily(
tfd.Categorical(**args), self.distributions
)
def sample(self, time, outputs, state, name=None):
"""
sample tokens for next step, notice the special form
of 'state'([decoded_ids, rnn_state])
"""
sample_method_sampler = \
tfpd.Categorical(probs=self._lambdas)
sample_method_id = sample_method_sampler.sample()
truth_feeding = lambda: tf.cond(
tf.less(time, tf.shape(self._ground_truth)[1]),
lambda: tf.cast(self._ground_truth[:, time], tf.int32),
lambda: tf.ones_like(self._ground_truth[:, 0],
dtype=tf.int32) * self._vocab.eos_token_id)
self_feeding = lambda: SampleEmbeddingHelper.sample(
self, time, outputs, state, name)
reward_feeding = lambda: self._sample_by_reward(time, state)
sample_ids = tf.cond(
tf.logical_or(tf.equal(time, 0), tf.equal(sample_method_id, 1)),
truth_feeding,
lambda: tf.cond(
tf.equal(sample_method_id, 2),
reward_feeding,
self_feeding))
return sample_ids
def get_distributions_from_tensor(t, dimension, num_mixes):
y_pred = tf.reshape(t, [-1, (2 * num_mixes * dimension + 1) + num_mixes],
name='reshape_ypreds')
out_e, out_pi, out_mus, out_stds = tf.split(y_pred,
num_or_size_splits=[
1, num_mixes,
num_mixes * dimension,
num_mixes * dimension
],
name='mdn_coef_split',
axis=-1)
cat = tfd.Categorical(logits=out_pi)
components_splits = [dimension] * num_mixes
mus = tf.split(out_mus, num_or_size_splits=components_splits, axis=1)
stds = tf.split(out_stds, num_or_size_splits=components_splits, axis=1)
components = [
tfd.MultivariateNormalDiag(loc=mu_i, scale_diag=std_i)
for mu_i, std_i in zip(mus, stds)
]
mix = tfd.Mixture(cat=cat, components=components)
stroke = tfd.Bernoulli(logits=out_e)
return mix, stroke
def make_mixture_prior(latent_size, mixture_components):
"""Creates the mixture of Gaussians prior distribution.
Args:
latent_size: The dimensionality of the latent representation.
mixture_components: Number of elements of the mixture.
Returns:
random_prior: A `tfd.Distribution` instance representing the distribution
over encodings in the absence of any evidence.
"""
if mixture_components == 1:
return tfd.MultivariateNormalDiag(
loc=tf.zeros([latent_size]), scale_identity_multiplier=1.0)
loc = tf.get_variable(name="loc", shape=[mixture_components, latent_size])
raw_scale_diag = tf.get_variable(
name="raw_scale_diag", shape=[mixture_components, latent_size])
mixture_logits = tf.get_variable(
name="mixture_logits", shape=[mixture_components])
return tfd.MixtureSameFamily(
components_distribution=tfd.MultivariateNormalDiag(
loc=loc, scale_diag=tf.nn.softplus(raw_scale_diag)),
mixture_distribution=tfd.Categorical(logits=mixture_logits),
name="prior")
def get_steering(preds):
alpha, mu, sigma = slice_parameter_vectors(preds.numpy(), components)
# print(alpha)
max_prob = np.max(alpha, axis=-1)
if max_prob > 0.9995:
index = np.argmax(alpha, axis=-1)
angle = mu[:, index[0]]
else:
angle = np.multiply(alpha, mu).sum(axis=-1)
# print(angle)
gm = tfd.MixtureSameFamily(
mixture_distribution=tfd.Categorical(probs=alpha),
components_distribution=tfd.Normal(loc=mu, scale=sigma),
)
x = np.linspace(-1, 1, int(1e3))
pyx = gm.prob(x)
plot = cv2.plot.Plot2d_create(
np.array(x).astype(np.float64),
np.array(pyx).astype(np.float64))
plot.setPlotBackgroundColor((255, 255, 255))
plot.setInvertOrientation(True)
plot.setPlotLineColor(0)
plot = plot.render()
cv2.imshow("Distribution", plot)
return angle[0]
