import tensorflow as tf

                print('Using ranking trajectory loss ', config.contra_horizon)
                contra_loss, cstr_pct = contra_traj_lossV9(
                    pred, target[name], horizon=config.contra_horizon, margin=config.margin)

            objectives.append((Objective(name, contra_loss, min, include, exclude)))
        elif name == 'reward' and config.r_loss == 'l2':
            pred = heads[name](features)
            l2_loss = tf.compat.v1.losses.mean_squared_error(target[name], pred)
            # l2_loss = tf.nn.l2_loss(pred - target[name])
            objectives.append((Objective(name, l2_loss, min, include, exclude)))
        else:
            if not config.aug_same and config.aug:
                recon_feat = tf.concat([features, target['aug']], -1)
                print('Use recon feature ', name, recon_feat)
                logprob = heads[name](recon_feat).log_prob(target[name])
                # logprob = heads[name](features).log_prob(target['ori_img'])
            else:
                logprob = heads[name](features).log_prob(target[name])
            objectives.append(Objective(name, logprob, max, include, exclude))

    objectives = [o._replace(value=tf.reduce_mean(o.value)) for o in objectives]

    return objectives, cstr_pct

