flow_copd/copd_loss.py

"""
Flow-CoPD loss: Contrastive On-Policy Distillation for flow-matching T2I.

See refine-logs/FINAL_PROPOSAL.md. One-line method:
  Flow-OPD / DiffusionOPD are POSITIVE-ONLY (pull student velocity toward a
  teacher velocity).  Flow-CoPD adds the missing CONTRASTIVE NEGATIVE: on the
  student's own LOW-reward on-policy trajectories, repel the student velocity
  from the (constant rectified-flow) velocity that regenerates those low-reward
  samples.  Asymmetric, data-source dependent (DistiLLM-2 port).

All functions are pure tensor ops so they can be unit-tested / reviewed in
isolation.  Shapes: v_* are (B, C, H, W) velocity-field predictions at one
denoising step; weight_t is a scalar or (B,1,1,1) per-step weight; advantages
is (B,) per-trajectory (already group-normalized, as in flow_grpo).

DESIGN NOTES / open questions flagged for cross-model review:
  [R1] Positive term = reverse-KL -> time-weighted velocity-L2 (DiffusionOPD
       ODE form).  This is NOT advantage-weighted (advantage-weighting the L2
       would collapse to plain loss-weighting; see TAD-OPD post-mortem).
  [R2] Negative term is SIGN-FLIPPED (repulsion), so it is NOT expressible as a
       positive reweighting of the positive L2 -> genuinely new signal, not a
       collapse.  It is CLAMPED (neg_clamp) so the unbounded -||.||^2 cannot
       diverge.
  [R3] v_neg_target default = constant rectified-flow velocity (noise - x0) that
       regenerates the self-generated low-reward sample.  Scheduler-convention
       sensitive -> verify against SD3 FlowMatchEulerDiscreteScheduler.
"""

import torch


def _mse(a, b):
    # per-sample mean squared error over all non-batch dims -> (B,)
    return ((a.float() - b.float()) ** 2).flatten(1).mean(dim=1)


def opd_positive_loss(v_student, v_teacher, weight_t=1.0):
    """A0 = positive-only OPD (Flow-OPD / DiffusionOPD reproduction).

    L_pos = E[ weight_t * || v_student - v_teacher ||^2 ].  v_teacher must be
    detached (computed under the frozen teacher adapter / no_grad).
    """
    per = _mse(v_student, v_teacher.detach())            # (B,)
    if torch.is_tensor(weight_t):
        per = weight_t.flatten() * per
    else:
        per = weight_t * per
    return per.mean()


def copd_loss(
    v_student,
    v_teacher,
    v_neg_target,
    advantages,
    weight_t=1.0,
    lambda_neg=1.0,
    neg_clamp=1.0,
    adv_thresh=0.0,
    detach_neg_target=True,
):
    """A2 = Flow-CoPD = positive OPD + contrastive negative.

    Args:
      v_student:    (B,C,H,W) trainable student velocity at this step.
      v_teacher:    (B,C,H,W) frozen teacher velocity at this step (detached).
      v_neg_target: (B,C,H,W) velocity that REGENERATES the self-generated
                    sample (constant rectified-flow target); we repel from it on
                    low-reward trajectories. Detached by default.
      advantages:   (B,) group-normalized per-trajectory advantage (reward
                    signal). >adv_thresh = "good", <-adv_thresh treated as "bad".
      weight_t:     per-step time weight (scalar or (B,1,1,1)).
      lambda_neg:   strength of the contrastive negative term.
      neg_clamp:    cap on the per-sample negative MSE (stability; prevents the
                    repulsive -||.||^2 from diverging).
      adv_thresh:   deadband; |adv|<=thresh contributes only to the positive term.

    Returns: (total_loss, dict_of_components_for_logging)
    """
    if detach_neg_target:
        v_neg_target = v_neg_target.detach()
    w = weight_t.flatten() if torch.is_tensor(weight_t) else weight_t

    # ---- positive term: pull every trajectory's velocity toward the teacher ----
    pos_per = _mse(v_student, v_teacher.detach())                       # (B,)
    pos_loss = (w * pos_per).mean() if torch.is_tensor(w) else (w * pos_per).mean()

    # ---- contrastive negative: repel ONLY low-reward trajectories from the
    #      velocity that regenerates them (clamped, sign-flipped) ----
    bad_mask = (advantages < -adv_thresh).float()                      # (B,)
    neg_per = _mse(v_student, v_neg_target)                            # (B,)
    neg_per = torch.clamp(neg_per, max=neg_clamp)                      # bound it
    # maximize distance on bad samples => subtract; weight by |advantage| so the
    # worst samples are repelled hardest (sign already handled by bad_mask).
    neg_weight = bad_mask * advantages.abs()
    if torch.is_tensor(w):
        neg_term = (w * neg_weight * neg_per).mean()
    else:
        neg_term = (w * neg_weight * neg_per).mean()
    neg_loss = -lambda_neg * neg_term

    total = pos_loss + neg_loss
    info = {
        "loss_pos": pos_loss.detach(),
        "loss_neg": neg_loss.detach(),
        "neg_frac": bad_mask.mean().detach(),
        "neg_mse_mean": neg_per.mean().detach(),
    }
    return total, info


def regen_velocity_per_step(x_t, x0, sigma_t, eps=1e-3):
    """Per-step model-output velocity that regenerates endpoint x0 from state x_t.

    SD3 FlowMatchEuler: x0 = x_t - sigma_t * v_model  =>  v_model = (x_t - x0)/sigma_t.
    So the velocity the student USED (in expectation) to drive x_t toward the
    (low-reward) endpoint x0 at this step is (x_t - x0)/sigma_t.  Repelling the
    student from THIS per-step target unlearns the path to bad samples.

    Fixes CRITICAL#1 from review: the old constant (noise - x0) target is only
    valid for a straight deterministic path, NOT the Flow-GRPO SDE rollout.

    Args:
      x_t:     (B,C,H,W) state before this step  (sample["latents"][:, j]).
      x0:      (B,C,H,W) low-reward endpoint     (sample["next_latents"][:, -1]).
      sigma_t: scalar or (B,) noise level at this step (from scheduler.sigmas[j]).
    """
    if not torch.is_tensor(sigma_t):
        sigma_t = torch.as_tensor(sigma_t, dtype=x_t.dtype, device=x_t.device)
    sigma_t = sigma_t.clamp(min=eps).reshape(-1, 1, 1, 1)
    return ((x_t - x0) / sigma_t).detach()


def rectified_flow_regen_velocity(latents_first, latents_last):
    """DEPRECATED (review CRITICAL#1): constant noise->x0 velocity. Wrong for SDE
    rollouts; kept only so old configs import-error loudly. Use regen_velocity_per_step."""
    raise DeprecationWarning("use regen_velocity_per_step(x_t, x0, sigma_t)")


def sde_time_weight(sigma_t, scheme="snr", eps=1e-3):
    """Per-step OPD weight (review MAJOR#5). 'uniform'->1; 'snr'-> 1/sigma^2-ish so
    low-noise (late) steps that decide detail are weighted up. Returns (B,1,1,1) or scalar."""
    if scheme == "uniform":
        return 1.0
    if not torch.is_tensor(sigma_t):
        sigma_t = torch.as_tensor(sigma_t)
    sigma_t = sigma_t.clamp(min=eps)
    w = 1.0 / (sigma_t ** 2 + 1.0)        # bounded SNR-like weight in (0,1]
    return w.reshape(-1, 1, 1, 1)