From 125b49dfc175161a2dbfbbb66e5f075b3729a618 Mon Sep 17 00:00:00 2001 From: RileyChen23 <2017009663@qq.com> Date: Sun, 12 Jul 2026 17:47:36 +0800 Subject: [PATCH] feat: integrate verifiable rewards into GRPO --- code_agent/__init__.py | 6 +++- code_agent/dataset.py | 58 +++++++++++++++++++++++++++++++++ code_agent/reward.py | 31 +++++++++++++++++- tests/test_code_reward_batch.py | 25 ++++++++++++++ tests/test_code_rl_dataset.py | 31 ++++++++++++++++++ trainer/train_grpo.py | 37 +++++++++++++++++---- 6 files changed, 180 insertions(+), 8 deletions(-) create mode 100644 code_agent/dataset.py create mode 100644 tests/test_code_reward_batch.py create mode 100644 tests/test_code_rl_dataset.py diff --git a/code_agent/__init__.py b/code_agent/__init__.py index 96ec0fa..ba0b72d 100644 --- a/code_agent/__init__.py +++ b/code_agent/__init__.py @@ -1,7 +1,8 @@ """Utilities for execution-feedback code generation experiments.""" from .agent import AgentResult, AttemptRecord, ExecutionFeedbackAgent -from .reward import CodeReward, score_verification +from .dataset import CodeRLDataset, collate_code_rl +from .reward import CodeReward, score_code_batch, score_verification from .evaluation import ( EvaluationReport, aggregate_reports, @@ -21,6 +22,7 @@ __all__ = [ "AgentResult", "AttemptRecord", "CodeReward", + "CodeRLDataset", "CodeTask", "EvaluationReport", "ExecutionFeedbackAgent", @@ -28,9 +30,11 @@ __all__ = [ "VerificationResult", "VerificationStatus", "aggregate_reports", + "collate_code_rl", "estimate_pass_at_k", "evaluate_candidates", "extract_python_code", "score_verification", + "score_code_batch", "verify_code", ] diff --git a/code_agent/dataset.py b/code_agent/dataset.py new file mode 100644 index 0000000..c6a8d06 --- /dev/null +++ b/code_agent/dataset.py @@ -0,0 +1,58 @@ +"""Dataset adapter for execution-verifiable GRPO code tasks.""" + +from __future__ import annotations + +import json +import random +from pathlib import Path + +from torch.utils.data import Dataset + +from .verifier import CodeTask + + +class CodeRLDataset(Dataset): + def __init__(self, jsonl_path, tokenizer, thinking_ratio: float = 0.5): + if not 0.0 <= thinking_ratio <= 1.0: + raise ValueError("thinking_ratio must be between 0 and 1") + self.tokenizer = tokenizer + self.thinking_ratio = thinking_ratio + self.tasks: list[CodeTask] = [] + with Path(jsonl_path).open("r", encoding="utf-8") as handle: + for line_number, line in enumerate(handle, start=1): + if not line.strip(): + continue + try: + self.tasks.append(CodeTask.from_dict(json.loads(line))) + except (KeyError, TypeError, ValueError, json.JSONDecodeError) as exc: + raise ValueError(f"Invalid code task at {jsonl_path}:{line_number}: {exc}") from exc + if not self.tasks: + raise ValueError("code RL dataset is empty") + + def __len__(self): + return len(self.tasks) + + def __getitem__(self, index): + task = self.tasks[index] + message = { + "role": "user", + "content": ( + "Solve the Python algorithm task below. Return one complete ```python``` code block.\n" + f"Required entry point: {task.entry_point}\n" + f"Task:\n{task.prompt}" + ), + } + prompt = self.tokenizer.apply_chat_template( + [message], + tokenize=False, + open_thinking=random.random() < self.thinking_ratio, + add_generation_prompt=True, + ) + return {"prompt": prompt, "task": task} + + +def collate_code_rl(samples): + return { + "prompt": [sample["prompt"] for sample in samples], + "tasks": [sample["task"] for sample in samples], + } diff --git a/code_agent/reward.py b/code_agent/reward.py index 630368d..435e61e 100644 --- a/code_agent/reward.py +++ b/code_agent/reward.py @@ -3,8 +3,9 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Sequence -from .verifier import VerificationResult, VerificationStatus +from .verifier import CodeTask, VerificationResult, VerificationStatus, verify_code @dataclass(frozen=True) @@ -34,3 +35,31 @@ def score_verification(result: VerificationResult) -> CodeReward: if status in {VerificationStatus.NO_CODE, VerificationStatus.POLICY_VIOLATION}: return CodeReward(-1.0, 0.0, 0.0, -1.0, status.value) return CodeReward(-1.0, 0.0, -1.0, 0.0, "verifier failure") + + +def score_code_batch( + tasks: Sequence[CodeTask], + responses: Sequence[str], + *, + num_generations: int, + timeout_seconds: float = 2.0, + memory_mb: int = 256, +) -> list[CodeReward]: + """Score grouped generations in task-major order for GRPO/CISPO.""" + if num_generations < 1: + raise ValueError("num_generations must be at least 1") + expected = len(tasks) * num_generations + if len(responses) != expected: + raise ValueError(f"expected {expected} responses for {len(tasks)} tasks, got {len(responses)}") + rewards: list[CodeReward] = [] + for task_index, task in enumerate(tasks): + start = task_index * num_generations + for response in responses[start:start + num_generations]: + result = verify_code( + response, + task, + timeout_seconds=timeout_seconds, + memory_mb=memory_mb, + ) + rewards.append(score_verification(result)) + return rewards diff --git a/tests/test_code_reward_batch.py b/tests/test_code_reward_batch.py new file mode 100644 index 0000000..798e833 --- /dev/null +++ b/tests/test_code_reward_batch.py @@ -0,0 +1,25 @@ +import pytest + +from code_agent import CodeTask, TestCase as CodeTestCase, score_code_batch + + +TASK = CodeTask( + task_id="identity", + prompt="Implement identity(x).", + entry_point="identity", + tests=(CodeTestCase(args=(7,), expected=7),), +) + + +def test_scores_task_major_grouped_responses(): + rewards = score_code_batch( + [TASK], + ["def identity(x):\n return x", "def identity(x):\n return 0"], + num_generations=2, + ) + assert [reward.total for reward in rewards] == [1.0, 0.0] + + +def test_rejects_misaligned_group_size(): + with pytest.raises(ValueError, match="expected 2 responses"): + score_code_batch([TASK], ["def identity(x): return x"], num_generations=2) diff --git a/tests/test_code_rl_dataset.py b/tests/test_code_rl_dataset.py new file mode 100644 index 0000000..9c19ab8 --- /dev/null +++ b/tests/test_code_rl_dataset.py @@ -0,0 +1,31 @@ +import json + +from code_agent import CodeRLDataset, collate_code_rl + + +class FakeTokenizer: + def apply_chat_template(self, messages, **kwargs): + return messages[0]["content"] + f"\nthinking={kwargs['open_thinking']}" + + +def test_code_rl_dataset_builds_prompt_and_preserves_tests(tmp_path): + path = tmp_path / "tasks.jsonl" + path.write_text( + json.dumps( + { + "task_id": "inc", + "prompt": "Implement inc(x).", + "entry_point": "inc", + "tests": [{"args": [1], "expected": 2}], + } + ) + + "\n", + encoding="utf-8", + ) + dataset = CodeRLDataset(path, FakeTokenizer(), thinking_ratio=0.0) + sample = dataset[0] + + assert "Required entry point: inc" in sample["prompt"] + assert sample["task"].tests[0].expected == 2 + batch = collate_code_rl([sample]) + assert batch["tasks"][0].task_id == "inc" diff --git a/trainer/train_grpo.py b/trainer/train_grpo.py index 200adb6..f4de0a7 100755 --- a/trainer/train_grpo.py +++ b/trainer/train_grpo.py @@ -24,6 +24,7 @@ from model.model_minimind import MiniMindConfig, MiniMindForCausalLM from dataset.lm_dataset import RLAIFDataset from trainer.trainer_utils import Logger, is_main_process, lm_checkpoint, init_distributed_mode, setup_seed, SkipBatchSampler, init_model, LMForRewardModel from trainer.rollout_engine import create_rollout_engine +from code_agent import CodeRLDataset, collate_code_rl, score_code_batch warnings.filterwarnings('ignore') @@ -34,7 +35,17 @@ def rep_penalty(text, n=3, cap=0.5): return min(cap, (len(grams) - len(set(grams))) * cap * 2 / len(grams)) if grams else 0.0 -def calculate_rewards(prompts, responses, reward_model): +def calculate_rewards(prompts, responses, reward_model, tasks=None): + if args.reward_type == "code": + code_rewards = score_code_batch( + tasks, + responses, + num_generations=args.num_generations, + timeout_seconds=args.code_timeout, + memory_mb=args.code_memory_mb, + ) + return torch.tensor([reward.total for reward in code_rewards], device=args.device) + rewards = torch.zeros(len(responses), device=args.device) with torch.no_grad(): @@ -92,7 +103,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod full_mask = (outputs != tokenizer.pad_token_id).long() logp_pos = prompt_lens.unsqueeze(1) - 1 + torch.arange(completion_ids.size(1), device=args.device).unsqueeze(0) - rewards = calculate_rewards(prompts, completions, reward_model).to(args.device) # [B*num_gen] + rewards = calculate_rewards(prompts, completions, reward_model, batch.get('tasks')).to(args.device) # [B*num_gen] model_unwrapped = model.module if isinstance(model, DistributedDataParallel) else model with autocast_ctx: @@ -119,6 +130,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod Logger('='*100) grouped_rewards = rewards.view(-1, args.num_generations) # [B, num_gen] + degenerate_group_rate = (grouped_rewards.std(dim=1, unbiased=False) < 1e-6).float().mean() mean_r = grouped_rewards.mean(dim=1).repeat_interleave(args.num_generations) # [B*num_gen] std_r = grouped_rewards.std(dim=1, unbiased=False).repeat_interleave(args.num_generations) # [B*num_gen] advantages = (rewards - mean_r) / (std_r + 1e-4) # [B*num_gen] @@ -163,6 +175,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod Logger(f'Epoch:[{epoch + 1}/{args.epochs}]({step}/{iters}), ' f'Reward: {avg_reward_val:.4f}, KL_ref: {kl_ref_val:.4f}, ' + f'Degenerate Groups: {degenerate_group_rate.item():.4f}, ' f'Adv Std: {advantages_std_val:.4f}, Adv Mean: {advantages_mean_val:.4f}, ' f'Actor Loss: {policy_loss_val:.4f}, Avg Response Len: {avg_len_val:.2f}, Learning Rate: {current_lr:.8f}') @@ -174,6 +187,7 @@ def grpo_train_epoch(epoch, loader, iters, rollout_engine, ref_model, reward_mod "advantages_mean": advantages_mean_val, "policy_loss": policy_loss_val, "avg_response_len": avg_len_val, + "degenerate_group_rate": degenerate_group_rate.item(), "learning_rate": current_lr }) @@ -230,6 +244,9 @@ if __name__ == "__main__": parser.add_argument("--epsilon_high", type=float, default=5.0, help="epsilon上界") parser.add_argument('--from_weight', default='full_sft', type=str, help="基于哪个权重训练") parser.add_argument("--reward_model_path", type=str, default="../../internlm2-1_8b-reward", help="Reward模型路径") + parser.add_argument("--reward_type", type=str, default="reward_model", choices=["reward_model", "code"], help="Reward source") + parser.add_argument("--code_timeout", type=float, default=2.0, help="Code execution timeout per candidate") + parser.add_argument("--code_memory_mb", type=int, default=256, help="POSIX memory limit for each code worker") parser.add_argument('--from_resume', default=0, type=int, choices=[0, 1], help="是否自动检测&续训（0=否，1=是）") parser.add_argument("--use_wandb", action="store_true", help="是否使用wandb") parser.add_argument("--wandb_project", type=str, default="MiniMind-GRPO", help="wandb项目名") @@ -276,7 +293,10 @@ if __name__ == "__main__": ref_model, _ = init_model(lm_config, base_weight, device=args.device) ref_model = ref_model.eval().requires_grad_(False) # Reward模型 - reward_model = LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16) + reward_model = ( + LMForRewardModel(args.reward_model_path, device=args.device, dtype=torch.float16) + if args.reward_type == "reward_model" else None + ) # Rollout引擎（可插拔替换，只负责 policy 推理） rollout_engine = create_rollout_engine( engine_type=args.rollout_engine, @@ -289,7 +309,11 @@ if __name__ == "__main__": sglang_shared_path=args.sglang_shared_path, ) # 数据和优化器 - train_ds = RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len, thinking_ratio=args.thinking_ratio) + train_ds = ( + CodeRLDataset(args.data_path, tokenizer, thinking_ratio=args.thinking_ratio) + if args.reward_type == "code" + else RLAIFDataset(args.data_path, tokenizer, max_length=lm_config.max_seq_len, thinking_ratio=args.thinking_ratio) + ) train_sampler = DistributedSampler(train_ds) if dist.is_initialized() else None optimizer = optim.AdamW(model.parameters(), lr=args.learning_rate) loader_for_count = DataLoader(train_ds, batch_size=args.batch_size, sampler=train_sampler) @@ -321,7 +345,8 @@ if __name__ == "__main__": setup_seed(42 + epoch); indices = torch.randperm(len(train_ds)).tolist() skip = start_step if (epoch == start_epoch and start_step > 0) else 0 batch_sampler = SkipBatchSampler(train_sampler or indices, args.batch_size, skip) - loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True) + collate_fn = collate_code_rl if args.reward_type == "code" else None + loader = DataLoader(train_ds, batch_sampler=batch_sampler, num_workers=args.num_workers, pin_memory=True, collate_fn=collate_fn) if skip > 0: Logger(f'Epoch [{epoch + 1}/{args.epochs}]: 跳过前{start_step}个step，从step {start_step + 1}开始') grpo_train_epoch(epoch, loader, len(loader) + skip, rollout_engine, ref_model, reward_model, start_step, wandb, use_sglang = (args.rollout_engine == "sglang")) @@ -331,4 +356,4 @@ if __name__ == "__main__": # ========== 9. 清理分布进程 ========== if dist.is_initialized(): dist.barrier() - dist.destroy_process_group() \ No newline at end of file + dist.destroy_process_group()