feat: integrate verifiable rewards into GRPO

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",
 ]

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],
+    }

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

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)

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"

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()
+        dist.destroy_process_group()