//  [python]
def compute_losses(self, object_states, isvalid_by_type_o1o2,
                       encoded_h_pre,
                       encoded_h_post_gt,
                       encoded_h_post_pred,
                       affordance_pred_by_type,
                       gt_affordances_decoder,
                       isvalid_by_type):
        """

        :param object_states: [batch_size, 4, len(names_and_arities)
        :param isvalid_by_type_o1o2: first two objs whteher they're valid [batch_size, 2]
        :return:
        """
        batch_size, num_duplicates_x_num_objs, nlen_names_and_arities = get_shape_list(object_states, 3)

        # MAGNITUDE LOSSES
        ###################
        # Check if anything changed
        norms = {}
        losses = {}
        pre_states, post_states = tf.unstack(
            tf.reshape(object_states, [batch_size, 2, 2, len(names_and_arities)]), axis=2)
        did_change = tf.not_equal(pre_states, post_states)
        didchange_weight = tf.cast(tf.reduce_any(did_change, -1), dtype=tf.float32) * isvalid_by_type_o1o2
        nochange_weight = (1.0 - tf.cast(tf.reduce_any(did_change, -1), dtype=tf.float32)) * isvalid_by_type_o1o2

        ### How much did things change
        ###############
        encoded_h_delta = encoded_h_post_pred - encoded_h_pre
        encoded_h_delta_l2 = tf.sqrt(tf.reduce_mean(tf.square(encoded_h_delta), -1))

        norms['didchange_hdelta_l2'] = tf.reduce_sum(encoded_h_delta_l2 * didchange_weight) / (tf.reduce_sum(
            didchange_weight) + 1e-5)
        norms['nochange_hdelta_l2'] = tf.reduce_sum(encoded_h_delta_l2 * nochange_weight) / (tf.reduce_sum(
            nochange_weight) + 1e-5)

        # Delta between pred and GT
        ###
        # gt_mu = tf.stop_gradient(encoded_h_post_gt[:, :, :self.hidden_size])
        # pred_mu = encoded_h_post_pred[:, :, :self.hidden_size]
        #
        # #########################################
        # # VAE loss
        # all_mu, all_logvar = tf.split(tf.reshape(tf.concat([encoded_h_pre,
        #                                                     encoded_h_post_gt,
        #                                                     encoded_h_post_pred], 1),
        #                                          [-1, self.hidden_size * 2]), [self.hidden_size, self.hidden_size],
        #                               axis=-1)
        # kld = -0.5 * tf.reduce_mean(1.0 + all_logvar - tf.square(all_mu) - tf.exp(all_logvar))
        # losses['kld'] = kld

        #########################################
        gt_stop = tf.stop_gradient(encoded_h_post_gt)
        hidden_state_diff_l2 = tf.sqrt(tf.reduce_mean(tf.square(encoded_h_post_pred - gt_stop), -1))
        hidden_state_diff_l1 = tf.reduce_mean(tf.abs(encoded_h_post_pred - gt_stop), -1)
        norms['hidden_state_diff_l2'] = tf.reduce_sum(hidden_state_diff_l2 * isvalid_by_type_o1o2) / (
                tf.reduce_sum(isvalid_by_type_o1o2) + 1e-5)
        norms['hidden_state_diff_l1'] = tf.reduce_sum(hidden_state_diff_l1 * isvalid_by_type_o1o2) / (
                tf.reduce_sum(isvalid_by_type_o1o2) + 1e-5)

        hidden_state_magn_l2 = tf.sqrt(tf.reduce_mean(tf.square(gt_stop), -1))
        norms['hidden_state_magn_l2'] = tf.reduce_sum(hidden_state_magn_l2 * isvalid_by_type_o1o2) / (
                tf.reduce_sum(isvalid_by_type_o1o2) + 1e-5)

        # Upweight changed losses
        # did change: [batch_size, num_objs, num_affordances]
        for i, (affordance_name, arity_) in enumerate(names_and_arities):
            arity = max(arity_, 2)

            losses[f'state/{affordance_name}_post'] = sequence_xe_loss(
                affordance_pred_by_type[:, 4:, i, :arity],
                gt_affordances_decoder[:, 4:, i],
                label_weights=isvalid_by_type[:, 4:],
            )

            losses[f'state/{affordance_name}_pre'] = sequence_xe_loss(
                affordance_pred_by_type[:, 0:2, i, :arity],
                gt_affordances_decoder[:, 0:2, i],
                label_weights=isvalid_by_type[:, 0:2],  # + tf.cast(did_change[:, :, i], dtype=tf.float32) * 100.0,
            )

            losses[f'state/{affordance_name}_postgt'] = sequence_xe_loss(
                affordance_pred_by_type[:, 2:4, i, :arity],
                gt_affordances_decoder[:, 2:4, i],
                label_weights=isvalid_by_type[:, 2:4],  # + tf.cast(did_change[:, :, i], dtype=tf.float32) * 100.0,
            )
        # # Another way for losses
        # losses_all = _sequence_xe_loss_noreduce(affordance_pred_by_type, gt_affordances_decoder)
        # loss_mask = tf.reshape(tf.tile(isvalid_by_type[:, :, None], [1, 1, len(names_and_arities)]), [-1])
        # losses['state/all'] = tf.reduce_sum(losses_all * loss_mask) / (tf.reduce_sum(loss_mask) + 1e-5)

        return losses, norms