Upload 3 files

grpo_dataloader.py

@@ -0,0 +1,191 @@
+"""
+GRPO专用数据加载器
+"""
+import torch
+from torch.utils.data import Dataset, DataLoader
+from datasets import load_dataset, interleave_datasets
+from typing import Optional, List
+import logging
+import os
+
+logger = logging.getLogger(__name__)
+
+from data_config import (
+    GRPO_DATASETS,
+    GRPO_PROMPT_MIX,
+    HF_CACHE_DIR
+)
+
+
+class GRPOPromptDataset(Dataset):
+    """
+    GRPO Prompt数据集 - 用于生成阶段
+    """
+    def __init__(
+        self,
+        mix_name: str = 'default',
+        tokenizer=None,
+        max_length: int = 512,
+        max_samples: Optional[int] = None
+    ):
+        super().__init__()
+        
+        if tokenizer is None:
+            raise ValueError("tokenizer cannot be None")
+        
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+        
+        # 获取混合配置
+        if mix_name not in GRPO_PROMPT_MIX:
+            raise ValueError(
+                f"Unknown mix: {mix_name}. "
+                f"Available: {list(GRPO_PROMPT_MIX.keys())}"
+            )
+        
+        mix_config = GRPO_PROMPT_MIX[mix_name]
+        dataset_names = mix_config.get('datasets', [])
+        weights = mix_config.get('weights', [])
+        
+        logger.info(f"Loading GRPO prompt mix: {mix_name}")
+        logger.info(f"  Datasets: {dataset_names}")
+        logger.info(f"  Weights: {weights}")
+        
+        # 加载数据集
+        all_datasets = []
+        
+        for name in dataset_names:
+            if name not in GRPO_DATASETS:
+                logger.warning(f"Dataset {name} not found")
+                continue
+            
+            config = GRPO_DATASETS[name]
+            
+            # 验证文件存在
+            data_file = config.get('data_files')
+            if data_file and not os.path.exists(data_file):
+                logger.error(f"Data file not found: {data_file}")
+                logger.error(f"请先运行 download_grpo_datasets.py 下载数据")
+                continue
+            
+            try:
+                load_kwargs = {
+                    'path': config['hf_path'],
+                    'split': config.get('split', 'train'),
+                    'cache_dir': HF_CACHE_DIR,
+                }
+                
+                if 'data_files' in config:
+                    load_kwargs['data_files'] = config['data_files']
+                
+                ds = load_dataset(**load_kwargs)
+                
+                # 限制样本数
+                if config.get('max_samples'):
+                    ds = ds.select(range(min(len(ds), config['max_samples'])))
+                
+                all_datasets.append(ds)
+                logger.info(f"  Loaded {name}: {len(ds)} samples")
+                
+            except Exception as e:
+                logger.error(f"Error loading {name}: {e}")
+                continue
+        
+        if not all_datasets:
+            raise ValueError("No datasets loaded successfully")
+        
+        # 合并数据集
+        if len(all_datasets) == 1:
+            self.dataset = all_datasets[0]
+        else:
+            probabilities = [w / sum(weights[:len(all_datasets)]) 
+                           for w in weights[:len(all_datasets)]]
+            self.dataset = interleave_datasets(
+                all_datasets,
+                probabilities=probabilities,
+                seed=42,
+                stopping_strategy='all_exhausted'
+            )
+        
+        # 限制总样本数
+        if max_samples and len(self.dataset) > max_samples:
+            self.dataset = self.dataset.select(range(max_samples))
+        
+        logger.info(f"Total prompts: {len(self.dataset)}")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        try:
+            sample = self.dataset[idx]
+            
+            # 提取prompt
+            prompt = sample.get('prompt', '')
+            
+            if not prompt:
+                logger.warning(f"Empty prompt at index {idx}")
+                return None
+            
+            # Tokenize (不添加EOS，因为这是prompt)
+            encoding = self.tokenizer(
+                prompt,
+                max_length=self.max_length,
+                truncation=True,
+                padding='max_length',
+                return_tensors='pt',
+                add_special_tokens=True
+            )
+            
+            return {
+                'input_ids': encoding['input_ids'].squeeze(0),
+                'attention_mask': encoding['attention_mask'].squeeze(0),
+                'prompt_text': prompt
+            }
+            
+        except Exception as e:
+            logger.debug(f"Error processing sample {idx}: {e}")
+            return None
+
+
+def grpo_collate_fn(batch):
+    """GRPO专用collate函数"""
+    # 过滤None
+    batch = [item for item in batch if item is not None]
+    
+    if not batch:
+        return None
+    
+    return {
+        'input_ids': torch.stack([item['input_ids'] for item in batch]),
+        'attention_mask': torch.stack([item['attention_mask'] for item in batch]),
+        'prompt_texts': [item['prompt_text'] for item in batch]
+    }
+
+
+def create_grpo_prompt_dataloader(
+    mix_name: str = 'default',
+    tokenizer=None,
+    batch_size: int = 4,
+    num_workers: int = 2,
+    max_length: int = 512,
+    max_samples: Optional[int] = None,
+    shuffle: bool = True
+):
+    """创建GRPO prompt数据加载器"""
+    dataset = GRPOPromptDataset(
+        mix_name=mix_name,
+        tokenizer=tokenizer,
+        max_length=max_length,
+        max_samples=max_samples
+    )
+    
+    return DataLoader(
+        dataset,
+        batch_size=batch_size,
+        shuffle=shuffle,
+        num_workers=num_workers,
+        collate_fn=grpo_collate_fn,
+        pin_memory=True,
+        drop_last=False
+    )

grpo_r1_train.py

@@ -0,0 +1,320 @@
+import os
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from transformers import AutoTokenizer
+from torch.utils.data import DataLoader, Dataset
+import json
+import logging
+from tqdm import tqdm
+import glob
+from datetime import datetime
+import gc
+from model import MultiModalDenseTransformer
+from grpo import GRPOZeroTrainer
+
+# ================= DDP 设置 =================
+def setup_distributed():
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        dist.init_process_group(backend="nccl")
+        rank = int(os.environ["RANK"])
+        local_rank = int(os.environ["LOCAL_RANK"])
+        world_size = int(os.environ["WORLD_SIZE"])
+        torch.cuda.set_device(local_rank)
+        print(f"Initialized DDP: Rank {rank}/{world_size}")
+        return rank, local_rank, world_size
+    else:
+        print("Initialized Single GPU Mode")
+        return 0, 0, 1
+
+RANK, LOCAL_RANK, WORLD_SIZE = setup_distributed()
+IS_MAIN = RANK == 0
+
+logging.basicConfig(
+    level=logging.INFO if IS_MAIN else logging.WARNING,
+    format=f'%(asctime)s - [Rank {RANK}] - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# ================= 数据集 =================
+class MathDataset(Dataset):
+    def __init__(self, path):
+        self.data = []
+        with open(path, 'r', encoding='utf-8') as f:
+            for line in f:
+                if line.strip():
+                    self.data.append(json.loads(line))
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def __getitem__(self, idx):
+        return self.data[idx]
+
+def math_collate(batch):
+    return {
+        'prompt': [item['prompt'] for item in batch],
+        'ground_truth': [item['ground_truth'] for item in batch]
+    }
+
+# ================= 主函数 =================
+def main():
+    # ------------------ 配置区域 ------------------
+    CONFIG = {
+        # 基础模型路径
+        'sft_checkpoint': '/root/checkpoints/dcpo_posttrain_round3/step_2600.pt', 
+        'data_path': '/root/dataset/r1_zero_math.jsonl',
+        'save_dir': '/root/checkpoints/r1_zero_reproduction',
+        'resume_from': None, # 或者具体路径
+        
+        # 模型参数 (需确保与 Checkpoint 一致)
+        'model_dim': 1536,
+        'n_layers': 12,
+        'n_heads': 12,
+        'n_kv_heads': 4,
+        
+        # 训练参数
+        'group_size': 4,      
+        'batch_size': 1,      # Prompt Batch Size
+        'learning_rate': 2e-6,
+        'max_steps': 190000,
+        'max_gen_len': 512,
+        'save_interval': 300,
+        
+        # 【新增】累积更新参数
+        # 实际 Update Batch = batch_size * group_size * accum_steps
+        # 例如: 1 * 4 * 8 = 32
+        'gradient_accumulation_steps': 8, 
+        'inner_batch_size': 4 # PPO Update 时的显存计算 Batch
+    }
+    # ---------------------------------------------
+    
+    if IS_MAIN:
+        os.makedirs(CONFIG['save_dir'], exist_ok=True)
+        current_time = datetime.now().strftime('%Y%m%d_%H%M%S')
+        log_file = os.path.join(CONFIG['save_dir'], f"train_{current_time}.log")
+        file_handler = logging.FileHandler(log_file, encoding='utf-8')
+        file_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))
+        logger.addHandler(file_handler)
+        logger.info(f"Configuration: {json.dumps(CONFIG, indent=2)}")
+    
+    # 1. 加载 Tokenizer
+    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-7B-Instruct", trust_remote_code=True)
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+        tokenizer.pad_token_id = tokenizer.eos_token_id
+        
+    # 2. 初始化模型
+    def create_model():
+        return MultiModalDenseTransformer(
+            model_dim=CONFIG['model_dim'],
+            vocab_size=len(tokenizer),
+            n_layers=CONFIG['n_layers'],
+            n_heads=CONFIG['n_heads'],
+            n_kv_heads=CONFIG['n_kv_heads'],
+            max_seq_len=2048,
+            use_gradient_checkpointing=True
+        )
+        
+    device = torch.device(f"cuda:{LOCAL_RANK}")
+    
+    logger.info("Initializing Actor Model...")
+    actor = create_model().to(device)
+    
+    logger.info("Initializing Ref Model...")
+    ref = create_model().to(device)
+    ref.eval()
+    ref.requires_grad_(False)
+    
+    # 3. 初始化训练器 (传入累积参数)
+    trainer = GRPOZeroTrainer(
+        actor_model=actor,
+        ref_model=ref,
+        tokenizer=tokenizer,
+        learning_rate=CONFIG['learning_rate'],
+        group_size=CONFIG['group_size'],
+        use_amp=True,
+        gradient_accumulation_steps=CONFIG['gradient_accumulation_steps'],
+        inner_batch_size=CONFIG['inner_batch_size']
+    )
+    
+    # 4. 加载权重与恢复
+    start_step = 0
+    samples_seen = 0
+    
+    if CONFIG['resume_from']:
+        resume_path = CONFIG['resume_from']
+        logger.info(f"Resuming from: {resume_path}")
+        checkpoint = torch.load(resume_path, map_location='cpu')
+        
+        actor.load_state_dict(checkpoint['model_state_dict'])
+        # 恢复优化器
+        if 'optimizer_state_dict' in checkpoint:
+            try:
+                trainer.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+            except Exception as e:
+                logger.warning(f"Optimizer load failed (param mismatch?): {e}")
+        
+        ref.load_state_dict(checkpoint['model_state_dict']) # Ref 与 Actor 同步起点
+        
+        start_step = checkpoint.get('step', 0) + 1
+        samples_seen = checkpoint.get('samples_seen', start_step * CONFIG['batch_size'] * WORLD_SIZE)
+        
+        del checkpoint
+        gc.collect()
+        torch.cuda.empty_cache()
+    else:
+        logger.info(f"Loading SFT checkpoint: {CONFIG['sft_checkpoint']}")
+        checkpoint = torch.load(CONFIG['sft_checkpoint'], map_location='cpu')
+        state_dict = checkpoint['model_state_dict'] if 'model_state_dict' in checkpoint else checkpoint
+        # 去除 module. 前缀
+        new_state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
+        
+        actor.load_state_dict(new_state_dict)
+        ref.load_state_dict(new_state_dict)
+        del checkpoint, state_dict, new_state_dict
+        gc.collect()
+        torch.cuda.empty_cache()
+    
+    if WORLD_SIZE > 1:
+        actor = DDP(actor, device_ids=[LOCAL_RANK], output_device=LOCAL_RANK)
+        
+    # 5. 数据加载
+    dataset = MathDataset(CONFIG['data_path'])
+    if WORLD_SIZE > 1:
+        sampler = torch.utils.data.DistributedSampler(
+            dataset, num_replicas=WORLD_SIZE, rank=RANK, shuffle=True, seed=42
+        )
+    else:
+        sampler = None
+    
+    dataloader = DataLoader(
+        dataset, batch_size=CONFIG['batch_size'], 
+        collate_fn=math_collate, sampler=sampler, shuffle=(sampler is None)
+    )
+    
+    # 6. 训练循环
+    logger.info(f"Starting Training from step {start_step}...")
+    
+    if sampler:
+        epoch = samples_seen // len(dataset)
+        sampler.set_epoch(epoch)
+    
+    data_iter = iter(dataloader)
+    
+    # 简单的跳过逻辑
+    if samples_seen > 0:
+        skip_batches = samples_seen // (CONFIG['batch_size'] * WORLD_SIZE)
+        logger.info(f"Skipping {skip_batches} batches...")
+        for _ in range(skip_batches):
+            try:
+                next(data_iter)
+            except StopIteration:
+                if sampler: sampler.set_epoch(sampler.epoch + 1)
+                data_iter = iter(dataloader)
+                next(data_iter)
+
+    progress_bar = tqdm(range(start_step, CONFIG['max_steps']), disable=not IS_MAIN, initial=start_step, total=CONFIG['max_steps'])
+    
+    # 状态追踪
+    current_samples = samples_seen
+    running_reward = 0.0
+    running_loss = 0.0
+    
+    for step in progress_bar:
+        try:
+            try:
+                batch = next(data_iter)
+            except StopIteration:
+                if sampler:
+                    epoch = current_samples // len(dataset)
+                    sampler.set_epoch(epoch)
+                data_iter = iter(dataloader)
+                batch = next(data_iter)
+            
+            current_samples += CONFIG['batch_size'] * WORLD_SIZE
+            
+            # 生成阶段
+            experience = trainer.generate_and_score(
+                batch, 
+                max_gen_len=CONFIG['max_gen_len']
+            )
+            
+            # 记录 Reward (平滑)
+            step_reward = experience['avg_reward']
+            if running_reward == 0: running_reward = step_reward
+            else: running_reward = 0.95 * running_reward + 0.05 * step_reward
+            
+            # 训练阶段 (可能返回 None)
+            loss = trainer.train_step(experience)
+            
+            # 日志与显示逻辑
+            status_dict = {"R": f"{running_reward:.3f}"}
+            
+            if loss is not None:
+                # 发生了权重更新
+                if running_loss == 0: running_loss = loss
+                else: running_loss = 0.9 * running_loss + 0.1 * loss
+                status_dict["L"] = f"{running_loss:.3f}"
+                
+                if IS_MAIN:
+                    # 写入 Metrics
+                    current_lr = trainer.optimizer.param_groups[0]['lr']
+                    metrics_data = {
+                        "step": step,
+                        "reward": float(step_reward), # 记录当前步的 reward
+                        "loss": float(loss),
+                        "lr": float(current_lr),
+                        "samples_seen": current_samples,
+                        "timestamp": datetime.now().isoformat()
+                    }
+                    with open(os.path.join(CONFIG['save_dir'], "metrics.jsonl"), "a") as f:
+                        f.write(json.dumps(metrics_data) + "\n")
+                        
+                    if step % 10 == 0:
+                        logger.info(f"Step {step} | Reward: {step_reward:.4f} | Loss: {loss:.4f} | LR: {current_lr:.2e}")
+            else:
+                # 正在累积
+                status_dict["State"] = "Acc"
+            
+            progress_bar.set_description(f"{' '.join([f'{k}:{v}' for k,v in status_dict.items()])}")
+            
+            # 保存逻辑
+            if step > 0 and step % CONFIG['save_interval'] == 0 and IS_MAIN:
+                save_path = f"{CONFIG['save_dir']}/step_{step}.pt"
+                model_to_save = actor.module if hasattr(actor, 'module') else actor
+                torch.save({
+                    'step': step,
+                    'samples_seen': current_samples,
+                    'model_state_dict': model_to_save.state_dict(),
+                    'optimizer_state_dict': trainer.optimizer.state_dict(),
+                }, save_path)
+                logger.info(f"Checkpoint saved: {save_path}")
+            
+            # 显存清理
+            del experience
+            del batch
+            # 这里的 empty_cache 是可选的，如果显存非常紧张建议开启
+            # torch.cuda.empty_cache() 
+            
+        except Exception as e:
+            logger.error(f"Step {step} Error: {e}")
+            import traceback
+            traceback.print_exc()
+            continue
+
+    # 结束保存
+    if IS_MAIN:
+        final_path = f"{CONFIG['save_dir']}/final_r1_zero.pt"
+        model_to_save = actor.module if hasattr(actor, 'module') else actor
+        torch.save({
+            'step': CONFIG['max_steps'],
+            'model_state_dict': model_to_save.state_dict(),
+        }, final_path)
+        logger.info("Training Finished.")
+        
+    if WORLD_SIZE > 1:
+        dist.destroy_process_group()
+
+if __name__ == "__main__":
+    main()

math_verifier.py

@@ -0,0 +1,270 @@
+import re
+import math
+import logging
+from difflib import SequenceMatcher
+
+logger = logging.getLogger(__name__)
+
+class MathReward:
+    def __init__(self, use_reference_comparison=True):
+        """
+        Args:
+            use_reference_comparison: 是否使用参考答案进行推理过程比较
+        """
+        # 编译正则表达式，强制要求 <think> 在前，<answer> 在后
+        self.format_pattern = re.compile(r"<think>(.*?)</think>\s*<answer>(.*?)</answer>", re.DOTALL)
+        self.use_reference_comparison = use_reference_comparison
+        
+        # 推理关键词（用于检查推理质量）
+        self.reasoning_keywords = [
+            '计算', '因为', '所以', '首先', '然后', '接着', '最后', '根据',
+            '第一步', '第二步', '第三步', '第', '步', '得到', '等于', 
+            '加', '减', '乘', '除', '=', '+', '-', '*', '/', '÷', '×'
+        ]
+        
+    def parse_number(self, text):
+        """
+        从文本中解析数值。
+        支持：整数、小数、分数(1/5)、百分数(20%)、带逗号的数字(1,000)
+        """
+        if not text:
+            return None
+        
+        # 预处理：移除空格、货币符号、常见的中文单位
+        text = text.strip()
+        clean_text = text.replace(" ", "").replace(",", "").replace("¥", "").replace("$", "")
+        clean_text = clean_text.replace("千克", "").replace("元", "").replace("个", "").replace("只", "")
+        clean_text = clean_text.replace("本", "").replace("米", "").replace("人", "")
+        
+        try:
+            # 1. 处理百分数 (e.g., "20%")
+            if "%" in clean_text:
+                return float(clean_text.replace("%", "")) / 100
+            
+            # 2. 处理分数 (e.g., "1/5" 或 "42/5")
+            if "/" in clean_text:
+                parts = clean_text.split("/")
+                if len(parts) == 2:
+                    try:
+                        return float(parts[0]) / float(parts[1])
+                    except:
+                        pass
+            
+            # 3. 处理科学记数法 (e.g., "1.5e-3")
+            if "e" in clean_text.lower() or "E" in clean_text:
+                return float(clean_text)
+            
+            # 4. 提取所有匹配的数字格式
+            # 匹配 浮点数 或 整数，忽略可能混杂的文字
+            matches = re.findall(r"[-+]?\d*\.\d+|\d+", clean_text)
+            if matches:
+                # 取最后一个作为最终答案（通常答案在最后）
+                return float(matches[-1])
+                
+        except Exception as e:
+            logger.debug(f"解析数字失败: {text}, 错误: {e}")
+            
+        return None
+
+    def check_reasoning_quality(self, think_content):
+        """
+        检查推理过程的质量
+        
+        返回质量评分 (0.0 - 1.0)
+        """
+        if not think_content:
+            return 0.0
+        
+        quality_score = 0.0
+        
+        # 1. 长度检查（基础）
+        length = len(think_content)
+        if length >= 100:
+            quality_score += 0.3
+        elif length >= 50:
+            quality_score += 0.15
+        
+        # 2. 关键词检查（推理步骤标识）
+        keyword_count = sum(1 for kw in self.reasoning_keywords if kw in think_content)
+        # 每出现一个关键词加分，最多加0.3分
+        quality_score += min(keyword_count * 0.05, 0.3)
+        
+        # 3. 数学表达式检查（是否包含计算过程）
+        # 匹配数学运算符或等式
+        math_expressions = re.findall(r'\d+\s*[+\-*/×÷=]\s*\d+', think_content)
+        if len(math_expressions) > 0:
+            quality_score += 0.2
+            # 多个表达式说明推理更详细
+            if len(math_expressions) >= 3:
+                quality_score += 0.1
+        
+        # 4. 结构检查（是否有步骤分隔）
+        has_steps = bool(re.search(r'第\d+步|步骤\d+|^\d+[.、]', think_content, re.MULTILINE))
+        if has_steps:
+            quality_score += 0.1
+        
+        return min(quality_score, 1.0)
+
+    def compute_reasoning_similarity(self, generated_reasoning, reference_reasoning):
+        """
+        计算生成的推理过程与参考推理过程的相似度
+        
+        使用序列匹配算法（考虑顺序）
+        返回相似度分数 (0.0 - 1.0)
+        """
+        if not generated_reasoning or not reference_reasoning:
+            return 0.0
+        
+        # 使用 difflib 的 SequenceMatcher 计算相似度
+        similarity = SequenceMatcher(None, generated_reasoning, reference_reasoning).ratio()
+        
+        return similarity
+
+    def compute_rewards(self, completions, ground_truths):
+        """
+        计算奖励
+        
+        Args:
+            completions: List[str] 模型生成的完整文本
+            ground_truths: List[dict] 对应的真值 
+                必须包含: 'answer_val': float
+                可选包含: 'reasoning': str, 'reference_completion': str
+        
+        Returns:
+            rewards: List[float]
+        """
+        rewards = []
+        
+        for completion, gt in zip(completions, ground_truths):
+            total_reward = 0.0
+            
+            # --- 1. 格式与结构检查 ---
+            match = self.format_pattern.search(completion)
+            
+            # 如果没有匹配到 <think>...</think><answer>...</answer> 结构
+            if match is None:
+                # 格式严重错误，给予重罚
+                rewards.append(-2.0)
+                continue
+            
+            # 提取内容
+            think_content = match.group(1).strip()
+            answer_content = match.group(2).strip()
+            
+            # 格式正确的基础分
+            total_reward += 0.6
+            
+            # --- 2. 思考过程质量检查 ---
+            reasoning_quality = self.check_reasoning_quality(think_content)
+            
+            if reasoning_quality < 0.3:
+                # 推理过程质量太低（可能是敷衍或格式化）
+                total_reward -= 0.5
+            else:
+                # 推理质量越高，奖励越多
+                total_reward += reasoning_quality * 1.0  # 最多1.0分
+            
+            # --- 3. 推理过程与参考对比（如果有参考） ---
+            if self.use_reference_comparison and 'reasoning' in gt:
+                reference_reasoning = gt['reasoning']
+                similarity = self.compute_reasoning_similarity(think_content, reference_reasoning)
+                
+                # 相似度奖励（最多0.5分）
+                # 注意：不要求完全一致，因为可能有多种正确推理方式
+                if similarity > 0.3:
+                    total_reward += similarity * 0.5
+            
+            # --- 4. 答案准确性检查（最重要） ---
+            pred_val = self.parse_number(answer_content)
+            gt_val = gt['answer_val']
+            
+            if pred_val is not None:
+                # 数值比较，允许 float 精度误差
+                if math.isclose(pred_val, gt_val, rel_tol=1e-4, abs_tol=1e-4):
+                    # 答对给予最高奖励
+                    total_reward += 3.0
+                else:
+                    # 答错扣分
+                    # 根据误差大小调整惩罚
+                    try:
+                        relative_error = abs(pred_val - gt_val) / (abs(gt_val) + 1e-8)
+                        if relative_error < 0.1:
+                            # 接近正确答案，轻微惩罚
+                            total_reward -= 0.3
+                        elif relative_error < 0.5:
+                            # 有一定误差
+                            total_reward -= 0.8
+                        else:
+                            # 完全错误
+                            total_reward -= 1.5
+                    except:
+                        total_reward -= 1.5
+            else:
+                # <answer> 标签内提取不到有效数字
+                total_reward -= 1.0
+            
+            # --- 5. 一致性检查：推理过程中的数字应该与答案相关 ---
+            # 提取推理过程中出现的所有数字
+            reasoning_numbers = re.findall(r'[-+]?\d*\.\d+|\d+', think_content)
+            if reasoning_numbers and pred_val is not None:
+                # 检查答案是否出现在推理过程中
+                answer_in_reasoning = any(
+                    math.isclose(float(num), pred_val, rel_tol=1e-3, abs_tol=1e-3)
+                    for num in reasoning_numbers
+                )
+                if answer_in_reasoning:
+                    total_reward += 0.2
+            
+            rewards.append(total_reward)
+        
+        return rewards
+    
+    def compute_metrics(self, completions, ground_truths):
+        """
+        计算详细的评估指标（用于分析）
+        
+        Returns:
+            dict: 包含各种指标的字典
+        """
+        metrics = {
+            'format_correct': 0,
+            'reasoning_quality_avg': 0.0,
+            'answer_correct': 0,
+            'answer_close': 0,  # 答案接近但不完全正确
+            'total': len(completions)
+        }
+        
+        quality_scores = []
+        
+        for completion, gt in zip(completions, ground_truths):
+            match = self.format_pattern.search(completion)
+            
+            if match:
+                metrics['format_correct'] += 1
+                
+                think_content = match.group(1).strip()
+                answer_content = match.group(2).strip()
+                
+                # 推理质量
+                quality = self.check_reasoning_quality(think_content)
+                quality_scores.append(quality)
+                
+                # 答案准确性
+                pred_val = self.parse_number(answer_content)
+                gt_val = gt['answer_val']
+                
+                if pred_val is not None and gt_val is not None:
+                    if math.isclose(pred_val, gt_val, rel_tol=1e-4, abs_tol=1e-4):
+                        metrics['answer_correct'] += 1
+                    elif math.isclose(pred_val, gt_val, rel_tol=0.1, abs_tol=0.1):
+                        metrics['answer_close'] += 1
+        
+        if quality_scores:
+            metrics['reasoning_quality_avg'] = sum(quality_scores) / len(quality_scores)
+        
+        # 计算百分比
+        metrics['format_correct_pct'] = metrics['format_correct'] / metrics['total'] * 100
+        metrics['answer_correct_pct'] = metrics['answer_correct'] / metrics['total'] * 100
+        metrics['answer_close_pct'] = metrics['answer_close'] / metrics['total'] * 100
+        
+        return metrics