def infer(self, obs, act):
action_probs = self.actor(obs)
dist = Categorical(probs=action_probs)
action_logprobs = dist.log_prob(act)
dist_entropy = dist.entropy()
q_value = self.critic(obs)
return action_logprobs, tf.squeeze(q_value), dist_entropy
def step(self, obs):
if obs.ndim < 2:
obs = obs[np.newaxis, :]
action_probs = self.actor(obs)
dist = Categorical(probs=action_probs)
action = dist.sample()
return action.numpy()[0], dist.log_prob(action)
def call(self, state):
if state.ndim < 2:
state = state[np.newaxis, :]
action_probs = self.net(state)
dist = Categorical(probs=action_probs)
action = dist.sample()
log_pi = dist.log_prob(action)
return action.numpy(), log_pi
def __init__(self,
data,
bandwidth=0.01,
kernel=Normal,
use_grid=False,
use_fft=False,
num_grid_points=1024,
reparameterize=False,
validate_args=False,
allow_nan_stats=True,
name='KernelDensityEstimation'):
components_distribution_generator = lambda loc, scale: Independent(
kernel(loc=loc, scale=scale))
with tf.name_scope(name) as name:
self._use_fft = use_fft
dtype = dtype_util.common_dtype([bandwidth, data], tf.float32)
self._bandwidth = tensor_util.convert_nonref_to_tensor(
bandwidth, name='bandwidth', dtype=dtype)
self._data = tensor_util.convert_nonref_to_tensor(data,
name='data',
dtype=dtype)
if (use_fft):
self._grid = self._generate_grid(num_grid_points)
self._grid_data = self._linear_binning()
mixture_distribution = Categorical(probs=self._grid_data)
components_distribution = components_distribution_generator(
loc=self._grid, scale=self._bandwidth)
elif (use_grid):
self._grid = self._generate_grid(num_grid_points)
self._grid_data = self._linear_binning()
mixture_distribution = Categorical(probs=self._grid_data)
components_distribution = components_distribution_generator(
loc=self._grid, scale=self._bandwidth)
else:
self._grid = None
self._grid_data = None
n = self._data.shape[0]
mixture_distribution = Categorical(probs=[1 / n] * n)
components_distribution = components_distribution_generator(
loc=self._data, scale=self._bandwidth)
super(KernelDensityEstimation, self).__init__(
mixture_distribution=mixture_distribution,
components_distribution=components_distribution,
reparameterize=reparameterize,
validate_args=validate_args,
allow_nan_stats=allow_nan_stats,
name=name)
def _mean(self, **kwargs):
params = self._params
if self.n_channels == 1:
component_logits, locs, scales = params
else:
# r ~ Logistic(loc_r, scale_r)
# g ~ Logistic(coef_rg * r + loc_g, scale_g)
# b ~ Logistic(coef_rb * r + coef_gb * g + loc_b, scale_b)
component_logits, locs, scales, coeffs = params
loc_tensors = tf.split(locs, self.n_channels, axis=-1)
coef_tensors = tf.split(coeffs, self.n_coeffs, axis=-1)
coef_count = 0
for i in range(self.n_channels):
for j in range(i):
loc_tensors[i] += loc_tensors[j] * coef_tensors[coef_count]
coef_count += 1
locs = tf.concat(loc_tensors, axis=-1)
## create the distrubtion
mixture_distribution = Categorical(logits=component_logits)
# Convert distribution parameters for pixel values in
# `[self._low, self._high]` for use with `QuantizedDistribution`
locs = self.low + 0.5 * (self.high - self.low) * (locs + 1.)
scales = scales * 0.5 * (self.high - self.low)
logistic_dist = TransformedDistribution(
distribution=Logistic(loc=locs, scale=scales),
bijector=Shift(shift=tf.cast(-0.5, self.dtype)))
dist = MixtureSameFamily(mixture_distribution=mixture_distribution,
components_distribution=Independent(
logistic_dist,
reinterpreted_batch_ndims=1))
mean = Independent(dist, reinterpreted_batch_ndims=2).mean()
## normalize the data back to the input domain
return _pixels_to(mean, self.inputs_domain, self.low, self.high)
def MixtureQLogistic(
locs: tf.Tensor,
scales: tf.Tensor,
logits: Optional[tf.Tensor] = None,
probs: Optional[tf.Tensor] = None,
batch_ndims: int = 0,
low: int = 0,
bits: int = 8,
name: str = 'MixtureQuantizedLogistic') -> MixtureSameFamily:
""" Mixture of quantized logistic distribution
Parameters
----------
locs : tf.Tensor
locs of all logistics components, shape `[batch_size, n_components, event_size]`
scales : tf.Tensor
scales of all logistics components, shape `[batch_size, n_components, event_size]`
logits, probs : tf.Tensor
probability for the mixture Categorical distribution, shape `[batch_size, n_components]`
low : int, optional
minimum quantized value, by default 0
bits : int, optional
number of bits for quantization, the maximum will be `2^bits - 1`, by default 8
name : str, optional
distribution name, by default 'MixtureQuantizedLogistic'
Returns
-------
MixtureSameFamily
the mixture of quantized logistic distribution
Example
-------
```
d = MixtureQLogistic(np.ones((12, 3, 8)).astype('float32'),
np.ones((12, 3, 8)).astype('float32'),
logits=np.random.rand(12, 3).astype('float32'),
batch_ndims=1)
```
Reference
---------
Salimans, T., Karpathy, A., Chen, X., Kingma, D.P., 2017.
PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture
Likelihood and Other Modifications. arXiv:1701.05517 [cs, stat].
"""
cats = Categorical(probs=probs, logits=logits)
dists = Logistic(loc=locs, scale=scales)
dists = TransformedDistribution(
distribution=dists, bijector=Shift(shift=tf.cast(-0.5, dists.dtype)))
dists = QuantizedDistribution(dists, low=low, high=2**bits - 1.)
dists = Independent(dists, reinterpreted_batch_ndims=batch_ndims)
dists = MixtureSameFamily(mixture_distribution=cats,
components_distribution=dists,
name=name)
return dists
def _log_prob(self, value: tf.Tensor):
""" expect `value` is output from ELU function """
params = self._params
transformed_value, value = _switch_domain(
value,
inputs_domain=self.inputs_domain,
low=self.low,
high=self.high)
## prepare the parameters
if self.n_channels == 1:
component_logits, locs, scales = params
else:
channel_tensors = tf.split(transformed_value,
self.n_channels,
axis=-1)
# If there is more than one channel, we create a linear autoregressive
# dependency among the location parameters of the channels of a single
# pixel (the scale parameters within a pixel are independent). For a pixel
# with R/G/B channels, the `r`, `g`, and `b` saturation values are
# distributed as:
#
# r ~ Logistic(loc_r, scale_r)
# g ~ Logistic(coef_rg * r + loc_g, scale_g)
# b ~ Logistic(coef_rb * r + coef_gb * g + loc_b, scale_b)
component_logits, locs, scales, coeffs = params
loc_tensors = tf.split(locs, self.n_channels, axis=-1)
coef_tensors = tf.split(coeffs, self.n_coeffs, axis=-1)
coef_count = 0
for i in range(self.n_channels):
channel_tensors[i] = channel_tensors[i][..., tf.newaxis, :]
for j in range(i):
loc_tensors[
i] += channel_tensors[j] * coef_tensors[coef_count]
coef_count += 1
locs = tf.concat(loc_tensors, axis=-1)
## create the distrubtion
mixture_distribution = Categorical(logits=component_logits)
# Convert distribution parameters for pixel values in
# `[self._low, self._high]` for use with `QuantizedDistribution`
locs = self.low + 0.5 * (self.high - self.low) * (locs + 1.)
scales = scales * 0.5 * (self.high - self.low)
logistic_dist = QuantizedDistribution(
distribution=TransformedDistribution(
distribution=Logistic(loc=locs, scale=scales),
bijector=Shift(shift=tf.cast(-0.5, self.dtype))),
low=self.low,
high=self.high,
)
dist = MixtureSameFamily(mixture_distribution=mixture_distribution,
components_distribution=Independent(
logistic_dist,
reinterpreted_batch_ndims=1))
dist = Independent(dist, reinterpreted_batch_ndims=2)
return dist.log_prob(value)
def test_pad_mixture_dimensions_mixture_same_family(self):
gm = MixtureSameFamily(
mixture_distribution=Categorical(probs=[0.3, 0.7]),
components_distribution=MultivariateNormalDiag(
loc=[[-1., 1], [1, -1]], scale_identity_multiplier=[1.0, 0.5]))
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
x_pad = distribution_util.pad_mixture_dimensions(
x, gm, gm.mixture_distribution, tensorshape_util.rank(gm.event_shape))
x_out, x_pad_out = self.evaluate([x, x_pad])
self.assertAllEqual(x_pad_out.shape, [2, 2, 1])
self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1]))
def test_pad_mixture_dimensions_mixture(self):
gm = Mixture(cat=Categorical(probs=[[0.3, 0.7]]),
components=[
Normal(loc=[-1.0], scale=[1.0]),
Normal(loc=[1.0], scale=[0.5])
])
x = tf.constant([[1.0, 2.0], [3.0, 4.0]])
x_pad = distribution_util.pad_mixture_dimensions(
x, gm, gm.cat, tensorshape_util.rank(gm.event_shape))
x_out, x_pad_out = self.evaluate([x, x_pad])
self.assertAllEqual(x_pad_out.shape, [2, 2])
self.assertAllEqual(x_out.reshape([-1]), x_pad_out.reshape([-1]))
def new(params,
probs_input=False,
dtype=None,
validate_args=False,
name='CategoricalLayer'):
"""Create the distribution instance from a `params` vector."""
params = tf.convert_to_tensor(value=params, name='params')
return Categorical(
logits=params if not probs_input else None,
probs=tf.clip_by_value(params, 1e-8, 1 - 1e-8) \
if probs_input else None,
dtype=dtype or params.dtype,
validate_args=validate_args,
name=name)
def set_prior(self,
loc=0.,
log_scale=np.log(np.expm1(1)),
mixture_logits=None):
r""" Set the prior for mixture density network
loc : Scalar or Tensor with shape `[n_components, event_size]`
log_scale : Scalar or Tensor with shape
`[n_components, event_size]` for 'none' and 'diag' component, and
`[n_components, event_size*(event_size +1)//2]` for 'full' component.
mixture_logits : Scalar or Tensor with shape `[n_components]`
"""
event_size = self.event_size
if self.covariance == 'diag':
scale_shape = [self.n_components, event_size]
fn = lambda l, s: MultivariateNormalDiag(
loc=l, scale_diag=tf.nn.softplus(s))
elif self.covariance == 'none':
scale_shape = [self.n_components, event_size]
fn = lambda l, s: Independent(
Normal(loc=l, scale=tf.math.softplus(s)), 1)
elif self.covariance == 'full':
scale_shape = [
self.n_components, event_size * (event_size + 1) // 2
]
fn = lambda l, s: MultivariateNormalTriL(
loc=l,
scale_tril=FillScaleTriL(diag_shift=1e-5)(tf.math.softplus(s)))
#
if isinstance(log_scale, Number) or tf.rank(log_scale) == 0:
loc = tf.fill([self.n_components, self.event_size], loc)
#
if isinstance(log_scale, Number) or tf.rank(log_scale) == 0:
log_scale = tf.fill(scale_shape, log_scale)
#
if mixture_logits is None:
p = 1. / self.n_components
mixture_logits = np.log(p / (1. - p))
if isinstance(mixture_logits, Number) or tf.rank(mixture_logits) == 0:
mixture_logits = tf.fill([self.n_components], mixture_logits)
#
loc = tf.cast(loc, self.dtype)
log_scale = tf.cast(log_scale, self.dtype)
mixture_logits = tf.cast(mixture_logits, self.dtype)
self._prior = MixtureSameFamily(
components_distribution=fn(loc, log_scale),
mixture_distribution=Categorical(logits=mixture_logits),
name="prior")
return self
def model_gmmprior(args: Arguments):
nets = get_networks(args.ds, zdim=args.zdim, is_hierarchical=False,
is_semi_supervised=False)
latent_size = np.prod(nets['latents'].event_shape)
n_components = 100
loc = tf.compat.v1.get_variable(name="loc", shape=[n_components, latent_size])
raw_scale_diag = tf.compat.v1.get_variable(
name="raw_scale_diag", shape=[n_components, latent_size])
mixture_logits = tf.compat.v1.get_variable(
name="mixture_logits", shape=[n_components])
nets['latents'].prior = MixtureSameFamily(
components_distribution=MultivariateNormalDiag(
loc=loc,
scale_diag=tf.nn.softplus(raw_scale_diag) + tf.math.exp(-7.)),
mixture_distribution=Categorical(logits=mixture_logits),
name="prior")
return VariationalAutoencoder(**nets, name='GMMPrior')
def model_fullcovgmm(args: Arguments):
nets = get_networks(args.ds, zdim=args.zdim, is_hierarchical=False,
is_semi_supervised=False)
latent_size = int(np.prod(nets['latents'].event_shape))
n_components = 100
loc = tf.compat.v1.get_variable(name="loc", shape=[n_components, latent_size])
raw_scale_diag = tf.compat.v1.get_variable(
name="raw_scale_diag", shape=[n_components, latent_size])
mixture_logits = tf.compat.v1.get_variable(
name="mixture_logits", shape=[n_components])
nets['latents'] = RVconf(
event_shape=latent_size,
projection=True,
posterior='mvntril',
prior=MixtureSameFamily(
components_distribution=MultivariateNormalDiag(
loc=loc,
scale_diag=tf.nn.softplus(raw_scale_diag) + tf.math.exp(-7.)),
mixture_distribution=Categorical(logits=mixture_logits),
name="prior"),
name='latents').create_posterior()
return VariationalAutoencoder(**nets, name='FullCov')
def __init__(self,
loc,
scale,
logits=None,
probs=None,
covariance_type='diag',
trainable=False,
validate_args=False,
allow_nan_stats=True,
name=None):
kw = dict(validate_args=validate_args, allow_nan_stats=allow_nan_stats)
self._trainable = bool(trainable)
self._llk_history = []
if trainable:
loc = tf.Variable(loc, trainable=True, name='loc')
scale = tf.Variable(scale, trainable=True, name='scale')
if logits is not None:
logits = tf.Variable(logits, trainable=True, name='logits')
if probs is not None:
probs = tf.Variable(probs, trainable=True, name='probs')
### initialize mixture Categorical
mixture = Categorical(logits=logits,
probs=probs,
name="MixtureWeights",
**kw)
n_components = mixture._num_categories()
### initialize Gaussian components
covariance_type = str(covariance_type).lower().strip()
if name is None:
name = 'Mixture%sGaussian' % \
(covariance_type.capitalize() if covariance_type != 'none' else
'Independent')
## create the components
if covariance_type == 'diag':
if tf.rank(scale) == 0: # scalar
extra_kw = dict(scale_identity_multiplier=scale)
else: # a tensor
extra_kw = dict(scale_diag=scale)
components = MultivariateNormalDiag(loc=loc,
name=name,
**kw,
**extra_kw)
elif covariance_type in ('tril', 'full'):
if tf.rank(scale) == 1 or \
(scale.shape[-1] != scale.shape[-2]):
scale_tril = FillScaleTriL(diag_shift=np.array(
1e-5,
tf.convert_to_tensor(scale).dtype.as_numpy_dtype()))
scale = scale_tril(scale)
components = MultivariateNormalTriL(loc=loc,
scale_tril=scale,
name=name,
**kw)
elif covariance_type == 'none':
components = Independent(distribution=Normal(loc=loc,
scale=scale,
**kw),
reinterpreted_batch_ndims=1,
name=name)
else:
raise ValueError("No support for covariance_type: '%s'" %
covariance_type)
### validate the n_components
assert (components.batch_shape[-1] == int(n_components)), \
"Number of components mismatch, given:%d, mixture:%d, components:%d" % \
(mixture.event_shape[-1], components.batch_shape[-1], int(n_components))
super().__init__(mixture_distribution=mixture,
components_distribution=components,
name=name,
**kw)
def train():
# Load MNIST data.
data = input_data.read_data_sets(FLAGS.data_dir + '/MNIST', one_hot=True)
# Create encoder graph.
with tf.variable_scope("encoder"):
inputs = tf.placeholder(tf.float32,
shape=[None, 28 * 28],
name='inputs')
tau = tf.placeholder(tf.float32, shape=[], name='temperature')
logits = encoder(inputs)
z = gumbel_softmax(
logits, tau,
hard=False) # (batch_size, num_cat_dists, num_classes)
# Create decoder graph.
with tf.variable_scope("decoder"):
p_x_given_z = decoder(z)
with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
categorical = Categorical(probs=np.ones(FLAGS.num_classes) /
FLAGS.num_classes)
z = categorical.sample(
sample_shape=[FLAGS.batch_size, FLAGS.num_cat_dists])
z = tf.one_hot(z, depth=FLAGS.num_classes)
p_x_given_z_eval = decoder(z)
# Define loss function and train opeator.
# NOTE: Categorically uniform prior p(z) is assumed.
# NOTE: Also, in this case, KL becomes negative entropy.
# NOTE: Summation becomes KLD over whole distribution q(z|x) since z is assumed to be elementwise independent.
KL = -tf.reduce_sum(tf.nn.softmax_cross_entropy_with_logits_v2(
labels=tf.nn.softmax(logits), logits=logits),
axis=1)
ELBO = tf.reduce_sum(p_x_given_z.log_prob(inputs), axis=1) - KL
loss = tf.reduce_mean(-ELBO)
train_op = tf.train.AdamOptimizer(
learning_rate=FLAGS.learning_rate).minimize(loss)
with tf.Session() as sess:
sess.run(tf.global_variables_initializer())
temperature = 1.0
for i in tqdm(range(1, FLAGS.num_iters)):
np_x, np_y = data.train.next_batch(FLAGS.batch_size)
_, np_loss = sess.run([train_op, loss], {
inputs: np_x,
tau: temperature
})
if i % 1000 == 0:
temperature = np.maximum(FLAGS.min_temp,
np.exp(-FLAGS.anneal_rate * i))
print('Temperature updated to {}\n'.format(temperature))
if i % 5000 == 1:
print('Iteration {}\nELBO: {}\n'.format(i, -np_loss))
# Plot results.
x_mean = p_x_given_z.mean()
batch = data.test.next_batch(FLAGS.batch_size)
np_x = sess.run(x_mean, {inputs: batch[0], tau: FLAGS.min_temp})
x_mean_eval = p_x_given_z_eval.mean()
np_x_eval = sess.run(x_mean_eval)
plot_squares(batch[0], np_x, np_x_eval, 8)