别再纠结RLHF了!用Python手把手教你实现一个极简RLVR训练脚本(附GRPO代码)
用Python实现RLVR训练从零构建可验证奖励的强化学习系统在当今AI领域强化学习从人类反馈(RLHF)和AI反馈(RLAIF)已成为模型对齐的主流方法。然而这些方法面临着标注成本高、奖励模型可能被欺骗等挑战。RLVR(Reinforcement Learning from Verifiable Rewards)提供了一种更直接、可扩展的替代方案——它使用可自动验证的程序作为奖励信号特别适合数学推理、代码生成等结构化任务。1. RLVR核心概念与优势RLVR的核心思想是将奖励设计成一个可被程序自动验证的函数完全摒弃了传统RLHF中需要人工标注或训练奖励模型的环节。这种方法在DeepSeek-R1、OpenAI的某些模型系列中已经得到成功应用。RLVR与传统方法的对比特性RLHFRLAIFRLVR奖励来源人类标注AI模型标注自动验证程序成本高昂中等极低可扩展性有限较好极佳适用场景开放任务(如对话)价值观对齐结构化任务(如数学、代码)奖励客观性主观依赖AI模型完全客观RLVR特别适合以下场景数学问题求解可以直接验证答案的正确性代码生成通过单元测试验证代码功能多选题回答对比标准答案医学问答验证关键事实的准确性提示RLVR的关键优势在于奖励函数的透明性和可重复性这使得模型优化过程更加可控也更容易诊断问题。2. 构建RLVR训练环境要实现一个完整的RLVR训练系统我们需要准备以下几个核心组件2.1 数据准备对于演示目的我们创建一个简单的数学运算数据集import random def generate_math_dataset(num_samples): dataset [] operators [, -, *, /] for _ in range(num_samples): a random.randint(1, 100) b random.randint(1, 100) op random.choice(operators) if op /: # 确保除法结果为整数 b random.randint(1, 10) a b * random.randint(1, 10) question f{a} {op} {b} answer eval(question) # 安全警告实际应用中应避免直接eval dataset.append({ prompt: f计算: {question}, answer: answer }) return dataset math_dataset generate_math_dataset(1000)2.2 可验证奖励函数奖励函数是RLVR的核心它需要能够自动判断模型输出的正确性import re def extract_last_float(text): 从文本中提取最后一个浮点数 numbers re.findall(r-?\d\.?\d*, text) return float(numbers[-1]) if numbers else None def math_verifier(prompt, completion, gold_answer, tolerance1e-6): 数学问题验证器 返回: 1.0(正确) 或 0.0(错误) pred extract_last_float(completion) if pred is None: return 0.0 return 1.0 if abs(pred - gold_answer) tolerance else 0.02.3 策略模型包装我们需要一个包装类来处理语言模型的采样和概率计算import torch from transformers import AutoModelForCausalLM, AutoTokenizer class PolicyLM: def __init__(self, model_name): self.device cuda if torch.cuda.is_available() else cpu self.tokenizer AutoTokenizer.from_pretrained(model_name) self.model AutoModelForCausalLM.from_pretrained(model_name).to(self.device) if not self.tokenizer.pad_token: self.tokenizer.pad_token self.tokenizer.eos_token def generate(self, prompts, num_samples1, max_length50): 为每个prompt生成num_samples个补全 inputs self.tokenizer(prompts, return_tensorspt, paddingTrue).to(self.device) with torch.no_grad(): outputs self.model.generate( **inputs, max_lengthmax_length, num_return_sequencesnum_samples, do_sampleTrue, pad_token_idself.tokenizer.pad_token_id ) completions [] for i in range(len(prompts)): group [] for j in range(num_samples): idx i * num_samples j completion outputs[idx][len(inputs[input_ids][i]):] group.append(self.tokenizer.decode(completion, skip_special_tokensTrue)) completions.append(group) return completions def logprobs(self, prompts, completions): 计算给定prompt和completion的对数概率 batch_logprobs [] for prompt, comp_group in zip(prompts, completions): group_logprobs [] for comp in comp_group: full_text prompt comp inputs self.tokenizer(full_text, return_tensorspt).to(self.device) with torch.no_grad(): outputs self.model(**inputs) logits outputs.logits logprobs torch.log_softmax(logits, dim-1)[0] # 计算completion部分的logprob总和 prompt_len len(self.tokenizer(prompt, return_tensorspt)[input_ids][0]) comp_ids inputs[input_ids][0][prompt_len:] comp_logprob 0.0 for i, token_id in enumerate(comp_ids): comp_logprob logprobs[i prompt_len - 1, token_id].item() group_logprobs.append(comp_logprob) batch_logprobs.append(group_logprobs) return torch.tensor(batch_logprobs, dtypetorch.float32)3. GRPO算法实现GRPO(Group Relative Policy Optimization)是RLVR中常用的优化算法它通过在同一个提示的多个响应之间进行相对比较来降低方差。3.1 优势计算def compute_advantages(rewards): 计算GRPO风格的优势 rewards: [batch_size, group_size]的张量 返回: 标准化后的优势相同形状 mean rewards.mean(dim1, keepdimTrue) std rewards.std(dim1, keepdimTrue) 1e-8 advantages (rewards - mean) / std return advantages3.2 完整训练循环from torch.optim import AdamW import random def train_rlvr( policy, ref_policy, dataset, verifier, epochs10, batch_size8, group_size4, kl_coef0.1, lr5e-6 ): optimizer AdamW(policy.model.parameters(), lrlr) policy.model.train() ref_policy.model.eval() for epoch in range(epochs): random.shuffle(dataset) total_loss 0 total_reward 0 num_batches len(dataset) // batch_size for batch_idx in range(num_batches): batch dataset[batch_idx*batch_size : (batch_idx1)*batch_size] prompts [item[prompt] for item in batch] gold_answers [item[answer] for item in batch] # 1. 采样生成 completions policy.generate(prompts, num_samplesgroup_size) # 2. 计算奖励 rewards [] for i, prompt in enumerate(prompts): group_rewards [] for comp in completions[i]: r verifier(prompt, comp, gold_answers[i]) group_rewards.append(r) rewards.append(group_rewards) rewards torch.tensor(rewards, dtypetorch.float32) total_reward rewards.mean().item() # 3. 计算优势 advantages compute_advantages(rewards) # 4. 计算当前策略和参考策略的logprob logprobs policy.logprobs(prompts, completions) with torch.no_grad(): ref_logprobs ref_policy.logprobs(prompts, completions) # 5. 计算KL散度和调整后的优势 kl logprobs - ref_logprobs adjusted_adv advantages - kl_coef * kl # 6. 计算策略梯度损失 loss -(adjusted_adv * logprobs).mean() total_loss loss.item() # 7. 反向传播 optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(policy.model.parameters(), 1.0) optimizer.step() avg_loss total_loss / num_batches avg_reward total_reward / num_batches print(fEpoch {epoch1}: Loss{avg_loss:.4f}, Avg Reward{avg_reward:.4f})4. 实战训练与优化4.1 初始化模型# 初始化策略模型和参考模型 base_model_name gpt2 # 实际应用中应使用更大的模型 policy PolicyLM(base_model_name) ref_policy PolicyLM(base_model_name) # 通常参考模型是冻结的 # 生成训练数据 train_dataset generate_math_dataset(5000)4.2 运行训练train_rlvr( policypolicy, ref_policyref_policy, datasettrain_dataset, verifiermath_verifier, epochs5, batch_size8, group_size4, kl_coef0.1, lr5e-6 )4.3 训练优化技巧课程学习从简单问题开始逐步增加难度先训练单运算符问题再引入混合运算逐步扩大数字范围和问题复杂度奖励塑形除了最终答案正确性还可以考虑def enhanced_verifier(prompt, completion, gold_answer): base_reward math_verifier(prompt, completion, gold_answer) # 额外奖励包含计算步骤 step_bonus 0.2 if in completion else 0 # 惩罚过长回答 length_penalty -0.1 * (len(completion.split()) / 10) return base_reward step_bonus length_penalty动态KL系数根据训练情况调整KL约束强度def dynamic_kl_coef(current_reward, target_reward0.8): 根据当前奖励动态调整KL系数 if current_reward target_reward * 0.5: return 0.05 # 弱约束 elif current_reward target_reward: return 0.1 # 中等约束 else: return 0.2 # 强约束以保持稳定性混合训练结合RLVR和RLHF对可验证部分使用RLVR对风格、安全性等使用RLHF/RLAIF5. 评估与应用5.1 评估指标除了简单的奖励值还应监控def evaluate_model(policy, eval_dataset, verifier): total len(eval_dataset) correct 0 exact_match 0 for item in eval_dataset: completions policy.generate([item[prompt]], num_samples1) comp completions[0][0] reward verifier(item[prompt], comp, item[answer]) correct reward # 检查是否精确匹配标准答案 pred extract_last_float(comp) if pred is not None and abs(pred - item[answer]) 1e-6: exact_match 1 accuracy correct / total exact_accuracy exact_match / total return accuracy, exact_accuracy5.2 应用到代码生成RLVR同样适用于代码生成任务只需修改验证器import subprocess def code_verifier(prompt, completion, test_cases): 代码验证器 test_cases: 输入输出对的列表 try: # 提取代码块 code extract_code_block(completion) if not code: return 0.0 # 为每个测试用例运行代码 correct 0 for input_case, expected_output in test_cases: try: process subprocess.run( [python, -c, code], inputinput_case, textTrue, capture_outputTrue, timeout5 ) output process.stdout.strip() if output expected_output: correct 1 except: continue return correct / len(test_cases) except: return 0.05.3 实际部署考虑安全沙箱代码执行需要在隔离环境中进行错误处理完善的超时和资源限制渐进式验证先进行静态分析再执行测试缓存机制避免重复执行相同测试6. 扩展与进阶方向6.1 多模态RLVR将RLVR扩展到图像生成等任务def image_verifier(prompt, generated_image, conditions): 图像验证器示例 conditions: 需要满足的条件列表 score 0 # 使用CLIP评估文本-图像对齐 clip_score calculate_clip_similarity(prompt, generated_image) if clip_score 0.8: score 0.5 # 使用分类器检查特定属性 for attr, value in conditions.items(): if check_image_attribute(generated_image, attr, value): score 0.5/len(conditions) return score6.2 分布式RLVR训练对于大规模任务可以使用Ray等框架实现分布式训练import ray ray.remote class RolloutWorker: def __init__(self, model_name): self.policy PolicyLM(model_name) def generate_batch(self, prompts, num_samples): return self.policy.generate(prompts, num_samples) def compute_logprobs(self, prompts, completions): return self.policy.logprobs(prompts, completions) # 初始化多个worker workers [RolloutWorker.remote(base_model_name) for _ in range(4)]6.3 混合奖励系统结合可验证奖励和学习奖励模型class HybridReward: def __init__(self, verifier, reward_model): self.verifier verifier self.reward_model reward_model def __call__(self, prompt, completion, gold_answerNone): verifiable_reward self.verifier(prompt, completion, gold_answer) learned_reward self.reward_model(prompt, completion) return 0.7 * verifiable_reward 0.3 * learned_reward7. 调试与问题解决7.1 常见问题及解决方案问题1奖励始终为0检查验证器是否能正确解析模型输出确认基础模型具备基本任务能力尝试简化任务难度问题2训练不稳定增加group_size以获得更稳定的优势估计调整KL系数平衡创新和稳定性尝试较小的学习率问题3模型走捷径在验证器中添加反作弊规则多样化奖励信号(如步骤分、格式分)加强KL约束防止偏离参考策略太远7.2 可视化监控实现训练过程可视化有助于调试import matplotlib.pyplot as plt def plot_training_history(history): plt.figure(figsize(12, 4)) plt.subplot(1, 2, 1) plt.plot(history[reward], labelAverage Reward) plt.xlabel(Epoch) plt.ylabel(Reward) plt.legend() plt.subplot(1, 2, 2) plt.plot(history[kl], labelKL Divergence) plt.xlabel(Epoch) plt.ylabel(KL) plt.legend() plt.tight_layout() plt.show()8. 实际应用案例8.1 数学解题助手经过RLVR训练的数学助手可以理解自然语言描述的问题生成分步解决方案最终给出精确答案math_prompt 小明有12个苹果他给了小红5个又买了8个新的他现在有多少个苹果 completion policy.generate([math_prompt])[0][0] 让我们一步步计算 1. 最初有12个苹果 2. 给小红5个后剩下12 - 5 7个 3. 买了8个新的后总数7 8 15个 最终答案是15.0 8.2 代码补全工具针对Python开发的RLVR代码补全code_prompt 实现一个函数计算斐波那契数列的第n项 completion policy.generate([code_prompt])[0][0] def fibonacci(n): if n 0: return 0 elif n 1: return 1 a, b 0, 1 for _ in range(n-1): a, b b, a b return b 8.3 科学问题解答结合知识检索的验证式问答science_prompt 光的传播速度是多少单位是m/s completion policy.generate([science_prompt])[0][0] 真空中的光速是299,792,458米/秒这是物理学中的一个基本常数。 9. 性能优化技巧9.1 内存高效训练使用梯度检查点和混合精度训练from torch.cuda.amp import GradScaler, autocast scaler GradScaler() for batch in dataloader: optimizer.zero_grad() with autocast(): loss compute_loss(batch) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()9.2 并行采样利用多GPU加速生成过程from torch.nn.parallel import DataParallel parallel_model DataParallel(policy.model) def parallel_generate(prompts, num_samples): inputs policy.tokenizer(prompts, return_tensorspt, paddingTrue).to(policy.device) with torch.no_grad(): outputs parallel_model.generate( **inputs, max_length50, num_return_sequencesnum_samples, do_sampleTrue ) # 后续处理与单GPU相同9.3 缓存机制缓存验证结果避免重复计算from functools import lru_cache lru_cache(maxsize10000) def cached_verifier(prompt, completion, gold_answer): return math_verifier(prompt, completion, gold_answer)10. 未来发展方向RLVR技术仍在快速发展几个值得关注的方向包括更强大的验证器结合符号推理和神经网络的混合验证系统多模态验证同时验证文本、代码、数学表达式和图像生成自动化验证器学习从少量标注中学习验证规则安全验证框架确保验证过程本身的安全性和鲁棒性分布式验证网络构建可扩展的验证基础设施在实际项目中成功应用RLVR的关键是设计好的验证函数和合理的奖励结构。与传统的RLHF相比RLVR提供了更透明、更可控的优化过程特别适合那些能够定义明确成功标准的任务。