def _model():
p = yield Root(tfd.Beta(dtype(1), dtype(1), name="p"))
gamma_C = yield Root(tfd.Beta(dtype(1), dtype(1), name="gamma_C"))
gamma_T = yield Root(tfd.Beta(dtype(1), dtype(1), name="gamma_T"))
eta_C = yield Root(tfd.Dirichlet(np.ones(K, dtype=dtype) / K,
name="eta_C"))
eta_T = yield Root(tfd.Dirichlet(np.ones(K, dtype=dtype) / K,
name="eta_T"))
loc = yield Root(tfd.Sample(tfd.Normal(dtype(0), dtype(1)),
sample_shape=K, name="loc"))
nu = yield Root(tfd.Sample(tfd.Uniform(dtype(10), dtype(50)),
sample_shape=K, name="nu"))
phi = yield Root(tfd.Sample(tfd.Normal(dtype(m_phi), dtype(s_phi)),
sample_shape=K, name="phi"))
sigma_sq = yield Root(tfd.Sample(tfd.InverseGamma(dtype(3), dtype(2)),
sample_shape=K,
name="sigma_sq"))
scale = np.sqrt(sigma_sq)
gamma_T_star = compute_gamma_T_star(gamma_C, gamma_T, p)
eta_T_star = compute_eta_T_star(gamma_C[..., tf.newaxis],
gamma_T[..., tf.newaxis],
eta_C, eta_T,
p[..., tf.newaxis],
gamma_T_star[..., tf.newaxis])
# likelihood
y_C = yield mix(nC, eta_C, loc, scale, name="y_C")
n0C = yield tfd.Binomial(nC, gamma_C, name="n0C")
y_T = yield mix(nT, eta_T_star, loc, scale, name="y_T")
n0T = yield tfd.Binomial(nT, gamma_T_star, name="n0T")
def _init_distribution(conditions, **kwargs):
return tfd.Mixture(
cat=tfd.Categorical(
probs=[1 - conditions["psi"], conditions["psi"]]),
components=[
tfd.Deterministic(loc=tf.zeros_like(conditions["n"])),
tfd.Binomial(total_count=conditions["n"],
probs=conditions["p"]),
],
**kwargs,
)
def create_model(npos_C, npos_T, n0_C, n0_T, K, dtype=np.float64):
n_C = npos_C + n0_C
n_T = npos_T + n0_T
return tfd.JointDistributionNamed(
dict(p=tfd.Beta(dtype(1), dtype(1)),
gamma_C=tfd.Beta(dtype(1), dtype(1)),
gamma_T=tfd.Beta(dtype(1), dtype(1)),
eta_C=tfd.Dirichlet(tf.ones(K, dtype=dtype) / K),
eta_T=tfd.Dirichlet(tf.ones(K, dtype=dtype) / K),
loc=tfd.Sample(tfd.Normal(dtype(0), dtype(1)), sample_shape=K),
sigma_sq=tfd.Sample(tfd.InverseGamma(dtype(3), dtype(2)),
sample_shape=K),
n0_C=lambda gamma_C: tfd.Binomial(n_C, gamma_C),
n0_T=lambda gamma_T: tfd.Binomial(n_C, gamma_C),
y_C=lambda gamma_C, eta_C, loc, sigma_sq: tfd.Sample(
mix(gamma_C, eta_C, loc, tf.sqrt(sigma_sq), dtype(neg_inf)),
n_C),
y_T=lambda gamma_C, gamma_T, eta_C, eta_T, p, loc, sigma_sq: tfd.
Sample(
mix_T(gamma_C, gamma_T, eta_C, eta_T, p, loc, tf.sqrt(
sigma_sq), dtype(neg_inf)), n_T)))
def _base_dist(self, *args, **kwargs):
"""
Zero-inflated binomial base distribution.
A ZeroInflatedBinomial is a mixture between a deterministic
distribution and a Binomial distribution.
"""
mix = kwargs.pop("mix")
return tfd.Mixture(
cat=tfd.Categorical(probs=[mix, 1.0 - mix]),
components=[tfd.Deterministic(0.0),
tfd.Binomial(*args, **kwargs)],
name="ZeroInflatedBinomial",
)
def rnn_sim(rnn: MDNRNN, z, states, a):
# Make one LSTM step
z = tf.reshape(tf.cast(z, dtype=tf.float32), (1, 1, rnn.args.z_size))
a = tf.reshape(tf.cast(a, dtype=tf.float32), (1, 1, rnn.args.a_width))
input_x = tf.concat((z, a), axis=2)
rnn_out, h, c = rnn.inference_base(input_x, initial_state=states)
rnn_state = [h, c]
rnn_out = tf.reshape(rnn_out, [-1, rnn.args.rnn_size])
# Predict z
mdnrnn_params = rnn.predict_z(rnn_out)
mdnrnn_params = tf.reshape(mdnrnn_params,
[-1, 3 * rnn.args.rnn_num_mixture])
mu, logstd, logpi = tf.split(mdnrnn_params, num_or_size_splits=3, axis=1)
logpi = logpi - tf.reduce_logsumexp(
input_tensor=logpi, axis=1, keepdims=True) # normalize
cat = tfd.Categorical(logits=logpi)
component_splits = [1] * rnn.args.rnn_num_mixture
mus = tf.split(mu, num_or_size_splits=component_splits, axis=1)
sigs = tf.split(tf.exp(logstd),
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)
z = tf.reshape(mixture.sample(), shape=(-1, rnn.args.z_size))
# Predict done
if rnn.args.rnn_predict_done:
d_distr = rnn.predict_done(rnn_out)
done_logit = tfd.Normal(d_distr[0][0], d_distr[0][1]).sample()
done_dist = tfd.Binomial(total_count=1, logits=done_logit)
done = tf.squeeze(done_dist.sample()) == 1.0
else:
done = False
# Predict reward
if rnn.args.rnn_predict_reward:
r_distr = rnn.predict_reward(rnn_out)
reward = tfd.Normal(r_distr[0][0], r_distr[0][1]).sample()
else:
reward = 1.0
return rnn_state, z, reward, done
def rnn_sim(rnn, z, states, a, training=True):
z = tf.reshape(tf.cast(z, dtype=tf.float32), (1, 1, rnn.args.z_size))
a = tf.reshape(tf.cast(a, dtype=tf.float32), (1, 1, rnn.args.a_width))
input_x = tf.concat((z, a), axis=2)
rnn_out, h, c = rnn.inference_base(
input_x, initial_state=states,
training=training) # set training True to use Dropout
rnn_state = [h, c]
rnn_out = tf.reshape(rnn_out, [-1, rnn.args.rnn_size])
out = rnn.out_net(rnn_out)
mdnrnn_params, r, d_logits = rnn.parse_rnn_out(out)
mdnrnn_params = tf.reshape(mdnrnn_params,
[-1, 3 * rnn.args.rnn_num_mixture])
mu, logstd, logpi = tf.split(mdnrnn_params, num_or_size_splits=3, axis=1)
logpi = logpi / rnn.args.rnn_temperature # temperature
logpi = logpi - tf.reduce_logsumexp(
input_tensor=logpi, axis=1, keepdims=True) # normalize
d_dist = tfd.Binomial(total_count=1, logits=d_logits)
d = tf.squeeze(d_dist.sample()) == 1.0
cat = tfd.Categorical(logits=logpi)
component_splits = [1] * rnn.args.rnn_num_mixture
mus = tf.split(mu, num_or_size_splits=component_splits, axis=1)
# temperature
sigs = tf.split(tf.exp(logstd) * tf.sqrt(rnn.args.rnn_temperature),
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)
z = tf.reshape(mixture.sample(), shape=(-1, rnn.args.z_size))
if rnn.args.rnn_r_pred == 0:
r = 1.0 # For Doom Reward is always 1.0 if the agent is alive
return rnn_state, z, r, d
def rnn_sim(rnn, z, states, a):
if rnn.args.env_name == 'CarRacing-v0':
raise ValueError('Not implemented yet for CarRacing')
z = tf.reshape(tf.cast(z, dtype=tf.float32), (1, 1, rnn.args.z_size))
a = tf.reshape(tf.cast(a, dtype=tf.float32), (1, 1, rnn.args.a_width))
input_x = tf.concat((z, a), axis=2)
rnn_out, h, c = rnn.inference_base(input_x, initial_state=states)
rnn_state = [h, c]
rnn_out = tf.reshape(rnn_out, [-1, rnn.args.rnn_size])
out = rnn.out_net(rnn_out)
mdnrnn_params, d_logits = out[:, :-1], out[:, -1:]
mdnrnn_params = tf.reshape(mdnrnn_params,
[-1, 3 * rnn.args.rnn_num_mixture])
mu, logstd, logpi = tf.split(mdnrnn_params, num_or_size_splits=3, axis=1)
logpi = logpi - tf.reduce_logsumexp(
input_tensor=logpi, axis=1, keepdims=True) # normalize
d_dist = tfd.Binomial(total_count=1, logits=d_logits)
d = tf.squeeze(d_dist.sample()) == 1.0
cat = tfd.Categorical(logits=logpi)
component_splits = [1] * rnn.args.rnn_num_mixture
mus = tf.split(mu, num_or_size_splits=component_splits, axis=1)
sigs = tf.split(tf.exp(logstd),
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)
z = tf.reshape(mixture.sample(), shape=(-1, rnn.args.z_size))
r = 1.0 # For Doom Reward is always 1.0 if the agent is alive
return rnn_state, z, r, d
def _init_distribution(conditions):
n, p = conditions["n"], conditions["p"]
return tfd.Binomial(total_count=n, probs=p)
def _init_distribution(conditions):
total_count, probs = conditions["total_count"], conditions["probs"]
return tfd.Binomial(total_count=total_count, probs=probs)
def calc_new_n_produced():
n_produced_float = tf.cast(n_produced, dtype=ztypes.float)
binomial = tfd.Binomial(n_produced_float,
probs=old_scaling / weights_scaling)
return tf.cast(tf.round(binomial.sample()), dtype=tf.int32)
def _base_dist(self, n: TensorLike, p: TensorLike, *args, **kwargs):
return tfd.Binomial(total_count=n, probs=p, *args, **kwargs)