def loss(self, model_output, inputs):
losses = AttrDict()
# reconstruction loss
losses.nll = NLL()(model_output.pred_act, self._regression_targets(inputs))
losses.total = self._compute_total_loss(losses)
return losses
def loss(self, model_output, inputs):
"""Loss computation of the SPIRL model.
:arg model_output: output of SPIRL model forward pass
:arg inputs: dict with 'states', 'actions', 'images' keys from data loader
"""
losses = AttrDict()
# reconstruction loss, assume unit variance model output Gaussian
losses.rec_mse = NLL(self._hp.reconstruction_mse_weight) \
(Gaussian(model_output.reconstruction, torch.zeros_like(model_output.reconstruction)),
self._regression_targets(inputs))
# KL loss
losses.kl_loss = KLDivLoss(self._hp.kl_div_weight)(model_output.q, model_output.p)
# learned skill prior net loss
losses.q_hat_loss = self._compute_learned_prior_loss(model_output)
losses.total = self._compute_total_loss(losses)
return losses
