from __future__ import annotations
def masked_sum(values: torch.Tensor, mask: torch.Tensor, axis: int | tuple[int, ...] | None=None) -> torch.Tensor:
    """Compute sum of tensor values where mask is True.

    NaN values outside the mask are replaced with zeros to prevent
    contaminating the sum.

    Args:
        values: Input tensor containing values to sum.
        mask: Boolean or numeric mask tensor (same shape as values).
            Non-zero values indicate elements to include.
        axis: Dimension(s) along which to sum. None sums all elements.

    Returns:
        torch.Tensor: Sum of masked values, reduced along specified axis.
    """
    valid_values = torch.where(mask.bool(), values, 0.0)
    return (valid_values * mask).sum(axis=axis)
from __future__ import annotations
def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str, dp_size: int=1, batch_num_tokens: Optional[int]=None, global_batch_size: Optional[int]=None, loss_scale_factor: Optional[int]=None):
    """
    Aggregate the loss across global batch to ensure the loss is invariant to fsdp/megatron parallelism.

    NOTE: The returned loss has different behaviors for different backend:
    - FSDP: the loss is directly used for backward.
    - Megatron: the loss should be scaled by `num_microbatches` and `cp_size` for pp schedule.

    Args:
        loss_mat: micro batch loss matrix, (bs, response_length)
        loss_mask: micro batch loss mask, (bs, response_length)
        loss_agg_mode: method to aggregate the loss matrix into a scalar
        dp_size: data parallel size
        batch_num_tokens: number of valid tokens in global batch
        global_batch_size: global batch size
        loss_scale_factor: scale factor for "seq-mean-token-sum-norm" mode. If None, uses loss_mask.shape[-1].
            Set this to a constant value to ensure consistent normalization throughout training.

    Returns:
        loss: `a scalar torch.Tensor`
            aggregated loss
    """
    if loss_agg_mode == 'token-mean':
        if batch_num_tokens is None:
            batch_num_tokens = loss_mask.sum()
        loss = verl_F.masked_sum(loss_mat, loss_mask) / batch_num_tokens * dp_size
    elif loss_agg_mode == 'seq-mean-token-sum':
        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1)
        seq_mask = (torch.sum(loss_mask, dim=-1) > 0).float()
        if global_batch_size is None:
            global_batch_size = seq_mask.sum()
        loss = verl_F.masked_sum(seq_losses, seq_mask) / global_batch_size * dp_size
    elif loss_agg_mode == 'seq-mean-token-mean':
        seq_mask = torch.sum(loss_mask, dim=-1)
        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / (seq_mask + 1e-08)
        seq_mask = (seq_mask > 0).float()
        if global_batch_size is None:
            global_batch_size = seq_mask.sum()
        loss = verl_F.masked_sum(seq_losses, seq_mask) / global_batch_size * dp_size
    elif loss_agg_mode == 'seq-mean-token-sum-norm':
        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1)
        if loss_scale_factor is None:
            loss_scale_factor = loss_mask.shape[-1]
        loss = torch.sum(seq_losses) / loss_scale_factor
    else:
        raise ValueError(f'Invalid loss_agg_mode: {loss_agg_mode}')
    return loss

def compute_self_distillation_loss(student_log_probs: torch.Tensor, teacher_log_probs: torch.Tensor, response_mask: torch.Tensor, self_distillation_config: Any, old_log_probs: Optional[torch.Tensor]=None, student_all_log_probs: Optional[torch.Tensor]=None, teacher_all_log_probs: Optional[torch.Tensor]=None, student_topk_log_probs: Optional[torch.Tensor]=None, teacher_topk_log_probs: Optional[torch.Tensor]=None, self_distillation_mask: Optional[torch.Tensor]=None, loss_agg_mode: str='token-mean', rollout_is_weights: Optional[torch.Tensor]=None) -> tuple[torch.Tensor, dict[str, Any]]:
    metrics = {}
    loss_mask = response_mask
    if self_distillation_mask is not None:
        loss_mask = loss_mask * self_distillation_mask.unsqueeze(1)
    if self_distillation_config.full_logit_distillation:
        use_topk = self_distillation_config.distillation_topk is not None
        if use_topk:
            if student_topk_log_probs is None or teacher_topk_log_probs is None:
                raise ValueError('top-k distillation requires student_topk_log_probs and teacher_topk_log_probs.')

            def add_tail(log_probs: torch.Tensor) -> torch.Tensor:
                log_s = torch.logsumexp(log_probs, dim=-1, keepdim=True)
                log_s = torch.clamp(log_s, max=-1e-07)
                tail_log = torch.log(-torch.expm1(log_s))
                return torch.cat([log_probs, tail_log], dim=-1)

            def renorm_topk_log_probs(logp: torch.Tensor) -> torch.Tensor:
                logZ = torch.logsumexp(logp, dim=-1, keepdim=True)
                return logp - logZ
            student_distill_log_probs = student_topk_log_probs
            teacher_distill_log_probs = teacher_topk_log_probs
            if self_distillation_config.distillation_add_tail:
                student_distill_log_probs = add_tail(student_distill_log_probs)
                teacher_distill_log_probs = add_tail(teacher_distill_log_probs)
            else:
                student_distill_log_probs = renorm_topk_log_probs(student_distill_log_probs)
                teacher_distill_log_probs = renorm_topk_log_probs(teacher_distill_log_probs)
        else:
            if student_all_log_probs is None or teacher_all_log_probs is None:
                raise ValueError('full_logit_distillation requires student_all_log_probs and teacher_all_log_probs.')
            student_distill_log_probs = student_all_log_probs
            teacher_distill_log_probs = teacher_all_log_probs
        if self_distillation_config.alpha == 0.0:
            kl_loss = F.kl_div(student_distill_log_probs, teacher_distill_log_probs, reduction='none', log_target=True)
        elif self_distillation_config.alpha == 1.0:
            kl_loss = F.kl_div(teacher_distill_log_probs, student_distill_log_probs, reduction='none', log_target=True)
        else:
            alpha = torch.tensor(self_distillation_config.alpha, dtype=student_distill_log_probs.dtype, device=student_distill_log_probs.device)
            mixture_log_probs = torch.logsumexp(torch.stack([student_distill_log_probs + torch.log(1 - alpha), teacher_distill_log_probs + torch.log(alpha)]), dim=0)
            kl_teacher = F.kl_div(mixture_log_probs, teacher_distill_log_probs, reduction='none', log_target=True)
            kl_student = F.kl_div(mixture_log_probs, student_distill_log_probs, reduction='none', log_target=True)
            kl_loss = torch.lerp(kl_student, kl_teacher, alpha)
        per_token_loss = kl_loss.sum(-1)
    else:
        assert self_distillation_config.alpha == 1.0, 'Only reverse KL is supported for non-full-logit distillation'
        log_ratio = student_log_probs - teacher_log_probs
        per_token_loss = log_ratio.detach() * student_log_probs
    is_clip = self_distillation_config.is_clip
    if is_clip is not None:
        if old_log_probs is None:
            raise ValueError('old_log_probs is required for distillation IS ratio.')
        negative_approx_kl = (student_log_probs - old_log_probs).detach()
        negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)
        ratio = torch.exp(negative_approx_kl).clamp(max=is_clip)
        per_token_loss = per_token_loss * ratio
    if rollout_is_weights is not None:
        per_token_loss = per_token_loss * rollout_is_weights
    loss = agg_loss(loss_mat=per_token_loss, loss_mask=loss_mask, loss_agg_mode=loss_agg_mode, batch_num_tokens=loss_mask.sum().clamp(min=1.0))
    return (loss, metrics)
from __future__ import annotations
def _loss(logprobs: torch.Tensor, batch: dict[str, Any]) -> tuple[torch.Tensor, torch.Tensor]:
    mask = batch['loss_mask'].to(logprobs.device).reshape(-1)
    loss_name = batch['loss_name']
    if loss_name == 'cross_entropy':
        loss = -(logprobs * mask).sum()
    elif loss_name in RL_LOSSES:
        sampling_logprobs = batch['sampling_logprobs'].to(logprobs.device).reshape(-1)
        advantages = batch['advantages'].to(logprobs.device).reshape(-1)
        probability_ratio = torch.exp(logprobs - sampling_logprobs)
        if loss_name == 'importance_sampling':
            objective = probability_ratio * advantages
        elif loss_name == 'dppo':
            threshold = batch['loss_config'].get('tv_threshold', 0.1)
            ratio = probability_ratio.detach()
            divergence = (ratio * sampling_logprobs.exp() - sampling_logprobs.exp()).abs()
            leaving = (advantages > 0) & (ratio > 1) | (advantages < 0) & (ratio < 1)
            blocked = leaving & (divergence > threshold)
            objective = probability_ratio * advantages * (~blocked).to(logprobs.dtype)
        elif loss_name == 'ppo':
            low = batch['loss_config'].get('clip_low_threshold', 0.8)
            high = batch['loss_config'].get('clip_high_threshold', 1.2)
            clipped = torch.clamp(probability_ratio, low, high)
            objective = torch.minimum(probability_ratio * advantages, clipped * advantages)
        elif loss_name == 'cispo':
            low = batch['loss_config'].get('clip_low_threshold', 0.0)
            high = batch['loss_config'].get('clip_high_threshold', 4.0)
            objective = torch.clamp(probability_ratio, low, high).detach() * logprobs * advantages
        else:
            beta = batch['loss_config'].get('beta', 0.05)
            objective = logprobs * advantages - 0.5 * beta * (logprobs - sampling_logprobs).square()
        loss = -(objective * mask).sum()
    return (loss, mask.count_nonzero())
from __future__ import annotations
def compute_advantages(trajectory_groups_P: list[TrajectoryGroup]) -> list[torch.Tensor]:
    """Compute advantages for each trajectory, centered within groups.

    Args:
        trajectory_groups_P (list[TrajectoryGroup]): Groups of trajectories,
            where each group's rewards are centered independently.

    Returns:
        list[torch.Tensor]: Per-group advantage tensors of shape ``(G,)``,
            where ``G`` is the number of trajectories in each group.
    """
    advantages_P: list[torch.Tensor] = []
    for traj_group in trajectory_groups_P:
        rewards_G = torch.tensor(traj_group.get_total_rewards())
        advantages_G = rewards_G - rewards_G.mean()
        advantages_P.append(advantages_G)
    return advantages_P