# sync_vs_async_comparison.py - 同步RL vs 异步RL效率对比

import numpy as np
import time
from typing import List

class SyncRLPipeline:
    """同步RL流水线模拟"""
    
    def __init__(self, group_size: int = 8):
        self.group_size = group_size
    
    def simulate_rollout(self, task_difficulty: str) -> float:
        """模拟一个rollout的生成时间"""
        time_map = {
            "simple": np.random.uniform(0.5, 1.5),
            "medium": np.random.uniform(1.5, 4.0),
            "complex": np.random.uniform(4.0, 10.0),
            "agent": np.random.uniform(8.0, 30.0),
        }
        return time_map.get(task_difficulty, 2.0)
    
    def run_batch(self, tasks: List[str]) -> float:
        """同步执行一批任务，等待最慢的完成"""
        start = time.time()
        
        times = []
        for task in tasks:
            t = self.simulate_rollout(task)
            times.append(t)
            time.sleep(t * 0.01)  # 模拟
        
        # 同步等待：取最慢的
        batch_time = max(times)
        total_wait = sum(times)
        idle_time = batch_time * len(times) - total_wait
        
        return {
            "batch_time": batch_time,
            "total_compute": total_wait,
            "idle_time": idle_time,
            "gpu_utilization": total_wait / (batch_time * len(times)) * 100,
        }

class AsyncRLPipeline:
    """异步RL流水线模拟"""
    
    def simulate_rollout(self, task_difficulty: str) -> float:
        time_map = {
            "simple": np.random.uniform(0.5, 1.5),
            "medium": np.random.uniform(1.5, 4.0),
            "complex": np.random.uniform(4.0, 10.0),
            "agent": np.random.uniform(8.0, 30.0),
        }
        return time_map.get(task_difficulty, 2.0)
    
    def run_batch(self, tasks: List[str]) -> float:
        """异步执行：每个rollout完成即训练，无需等待"""
        start = time.time()
        
        times = []
        for task in tasks:
            t = self.simulate_rollout(task)
            times.append(t)
            time.sleep(t * 0.01)
        
        # 异步：各rollout独立完成，立即更新
        # 总时间 = 最长rollout时间
        total_time = max(times)
        total_compute = sum(times)
        
        # 异步节省的时间
        sync_time = total_time * len(times)
        async_time = total_time
        
        return {
            "total_time": total_time,
            "total_compute": total_compute,
            "sync_time": sync_time,
            "async_time": async_time,
            "speedup": sync_time / async_time,
            "gpu_utilization": 100.0,  # 异步无等待
        }

# 对比测试
np.random.seed(42)

# Agent任务混合
agent_tasks = ["agent"] * 2 + ["complex"] * 3 + ["medium"] * 3
simple_tasks = ["simple"] * 8

for name, tasks in [("Agent混合任务", agent_tasks), ("简单任务", simple_tasks)]:
    sync = SyncRLPipeline(group_size=8)
    async_p = AsyncRLPipeline()
    
    sync_result = sync.run_batch(tasks)
    async_result = async_p.run_batch(tasks)
    
    print(f"\n=== {name} ===")
    print(f"同步RL: 耗时={sync_result['batch_time']:.1f}s, GPU利用率={sync_result['gpu_utilization']:.0f}%")
    print(f"异步RL: 耗时={async_result['total_time']:.1f}s, 加速比={async_result['speedup']:.1f}x")

# sao_dis.py - SAO直接双边重要性采样实现

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class DirectImportanceSampling:
    """
    直接双边重要性采样（DIS）
    核心：用rollout引擎记录的log-probabilities，计算当前策略与行为策略的概率比率
    
    与标准PPO的差异：
    - PPO：clip(ratio, 1-ε, 1+ε) 截断到边界
    - SAO-DIS：超出信任域的token直接掩码丢弃，不参与梯度计算
    """
    
    def __init__(self, 
                 clip_lower: float = 0.2,
                 clip_upper: float = 5.0,
                 epsilon: float = 1e-8):
        """
        Args:
            clip_lower: 下界阈值（概率比率低于此值直接掩码）
            clip_upper: 上界阈值（概率比率高于此值直接掩码）
            epsilon: 防止除零
        """
        self.clip_lower = clip_lower
        self.clip_upper = clip_upper
        self.epsilon = epsilon
        
        # 统计
        self.stats = {
            "total_tokens": 0,
            "masked_low": 0,
            "masked_high": 0,
            "kept_tokens": 0,
        }
    
    def compute_importance_ratio(self,
                                  current_log_probs: torch.Tensor,
                                  behavior_log_probs: torch.Tensor) -> torch.Tensor:
        """
        计算重要性采样比率
        ratio = exp(current_log_prob - behavior_log_prob)
        
        Args:
            current_log_probs: 当前策略的log概率 [batch, seq_len]
            behavior_log_probs: rollout引擎记录的log概率 [batch, seq_len]
        Returns:
            ratios: 重要性采样比率 [batch, seq_len]
        """
        log_ratio = current_log_probs - behavior_log_probs
        ratios = torch.exp(log_ratio)
        return ratios
    
    def apply_clipping(self, 
                        ratios: torch.Tensor,
                        advantages: torch.Tensor) -> torch.Tensor:
        """
        应用双侧Token级裁剪
        超出[clip_lower, clip_upper]区间的token被完全掩码
        
        Args:
            ratios: 重要性采样比率 [batch, seq_len]
            advantages: 优势值 [batch, seq_len]
        Returns:
            clipped_objective: 裁剪后的策略目标 [batch, seq_len]
        """
        # 1. 计算未裁剪的目标
        surr1 = ratios * advantages
        
        # 2. 检测超出信任域的token
        mask_low = ratios < self.clip_lower
        mask_high = ratios > self.clip_upper
        
        # 3. 被掩码的token：目标设为0（不参与梯度计算）
        clipped = torch.where(
            mask_low | mask_high,
            torch.zeros_like(surr1),
            surr1
        )
        
        # 更新统计
        self.stats["total_tokens"] += ratios.numel()
        self.stats["masked_low"] += mask_low.sum().item()
        self.stats["masked_high"] += mask_high.sum().item()
        self.stats["kept_tokens"] += (~(mask_low | mask_high)).sum().item()
        
        return clipped
    
    def get_mask_rate(self) -> dict:
        """获取掩码率统计"""
        total = self.stats["total_tokens"]
        return {
            "total_tokens": total,
            "low_mask_rate": self.stats["masked_low"] / max(total, 1),
            "high_mask_rate": self.stats["masked_high"] / max(total, 1),
            "kept_rate": self.stats["kept_tokens"] / max(total, 1),
        }

class SAOPolicyLoss(nn.Module):
    """SAO策略损失函数"""
    
    def __init__(self, clip_lower=0.2, clip_upper=5.0):
        super().__init__()
        self.dis = DirectImportanceSampling(clip_lower, clip_upper)
    
    def forward(self, 
                current_log_probs: torch.Tensor,
                behavior_log_probs: torch.Tensor,
                advantages: torch.Tensor) -> torch.Tensor:
        """
        SAO策略损失
        L = -E[clip(ratio, low, high) * advantage]
        """
        ratios = self.dis.compute_importance_ratio(
            current_log_probs, behavior_log_probs
        )
        clipped = self.dis.apply_clipping(ratios, advantages)
        
        # 负号：梯度上升（最大化目标）
        loss = -clipped.mean()
        
        return loss

# 模拟SAO vs GRPO的训练稳定性
def simulate_training_stability():
    """模拟SAO和GRPO的训练稳定性对比"""
    np.random.seed(42)
    
    n_steps = 1000
    
    # SAO: 稳定提升，不崩溃
    sao_performance = 10 + 80 * (1 - np.exp(-np.arange(n_steps) / 300))
    sao_performance += np.random.randn(n_steps) * 2  # 小幅噪声
    
    # GRPO: 约160步后崩溃
    grpo_performance = 10 + 80 * (1 - np.exp(-np.arange(n_steps) / 200))
    collapse_point = 160
    if collapse_point < n_steps:
        grpo_performance[collapse_point:] = (
            grpo_performance[collapse_point] * 
            np.exp(-(np.arange(n_steps - collapse_point)) / 50)
        )
    grpo_performance += np.random.randn(n_steps) * 3
    
    # 打印关键点
    for step in [0, 100, 200, 400, 600, 800, 1000]:
        if step < n_steps:
            print(f"Step {step:4d}: SAO={sao_performance[step]:.1f}  GRPO={grpo_performance[step]:.1f}")

simulate_training_stability()

// sao_critic.go - SAO Critic网络优化实现

package main

import (
	"fmt"
	"math"
)

// SAOCriticConfig SAO Critic网络配置
type SAOCriticConfig struct {
	// 更新频率
	ActorUpdateSteps  int     // 策略每n步更新一次
	CriticUpdateSteps int     // 价值每n步更新一次
	
	// 冻结策略
	FreezeAttention  bool    // 是否冻结注意力层
	FreezeMLP        bool    // 是否冻结MLP层
	TrainMoEOnly     bool    // 是否仅训练MoE投影层
	
	// GAE参数
	Gamma            float64 // 折扣因子
	Lambda           float64 // GAE lambda
	SkipObservation  bool    // 是否跳过观测token
}

// SAOCritic SAO价值网络
type SAOCritic struct {
	Config SAOCriticConfig
	
	// 模拟参数
	attentionParams int64
	moeParams       int64
	totalParams     int64
}

// NewSAOCritic 创建SAO Critic
func NewSAOCritic(cfg SAOCriticConfig) *SAOCritic {
	return &SAOCritic{
		Config: cfg,
		// 模拟GLM-5.2 Critic网络参数分布
		attentionParams: 80 * 8192 * 8192,     // 80层注意力
		moeParams:       80 * 64 * 8192 * 16384, // MoE专家
		totalParams:     750_000_000_000,        // 750B
	}
}

// CalculateTrainableParams 计算可训练参数量
func (c *SAOCritic) CalculateTrainableParams() map[string]int64 {
	result := make(map[string]int64)
	
	if c.Config.FreezeAttention {
		result["attention_frozen"] = c.attentionParams
		result["attention_trainable"] = 0
	} else {
		result["attention_frozen"] = 0
		result["attention_trainable"] = c.attentionParams
	}
	
	if c.Config.TrainMoEOnly {
		result["moe_trainable"] = c.moeParams
		result["other_trainable"] = 0
	} else {
		result["moe_trainable"] = c.moeParams
		result["other_trainable"] = c.totalParams - c.attentionParams - c.moeParams
	}
	
	result["total_trainable"] = result["attention_trainable"] + 
		result["moe_trainable"] + result["other_trainable"]
	result["frozen_ratio"] = c.totalParams - result["total_trainable"]
	
	return result
}

// GetUpdateRatio 获取更新频率比
func (c *SAOCritic) GetUpdateRatio() float64 {
	return float64(c.CriticUpdateSteps) / float64(c.ActorUpdateSteps)
}

// SimulateValueStability 模拟价值模型训练稳定性
func (c *SAOCritic) SimulateValueStability() map[string]float64 {
	results := make(map[string]float64)
	
	// 模拟不同配置下的价值模型方差
	configs := []struct {
		name        string
		freezeAttn  bool
		criticRatio float64
		skipObs     bool
	}{
		{"SAO完整", true, 2.0, true},
		{"无冻结注意力", false, 2.0, true},
		{"等频更新", true, 1.0, true},
		{"无跳过观测", true, 2.0, false},
	}
	
	for _, cfg := range configs {
		// 模拟价值模型方差
		baseVariance := 0.5
		if !cfg.freezeAttn {
			baseVariance *= 1.8
		}
		if cfg.criticRatio < 1.5 {
			baseVariance *= 1.3
		}
		if !cfg.skipObs {
			baseVariance *= 1.4
		}
		
		results[cfg.name] = baseVariance
	}
	
	return results
}

func main() {
	cfg := SAOCriticConfig{
		ActorUpdateSteps:  1,
		CriticUpdateSteps: 2,  // 策略每更新1次，价值更新2次
		FreezeAttention:   true,
		TrainMoEOnly:      true,
		Gamma:             0.99,
		Lambda:            0.95,
		SkipObservation:   true,
	}
	
	critic := NewSAOCritic(cfg)
	
	// 计算可训练参数
	params := critic.CalculateTrainableParams()
	fmt.Println("=== SAO Critic参数量 ===")
	fmt.Printf("注意力层冻结: %d (节省%.1f%%参数)\n", 
		params["attention_frozen"],
		float64(params["attention_frozen"])/float64(critic.totalParams)*100)
	fmt.Printf("MoE训练: %d\n", params["moe_trainable"])
	fmt.Printf("总可训练参数: %d (%.1f%%)\n", 
		params["total_trainable"],
		float64(params["total_trainable"])/float64(critic.totalParams)*100)
	
	// 更新频率比
	fmt.Printf("\n更新频率比 (Critic/Actor): %.0f:1\n", critic.GetUpdateRatio())
	
	// 稳定性对比
	stability := critic.SimulateValueStability()
	fmt.Println("\n=== 价值模型训练方差对比 ===")
	for name, variance := range stability {
		fmt.Printf("%-20s 方差: %.3f\n", name, variance)
	}
}

# sao_gae.py - 跳过观测的Token级GAE实现

import numpy as np
from typing import List, Tuple, Optional

class SkipObservationGAE:
    """
    跳过观测的Token级GAE
    核心：在计算优势时跳过环境反馈token，只连接模型动作token
    
    典型Agent轨迹：
    [Action] "search_web(\"AI news\")"
    [Observation] "搜索结果: 2026年7月..."
    [Action] "summarize(results)"
    [Observation] "摘要已完成..."
    [Reward] 0.85
    """
    
    def __init__(self, 
                 gamma: float = 0.99,
                 lambda_: float = 0.95,
                 action_token_ids: Optional[List[int]] = None,
                 obs_token_ids: Optional[List[int]] = None):
        """
        Args:
            gamma: 折扣因子
            lambda_: GAE lambda参数
            action_token_ids: 模型动作token的ID列表
            obs_token_ids: 环境观测token的ID列表
        """
        self.gamma = gamma
        self.lambda_ = lambda_
        self.action_token_ids = action_token_ids or [1, 2, 3]  # 示例
        self.obs_token_ids = obs_token_ids or [4, 5, 6]  # 示例
    
    def detect_action_tokens(self, 
                              token_ids: np.ndarray,
                              token_types: Optional[np.ndarray] = None) -> np.ndarray:
        """
        检测哪些token是模型动作token
        Args:
            token_ids: token序列 [seq_len]
            token_types: token类型标记 0=模型动作, 1=环境观测
        Returns:
            action_mask: 布尔掩码，True表示动作token
        """
        if token_types is not None:
            return token_types == 0
        
        # 基于token ID判断
        action_mask = np.isin(token_ids, self.action_token_ids)
        return action_mask
    
    def compute_skip_observation_gae(
        self,
        rewards: np.ndarray,
        values: np.ndarray,
        token_types: np.ndarray,
        dones: Optional[np.ndarray] = None,
    ) -> Tuple[np.ndarray, dict]:
        """
        计算跳过观测的GAE
        
        Args:
            rewards: 奖励序列 [seq_len]
            values: 价值序列 [seq_len]
            token_types: token类型 0=动作, 1=观测
            dones: 终止标记 [seq_len]
        Returns:
            advantages: 优势值 [seq_len]
            stats: 统计信息
        """
        seq_len = len(rewards)
        advantages = np.zeros(seq_len)
        
        if dones is None:
            dones = np.zeros(seq_len)
        
        last_gae = 0.0
        action_count = 0
        obs_count = 0
        
        # 反向遍历
        for t in reversed(range(seq_len)):
            next_value = values[t + 1] if t + 1 < seq_len else 0.0
            
            if token_types[t] == 0:  # 动作token
                # 正常GAE计算
                delta = rewards[t] + self.gamma * next_value * (1 - dones[t]) - values[t]
                last_gae = delta + self.gamma * self.lambda_ * last_gae * (1 - dones[t])
                advantages[t] = last_gae
                action_count += 1
            else:  # 观测token
                # 跳过观测：直接连接上一个动作的value到下一个动作
                # 即：delta = 0（观测本身不产生奖励），价值传递不受影响
                # 但实际上观测token不参与优势计算
                advantages[t] = 0.0
                obs_count += 1
                # 注意：last_gae保持不变，观测token不打断GAE传播
        
        stats = {
            "action_tokens": action_count,
            "observation_tokens": obs_count,
            "skip_ratio": obs_count / max(seq_len, 1),
            "mean_advantage": float(np.mean(advantages[token_types == 0])),
            "std_advantage": float(np.std(advantages[token_types == 0])),
        }
        
        return advantages, stats
    
    def compute_standard_gae(
        self,
        rewards: np.ndarray,
        values: np.ndarray,
        dones: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """标准GAE（不跳过观测，用于对比）"""
        seq_len = len(rewards)
        advantages = np.zeros(seq_len)
        
        if dones is None:
            dones = np.zeros(seq_len)
        
        last_gae = 0.0
        for t in reversed(range(seq_len)):
            next_value = values[t + 1] if t + 1 < seq_len else 0.0
            delta = rewards[t] + self.gamma * next_value * (1 - dones[t]) - values[t]
            last_gae = delta + self.gamma * self.lambda_ * last_gae * (1 - dones[t])
            advantages[t] = last_gae
        
        return advantages

# 模拟Agent轨迹
def simulate_agent_trajectory():
    """模拟一个多轮Agent交互轨迹"""
    np.random.seed(42)
    
    # 构建轨迹：动作-观测交替
    seq_len = 100
    token_types = np.zeros(seq_len)
    
    # 动作-观测交替模式
    # 第1轮
    token_types[0:10] = 0   # 动作
    token_types[10:20] = 1  # 观测
    # 第2轮
    token_types[20:35] = 0
    token_types[35:45] = 1
    # 第3轮
    token_types[45:55] = 0
    token_types[55:60] = 1
    # 第4轮
    token_types[60:70] = 0
    token_types[70:80] = 1
    # 第5轮
    token_types[80:90] = 0
    token_types[90:100] = 1
    
    # 奖励：最终奖励集中在最后
    rewards = np.zeros(seq_len)
    rewards[-1] = 1.0  # 最终成功奖励
    
    # 价值预测
    values = np.random.randn(seq_len) * 0.1 + 0.5
    
    return token_types, rewards, values

# 对比测试
gae = SkipObservationGAE(gamma=0.99, lambda_=0.95)
token_types, rewards, values = simulate_agent_trajectory()

# 跳过观测的GAE
advantages_skip, stats = gae.compute_skip_observation_gae(
    rewards, values, token_types
)

# 标准GAE
advantages_standard = gae.compute_standard_gae(rewards, values)

# 对比
action_mask = token_types == 0
print("=== 跳过观测GAE vs 标准GAE ===")
print(f"动作token数: {stats['action_tokens']}, 观测token数: {stats['observation_tokens']}")
print(f"跳过比例: {stats['skip_ratio']:.1%}")
print(f"\n动作token优势值对比:")
print(f"  SAO跳过观测: 均值={np.mean(advantages_skip[action_mask]):.4f}, "
      f"标准差={np.std(advantages_skip[action_mask]):.4f}")
print(f"  标准GAE:      均值={np.mean(advantages_standard[action_mask]):.4f}, "
      f"标准差={np.std(advantages_standard[action_mask]):.4f}")

# 优势值的信噪比
snr_skip = np.mean(advantages_skip[action_mask]) / max(np.std(advantages_skip[action_mask]), 1e-10)
snr_standard = np.mean(advantages_standard[action_mask]) / max(np.std(advantages_standard[action_mask]), 1e-10)
print(f"\n信噪比 (SNR):")
print(f"  SAO跳过观测: {snr_skip:.2f}")
print(f"  标准GAE:      {snr_standard:.2f}")
print(f"  提升:         {(snr_skip/snr_standard - 1)*100:.1f}%")

# glm52_sao_pipeline.py - GLM-5.2 SAO训练管线

import numpy as np
from typing import List, Dict, Any, Optional
from dataclasses import dataclass
import time

@dataclass
class SAOTrainingConfig:
    """SAO训练配置"""
    # 模型配置
    model_size: str = "GLM-5.2-750B"
    batch_size: int = 64
    max_seq_len: int = 8192
    
    # SAO特定参数
    critic_update_ratio: float = 2.0  # Critic:Actor更新比
    freeze_attention: bool = True     # 冻结注意力
    skip_observation: bool = True     # 跳过观测
    clip_lower: float = 0.2           # DIS下界
    clip_upper: float = 5.0           # DIS上界
    gamma: float = 0.99               # 折扣因子
    lambda_: float = 0.95             # GAE lambda
    
    # 训练参数
    learning_rate: float = 1e-5
    max_steps: int = 1000
    warmup_steps: int = 50
    
    # 任务配置
    task_types: List[str] = None
    
    def __post_init__(self):
        if self.task_types is None:
            self.task_types = ["coding", "reasoning", "agent"]

class SAOTrainer:
    """SAO训练器"""
    
    def __init__(self, config: SAOTrainingConfig):
        self.config = config
        self.step = 0
        self.critic_step = 0
        
        self.training_stats = {
            "policy_loss": [],
            "value_loss": [],
            "explained_variance": [],
            "acceptance_rate": [],
            "kl_divergence": [],
        }
    
    def train_step(self, 
                   batch: Dict[str, np.ndarray],
                   is_critic_step: bool = False) -> Dict[str, float]:
        """
        单步训练
        Args:
            batch: 训练数据
            is_critic_step: 是否是Critic更新步
        Returns:
            metrics: 训练指标
        """
        # 模拟前向传播
        rollout_log_probs = batch.get("log_probs", np.random.randn(64, 128))
        current_log_probs = rollout_log_probs + np.random.randn(*rollout_log_probs.shape) * 0.1
        
        # 模拟优势计算
        advantages = np.random.randn(*rollout_log_probs.shape) * 0.2 + 0.1
        
        # 计算策略损失
        ratio = np.exp(current_log_probs - rollout_log_probs)
        clipped = np.clip(ratio, self.config.clip_lower, self.config.clip_upper)
        policy_loss = -np.mean(clipped * advantages)
        
        # 计算价值损失（仅Critic步）
        value_loss = 0.0
        if is_critic_step:
            predicted_values = np.random.randn(64) * 0.1
            target_values = np.random.randn(64) * 0.1 + 0.5
            value_loss = np.mean((predicted_values - target_values) ** 2)
        
        # 更新统计
        self.training_stats["policy_loss"].append(policy_loss)
        if is_critic_step:
            self.training_stats["value_loss"].append(value_loss)
        
        return {
            "policy_loss": float(policy_loss),
            "value_loss": float(value_loss),
            "is_critic_step": is_critic_step,
        }
    
    def train(self, num_steps: int) -> Dict[str, List[float]]:
        """
        执行SAO训练
        """
        start_time = time.time()
        
        for step in range(num_steps):
            # 判断是否为Critic更新步
            is_critic = (step % int(1 / self.config.critic_update_ratio) == 0)
            
            # 模拟batch数据
            batch = {
                "log_probs": np.random.randn(64, 128) * 0.1,
                "advantages": np.random.randn(64, 128) * 0.2,
                "rewards": np.random.randn(64) * 0.1 + 0.5,
                "values": np.random.randn(64) * 0.1,
            }
            
            # 训练步
            metrics = self.train_step(batch, is_critic)
            
            self.step += 1
            if is_critic:
                self.critic_step += 1
        
        elapsed = time.time() - start_time
        
        return {
            "total_steps": self.step,
            "critic_steps": self.critic_step,
            "actor_steps": self.step - self.critic_step,
            "training_time_s": elapsed,
            "avg_policy_loss": np.mean(self.training_stats["policy_loss"][-100:]),
            "avg_value_loss": np.mean(self.training_stats["value_loss"][-100:]) 
                if self.training_stats["value_loss"] else 0,
        }

# 运行训练模拟
config = SAOTrainingConfig(
    model_size="GLM-5.2-750B",
    batch_size=64,
    critic_update_ratio=2.0,
    freeze_attention=True,
    skip_observation=True,
    max_steps=1000,
)
trainer = SAOTrainer(config)
results = trainer.train(1000)

print("=== GLM-5.2 SAO训练结果 ===")
print(f"总步数: {results['total_steps']}")
print(f"Actor步数: {results['actor_steps']}")
print(f"Critic步数: {results['critic_steps']}")
print(f"训练时间: {results['training_time_s']:.1f}s")
print(f"平均策略损失: {results['avg_policy_loss']:.4f}")
print(f"平均价值损失: {results['avg_value_loss']:.4f}")

# ablation_study.py - SAO消融实验

class SAOAblationStudy:
    """SAO消融实验"""
    
    def __init__(self):
        self.results = {}
    
    def run_ablation(self, 
                     name: str,
                     use_dis: bool,
                     critic_ratio: float,
                     freeze_attn: bool,
                     skip_obs: bool) -> dict:
        """运行单次消融实验"""
        
        # 模拟各基准测试结果
        base_score = {
            "AIME2025": 72.1,
            "SWE-Bench": 23.0,
            "BeyondAIME": 58.3,
        }
        
        # 各组件贡献
        score = base_score.copy()
        
        if use_dis:
            for k in score:
                score[k] += {
                    "AIME2025": 5.0,
                    "SWE-Bench": 2.0,
                    "BeyondAIME": 3.5,
                }[k]
        
        if critic_ratio >= 2.0:
            for k in score:
                score[k] += {
                    "AIME2025": 8.0,
                    "SWE-Bench": 3.0,
                    "BeyondAIME": 6.0,
                }[k]
        elif critic_ratio >= 1.0:
            for k in score:
                score[k] += {
                    "AIME2025": 4.0,
                    "SWE-Bench": 1.5,
                    "BeyondAIME": 3.0,
                }[k]
        
        if freeze_attn:
            for k in score:
                score[k] += {
                    "AIME2025": 5.0,
                    "SWE-Bench": 1.0,
                    "BeyondAIME": 3.0,
                }[k]
        
        if skip_obs:
            for k in score:
                score[k] += {
                    "AIME2025": 3.0,
                    "SWE-Bench": 0.5,
                    "BeyondAIME": 2.0,
                }[k]
        
        # 完整SAO增益
        if use_dis and critic_ratio >= 2.0 and freeze_attn and skip_obs:
            for k in score:
                score[k] += 2.0  # 协同增益
        
        self.results[name] = score
        return score

ablation = SAOAblationStudy()

configs = [
    ("GRPO基线", False, 0.0, False, False),
    ("+DIS", True, 0.0, False, False),
    ("+DIS+Critic(1x)", True, 1.0, False, False),
    ("+DIS+Critic(2x)", True, 2.0, False, False),
    ("+DIS+Critic(2x)+冻结Attn", True, 2.0, True, False),
    ("SAO完整", True, 2.0, True, True),
]

print(f"{'配置':<30s} {'AIME2025':<10s} {'SWE-Bench':<10s} {'BeyondAIME':<10s}")
print("-" * 60)
for name, dis, ratio, freeze, skip in configs:
    result = ablation.run_ablation(name, dis, ratio, freeze, skip)
    print(f"{name:<30s} {result['AIME2025']:<10.1f} {result['SWE-Bench']:<10.1f} {result['BeyondAIME']:<10.1f}")