def loss(self, inputs, model_output):
losses = AttrDict()
# action prediction loss
n_actions = model_output.actions.shape[1]
losses.action_reconst = L2Loss(1.0)(model_output.actions, inputs.actions[:, :n_actions],
weights=broadcast_final(inputs.pad_mask[:, :n_actions], inputs.actions))
# compute total loss
#total_loss = torch.stack([loss[1].value * loss[1].weight for loss in losses.items()]).sum()
#losses.total = AttrDict(value=total_loss)
# losses.total = total_loss*torch.tensor(np.nan) # for checking if backprop works
return losses
def loss(self, inputs, outputs, add_total=True):
losses = AttrDict()
# subgoal reconstruction loss
n_action_output = outputs.actions.shape[1]
loss_weights = broadcast_final(
outputs.pad_mask[:, :n_action_output],
inputs.actions) if 'pad_mask' in outputs else 1
losses.action_reconst = L2Loss(self._hp.action_rec_weight)(
outputs.actions,
outputs.action_targets[:, :n_action_output],
weights=loss_weights)
if self._hp.pred_states:
losses.state_reconst = L2Loss(self._hp.state_rec_weight)(
outputs.states, outputs.state_targets)
return losses
