def _draw_learned_prior(self, ax, prior, color):
if isinstance(prior, GMM):
[
self._draw_gaussian(ax, component.tensor(), color,
ten2ar(weight))
for weight, component in prior
]
else:
self._draw_gaussian(ax, prior.tensor(), color)
def _draw_gaussian(ax, gauss_tensor, color, weight=None):
px, py, p_logsig_x, p_logsig_y = split_along_axis(ten2ar(gauss_tensor), axis=0)
def logsig2std(logsig):
return np.exp(logsig)
ell = Ellipse(xy=(px, py),
width=2*logsig2std(p_logsig_x), height=2*logsig2std(p_logsig_y),
angle=0, color=color) # this assumes diagonal gaussian
if weight is not None:
ell.set_alpha(weight)
else:
ell.set_facecolor('none')
ax.add_artist(ell)
def __call__(self, estimates, targets, per_datum=False):
"""
:param estimates:
:param targets:
:param per_datum: If this is True, return a tensor of shape: [batch_size], otherwise: [1]
:return:
"""
error = self.func(estimates, targets)
if isinstance(error, torch.Tensor): error = ten2ar(error)
if per_datum:
return np.mean(error, axis=get_dim_inds(error)[1:])
else:
return np.mean(error)
def psnr(estimates, targets, data_dims=3):
# NOTE: PSNR is not dimension-independent. The number of dimensions which are part of the metric has to be specified
# I.e 2 for grayscale, 3 for color images.
if isinstance(estimates, torch.Tensor): estimates = ten2ar(estimates)
if isinstance(targets, torch.Tensor): targets = ten2ar(targets)
estimates = (estimates + 1) / 2
targets = (targets + 1) / 2
max_pix_val = 1.0
tolerance = 0.001
assert (0 - tolerance) <= np.min(targets) and np.max(
targets) <= max_pix_val * (1 + tolerance)
assert (0 - tolerance) <= np.min(estimates) and np.max(
estimates) <= max_pix_val * (1 + tolerance)
mse = (np.square(estimates - targets))
mse = np.mean(mse, axis=get_dim_inds(mse)[-data_dims:])
psnr = 10 * np.log(max_pix_val / mse) / np.log(10)
if np.any(np.isinf(psnr)):
import pdb
pdb.set_trace()
return psnr
def step(self, action):
if isinstance(action, torch.Tensor): action = ten2ar(action)
try:
obs, reward, done, info = self._env.step(action)
reward = reward / self._hp.reward_norm
except self._mj_except:
# this can happen when agent drives simulation to unstable region (e.g. very fast speeds)
print("Catch env exception!")
obs = self.reset()
reward = self._hp.punish_reward # this avoids that the agent is going to these states again
done = np.array(
True
) # terminate episode (observation will get overwritten by env reset)
info = {}
return self._wrap_observation(obs), reward, np.array(done), info
def mse(estimates, targets):
if isinstance(estimates, torch.Tensor): estimates = ten2ar(estimates)
if isinstance(targets, torch.Tensor): targets = ten2ar(targets)
return np.square(estimates - targets)
def to_numpy(self):
"""Convert internal variables to numpy arrays."""
return type(self)(ten2ar(self.mu), ten2ar(self.log_sigma))
data_sample = data_dist.sample()
gmm_sample = gmm_dist.rsample()
# loss = gmm_dist.nll(pydata).mean()
# loss = (gmm_dist.log_prob(gmm_sample) - data_dist.log_prob(gmm_sample))
loss = (data_dist.log_prob(data_sample) -
gmm_dist.log_prob(data_sample))
# loss = (gmm_dist.log_prob(gmm_sample) - data_dist.log_prob(gmm_sample)) + \
# (data_dist.log_prob(data_sample) - gmm_dist.log_prob(data_sample))
loss_samples.append(loss)
loss = torch.cat(loss_samples).mean()
loss.backward()
optimizer.step()
if i % 100 == 0:
print(f"Iter: {i}\t" + f"NLL: {loss.mean().data:.2f}\t")
# visualize samples
samples = gmm_dist.sample().data.numpy()
fig = plt.figure()
ax = plt.subplot(111)
plt.xlim(-2, 2)
plt.ylim(-2, 2)
# plt.scatter(data[:, 0], data[:, 1], c='black', alpha=0.1)
# plt.scatter(samples[:, 0], samples[:, 1], c='green', alpha=0.5)
[
_draw_gaussian(ax, component.tensor(), 'green', ten2ar(weight))
for weight, component in gmm_dist[0]
]
plt.axis("equal")
plt.savefig("gmm_fit.png")
# plt.show()