Buckets:
| """Factor-wise clipped multimodal GRPO objective with reference JSD.""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| import torch | |
| from shared.distributions import symmetric_jsd_from_logits | |
| class TrajectoryPolicyInputs: | |
| advantage: torch.Tensor | |
| text_current_log_probs: torch.Tensor | |
| text_old_log_probs: torch.Tensor | |
| visual_current_factor_log_probs: torch.Tensor | |
| visual_old_factor_log_probs: torch.Tensor | |
| text_current_logits: torch.Tensor | None = None | |
| text_reference_logits: torch.Tensor | None = None | |
| visual_current_logits: torch.Tensor | None = None | |
| visual_reference_logits: torch.Tensor | None = None | |
| text_reference_jsd: torch.Tensor | None = None | |
| class GRPOLossOutput: | |
| loss: torch.Tensor | |
| surrogate: torch.Tensor | |
| reference_jsd: torch.Tensor | |
| statistics: dict[str, torch.Tensor] | |
| def _clip_objective( | |
| ratio: torch.Tensor, | |
| advantage: torch.Tensor, | |
| *, | |
| low: float, | |
| high: float, | |
| ) -> torch.Tensor: | |
| unclipped = ratio * advantage | |
| clipped = ratio.clamp(1.0 - low, 1.0 + high) * advantage | |
| return torch.minimum(unclipped, clipped) | |
| def _validate(item: TrajectoryPolicyInputs, reference_beta: float) -> None: | |
| if item.advantage.numel() != 1: | |
| raise ValueError("trajectory advantage must be scalar") | |
| if item.text_current_log_probs.shape != item.text_old_log_probs.shape: | |
| raise ValueError("text current/old log-probabilities differ") | |
| if ( | |
| item.visual_current_factor_log_probs.shape | |
| != item.visual_old_factor_log_probs.shape | |
| or item.visual_current_factor_log_probs.ndim != 2 | |
| ): | |
| raise ValueError("visual current/old factor log-probabilities differ") | |
| if reference_beta: | |
| if item.text_reference_jsd is not None: | |
| if item.text_reference_jsd.shape != item.text_current_log_probs.shape: | |
| raise ValueError("text reference JSD does not match actions") | |
| elif item.text_current_log_probs.numel() and ( | |
| item.text_current_logits is None | |
| or item.text_reference_logits is None | |
| or item.text_current_logits.shape != item.text_reference_logits.shape | |
| or item.text_current_logits.shape[0] | |
| != item.text_current_log_probs.shape[0] | |
| ): | |
| raise ValueError("text reference distributions do not match actions") | |
| if item.visual_current_factor_log_probs.numel() and ( | |
| item.visual_current_logits is None | |
| or item.visual_reference_logits is None | |
| or item.visual_current_logits.shape | |
| != item.visual_reference_logits.shape | |
| or item.visual_current_logits.shape[:-1] | |
| != item.visual_current_factor_log_probs.shape | |
| or item.visual_current_logits.shape[-1] != 2 | |
| ): | |
| raise ValueError("visual reference distributions do not match factors") | |
| def multimodal_grpo_loss( | |
| trajectories: list[TrajectoryPolicyInputs], | |
| *, | |
| clip_low: float, | |
| clip_high: float, | |
| reference_beta: float, | |
| ) -> GRPOLossOutput: | |
| if not trajectories: | |
| raise ValueError("GRPO objective requires sampled trajectories") | |
| if not 0 <= clip_low < 1 or clip_high < 0 or reference_beta < 0: | |
| raise ValueError("invalid clip or reference coefficient") | |
| trajectory_surrogates = [] | |
| trajectory_references = [] | |
| text_ratios = [] | |
| visual_ratios = [] | |
| text_clipped = [] | |
| visual_clipped = [] | |
| text_steps = visual_steps = 0 | |
| for item in trajectories: | |
| _validate(item, reference_beta) | |
| advantage = item.advantage.detach().float().reshape(()) | |
| text_log_ratio = ( | |
| item.text_current_log_probs.float() | |
| - item.text_old_log_probs.detach().to(item.text_current_log_probs).float() | |
| ) | |
| visual_log_ratio = ( | |
| item.visual_current_factor_log_probs.float() | |
| - item.visual_old_factor_log_probs.detach() | |
| .to(item.visual_current_factor_log_probs) | |
| .float() | |
| ) | |
| text_ratio = text_log_ratio.exp() | |
| visual_ratio = visual_log_ratio.exp() | |
| text_values = _clip_objective( | |
| text_ratio, advantage, low=clip_low, high=clip_high | |
| ) | |
| visual_factor_values = _clip_objective( | |
| visual_ratio, advantage, low=clip_low, high=clip_high | |
| ) | |
| visual_values = visual_factor_values.mean(-1) | |
| all_steps = torch.cat((text_values, visual_values)) | |
| anchor = ( | |
| item.text_current_log_probs.sum() | |
| + item.visual_current_factor_log_probs.sum() | |
| ) * 0.0 | |
| trajectory_surrogates.append(all_steps.mean() if all_steps.numel() else anchor) | |
| text_steps += text_values.numel() | |
| visual_steps += visual_values.numel() | |
| text_ratios.append(text_ratio.detach()) | |
| visual_ratios.append(visual_ratio.detach().flatten()) | |
| text_clipped.append( | |
| ((text_ratio < 1.0 - clip_low) | (text_ratio > 1.0 + clip_high)) | |
| .detach() | |
| .float() | |
| ) | |
| visual_clipped.append( | |
| ((visual_ratio < 1.0 - clip_low) | (visual_ratio > 1.0 + clip_high)) | |
| .detach() | |
| .float() | |
| .flatten() | |
| ) | |
| reference_steps = [] | |
| if reference_beta and item.text_current_log_probs.numel(): | |
| reference_steps.append( | |
| item.text_reference_jsd if item.text_reference_jsd is not None else symmetric_jsd_from_logits( | |
| item.text_reference_logits, | |
| item.text_current_logits, | |
| ) | |
| ) | |
| if reference_beta and item.visual_current_factor_log_probs.numel(): | |
| visual_jsd = symmetric_jsd_from_logits( | |
| item.visual_reference_logits, | |
| item.visual_current_logits, | |
| ).mean(-1) | |
| reference_steps.append(visual_jsd) | |
| trajectory_references.append( | |
| torch.cat(reference_steps).mean() | |
| if reference_steps | |
| else anchor | |
| ) | |
| surrogate = torch.stack(trajectory_surrogates).mean() | |
| reference_jsd = torch.stack(trajectory_references).mean() | |
| loss = -surrogate + reference_beta * reference_jsd | |
| def mean_or_zero(values: list[torch.Tensor]) -> torch.Tensor: | |
| nonempty = [value for value in values if value.numel()] | |
| return ( | |
| torch.cat(nonempty).mean() | |
| if nonempty | |
| else loss.detach().new_zeros(()) | |
| ) | |
| statistics = { | |
| "surrogate": surrogate.detach(), | |
| "reference_jsd": reference_jsd.detach(), | |
| "text_ratio_mean": mean_or_zero(text_ratios), | |
| "visual_factor_ratio_mean": mean_or_zero(visual_ratios), | |
| "text_clip_fraction": mean_or_zero(text_clipped), | |
| "visual_factor_clip_fraction": mean_or_zero(visual_clipped), | |
| "text_steps": loss.detach().new_tensor(text_steps), | |
| "visual_steps": loss.detach().new_tensor(visual_steps), | |
| "autonomous_ar_steps": loss.detach().new_tensor(text_steps + visual_steps), | |
| } | |
| return GRPOLossOutput(loss, surrogate, reference_jsd, statistics) | |
| __all__ = [ | |
| "GRPOLossOutput", | |
| "multimodal_grpo_loss", | |
| "TrajectoryPolicyInputs", | |
| ] | |
Xet Storage Details
- Size:
- 7.27 kB
- Xet hash:
- 424bf69db65fd1a8527281faa30272c49733895607808485a14cff2591498395
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.