From 58954c660bcf897036d3bb4fa09be5e2e0cf6969 Mon Sep 17 00:00:00 2001 From: Jiali Chen Date: Wed, 19 Nov 2025 06:15:23 +0000 Subject: [PATCH] implement Phase 1 backbone code --- fastvideo/fastvideo_args.py | 208 ++++++++++ fastvideo/pipelines/pipeline_batch_info.py | 24 ++ fastvideo/training/__init__.py | 21 +- fastvideo/training/reward_models.py | 419 +++++++++++++++++++ fastvideo/training/rl_pipeline.py | 448 +++++++++++++++++++++ fastvideo/training/rl_utils.py | 385 ++++++++++++++++++ 6 files changed, 1504 insertions(+), 1 deletion(-) create mode 100644 fastvideo/training/reward_models.py create mode 100644 fastvideo/training/rl_pipeline.py create mode 100644 fastvideo/training/rl_utils.py diff --git a/fastvideo/fastvideo_args.py b/fastvideo/fastvideo_args.py index c2afc9f3a..921d31d29 100644 --- a/fastvideo/fastvideo_args.py +++ b/fastvideo/fastvideo_args.py @@ -740,6 +740,67 @@ def get_current_fastvideo_args() -> FastVideoArgs: return _current_fastvideo_args +@dataclasses.dataclass +class RLArgs: + """ + Reinforcement Learning (RL) specific arguments + + #TODO(Jiali) refactoring to separate config classes when implementing multiple algorithms. + """ + # ============================================================================ + # SHARED RL CONFIGURATION + rl_mode: bool = False # Enable RL training mode + rl_algorithm: str = "grpo" # RL algorithm to use: "grpo", "ppo", "dpo" + + # Trajectory collection + rl_num_rollouts: int = 4 # Number of rollouts to collect per training step + rl_rollout_steps: str = "20,30" # Random intermediate steps for sampling (comma-separated) + rl_noise_injection_min: int = 10 # Minimum timestep for noise injection + rl_noise_injection_max: int = 40 # Maximum timestep for noise injection + rl_use_sde_sampling: bool = True # Use SDE sampling (Flow-GRPO-Fast) + rl_num_denoising_steps: int = 2 # Number of denoising steps per trajectory (1-2 for fast) + + # Advantage estimation + rl_gamma: float = 0.99 # Discount factor for returns + rl_lambda: float = 0.95 # GAE lambda parameter + rl_use_gae: bool = True # Use Generalized Advantage Estimation + rl_normalize_advantages: bool = True # Normalize advantages before policy update + + # Reward models + reward_model_paths: str = "" # Comma-separated paths to reward models + reward_weights: str = "" # Comma-separated weights for reward aggregation + reward_model_types: str = "" # Comma-separated reward types (pickscore,geneval,ocr,etc) + value_model_path: str = "" # Path to value model (can be empty to train from scratch) + value_model_share_backbone: bool = False # Share transformer backbone between policy and value + + # Training schedule + rl_policy_value_update_ratio: float = 1.0 # Ratio of policy:value updates (like distillation) + rl_warmup_steps: int = 1000 # Collect SFT-style data before starting RL + rl_collect_on_policy: bool = True # Collect fresh rollouts each step (on-policy) + + # ============================================================================ + # ALGORITHM-SPECIFIC CONFIGURATION + + # Policy optimization (GRPO/PPO shared) + rl_policy_clip_range: float = 0.2 # PPO-style clipping range for policy ratio + rl_value_clip_range: float = 0.2 # Value function clipping range + rl_num_policy_epochs: int = 1 # Number of policy update epochs (GRPO typically uses 1) + rl_num_value_epochs: int = 1 # Number of value function update epochs + rl_target_kl: float = 0.01 # Target KL divergence for early stopping + rl_entropy_coef: float = 0.0 # Entropy coefficient for exploration + rl_value_loss_coef: float = 0.5 # Value loss coefficient + + # GRPO-specific (GRPO-Guard safety mechanisms) + rl_use_grpo_guard: bool = True # Enable GRPO-Guard safety mechanisms + rl_ratio_norm_correction: bool = True # RatioNorm: correct importance ratio bias + rl_gradient_reweighting: bool = True # Reweight gradients across denoising steps + rl_max_importance_ratio: float = 10.0 # Clip importance ratios above this value + + # DPO-specific (not used yet, add when implementing DPO) + # dpo_beta: float = 0.1 # DPO regularization parameter + # dpo_label_smoothing: float = 0.0 # Label smoothing for DPO loss + + @dataclasses.dataclass class TrainingArgs(FastVideoArgs): """ @@ -862,6 +923,9 @@ class TrainingArgs(FastVideoArgs): last_step_only: bool = False # Only use the last timestep for training context_noise: int = 0 # Context noise level for cache updates + # Nested RL configuration + rl_args: RLArgs = dataclasses.field(default_factory=RLArgs) + @classmethod def from_cli_args(cls, args: argparse.Namespace) -> "TrainingArgs": provided_args = clean_cli_args(args) @@ -886,6 +950,20 @@ class TrainingArgs(FastVideoArgs): kwargs[attr] = WorkloadType.from_string( workload_type_value) if isinstance( workload_type_value, str) else workload_type_value + elif attr == 'rl_args': + # Construct nested RLArgs from CLI arguments + rl_kwargs = {} + for rl_field in dataclasses.fields(RLArgs): + rl_attr = rl_field.name + if hasattr(args, rl_attr): + rl_kwargs[rl_attr] = getattr(args, rl_attr) + else: + # Use default value from RLArgs + if rl_field.default_factory is not dataclasses.MISSING: + rl_kwargs[rl_attr] = rl_field.default_factory() + elif rl_field.default is not dataclasses.MISSING: + rl_kwargs[rl_attr] = rl_field.default + kwargs[attr] = RLArgs(**rl_kwargs) # Use getattr with default value from the dataclass for potentially missing attributes else: # Get the field to check its default value @@ -1284,6 +1362,136 @@ class TrainingArgs(FastVideoArgs): default=TrainingArgs.context_noise, help="Context noise level for cache updates") + # RL (Reinforcement Learning) arguments + parser.add_argument("--rl-mode", + action=StoreBoolean, + help="Enable RL training mode") + parser.add_argument("--rl-algorithm", + type=str, + default=RLArgs.rl_algorithm, + choices=["grpo", "ppo", "dpo"], + help="RL algorithm to use (grpo, ppo, dpo)") + + # Trajectory collection (Flow-GRPO-Fast) + parser.add_argument("--rl-num-rollouts", + type=int, + default=RLArgs.rl_num_rollouts, + help="Number of rollouts to collect per training step") + parser.add_argument("--rl-rollout-steps", + type=str, + default=RLArgs.rl_rollout_steps, + help="Random intermediate steps for sampling (comma-separated)") + parser.add_argument("--rl-noise-injection-min", + type=int, + default=RLArgs.rl_noise_injection_min, + help="Minimum timestep for noise injection") + parser.add_argument("--rl-noise-injection-max", + type=int, + default=RLArgs.rl_noise_injection_max, + help="Maximum timestep for noise injection") + parser.add_argument("--rl-use-sde-sampling", + action=StoreBoolean, + help="Use SDE sampling (Flow-GRPO-Fast)") + parser.add_argument("--rl-num-denoising-steps", + type=int, + default=RLArgs.rl_num_denoising_steps, + help="Number of denoising steps per trajectory (1-2 for fast)") + + # Advantage estimation + parser.add_argument("--rl-gamma", + type=float, + default=RLArgs.rl_gamma, + help="Discount factor for returns") + parser.add_argument("--rl-lambda", + type=float, + default=RLArgs.rl_lambda, + help="GAE lambda parameter") + parser.add_argument("--rl-use-gae", + action=StoreBoolean, + help="Use Generalized Advantage Estimation") + parser.add_argument("--rl-normalize-advantages", + action=StoreBoolean, + help="Normalize advantages before policy update") + + # Policy optimization (GRPO/PPO) + parser.add_argument("--rl-policy-clip-range", + type=float, + default=RLArgs.rl_policy_clip_range, + help="PPO-style clipping range for policy ratio") + parser.add_argument("--rl-value-clip-range", + type=float, + default=RLArgs.rl_value_clip_range, + help="Value function clipping range") + parser.add_argument("--rl-num-policy-epochs", + type=int, + default=RLArgs.rl_num_policy_epochs, + help="Number of policy update epochs (GRPO typically uses 1)") + parser.add_argument("--rl-num-value-epochs", + type=int, + default=RLArgs.rl_num_value_epochs, + help="Number of value function update epochs") + parser.add_argument("--rl-target-kl", + type=float, + default=RLArgs.rl_target_kl, + help="Target KL divergence for early stopping") + parser.add_argument("--rl-entropy-coef", + type=float, + default=RLArgs.rl_entropy_coef, + help="Entropy coefficient for exploration") + parser.add_argument("--rl-value-loss-coef", + type=float, + default=RLArgs.rl_value_loss_coef, + help="Value loss coefficient") + + # GRPO-Guard (safety mechanisms) + parser.add_argument("--rl-use-grpo-guard", + action=StoreBoolean, + help="Enable GRPO-Guard safety mechanisms") + parser.add_argument("--rl-ratio-norm-correction", + action=StoreBoolean, + help="RatioNorm: correct importance ratio bias") + parser.add_argument("--rl-gradient-reweighting", + action=StoreBoolean, + help="Reweight gradients across denoising steps") + parser.add_argument("--rl-max-importance-ratio", + type=float, + default=RLArgs.rl_max_importance_ratio, + help="Clip importance ratios above this value") + + # Reward models + parser.add_argument("--reward-model-paths", + type=str, + default=RLArgs.reward_model_paths, + help="Comma-separated paths to reward models") + parser.add_argument("--reward-weights", + type=str, + default=RLArgs.reward_weights, + help="Comma-separated weights for reward aggregation") + parser.add_argument("--reward-model-types", + type=str, + default=RLArgs.reward_model_types, + help="Comma-separated reward types (pickscore,geneval,ocr,etc)") + parser.add_argument("--value-model-path", + type=str, + default=RLArgs.value_model_path, + help="Path to value model (can be empty to train from scratch)") + parser.add_argument("--value-model-share-backbone", + action=StoreBoolean, + help="Share transformer backbone between policy and value") + + # Training schedule + parser.add_argument("--rl-policy-value-update-ratio", + type=float, + default=RLArgs.rl_policy_value_update_ratio, + help="Ratio of policy:value updates (like distillation)") + parser.add_argument("--rl-warmup-steps", + type=int, + default=RLArgs.rl_warmup_steps, + help="Collect SFT-style data before starting RL") + parser.add_argument("--rl-collect-on-policy", + action=StoreBoolean, + help="Collect fresh rollouts each step (on-policy)") + return parser diff --git a/fastvideo/pipelines/pipeline_batch_info.py b/fastvideo/pipelines/pipeline_batch_info.py index 32d15e860..47bc2454f 100644 --- a/fastvideo/pipelines/pipeline_batch_info.py +++ b/fastvideo/pipelines/pipeline_batch_info.py @@ -267,6 +267,30 @@ class TrainingBatch: latent_vis_dict: dict[str, Any] = field(default_factory=dict) fake_score_latent_vis_dict: dict[str, Any] = field(default_factory=dict) + # RL/GRPO-specific attributes + reward_scores: torch.Tensor | None = None # Computed rewards from reward models + log_probs: torch.Tensor | None = None # Current policy log probabilities + old_log_probs: torch.Tensor | None = None # Old policy log probs (for importance ratio) + advantages: torch.Tensor | None = None # GAE advantages + returns: torch.Tensor | None = None # TD returns (advantages + values) + values: torch.Tensor | None = None # Value function predictions + old_values: torch.Tensor | None = None # Old value predictions (for clipping) + + # RL loss components + policy_loss: float = 0.0 # GRPO/PPO policy loss + value_loss: float = 0.0 # Value function loss + kl_divergence: float = 0.0 # KL(new_policy || old_policy) + importance_ratio: float = 1.0 # exp(log_prob - old_log_prob) + clip_fraction: float = 0.0 # Fraction of ratios that were clipped + + # RL metrics + advantage_mean: float = 0.0 # Mean advantage (should be ~0 after normalization) + advantage_std: float = 1.0 # Std of advantages + reward_mean: float = 0.0 # Mean reward across batch + reward_std: float = 0.0 # Std of rewards + value_mean: float = 0.0 # Mean value prediction + entropy: float = 0.0 # Policy entropy (for exploration) + @dataclass class PreprocessBatch(ForwardBatch): diff --git a/fastvideo/training/__init__.py b/fastvideo/training/__init__.py index 2eeba4d18..ad78debde 100644 --- a/fastvideo/training/__init__.py +++ b/fastvideo/training/__init__.py @@ -1,5 +1,24 @@ from .distillation_pipeline import DistillationPipeline from .training_pipeline import TrainingPipeline from .wan_training_pipeline import WanTrainingPipeline +from .rl_pipeline import RLPipeline, create_rl_pipeline +from .reward_models import ( + BaseRewardModel, + MultiRewardAggregator, + ValueModel, + DummyRewardModel, + create_reward_models +) -__all__ = ["TrainingPipeline", "WanTrainingPipeline", "DistillationPipeline"] +__all__ = [ + "TrainingPipeline", + "WanTrainingPipeline", + "DistillationPipeline", + "RLPipeline", + "create_rl_pipeline", + "BaseRewardModel", + "MultiRewardAggregator", + "ValueModel", + "DummyRewardModel", + "create_reward_models", +] diff --git a/fastvideo/training/reward_models.py b/fastvideo/training/reward_models.py new file mode 100644 index 000000000..b8b68feb2 --- /dev/null +++ b/fastvideo/training/reward_models.py @@ -0,0 +1,419 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Base infrastructure for VIDEO reward models in RL/GRPO training. + +IMPORTANT: This module is designed exclusively for VIDEO generation models. +All reward models must operate on video sequences [B, T, C, H, W], not single frames. + +This module provides: +1. Abstract base class for VIDEO reward models +2. Multi-reward aggregation for video +3. Value model wrapper +4. Integration with FastVideo video generation infrastructure + +Scope: VIDEO ONLY - No image-only reward models +""" + +from abc import ABC, abstractmethod +from typing import Any + +import torch +import torch.nn as nn + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +class BaseRewardModel(ABC, nn.Module): + """ + Abstract base class for VIDEO reward models. + + All VIDEO reward models should inherit from this class and implement + the compute_reward() method. + + IMPORTANT: Reward models must process FULL VIDEO SEQUENCES, not individual frames. + Input shape is [B, T, C, H, W] where T is the temporal (frame) dimension. + + For video-specific rewards, consider: + - Temporal coherence across frames + - Motion quality and smoothness + - Video-text alignment (not just frame-text) + - Multi-frame aesthetic quality + """ + + def __init__(self, model_path: str | None = None, device: str = "cuda"): + super().__init__() + self.model_path = model_path + self.device = device + + @abstractmethod + def compute_reward( + self, + videos: torch.Tensor, # [B, T, C, H, W] decoded video sequences + prompts: list[str], # Text prompts describing videos + **kwargs: Any + ) -> torch.Tensor: + """ + Compute rewards for generated VIDEO sequences. + + IMPORTANT: This method must process the FULL temporal sequence [B, T, C, H, W]. + Do NOT evaluate individual frames independently and average. + + Args: + videos: Decoded video tensors [B, T, C, H, W] in range [0, 1] + B = batch size + T = number of frames (temporal dimension) + C = channels (typically 3 for RGB) + H, W = height, width + prompts: List of text prompts (length B) describing each video + **kwargs: Additional model-specific arguments + + Returns: + rewards: Tensor of shape [B] with reward scores for each video sequence + + Example: + >>> videos = torch.rand(4, 17, 3, 256, 256) # 4 videos, 17 frames each + >>> prompts = ["A cat jumping", "A dog running", ...] + >>> rewards = model.compute_reward(videos, prompts) + >>> rewards.shape # torch.Size([4]) + """ + raise NotImplementedError("Subclasses must implement compute_reward()") + + @abstractmethod + def load_model(self) -> None: + """Load the reward model from checkpoint.""" + raise NotImplementedError("Subclasses must implement load_model()") + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(model_path={self.model_path})" + + +class MultiRewardAggregator(nn.Module): + """ + Aggregates multiple reward models with configurable weights. + + This implements the multi-reward aggregation strategy from flow_grpo, + allowing combination of different reward signals (aesthetic quality, + text-video alignment, compositional understanding, etc.) + """ + + def __init__( + self, + reward_models: list[BaseRewardModel], + reward_weights: list[float] | None = None, + normalize_rewards: bool = True + ): + """ + Initialize multi-reward aggregator. + + Args: + reward_models: List of reward model instances + reward_weights: Weights for each reward model (default: uniform) + normalize_rewards: Whether to normalize rewards before aggregation + """ + super().__init__() + self.reward_models = nn.ModuleList(reward_models) + + if reward_weights is None: + reward_weights = [1.0 / len(reward_models)] * len(reward_models) + + assert len(reward_weights) == len(reward_models), \ + f"Number of weights ({len(reward_weights)}) must match number of models ({len(reward_models)})" + + assert abs(sum(reward_weights) - 1.0) < 1e-6, \ + f"Reward weights must sum to 1.0, got {sum(reward_weights)}" + + self.reward_weights = reward_weights + self.normalize_rewards = normalize_rewards + + logger.info( + "Initialized MultiRewardAggregator with %d models: %s", + len(reward_models), + [(type(m).__name__, w) for m, w in zip(reward_models, reward_weights, strict=False)] + ) + + def compute_reward( + self, + videos: torch.Tensor, + prompts: list[str], + return_individual: bool = False, + **kwargs: Any + ) -> torch.Tensor | dict[str, torch.Tensor]: + """ + Compute aggregated reward from multiple models. + + Args: + videos: Decoded video tensors [B, C, T, H, W] + prompts: List of text prompts + return_individual: If True, return dict with individual rewards + **kwargs: Additional arguments passed to reward models + + Returns: + If return_individual=False: aggregated_rewards [B] + If return_individual=True: dict with "aggregated" and individual model rewards + """ + batch_size = videos.shape[0] + individual_rewards: dict[str, torch.Tensor] = {} + + # Collect rewards from all models + all_rewards = [] + for i, (model, weight) in enumerate(zip(self.reward_models, self.reward_weights, strict=False)): + reward = model.compute_reward(videos, prompts, **kwargs) + assert reward.shape == (batch_size,), \ + f"Reward model {i} returned shape {reward.shape}, expected ({batch_size},)" + + # Optionally normalize individual rewards + if self.normalize_rewards: + reward = (reward - reward.mean()) / (reward.std() + 1e-8) + + individual_rewards[f"reward_{type(model).__name__}"] = reward + all_rewards.append(weight * reward) + + # Aggregate with weights + aggregated = sum(all_rewards) + + if return_individual: + individual_rewards["aggregated"] = aggregated + return individual_rewards + + return aggregated + + def __repr__(self) -> str: + models_str = ", ".join([ + f"{type(m).__name__}(w={w:.3f})" + for m, w in zip(self.reward_models, self.reward_weights, strict=False) + ]) + return f"MultiRewardAggregator({models_str})" + + +class ValueModel(nn.Module): + """ + Value function model wrapper for RL training. + + The value model can either: + 1. Share the transformer backbone with the policy (memory efficient) + 2. Use a separate transformer (more flexible) + + For now, this is a placeholder that will be expanded based on + the chosen architecture strategy. + """ + + def __init__( + self, + transformer: nn.Module, + share_backbone: bool = False, + hidden_size: int | None = None + ): + """ + Initialize value model. + + Args: + transformer: Transformer model (policy or separate) + share_backbone: Whether to share backbone with policy + hidden_size: Hidden size for value head (inferred if None) + """ + super().__init__() + self.transformer = transformer + self.share_backbone = share_backbone + + # Value head will be added later based on transformer architecture + # For now, just store the transformer reference + logger.info( + "Initialized ValueModel (share_backbone=%s)", + share_backbone + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + timestep: torch.Tensor, + **kwargs: Any + ) -> torch.Tensor: + """ + Forward pass to compute value predictions. + + Args: + hidden_states: Latent states [B, C, T, H, W] + encoder_hidden_states: Text embeddings [B, L, D] + timestep: Timesteps [B] + **kwargs: Additional transformer arguments + + Returns: + values: Value predictions [B] + """ + # TODO: Implement value prediction + # For now, return dummy values + batch_size = hidden_states.shape[0] + return torch.zeros(batch_size, device=hidden_states.device) + + +class DummyRewardModel(BaseRewardModel): + """ + Dummy VIDEO reward model for testing and development. + + Returns random rewards in the range [0, 1] for VIDEO inputs. + This is a placeholder for testing the RL pipeline before real video reward models + are implemented. + + NOTE: This does NOT actually evaluate video quality - it's just for testing! + """ + + def __init__(self, mean: float = 0.5, std: float = 0.1): + super().__init__(model_path=None) + self.mean = mean + self.std = std + logger.info("Initialized DummyRewardModel (VIDEO) - mean=%.2f, std=%.2f", mean, std) + logger.warning( + "DummyRewardModel is for TESTING ONLY - does not evaluate actual video quality!" + ) + + def compute_reward( + self, + videos: torch.Tensor, # [B, T, C, H, W] + prompts: list[str], + **kwargs: Any + ) -> torch.Tensor: + """ + Return random rewards for testing. + + Args: + videos: Video sequences [B, T, C, H, W] + prompts: Text prompts + + Returns: + Random rewards [B] in range [0, 1] + """ + batch_size = videos.shape[0] + num_frames = videos.shape[1] + + logger.debug( + "DummyRewardModel processing %d videos with %d frames each", + batch_size, + num_frames + ) + + # Generate random rewards (not based on actual video content!) + rewards = torch.randn(batch_size, device=videos.device) * self.std + self.mean + return rewards.clamp(0.0, 1.0) + + def load_model(self) -> None: + """No model to load for dummy.""" + pass + + +def create_reward_models( + reward_model_paths: str, + reward_weights: str, + reward_model_types: str, + device: str = "cuda" +) -> MultiRewardAggregator: + """ + Factory function to create VIDEO reward models from configuration strings. + + IMPORTANT: Only creates VIDEO reward models. Image-only reward models + (PickScore, ImageReward, GenEval, etc.) are NOT supported. + + Args: + reward_model_paths: Comma-separated paths to VIDEO reward models + reward_weights: Comma-separated weights for aggregation + reward_model_types: Comma-separated VIDEO reward types + device: Device to load models on + + Returns: + MultiRewardAggregator with loaded VIDEO reward models + + Supported VIDEO Reward Types (Phase 2+): + - "video_score": Video aesthetic quality (multi-frame) + - "video_text_alignment": CLIP-based video-text similarity + - "temporal_coherence": Frame-to-frame consistency + - "motion_quality": Motion smoothness and realism + - "dummy": Random rewards for testing (VIDEO-aware) + + NOT Supported (Image-Only): + - "pickscore": Image aesthetic (use "video_score" instead) + - "imagereward": Image quality (use "video_score" instead) + - "geneval": Image compositional (no video equivalent yet) + - Any single-frame reward models + + Example: + >>> models = create_reward_models( + ... reward_model_paths="/path/to/video_score,/path/to/video_clip", + ... reward_weights="0.5,0.5", + ... reward_model_types="video_score,video_text_alignment", + ... device="cuda" + ... ) + """ + if not reward_model_paths: + logger.warning("No reward models specified, using DummyRewardModel (VIDEO)") + return MultiRewardAggregator([DummyRewardModel()], [1.0]) + + paths = [p.strip() for p in reward_model_paths.split(",")] + types = [t.strip() for t in reward_model_types.split(",")] + + if reward_weights: + weights = [float(w.strip()) for w in reward_weights.split(",")] + else: + weights = [1.0 / len(paths)] * len(paths) + + assert len(paths) == len(types), \ + f"Number of paths ({len(paths)}) must match number of types ({len(types)})" + + assert len(paths) == len(weights), \ + f"Number of paths ({len(paths)}) must match number of weights ({len(weights)})" + + # Validate that no image-only reward types are specified + image_only_types = {"pickscore", "imagereward", "geneval", "aesthetic_predictor"} + for reward_type in types: + if reward_type.lower() in image_only_types: + raise ValueError( + f"Image-only reward type '{reward_type}' is not supported. " + f"This RL implementation is VIDEO-ONLY. " + f"Please use video-specific reward models instead." + ) + + # Create reward models based on types + reward_models: list[BaseRewardModel] = [] + for path, reward_type in zip(paths, types, strict=False): + if reward_type == "dummy": + model = DummyRewardModel() + elif reward_type == "video_score": + # TODO: Implement VideoScore reward model (Phase 2) + logger.warning( + "VideoScore reward not implemented yet, using DummyRewardModel" + ) + model = DummyRewardModel() + elif reward_type == "video_text_alignment": + # TODO: Implement VideoTextAlignment reward model (Phase 2) + logger.warning( + "VideoTextAlignment reward not implemented yet, using DummyRewardModel" + ) + model = DummyRewardModel() + elif reward_type == "temporal_coherence": + # TODO: Implement TemporalCoherence reward model (Phase 2) + logger.warning( + "TemporalCoherence reward not implemented yet, using DummyRewardModel" + ) + model = DummyRewardModel() + elif reward_type == "motion_quality": + # TODO: Implement MotionQuality reward model (Phase 2) + logger.warning( + "MotionQuality reward not implemented yet, using DummyRewardModel" + ) + model = DummyRewardModel() + else: + logger.warning( + "Unknown VIDEO reward type '%s', using DummyRewardModel", + reward_type + ) + model = DummyRewardModel() + + reward_models.append(model) + + logger.info( + "Created MultiRewardAggregator with %d VIDEO reward models", + len(reward_models) + ) + + return MultiRewardAggregator(reward_models, weights, normalize_rewards=True) diff --git a/fastvideo/training/rl_pipeline.py b/fastvideo/training/rl_pipeline.py new file mode 100644 index 000000000..2431ec1dc --- /dev/null +++ b/fastvideo/training/rl_pipeline.py @@ -0,0 +1,448 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +RL training pipeline for FastVideo. + +This module currently implements reinforcement learning training using GRPO (Group Relative Policy Optimization) +and related algorithms. It extends the base TrainingPipeline with RL-specific functionality: + +Reference: + Flow-GRPO: https://github.com/yifan123/flow_grpo +""" + +import torch +import torch.nn as nn +from typing import Any + +from fastvideo.distributed import get_local_torch_device, get_sp_group, get_world_group +from fastvideo.fastvideo_args import TrainingArgs +from fastvideo.forward_context import set_forward_context +from fastvideo.logger import init_logger +from fastvideo.pipelines import TrainingBatch +from fastvideo.training.training_pipeline import TrainingPipeline +from fastvideo.training.reward_models import ( + create_reward_models, + MultiRewardAggregator, + ValueModel +) +from fastvideo.training.rl_utils import ( + compute_gae, + normalize_advantages, + compute_grpo_policy_loss, + compute_value_loss, + compute_policy_entropy, + sample_random_timesteps, + compute_reward_statistics, + check_early_stopping, + compute_log_probs_from_model_output, + check_for_nan_inf +) +from fastvideo.training.training_utils import ( + get_scheduler, + save_checkpoint, + count_trainable +) + +logger = init_logger(__name__) + + +class RLPipeline(TrainingPipeline): + """ + RL training pipeline using GRPO for flow matching models. + + This pipeline implements online reinforcement learning for video generation models. + It follows the Flow-GRPO approach with: + - Fast trajectory collection (1-2 denoising steps) + - Multi-reward aggregation + - GRPO policy optimization + - GRPO-Guard safety mechanisms + """ + + def __init__( + self, + model_path: str, + fastvideo_args: TrainingArgs, + required_config_modules: list[str] | None = None, + loaded_modules: dict[str, nn.Module] | None = None + ) -> None: + """Initialize RL pipeline.""" + if not fastvideo_args.rl_args.rl_mode: + logger.warning( + "rl_mode is False, but RLPipeline is being initialized. " + "Setting rl_mode=True." + ) + fastvideo_args.rl_args.rl_mode = True + + super().__init__(model_path, fastvideo_args, required_config_modules, loaded_modules) + + # RL-specific components (will be initialized in initialize_training_pipeline) + self.reward_models: MultiRewardAggregator | None = None + self.value_model: ValueModel | None = None + self.value_optimizer: torch.optim.Optimizer | None = None + self.value_scheduler: Any | None = None + + logger.info( + "Initialized RLPipeline with algorithm: %s", + fastvideo_args.rl_args.rl_algorithm + ) + + def initialize_training_pipeline(self, training_args: TrainingArgs): + """Initialize the RL training pipeline with reward and value models.""" + # First initialize base pipeline (transformer, optimizer, dataloader, etc.) + super().initialize_training_pipeline(training_args) + + logger.info("Initializing RL-specific components...") + + # Initialize reward models + self.reward_models = create_reward_models( + reward_model_paths=training_args.rl_args.reward_model_paths, + reward_weights=training_args.rl_args.reward_weights, + reward_model_types=training_args.rl_args.reward_model_types, + device=str(self.device) + ) + logger.info("Loaded reward models: %s", self.reward_models) + + # Initialize value model + if training_args.rl_args.value_model_share_backbone: + # Share transformer backbone with policy + logger.info("Value model will share backbone with policy transformer") + self.value_model = ValueModel( + self.transformer, + share_backbone=True + ) + else: + # Separate value model (clone transformer architecture) + logger.info("Creating separate value model") + # TODO: Implement separate value model initialization + # For now, use shared backbone + self.value_model = ValueModel( + self.transformer, + share_backbone=True + ) + + # Create optimizer and scheduler for value model + if not training_args.rl_args.value_model_share_backbone: + value_params = list(self.value_model.parameters()) + self.value_optimizer = torch.optim.AdamW( + value_params, + lr=training_args.learning_rate, + betas=(0.9, 0.999), + weight_decay=training_args.weight_decay, + eps=1e-8, + ) + + self.value_scheduler = get_scheduler( + training_args.lr_scheduler, + optimizer=self.value_optimizer, + num_warmup_steps=training_args.lr_warmup_steps, + num_training_steps=training_args.max_train_steps, + num_cycles=training_args.lr_num_cycles, + power=training_args.lr_power, + min_lr_ratio=training_args.min_lr_ratio, + last_epoch=self.init_steps - 1, + ) + + logger.info("Created separate optimizer and scheduler for value model") + logger.info("Value model trainable parameters: %s B", + round(count_trainable(self.value_model) / 1e9, 3)) + + logger.info("RL pipeline initialization complete") + + def initialize_validation_pipeline(self, training_args: TrainingArgs): + """Initialize validation pipeline for RL (similar to base implementation).""" + # For now, reuse the base implementation + # In the future, we can add RL-specific validation (e.g., compare old vs new policy) + logger.info("RL validation pipeline will be implemented based on task requirements") + # Set validation_pipeline to None for now + self.validation_pipeline = None + + def collect_trajectories(self, training_batch: TrainingBatch) -> TrainingBatch: + """ + Collect on-policy trajectories using Flow-GRPO-Fast approach. + + This implements fast trajectory collection by: + 1. Sampling random intermediate timesteps for noise injection + 2. Running 1-2 denoising steps (not full trajectory) + 3. Computing log probabilities at sampled steps + + Args: + training_batch: Current training batch + + Returns: + Updated training_batch with trajectory information + """ + logger.debug("Collecting trajectories with Flow-GRPO-Fast") + + # Parse rollout steps from config + rollout_steps_str = self.training_args.rl_args.rl_rollout_steps + rollout_steps = [int(s.strip()) for s in rollout_steps_str.split(",")] + + # Sample random timesteps for noise injection + batch_size = training_batch.latents.shape[0] + timesteps = sample_random_timesteps( + batch_size=batch_size, + min_timestep=self.training_args.rl_args.rl_noise_injection_min, + max_timestep=self.training_args.rl_args.rl_noise_injection_max, + device=training_batch.latents.device, + generator=self.noise_random_generator + ) + + # TODO: Implement actual trajectory collection + # use existing noisy inputs and compute log probs for now + training_batch.timesteps = timesteps + + # Store old log probs for importance ratio + training_batch.old_log_probs = training_batch.log_probs + + logger.debug("Trajectory collection complete") + return training_batch + + def compute_rewards(self, training_batch: TrainingBatch) -> TrainingBatch: + """ + Compute rewards using multi-reward aggregation. + + Args: + training_batch: Training batch with generated videos + + Returns: + Updated training_batch with reward scores + """ + logger.debug("Computing rewards from reward models") + + # TODO: Implement actual reward computation + # This requires decoding latents to videos and running reward models + # For now, use dummy rewards + batch_size = training_batch.latents.shape[0] + dummy_rewards = torch.randn(batch_size, device=training_batch.latents.device) * 0.1 + 0.5 + training_batch.reward_scores = dummy_rewards.clamp(0.0, 1.0) + + # Compute reward statistics + reward_stats = compute_reward_statistics(training_batch.reward_scores) + training_batch.reward_mean = reward_stats["reward_mean"] + training_batch.reward_std = reward_stats["reward_std"] + + logger.debug("Rewards computed: mean=%.3f, std=%.3f", + reward_stats["reward_mean"], + reward_stats["reward_std"]) + + return training_batch + + def compute_values(self, training_batch: TrainingBatch) -> TrainingBatch: + """ + Compute value predictions using value model. + + Args: + training_batch: Training batch + + Returns: + Updated training_batch with value predictions + """ + logger.debug("Computing value predictions") + + # TODO: Implement actual value computation + # For now, use dummy values + batch_size = training_batch.latents.shape[0] + dummy_values = torch.randn(batch_size, device=training_batch.latents.device) * 0.1 + 0.5 + training_batch.values = dummy_values + training_batch.old_values = training_batch.values.clone() + + training_batch.value_mean = training_batch.values.mean().item() + + logger.debug("Values computed: mean=%.3f", training_batch.value_mean) + return training_batch + + def compute_advantages(self, training_batch: TrainingBatch) -> TrainingBatch: + """ + Compute advantages using GAE-lambda. + + Args: + training_batch: Training batch with rewards and values + + Returns: + Updated training_batch with advantages and returns + """ + logger.debug("Computing advantages using GAE") + + # For single-step (Flow-GRPO-Fast), we have simple advantage computation + # advantages = rewards - values + # In multi-step, we'd use GAE with gamma and lambda + + if training_batch.reward_scores is None or training_batch.values is None: + raise ValueError("Rewards and values must be computed before advantages") + + # Simple advantage for single-step + # For multi-step, use compute_gae() + rewards = training_batch.reward_scores + values = training_batch.values + next_values = values # For single-step, next_value = current_value + + advantages, returns = compute_gae( + rewards=rewards, + values=values, + next_values=next_values, + gamma=self.training_args.rl_args.rl_gamma, + lambda_=self.training_args.rl_args.rl_lambda + ) + + # Normalize advantages + if self.training_args.rl_args.rl_normalize_advantages: + advantages = normalize_advantages(advantages) + + training_batch.advantages = advantages + training_batch.returns = returns + training_batch.advantage_mean = advantages.mean().item() + training_batch.advantage_std = advantages.std().item() + + logger.debug("Advantages computed: mean=%.3f, std=%.3f", + training_batch.advantage_mean, + training_batch.advantage_std) + + return training_batch + + def train_one_step(self, training_batch: TrainingBatch) -> TrainingBatch: + """ + Train one step using RL/GRPO algorithm. + + This overrides the base class train_one_step to implement RL-specific logic: + 1. Collect trajectories with current policy + 2. Compute rewards + 3. Compute values + 4. Compute advantages + 5. Update policy with GRPO loss + 6. Update value function + + Args: + training_batch: Current training batch + + Returns: + Updated training_batch with loss and metrics + """ + # Check if we're in warmup phase (do SFT instead of RL) + if training_batch.current_timestep < self.training_args.rl_args.rl_warmup_steps: + logger.debug("In warmup phase, using standard SFT training") + return super().train_one_step(training_batch) + + training_batch = self._prepare_training(training_batch) + + # Gradient accumulation loop + for _ in range(self.training_args.gradient_accumulation_steps): + training_batch = self._get_next_batch(training_batch) + training_batch = self._normalize_dit_input(training_batch) + training_batch = self._prepare_dit_inputs(training_batch) + + # === RL-specific steps === + + # 1. Collect trajectories (get log probs from current policy) + training_batch = self.collect_trajectories(training_batch) + + # 2. Compute rewards (requires running reward models) + training_batch = self.compute_rewards(training_batch) + + # 3. Compute value predictions + training_batch = self.compute_values(training_batch) + + # 4. Compute advantages using GAE + training_batch = self.compute_advantages(training_batch) + + # 5. Compute policy loss (GRPO) + if training_batch.log_probs is not None and training_batch.old_log_probs is not None: + policy_loss, policy_info = compute_grpo_policy_loss( + log_probs=training_batch.log_probs, + old_log_probs=training_batch.old_log_probs, + advantages=training_batch.advantages, + clip_range=self.training_args.rl_args.rl_policy_clip_range, + use_ratio_norm=self.training_args.rl_args.rl_ratio_norm_correction, + max_importance_ratio=self.training_args.rl_args.rl_max_importance_ratio + ) + + training_batch.policy_loss = policy_info["policy_loss"] + training_batch.kl_divergence = policy_info["kl_divergence"] + training_batch.importance_ratio = policy_info["importance_ratio_mean"] + training_batch.clip_fraction = policy_info["clip_fraction"] + + # Backward pass for policy + (policy_loss / self.training_args.gradient_accumulation_steps).backward() + + # 6. Compute value loss + if training_batch.values is not None and training_batch.returns is not None: + value_loss, value_info = compute_value_loss( + values=training_batch.values, + returns=training_batch.returns, + old_values=training_batch.old_values, + clip_range=self.training_args.rl_args.rl_value_clip_range, + use_clipping=True + ) + + training_batch.value_loss = value_info["value_loss"] + + # Backward pass for value + value_loss_scaled = value_loss * self.training_args.rl_args.rl_value_loss_coef + (value_loss_scaled / self.training_args.gradient_accumulation_steps).backward() + + # 7. Entropy bonus (optional) + if self.training_args.rl_args.rl_entropy_coef > 0.0 and training_batch.log_probs is not None: + entropy = compute_policy_entropy(training_batch.log_probs) + training_batch.entropy = entropy.item() + entropy_loss = -self.training_args.rl_args.rl_entropy_coef * entropy + (entropy_loss / self.training_args.gradient_accumulation_steps).backward() + + # Accumulate total loss + total_loss = training_batch.policy_loss + training_batch.value_loss + training_batch.total_loss += total_loss + + # Clip gradients + training_batch = self._clip_grad_norm(training_batch) + + # Optimizer step + with self.tracker.timed("timing/optimizer_step"): + self.optimizer.step() + self.lr_scheduler.step() + + if self.value_optimizer is not None: + self.value_optimizer.step() + self.value_scheduler.step() + + # Check for early stopping based on KL divergence + if check_early_stopping(training_batch.kl_divergence, self.training_args.rl_args.rl_target_kl): + logger.warning( + "Early stopping at step %d due to high KL divergence", + training_batch.current_timestep + ) + + return training_batch + + def set_trainable(self) -> None: + """Set which parameters should be trainable.""" + # Policy (transformer) is trainable + for param in self.transformer.parameters(): + param.requires_grad = True + + # Value model is trainable if not sharing backbone + if self.value_model is not None and not self.training_args.rl_args.value_model_share_backbone: + for param in self.value_model.parameters(): + param.requires_grad = True + + # Freeze reward models (they should not be trained) + if self.reward_models is not None: + for param in self.reward_models.parameters(): + param.requires_grad = False + + logger.info("Set trainable parameters for RL training") + + +def create_rl_pipeline( + model_path: str, + training_args: TrainingArgs +) -> RLPipeline: + """ + Factory function to create RL pipeline. + + Args: + model_path: Path to pretrained model + training_args: Training arguments with RL configuration + + Returns: + Initialized RLPipeline + """ + pipeline = RLPipeline(model_path, training_args) + pipeline.initialize_training_pipeline(training_args) + return pipeline diff --git a/fastvideo/training/rl_utils.py b/fastvideo/training/rl_utils.py new file mode 100644 index 000000000..fdb34e575 --- /dev/null +++ b/fastvideo/training/rl_utils.py @@ -0,0 +1,385 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Utility functions for RL/GRPO training. +""" + +from typing import Any + +import torch +import torch.nn.functional as F + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + + +def compute_gae( + rewards: torch.Tensor, + values: torch.Tensor, + next_values: torch.Tensor, + dones: torch.Tensor | None = None, + gamma: float = 0.99, + lambda_: float = 0.95 +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute Generalized Advantage Estimation (GAE-lambda). + + GAE reduces variance in advantage estimation while allowing some bias. + This is a key component of modern policy gradient methods like PPO and GRPO. + + Args: + rewards: Rewards at each step [B, T] or [B] + values: Value predictions at each step [B, T] or [B] + next_values: Value predictions at next step [B, T] or [B] + dones: Episode termination flags [B, T] or [B] (1 if done, 0 otherwise) + gamma: Discount factor + lambda_: GAE lambda parameter (0=TD(0), 1=Monte Carlo) + + Returns: + advantages: GAE advantages [B, T] or [B] + returns: TD(lambda) returns [B, T] or [B] + + Reference: + Schulman et al. "High-Dimensional Continuous Control Using Generalized Advantage Estimation" + https://arxiv.org/abs/1506.02438 + """ + if dones is None: + dones = torch.zeros_like(rewards) + + # Compute TD residuals: delta_t = r_t + gamma * V(s_{t+1}) - V(s_t) + deltas = rewards + gamma * next_values * (1.0 - dones) - values + + # If single step (no time dimension), return directly + if deltas.dim() == 1: + advantages = deltas + returns = advantages + values + return advantages, returns + + # Multi-step: compute GAE recursively + batch_size, num_steps = deltas.shape + advantages = torch.zeros_like(deltas) + gae = torch.zeros(batch_size, device=deltas.device) + + # Backward pass to compute GAE + for t in reversed(range(num_steps)): + gae = deltas[:, t] + gamma * lambda_ * (1.0 - dones[:, t]) * gae + advantages[:, t] = gae + + # Returns are advantages + values + returns = advantages + values + + return advantages, returns + + +def normalize_advantages( + advantages: torch.Tensor, + epsilon: float = 1e-8 +) -> torch.Tensor: + """ + Normalize advantages to have zero mean and unit variance. + + This is a common practice in PPO and GRPO to stabilize training. + + Args: + advantages: Raw advantages [B, ...] + epsilon: Small constant for numerical stability + + Returns: + normalized_advantages: Normalized advantages [B, ...] + """ + mean = advantages.mean() + std = advantages.std() + return (advantages - mean) / (std + epsilon) + + +def compute_grpo_policy_loss( + log_probs: torch.Tensor, + old_log_probs: torch.Tensor, + advantages: torch.Tensor, + clip_range: float = 0.2, + use_ratio_norm: bool = True, + max_importance_ratio: float = 10.0 +) -> tuple[torch.Tensor, dict[str, Any]]: + """ + Compute GRPO policy loss with importance sampling and clipping. + + This implements the core GRPO objective with safety mechanisms from GRPO-Guard: + - Importance ratio clipping (PPO-style) + - RatioNorm correction (GRPO-Guard) + - Ratio clamping for extreme values + + Args: + log_probs: Log probabilities from current policy [B] + old_log_probs: Log probabilities from old policy [B] + advantages: Advantages [B] + clip_range: Clipping range for importance ratios + use_ratio_norm: Apply RatioNorm correction (GRPO-Guard) + max_importance_ratio: Maximum importance ratio before clamping + + Returns: + loss: Policy loss (scalar) + info: Dictionary with diagnostic information + + Reference: + - PPO: Schulman et al. "Proximal Policy Optimization Algorithms" + - GRPO-Guard: RatioNorm and gradient reweighting + """ + # Compute importance ratio: r_t = pi_new(a|s) / pi_old(a|s) + log_ratio = log_probs - old_log_probs + ratio = torch.exp(log_ratio) + + # Clamp extreme ratios for numerical stability + ratio = torch.clamp(ratio, 1.0 / max_importance_ratio, max_importance_ratio) + + # RatioNorm correction (GRPO-Guard) + # Corrects bias in importance sampling when ratio >> 1 + if use_ratio_norm: + ratio_mean = ratio.mean() + ratio = ratio / (ratio_mean + 1e-8) + + # Clipped surrogate objective + ratio_clipped = torch.clamp(ratio, 1.0 - clip_range, 1.0 + clip_range) + surrogate1 = ratio * advantages + surrogate2 = ratio_clipped * advantages + policy_loss = -torch.min(surrogate1, surrogate2).mean() + + # Compute diagnostics + with torch.no_grad(): + # Clip fraction: how often ratios were clipped + clip_fraction = ((ratio < 1.0 - clip_range) | (ratio > 1.0 + clip_range)).float().mean() + + # KL divergence (approximate) + kl_div = log_ratio.mean() + + # Importance ratio stats + importance_ratio_mean = ratio.mean() + importance_ratio_std = ratio.std() + + info = { + "policy_loss": policy_loss.item(), + "clip_fraction": clip_fraction.item(), + "kl_divergence": kl_div.item(), + "importance_ratio_mean": importance_ratio_mean.item(), + "importance_ratio_std": importance_ratio_std.item(), + } + + return policy_loss, info + + +def compute_value_loss( + values: torch.Tensor, + returns: torch.Tensor, + old_values: torch.Tensor | None = None, + clip_range: float = 0.2, + use_clipping: bool = True +) -> tuple[torch.Tensor, dict[str, Any]]: + """ + Compute value function loss with optional clipping. + + Args: + values: Value predictions from current model [B] + returns: Target returns (from GAE) [B] + old_values: Value predictions from old model [B] (for clipping) + clip_range: Clipping range for value updates + use_clipping: Whether to use clipped value loss (PPO-style) + + Returns: + loss: Value loss (scalar) + info: Dictionary with diagnostic information + """ + # Standard MSE loss + value_loss_unclipped = F.mse_loss(values, returns, reduction="none") + + # Clipped value loss (PPO-style) + if use_clipping and old_values is not None: + values_clipped = old_values + torch.clamp( + values - old_values, + -clip_range, + clip_range + ) + value_loss_clipped = F.mse_loss(values_clipped, returns, reduction="none") + value_loss = torch.max(value_loss_unclipped, value_loss_clipped).mean() + else: + value_loss = value_loss_unclipped.mean() + + # Compute diagnostics + with torch.no_grad(): + explained_variance = 1.0 - (returns - values).var() / (returns.var() + 1e-8) + + info = { + "value_loss": value_loss.item(), + "explained_variance": explained_variance.item(), + "value_mean": values.mean().item(), + "value_std": values.std().item(), + } + + return value_loss, info + + +def compute_policy_entropy(log_probs: torch.Tensor) -> torch.Tensor: + """ + Compute policy entropy for exploration bonus. + + Args: + log_probs: Log probabilities [B] + + Returns: + entropy: Mean entropy across batch (scalar) + """ + # For continuous actions: H = -log_prob (assuming Gaussian) + # For discrete: H = -sum(p * log(p)) + # Here we use a simple approximation + entropy = -log_probs.mean() + return entropy + + +def apply_gradient_reweighting( + gradients: torch.Tensor, + timesteps: torch.Tensor, + num_train_timesteps: int = 1000 +) -> torch.Tensor: + """ + Apply GRPO-Guard gradient reweighting across denoising steps. + + This reweights gradients based on the timestep to balance learning + across different noise levels. + + Args: + gradients: Gradients to reweight [B, ...] + timesteps: Timesteps at which gradients were computed [B] + num_train_timesteps: Total number of training timesteps + + Returns: + reweighted_gradients: Reweighted gradients [B, ...] + """ + # Compute timestep weights (higher weight for later timesteps) + # This is a simple linear weighting, can be made more sophisticated + timestep_weights = 1.0 + (timesteps.float() / num_train_timesteps) + timestep_weights = timestep_weights.view(-1, *([1] * (gradients.dim() - 1))) + + return gradients * timestep_weights + + +def sample_random_timesteps( + batch_size: int, + min_timestep: int, + max_timestep: int, + device: torch.device, + generator: torch.Generator | None = None +) -> torch.Tensor: + """ + Sample random timesteps for noise injection (Flow-GRPO-Fast). + + Args: + batch_size: Number of samples + min_timestep: Minimum timestep + max_timestep: Maximum timestep + device: Device for tensor + generator: Random generator for reproducibility + + Returns: + timesteps: Random timesteps [B] + """ + if generator is not None: + timesteps = torch.randint( + min_timestep, + max_timestep + 1, + (batch_size,), + device=device, + generator=generator + ) + else: + timesteps = torch.randint( + min_timestep, + max_timestep + 1, + (batch_size,), + device=device + ) + + return timesteps + + +def compute_reward_statistics( + rewards: torch.Tensor +) -> dict[str, float]: + """ + Compute statistics for reward distribution. + + Args: + rewards: Reward values [B] + + Returns: + stats: Dictionary with mean, std, min, max + """ + return { + "reward_mean": rewards.mean().item(), + "reward_std": rewards.std().item(), + "reward_min": rewards.min().item(), + "reward_max": rewards.max().item(), + } + + +def check_early_stopping( + kl_divergence: float, + target_kl: float +) -> bool: + """ + Check if training should stop early based on KL divergence. + + Args: + kl_divergence: Current KL divergence + target_kl: Target KL threshold + + Returns: + should_stop: True if KL exceeds target + """ + if kl_divergence > target_kl: + logger.warning( + "Early stopping triggered: KL divergence %.4f > target %.4f", + kl_divergence, + target_kl + ) + return True + return False + + +def compute_log_probs_from_model_output( + model_output: torch.Tensor, + target: torch.Tensor, + noise_level: float = 0.1 +) -> torch.Tensor: + """ + Compute log probabilities from model predictions. + + For diffusion models, we approximate log probabilities using the + negative squared error (assuming Gaussian likelihood). + + Args: + model_output: Model predictions [B, C, T, H, W] + target: Target values [B, C, T, H, W] + noise_level: Assumed noise level (std) for Gaussian likelihood + + Returns: + log_probs: Log probabilities [B] + """ + # Compute mean squared error per sample + mse = ((model_output - target) ** 2).flatten(1).mean(dim=1) + + # Log probability under Gaussian: log p(x) = -0.5 * (x - mu)^2 / sigma^2 + const + log_probs = -0.5 * mse / (noise_level ** 2) + + return log_probs + + +def check_for_nan_inf(tensor: torch.Tensor, name: str) -> None: + """ + Check tensor for NaN or Inf values and raise error if found. + + Args: + tensor: Tensor to check + name: Name for error message + """ + if torch.isnan(tensor).any(): + raise ValueError(f"{name} contains NaN values") + if torch.isinf(tensor).any(): + raise ValueError(f"{name} contains Inf values")