def add_decoded_goal_to_path(self, path):
set_idx = path["full_observations"][0][self.set_index_key]
set = self.sets[set_idx]
sampled_data = set.example_dict[self.example_image_key]
posteriors = self.model.encoder(ptu.from_numpy(sampled_data))
learned_prior = set_vae_trainer.compute_prior(posteriors)
decoded = self.model.decoder(learned_prior.mean)
decoded_img = ptu.get_numpy(decoded.mean)[0]
for i_in_path, d in enumerate(path["full_observations"]):
d[self.decode_set_image_key] = decoded_img
def update_encodings(self):
means = []
covariances = []
for set in self.sets:
sampled_data = set.example_dict[self.data_key]
posteriors = self.vae.encoder(ptu.from_numpy(sampled_data))
learned_prior = set_vae_trainer.compute_prior(posteriors)
means.append(ptu.get_numpy(learned_prior.mean)[0])
covariances.append(ptu.get_numpy(learned_prior.variance)[0])
self.means = np.array(means)
self.covariances = np.array(covariances)