training/configs/config.py

#!/usr/bin/env python3
"""
Configuration Management for LKAlert SFT and DPO Training Pipeline

This module provides comprehensive configuration classes for:
- Model settings (VLM backbone, LoRA, heads)
- Data settings (NEXAR, DADA-2000 datasets)
- SFT training settings (TTA regression with uncertainty)
- DPO training settings (policy learning)
- Evaluation settings (benchmark metrics)
- Ablation study configurations

Key Data Format:
- Frame rate: 20Hz (0.05s per frame)
- accident_time: frame number when accident occurs
- risky_time: frame number when risk first becomes observable
- TTA = (accident_time - current_frame) * 0.05 seconds

Author: LKAlert Team
Version: 3.0 (Production)
"""

import json
import copy
from dataclasses import dataclass, field, asdict
from typing import Dict, List, Optional, Any, Tuple, Union
from pathlib import Path


# ============================================================================
# Data Configuration
# ============================================================================

@dataclass
class DataConfig:
    """
    Data configuration for NEXAR and DADA-2000 datasets
    
    NEXAR Structure:
        PROJECT_ROOT/NEXAR_COLLISION/dataset/
            train/ | test-public/ | test-private/
                positive/ | negative/
                    00000/
                        000.jpg, 001.jpg, ..., annotation.json
    
    DADA-2000 Structure:
        PROJECT_ROOT/DADA-2000/
            positive/ | negative/ | non-ego/
                0001/
                    frames: 000.jpg, 001.jpg, ...
                    annotation.json
    
    Label Format (annotation.json):
        - accident_time: frame number when accident occurs (e.g., 415 = 20.75s)
        - risky_time: frame number when risk first observable (e.g., 383 = 19.15s)
        - TTA = (accident_frame - current_frame) * 0.05
    """
    
    # Data roots
    nexar_root: str = "PROJECT_ROOT/NEXAR_COLLISION/dataset"
    dada_root: str = "PROJECT_ROOT/DADA-2000"
    
    # Frame settings (20Hz = 0.05s per frame)
    frame_rate: int = 20  # Hz
    frame_interval: float = 0.05  # seconds per frame
    
    # Observation window settings
    # Standard window: 2 seconds (40 frames @ 20Hz)
    # Extended window after OBSERVE: 3 seconds (60 frames @ 20Hz)
    window_size_frames: int = 40   # 2 seconds @ 20Hz
    extended_window_frames: int = 60  # 3 seconds @ 20Hz
    stride_frames: int = 10  # 0.5 seconds sliding window stride
    
    # Frame sampling for VLM (reduce token count)
    frame_sample_rate: int = 4  # Sample every N frames
    max_frames_per_sample: int = 10  # Maximum frames to feed VLM
    
    # TTA settings
    max_tta_seconds: float = 10.0  # Maximum TTA for training
    min_tta_seconds: float = 0.1   # Minimum TTA (clip values)
    
    # Dataset balance
    use_negatives: bool = True
    negative_ratio: float = 0.3  # Ratio of negative to positive samples
    
    # Data augmentation
    use_augmentation: bool = True
    augmentation_prob: float = 0.5
    
    # Image settings
    image_size: Tuple[int, int] = (384, 384)
    
    @property
    def window_size_seconds(self) -> float:
        """Standard observation window in seconds"""
        return self.window_size_frames / self.frame_rate
    
    @property
    def extended_window_seconds(self) -> float:
        """Extended window after OBSERVE action in seconds"""
        return self.extended_window_frames / self.frame_rate
    
    def frame_to_time(self, frame: int) -> float:
        """Convert frame number to time in seconds"""
        return frame * self.frame_interval
    
    def time_to_frame(self, time_s: float) -> int:
        """Convert time in seconds to frame number"""
        return int(time_s / self.frame_interval)


# ============================================================================
# Model Configuration
# ============================================================================

@dataclass
class ModelConfig:
    """
    Model configuration for BeliefActionVLM
    
    Architecture:
    - VLM Backbone: Qwen2.5-VL-3B/7B with LoRA fine-tuning
    - Belief Aggregator: Compresses hidden states to belief representation
    - TTA Head: Regresses time-to-accident with uncertainty
    - Policy Head: Selects actions (SILENT/OBSERVE/ALERT)
    """
    
    # VLM Backbone
    model_name: str = "Qwen/Qwen2.5-VL-3B-Instruct"
    hidden_dim: int = 2048  # 3B=2048, 7B=3584 (auto-detected)
    
    # Belief aggregation strategy
    # Options: "mean_pool", "last_token", "attention_pool"
    belief_aggregation: str = "mean_pool"
    belief_compression_dim: int = 256  # Optional compression
    use_belief_compression: bool = False
    
    # LoRA configuration for parameter-efficient fine-tuning
    use_lora: bool = True
    lora_r: int = 32
    lora_alpha: int = 64
    lora_dropout: float = 0.1
    lora_target_modules: List[str] = field(default_factory=lambda: [
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ])
    
    # TTA Head configuration
    tta_intermediate_dim: int = 512
    tta_dropout: float = 0.1
    
    # Policy Head configuration
    policy_intermediate_dim: int = 512
    policy_dropout: float = 0.1
    num_actions: int = 3  # SILENT, OBSERVE, ALERT
    
    # Pretrained checkpoints
    pretrained_vlm_path: Optional[str] = None
    pretrained_lora_path: str = "PROJECT_ROOT/checkpoints/pretrain/stage_b/Qwen2.5-VL-3B-Instruct"
    
    # Mixed precision
    use_bf16: bool = True


# ============================================================================
# SFT Training Configuration
# ============================================================================

@dataclass
class SFTTrainingConfig:
    """
    SFT (Supervised Fine-Tuning) training configuration
    
    Stage 1: Train VLM backbone and TTA head for multi-modal TTA regression
    
    Loss Function (Eq. 23):
        L_SFT = L_MSE + λ * L_NLL
        
        where:
        - L_MSE = (TTA_pred - TTA_true)²
        - L_NLL = 0.5 * (log(σ²) + (TTA_pred - TTA_true)² / σ²)
        - λ controls uncertainty calibration weight
    
    Curriculum Scheduled Sampling (Section 1.5):
        - Phase 0 (0-30%): Warmup with rule-based actions
        - Phase 1 (30-70%): Transition with mixed actions  
        - Phase 2 (70-100%): Full self-play with model actions
    """
    
    # Training epochs and batches
    num_epochs: int = 10
    batch_size: int = 4
    gradient_accumulation_steps: int = 4
    
    # Learning rate settings
    learning_rate: float = 1e-4
    tta_head_lr: float = 1e-3
    vlm_lr_multiplier: float = 0.1  # VLM gets lower LR
    min_lr: float = 1e-6
    weight_decay: float = 0.01
    
    # Learning rate schedule
    scheduler_type: str = "cosine"  # "cosine", "linear", "constant"
    warmup_ratio: float = 0.1
    warmup_steps: Optional[int] = None
    
    # Loss weights (Eq. 23)
    mse_weight: float = 1.0
    nll_weight: float = 0.1  # λ for uncertainty calibration
    
    # Curriculum learning settings
    use_curriculum: bool = True
    curriculum_phases: List[float] = field(default_factory=lambda: [0.3, 0.7, 1.0])
    # Phase 0: [0, 0.3) - Rule-based
    # Phase 1: [0.3, 0.7) - Mixed
    # Phase 2: [0.7, 1.0] - Self-play
    
    # Gradient settings
    max_grad_norm: float = 1.0
    
    # Checkpointing
    save_steps: int = 500
    eval_steps: int = 250
    logging_steps: int = 50
    save_total_limit: int = 3
    
    # Output
    output_dir: str = "PROJECT_ROOT/checkpoints/sft"
    experiment_name: str = "sft_default"
    
    # Mixed precision
    use_amp: bool = True
    
    # Debugging
    debug: bool = False
    debug_samples: int = 100
    
    # Wandb
    use_wandb: bool = True
    wandb_project: str = "lkalert-sft"


# ============================================================================
# DPO Training Configuration
# ============================================================================

@dataclass
class DPOTrainingConfig:
    """
    DPO (Direct Preference Optimization) training configuration
    
    Stage 2: Train policy head for action selection
    
    Loss Function (Bradley-Terry, Eq. 28):
        L_DPO = -log σ(β * (log π(τ⁺)/π_ref(τ⁺) - log π(τ⁻)/π_ref(τ⁻)))
        
        where:
        - τ⁺: Preferred trajectory (higher reward)
        - τ⁻: Dispreferred trajectory (lower reward)
        - β: Temperature parameter
    
    Reward Function (Eq. 27):
        R(τ) = Σ_t r(s_t, a_t, s_{t+1})
        
        where r() assigns:
        - +10: Timely alert (TTA ∈ [2, 5] seconds)
        - -20: Miss (no alert before accident)
        - -5: False alarm (alert when TTA > 5s)
        - +3: OBSERVE action that reduces uncertainty
    """
    
    # Training epochs and batches
    num_epochs: int = 5
    batch_size: int = 2
    gradient_accumulation_steps: int = 8
    
    # DPO hyperparameters
    beta: float = 0.1  # Temperature for preference learning
    reference_free: bool = False  # Use reference model
    
    # Learning rate settings
    learning_rate: float = 5e-5
    min_lr: float = 1e-6
    weight_decay: float = 0.01
    
    # Learning rate schedule
    scheduler_type: str = "cosine"
    warmup_ratio: float = 0.1
    
    # Reward function parameters (Eq. 27)
    reward_timely_alert: float = 10.0
    reward_miss: float = -20.0
    reward_false_alarm: float = -5.0
    reward_observe_uncertainty: float = 3.0
    
    # Alert thresholds
    min_alert_tta: float = 2.0  # Minimum TTA for valid alert
    max_alert_tta: float = 5.0  # Maximum TTA for timely alert
    
    # Preference pair generation
    min_reward_margin: float = 3.0  # Minimum margin between τ⁺ and τ⁻
    trajectories_per_video: int = 5  # Number of policy variants to generate
    
    # Gradient settings
    max_grad_norm: float = 1.0
    
    # Checkpointing
    save_steps: int = 200
    eval_steps: int = 100
    logging_steps: int = 20
    save_total_limit: int = 3
    
    # Output
    output_dir: str = "PROJECT_ROOT/checkpoints/dpo"
    experiment_name: str = "dpo_default"
    
    # SFT checkpoint to load
    sft_checkpoint: str = "PROJECT_ROOT/checkpoints/sft/best"
    
    # Mixed precision
    use_amp: bool = True
    
    # Debugging
    debug: bool = False
    debug_samples: int = 50
    
    # Wandb
    use_wandb: bool = True
    wandb_project: str = "lkalert-dpo"


# ============================================================================
# Evaluation Configuration
# ============================================================================

@dataclass
class EvaluationConfig:
    """
    Evaluation configuration for benchmark testing
    
    Metrics:
    - TTA Regression: MAE, RMSE, R², Calibration Error
    - Policy Performance: Precision, Recall, F1 for alerts
    - System Performance: Miss Rate, False Alarm Rate, Detection Time
    """
    
    # Test datasets
    test_nexar: bool = True  # Use NEXAR test-private
    test_dada: bool = True   # Use DADA-2000
    
    # Evaluation settings
    batch_size: int = 8
    num_workers: int = 4
    
    # Alert thresholds
    alert_tta_threshold: float = 2.0  # TTA below this triggers alert
    uncertainty_threshold: float = 0.5  # Uncertainty above this suggests OBSERVE
    
    # Calibration
    num_calibration_bins: int = 10
    
    # Output
    output_dir: str = "PROJECT_ROOT/evaluation_results"
    save_predictions: bool = True
    save_visualizations: bool = True
    
    # Visualization
    plot_format: str = "pdf"
    plot_dpi: int = 300
    use_latex: bool = True
    font_family: str = "Times New Roman"


# ============================================================================
# Full Configuration
# ============================================================================

@dataclass
class FullConfig:
    """
    Complete configuration combining all sub-configs
    """
    data: DataConfig = field(default_factory=DataConfig)
    model: ModelConfig = field(default_factory=ModelConfig)
    sft: SFTTrainingConfig = field(default_factory=SFTTrainingConfig)
    dpo: DPOTrainingConfig = field(default_factory=DPOTrainingConfig)
    evaluation: EvaluationConfig = field(default_factory=EvaluationConfig)
    
    def save(self, path: str):
        """Save configuration to JSON file"""
        path = Path(path)
        path.parent.mkdir(parents=True, exist_ok=True)
        
        # Convert to dict recursively
        def to_dict(obj):
            if hasattr(obj, '__dataclass_fields__'):
                return {k: to_dict(v) for k, v in asdict(obj).items()}
            elif isinstance(obj, (list, tuple)):
                return [to_dict(v) for v in obj]
            else:
                return obj
        
        config_dict = to_dict(self)
        
        with open(path, 'w') as f:
            json.dump(config_dict, f, indent=2)
        
        print(f"✅ Configuration saved to {path}")
    
    @classmethod
    def load(cls, path: str) -> 'FullConfig':
        """Load configuration from JSON file"""
        with open(path, 'r') as f:
            config_dict = json.load(f)
        
        return cls(
            data=DataConfig(**config_dict.get('data', {})),
            model=ModelConfig(**config_dict.get('model', {})),
            sft=SFTTrainingConfig(**config_dict.get('sft', {})),
            dpo=DPOTrainingConfig(**config_dict.get('dpo', {})),
            evaluation=EvaluationConfig(**config_dict.get('evaluation', {}))
        )


# ============================================================================
# Ablation Study Configurations
# ============================================================================

ABLATION_CONFIGS = {
    # Belief Aggregation Ablations
    "belief_mean_pool": {
        "model.belief_aggregation": "mean_pool"
    },
    "belief_last_token": {
        "model.belief_aggregation": "last_token"
    },
    "belief_attention_pool": {
        "model.belief_aggregation": "attention_pool"
    },
    
    # Curriculum Learning Ablations
    "no_curriculum": {
        "sft.use_curriculum": False
    },
    "with_curriculum": {
        "sft.use_curriculum": True
    },
    
    # Loss Weight Ablations
    "mse_only": {
        "sft.nll_weight": 0.0
    },
    "nll_heavy": {
        "sft.nll_weight": 0.5
    },
    "nll_light": {
        "sft.nll_weight": 0.05
    },
    
    # Window Size Ablations
    "window_1s": {
        "data.window_size_frames": 20,
        "data.extended_window_frames": 40
    },
    "window_3s": {
        "data.window_size_frames": 60,
        "data.extended_window_frames": 80
    },
    
    # Model Size Ablations
    "3B_model": {
        "model.model_name": "Qwen/Qwen2.5-VL-3B-Instruct",
        "model.hidden_dim": 2048,
        "model.pretrained_lora_path": "PROJECT_ROOT/checkpoints/pretrain/stage_b/Qwen2.5-VL-3B-Instruct"
    },
    "7B_model": {
        "model.model_name": "Qwen/Qwen2.5-VL-7B-Instruct",
        "model.hidden_dim": 3584,
        "model.pretrained_lora_path": "PROJECT_ROOT/checkpoints/pretrain/stage_b/Qwen2.5-VL-7B-Instruct"
    },
    
    # LoRA Rank Ablations
    "lora_r16": {
        "model.lora_r": 16,
        "model.lora_alpha": 32
    },
    "lora_r64": {
        "model.lora_r": 64,
        "model.lora_alpha": 128
    },
    
    # DPO Beta Ablations
    "beta_0.05": {
        "dpo.beta": 0.05
    },
    "beta_0.2": {
        "dpo.beta": 0.2
    },
    
    # Frame Sampling Ablations
    "dense_frames": {
        "data.frame_sample_rate": 2,
        "data.max_frames_per_sample": 20
    },
    "sparse_frames": {
        "data.frame_sample_rate": 8,
        "data.max_frames_per_sample": 5
    },
    
    # Negative Sample Ablations
    "no_negatives": {
        "data.use_negatives": False
    },
    "more_negatives": {
        "data.negative_ratio": 0.5
    }
}


def create_ablation_config(base_config: FullConfig, ablation_name: str) -> FullConfig:
    """
    Create an ablation configuration by modifying the base config
    
    Args:
        base_config: Base configuration to modify
        ablation_name: Name of ablation from ABLATION_CONFIGS
        
    Returns:
        Modified configuration
    """
    if ablation_name not in ABLATION_CONFIGS:
        raise ValueError(f"Unknown ablation: {ablation_name}. "
                        f"Available: {list(ABLATION_CONFIGS.keys())}")
    
    # Deep copy base config
    config = copy.deepcopy(base_config)
    
    # Apply modifications
    modifications = ABLATION_CONFIGS[ablation_name]
    
    for key, value in modifications.items():
        parts = key.split('.')
        obj = config
        
        # Navigate to the nested attribute
        for part in parts[:-1]:
            obj = getattr(obj, part)
        
        # Set the value
        setattr(obj, parts[-1], value)
    
    # Update experiment name
    config.sft.experiment_name = f"sft_{ablation_name}"
    config.dpo.experiment_name = f"dpo_{ablation_name}"
    
    return config


def get_debug_config() -> FullConfig:
    """Get configuration for debugging with small dataset"""
    config = FullConfig()
    
    # Enable debug mode
    config.sft.debug = True
    config.sft.debug_samples = 100
    config.sft.num_epochs = 2
    config.sft.batch_size = 2
    config.sft.eval_steps = 50
    config.sft.save_steps = 100
    config.sft.logging_steps = 10
    
    config.dpo.debug = True
    config.dpo.debug_samples = 50
    config.dpo.num_epochs = 1
    config.dpo.batch_size = 1
    config.dpo.eval_steps = 25
    config.dpo.save_steps = 50
    
    return config


def get_fast_config() -> FullConfig:
    """Get configuration for fast training (fewer epochs, larger batches)"""
    config = FullConfig()
    
    config.sft.num_epochs = 5
    config.sft.batch_size = 8
    config.sft.gradient_accumulation_steps = 2
    
    config.dpo.num_epochs = 3
    config.dpo.batch_size = 4
    config.dpo.gradient_accumulation_steps = 4
    
    return config


# ============================================================================
# __init__.py content
# ============================================================================

__all__ = [
    'DataConfig',
    'ModelConfig', 
    'SFTTrainingConfig',
    'DPOTrainingConfig',
    'EvaluationConfig',
    'FullConfig',
    'ABLATION_CONFIGS',
    'create_ablation_config',
    'get_debug_config',
    'get_fast_config'
]


if __name__ == "__main__":
    # Test configuration
    print("=" * 60)
    print("LKAlert Configuration Test")
    print("=" * 60)
    
    # Create default config
    config = FullConfig()
    print(f"\n📊 Default Configuration:")
    print(f"  Data:")
    print(f"    NEXAR root: {config.data.nexar_root}")
    print(f"    DADA root: {config.data.dada_root}")
    print(f"    Frame rate: {config.data.frame_rate} Hz")
    print(f"    Window size: {config.data.window_size_seconds}s ({config.data.window_size_frames} frames)")
    print(f"    Extended window: {config.data.extended_window_seconds}s ({config.data.extended_window_frames} frames)")
    
    print(f"\n  Model:")
    print(f"    VLM: {config.model.model_name}")
    print(f"    Belief aggregation: {config.model.belief_aggregation}")
    print(f"    LoRA: r={config.model.lora_r}, alpha={config.model.lora_alpha}")
    
    print(f"\n  SFT Training:")
    print(f"    Epochs: {config.sft.num_epochs}")
    print(f"    Batch size: {config.sft.batch_size}")
    print(f"    Learning rate: {config.sft.learning_rate}")
    print(f"    Curriculum learning: {config.sft.use_curriculum}")
    
    print(f"\n  DPO Training:")
    print(f"    Epochs: {config.dpo.num_epochs}")
    print(f"    Beta: {config.dpo.beta}")
    print(f"    Reward (timely alert): {config.dpo.reward_timely_alert}")
    print(f"    Reward (miss): {config.dpo.reward_miss}")
    
    # Test ablation creation
    print(f"\n📊 Available Ablations:")
    for name in ABLATION_CONFIGS:
        print(f"  - {name}")
    
    # Save test config
    config.save("/tmp/lkalert_config_test.json")
    
    # Load and verify
    loaded = FullConfig.load("/tmp/lkalert_config_test.json")
    print(f"\n✅ Config save/load test passed!")