def forward(self, input: torch.Tensor):
encoder_output = self.proposal_network(input)
encoded_shape = encoder_output.shape
encoded_resolution = encoded_shape[-1] * encoded_shape[-2]
# rearrange feature maps into "latent space"
rearranged_encoder_output = encoder_output \
.flatten(start_dim=-2) \
.transpose(-1, -2) \
.flatten(end_dim=1)
proposal_logits = self.proposal_logits_head(rearranged_encoder_output)
proposal_distribution = distributions.RelaxedOneHotCategorical(
self.temperature, logits=proposal_logits)
proposal_sample = proposal_distribution.rsample()
proposal_sample_copy = proposal_sample
proposal_sample = proposal_sample.reshape(-1, encoded_resolution, self.latent_dim) \
.transpose(-1, -2) \
.reshape(encoded_shape)
reconstruction = self.generative_network(proposal_sample)
return reconstruction, proposal_distribution, proposal_sample_copy
def forward(self, state_features):
x = self.feedforward_model(state_features)
if self.dist == 'tanh_normal':
mean, std = th.chunk(x, 2, -1)
mean = self.mean_scale * th.tanh(mean / self.mean_scale)
std = F.softplus(std + self.raw_init_std) + self.min_std
dist = td.Normal(mean, std)
# TODO: fix nan problem
dist = td.TransformedDistribution(dist,
td.TanhTransform(cache_size=1))
dist = td.Independent(dist, 1)
dist = SampleDist(dist)
elif self.dist == 'trunc_normal':
mean, std = th.chunk(x, 2, -1)
std = 2 * th.sigmoid((std + self.raw_init_std) / 2) + self.min_std
from rls.nn.dists.TruncatedNormal import \
TruncatedNormal as TruncNormalDist
dist = TruncNormalDist(th.tanh(mean), std, -1, 1)
dist = td.Independent(dist, 1)
elif self.dist == 'one_hot':
dist = td.OneHotCategoricalStraightThrough(logits=x)
elif self.dist == 'relaxed_one_hot':
dist = td.RelaxedOneHotCategorical(th.tensor(0.1), logits=x)
else:
raise NotImplementedError(f"{self.dist} is not implemented.")
return dist
def forward(self, input: torch.Tensor,
proposal: distributions.RelaxedOneHotCategorical,
proposal_sample: torch.Tensor,
reconstruction: torch.Tensor) -> torch.Tensor:
if self.likelihood == 'bernoulli':
likelihood = distributions.Bernoulli(probs=reconstruction)
else:
likelihood = distributions.Normal(reconstruction,
torch.ones_like(reconstruction))
likelihood = distributions.Independent(likelihood,
reinterpreted_batch_ndims=-1)
reconstruction_loss = likelihood.log_prob(input).mean()
assert proposal.logits.dim(
) == 2, "proposal.shape == [*, D], D is shape of isotopic gaussian"
prior = distributions.RelaxedOneHotCategorical(proposal.temperature,
logits=torch.ones_like(
proposal.logits))
regularization = (proposal.log_prob(proposal_sample) - prior.log_prob(proposal_sample)) \
.mean()
# evidence lower bound (maximize)
total_loss = reconstruction_loss - self.beta * regularization
return -total_loss, -reconstruction_loss, regularization
def sample_q(self, k, mode):
if mode == ModeKeys.TRAIN:
z_dist = td.RelaxedOneHotCategorical(self.temp,
logits=self.q_dist.logits)
z_NK = z_dist.rsample((k, ))
elif mode == ModeKeys.EVAL:
z_NK = self.q_dist.sample((k, ))
return torch.reshape(z_NK, (k, -1, self.z_dim))
def __init__(self, n_clusters, cluster_to_params_graph):
super().__init__()
self.cluster_logits = nn.Parameter(torch.randn(n_clusters))
self.cluster_to_params_graph = cluster_to_params_graph
self.cluster_temperature_logit = nn.Parameter(torch.tensor(1.))
self.cluster_distr = td.RelaxedOneHotCategorical(
temperature=F.softplus(self.cluster_temperature_logit),
logits=self.cluster_logits)
def _build_posterior(self, param_loc, param_spi_scale):
tn_posterior = TruncatedNormal(
loc=torch.clamp(param_loc, self.min_size - 1, self.max_size + 1),
scale=torch.clamp_min(nn.functional.softplus(param_spi_scale),
1e-3),
low=self.min_size,
high=self.max_size)
logits_posterior = tn_posterior.log_prob(self._categories)
dist_posterior = distributions.RelaxedOneHotCategorical(
logits=torch.clamp_min(logits_posterior, _LOG_EPSILON),
temperature=self.temperature)
return dist_posterior
def _build_prior(self, prior_loc, prior_spi_scale, prior_temperature):
tn_prior = TruncatedNormal(loc=prior_loc,
scale=torch.clamp_min(
nn.functional.softplus(prior_spi_scale),
1e-3),
low=self.min_size,
high=self.max_size)
logits_prior = tn_prior.log_prob(self._categories)
dist_prior = distributions.RelaxedOneHotCategorical(
logits=torch.clamp_min(logits_prior, _LOG_EPSILON),
temperature=torch.tensor(prior_temperature, dtype=torch.float32))
return dist_prior
def __init__(self, temp, latent_num, latent_dim):
super(Model, self).__init__()
if type(temp) != torch.Tensor:
temp = torch.tensor(temp)
self.__temp = temp
self.latent_num = latent_num
self.latent_dim = latent_dim
self.encoder = Encoder(latent_num=latent_num, latent_dim=latent_dim)
self.decoder = Decoder(latent_num=latent_num, latent_dim=latent_dim)
if 'ExpTDModel' in str(self.__class__):
self.prior = ExpRelaxedCategorical(temp, probs=torch.ones(latent_dim).cuda())
else:
self.prior = dist.RelaxedOneHotCategorical(temp, probs=torch.ones(latent_dim).cuda())
self.initialize()
self.softmax = nn.Softmax(dim=-1)
def forward(self, input: torch.Tensor):
encoder_output = self.proposal_network(input)
bottleneck_resolution = encoder_output.size()[-2:]
encoder_output = self.prenet(encoder_output)
assert encoder_output.size(-1) == self.latent_dim * self.num_latents
proposal_logits = self.proposal_logits_head(encoder_output) \
.reshape(-1, self.num_latents, self.latent_dim) \
.flatten(end_dim=-2)
proposal_distribution = distributions.RelaxedOneHotCategorical(
self.temperature, logits=proposal_logits)
proposal_sample = proposal_distribution.rsample()
proposal_sample_copy = proposal_sample
proposal_sample = proposal_sample.reshape(
-1, self.num_latents * self.latent_dim, 1, 1)
proposal_sample = F.interpolate(proposal_sample, bottleneck_resolution)
reconstruction = self.generative_network(proposal_sample)
return reconstruction, proposal_distribution, proposal_sample_copy
def _sample_categorical(self, probs):
if self.test_mode:
return dist.OneHotCategorical(probs=probs).sample()
else:
return dist.RelaxedOneHotCategorical(temperature=self.temperature,
probs=probs).sample()
def sample(self, log_alpha, temp):
v_dist = dist.RelaxedOneHotCategorical(temp, logits=log_alpha)
concrete = v_dist.rsample()
return concrete, v_dist
def temp(self, value):
self.__temp = value
if 'ExpTDModel' in str(self.__class__):
self.prior = ExpRelaxedCategorical(value, probs=torch.ones(self.latent_dim).cuda())
else:
self.prior = dist.RelaxedOneHotCategorical(value, probs=torch.ones(self.latent_dim).cuda())