def forward(self, x, compute_pi=True, compute_log_pi=True):
for layer in self.feature_layers:
x = F.relu(layer(x))
mu = self.mean_head(x)
logstd = self.logstd_head(x)
logstd = torch.tanh(logstd)
logstd = LOGSTD_MIN + 0.5 * (LOGSTD_MAX - LOGSTD_MIN) * (
logstd + 1)
dist = TransformedDistribution(Independent(Normal(mu, logstd.exp()), 1), [TanhTransform(cache_size=1)])
if compute_pi:
#std = logstd.exp()
#noise = torch.randn_like(mu)
#pi = mu + noise * std
pi = dist.rsample()
else:
pi = None
if compute_log_pi:
#log_pi = Independent(Normal(mu, logstd.exp()), 1).log_prob(pi).unsqueeze(-1)
#log_pi = gaussian_likelihood(noise, logstd)
log_pi = dist.log_prob(pi).unsqueeze(-1)
else:
log_pi = None
mu = torch.tanh(mu)
#if compute_pi:
# pi = torch.tanh(pi)
#if compute_log_pi:
# log_pi -= torch.log(F.relu(1 - pi.pow(2)) + 1e-6).sum(-1, keepdim=True)
#print(mu.shape, pi.shape, log_pi.shape)
#print(log_pi)
#breakpoint()
#mu, pi, log_pi = apply_squashing_func(mu, pi, log_pi)
return mu, pi, log_pi
# Then we have $$ D_\mathrm{KL} [ p_x^*(\mathbf{x}) \mathrel{\|} p_x(\mathbf{x} ; \boldsymbol{\theta}) ] = D_\mathrm{KL} [ p_u^*(\mathbf{u} ; \phi) \mathrel{\|} p_u(\mathbf{u} ; \psi) ] $$
#
# This means that fitting the model to the target using the forward KL (maximum likelihood) is equivalent to fitting the induced distribution $ p_u^*(\mathbf{u} ; \phi) $ to the base $ p_u(\mathbf{u} ; \psi) $ under the reverse KL.
#
# Reciprocally, there is also:
#
# $$ D_\mathrm{KL} [ p_x(\mathbf{x} ; \boldsymbol{\theta}) \mathrel{\|} p_x^*(\mathbf{x}) ] = D_\mathrm{KL} [ p_u(\mathbf{u} ; \psi) \mathrel{\|} p_u^*(\mathbf{u} ; \phi) ] $$
# %% [markdown]
# ## Base distribution $p_\mathrm{u}(\mathbf{u})$
# %%
base_mu, base_cov = torch.zeros(2), torch.eye(2)
# base_dist = MultivariateNormal(base_mu, base_cov)
prior = TransformedDistribution(Uniform(torch.zeros(2), torch.ones(2)), SigmoidTransform().inv) # Logistic distribution
U = prior.rsample(sample_shape=(512,))
plt.scatter(U[:, 0], U[:, 1])
plt.show()
# %% [markdown]
# ## Target distribution $p_\mathrm{x}(\mathbf{x})$
# %%
rng = default_rng()
random_state = 42
class DatasetMoons:
""" two half-moons """
def __init__(self, noise=0.05, random_state=42):
self.random_state = random_state
self.noise = noise
import torch
from torch.distributions import Independent, Normal, TransformedDistribution
from torch.distributions.transforms import TanhTransform
import numpy as np
batch_size = 400
torch.set_default_dtype(torch.float64)
n = 40
print(n)
done = False
i = 0
while not done:
mu = torch.as_tensor(np.random.random([batch_size, n]))
log_std = torch.as_tensor(np.random.random([batch_size, n]))
transform = TransformedDistribution(
Independent(Normal(mu, log_std.exp()), 1), TanhTransform())
input = transform.rsample()
output = transform.log_prob(input)
if torch.isnan(output).any().item():
done = True
if (input == -1).any() or (input == 1).any():
print("somethings wrong...")
print(output)
print("something was wrong")