def infer_posterior(self, context):
r"""Compute :math:`q(z \| c)` as a function of input context and sample new z.
Args:
context (torch.Tensor): Context values, with shape
:math:`(X, N, C)`. X is the number of tasks. N is batch size. C
is the combined size of observation, action, reward, and next
observation if next observation is used in context. Otherwise,
C is the combined size of observation, action, and reward.
"""
params = self._context_encoder.forward(context)
params = params.view(context.size(0), -1,
self._context_encoder.output_dim)
# with probabilistic z, predict mean and variance of q(z | c)
if self._use_information_bottleneck:
mu = params[..., :self._latent_dim]
sigma_squared = F.softplus(params[..., self._latent_dim:])
z_params = [
product_of_gaussians(m, s)
for m, s in zip(torch.unbind(mu), torch.unbind(sigma_squared))
]
self.z_means = torch.stack([p[0] for p in z_params])
self.z_vars = torch.stack([p[1] for p in z_params])
else:
self.z_means = torch.mean(params, dim=1)
self.sample_from_belief()
def test_product_of_gaussians():
"""Test computing mu, sigma of product of gaussians."""
size = 5
mu = torch.ones(size)
sigmas_squared = torch.ones(size)
output = product_of_gaussians(mu, sigmas_squared)
assert output[0] == 1
assert output[1] == 1 / size