implement Phase 1 backbone code

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
 
 

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):

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",
+]

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)

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

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")