RLVR: Reinforcement Learning with Verifiable Rewards

# Verifiable reward functions for different domains
import subprocess
import json
import re
from typing import Optional

def math_reward(
 model_answer: str,
 ground_truth: str,
 tolerance: float = 1e-6,
) -> float:
 """Binary reward for math: 1.0 if correct, 0.0 otherwise."""
 # Extract numerical answer from model output
 extracted = extract_final_answer(model_answer)
 if extracted is None:
 return 0.0

 try:
 model_val = float(extracted)
 truth_val = float(ground_truth)
 return 1.0 if abs(model_val - truth_val) < tolerance else 0.0
 except ValueError:
 # String comparison for symbolic answers
 return 1.0 if extracted.strip() == ground_truth.strip() else 0.0

def code_reward(
 generated_code: str,
 test_cases: list,
 timeout: int = 10,
) -> float:
 """Graded reward for code: fraction of tests passed."""
 passed = 0
 for test in test_cases:
 full_code = generated_code + "\n" + test["test_code"]
 try:
 result = subprocess.run(
 ["python", "-c", full_code],
 capture_output=True,
 timeout=timeout,
 text=True,
 )
 if result.returncode == 0:
 passed += 1
 except subprocess.TimeoutExpired:
 continue

 return passed / len(test_cases) if test_cases else 0.0

def proof_reward(
 proof_text: str,
 theorem_statement: str,
 lean_project_path: str,
) -> float:
 """Binary reward for Lean 4 proofs."""
 # Write proof to a temporary Lean file
 lean_code = f"{theorem_statement}\n{proof_text}"
 with open(f"{lean_project_path}/Temp.lean", "w") as f:
 f.write(lean_code)

 result = subprocess.run(
 ["lake", "build", "Temp"],
 capture_output=True,
 cwd=lean_project_path,
 timeout=60,
 )
 return 1.0 if result.returncode == 0 else 0.0

def extract_final_answer(text: str) -> Optional[str]:
 """Extract boxed answer from math response."""
 # Look for \boxed{...} or "The answer is ..."
 boxed = re.findall(r"\\boxed\{([^}]+)\}", text)
 if boxed:
 return boxed[-1]

 answer_pattern = re.findall(
 r"(?:the answer is|therefore|thus)[:\s]+([^\n.]+)",
 text, re.IGNORECASE,
 )
 if answer_pattern:
 return answer_pattern[-1].strip()

 return None

Let me think about this carefully...
Let the distance between the two cities be d km.
Time for A to B: d/60 hours
Time for B to A: d/40 hours
Total distance: 2d
Total time: d/60 + d/40 = d(2+3)/120 = 5d/120 = d/24
Average speed: 2d / (d/24) = 48 km/h

The average speed for the round trip is **48 km/h**.

# GRPO with Verifiable Rewards for Math Reasoning
import torch
from typing import List, Callable

def grpo_math_training_step(
 policy_model,
 ref_model,
 tokenizer,
 math_problems: List[dict], # {"prompt": ..., "answer": ...}
 group_size: int = 16,
 beta: float = 0.04,
 clip_range: float = 0.2,
):
 """
 One GRPO training step using verifiable math rewards.

 For each problem, generate group_size solutions,
 check correctness, normalize rewards within the group,
 and compute the policy gradient.
 """
 total_loss = 0.0
 total_correct = 0
 total_generated = 0

 for problem in math_problems:
 prompt = problem["prompt"]
 ground_truth = problem["answer"]
 input_ids = tokenizer.encode(prompt, return_tensors="pt")

 # Generate a group of solutions
 solutions = []
 for _ in range(group_size):
 output = policy_model.generate(
 input_ids,
 max_new_tokens=1024,
 do_sample=True,
 temperature=0.7,
 top_p=0.95,
 )
 solutions.append(output[0])

 # Compute verifiable rewards (binary: correct or not)
 rewards = torch.tensor([
 math_reward(tokenizer.decode(s), ground_truth)
 for s in solutions
 ])

 total_correct += rewards.sum().item()
 total_generated += group_size

 # Skip if all correct or all wrong (no learning signal)
 if rewards.sum() == 0 or rewards.sum() == group_size:
 continue

 # Group normalization (GRPO core idea)
 advantages = (rewards - rewards.mean()) / (rewards.std() + 1e-8)

 # Compute policy gradient for each solution
 for solution, advantage in zip(solutions, advantages):
 policy_logprobs = compute_logprobs(
 policy_model, input_ids, solution
 )
 with torch.no_grad():
 ref_logprobs = compute_logprobs(
 ref_model, input_ids, solution
 )

 # PPO-style clipped objective
 ratio = torch.exp(policy_logprobs - ref_logprobs)
 surr1 = ratio * advantage
 surr2 = torch.clamp(
 ratio, 1 - clip_range, 1 + clip_range
 ) * advantage
 policy_loss = -torch.min(surr1, surr2).mean()

 # KL penalty
 kl = (ref_logprobs - policy_logprobs).mean()
 total_loss += policy_loss + beta * kl

 accuracy = total_correct / total_generated if total_generated > 0 else 0
 return total_loss / len(math_problems), accuracy

# Simplified DeepSeek-R1 style training pipeline
from dataclasses import dataclass
from typing import List, Dict
import json

@dataclass
class RLVRPipelineConfig:
 """Configuration for a full RLVR training pipeline."""
 base_model: str = "deepseek-ai/DeepSeek-V3-Base"
 cold_start_data: str = "path/to/reasoning_examples.json"
 math_dataset: str = "hendrycks/MATH"
 code_dataset: str = "openai/humaneval"

 # GRPO hyperparameters
 group_size: int = 16
 beta: float = 0.04
 clip_range: float = 0.2
 learning_rate: float = 1e-6

 # Training schedule
 cold_start_epochs: int = 2
 rl_steps: int = 10000
 rejection_sampling_n: int = 64

def run_rlvr_pipeline(config: RLVRPipelineConfig):
 """Run the full RLVR pipeline (simplified)."""

 # Stage 1: Cold-start SFT
 print("Stage 1: Cold-start SFT on reasoning examples")
 cold_start_data = load_json(config.cold_start_data)
 model = sft_train(
 config.base_model,
 cold_start_data,
 epochs=config.cold_start_epochs,
 )

 # Stage 2: RLVR with GRPO
 print("Stage 2: RLVR training with verifiable rewards")
 math_problems = load_dataset(config.math_dataset)
 code_problems = load_dataset(config.code_dataset)

 for step in range(config.rl_steps):
 # Sample a batch of problems (mix of math and code)
 batch = sample_mixed_batch(math_problems, code_problems)

 # GRPO step with verifiable rewards
 loss, accuracy = grpo_math_training_step(
 policy_model=model,
 ref_model=ref_model,
 tokenizer=tokenizer,
 math_problems=batch,
 group_size=config.group_size,
 beta=config.beta,
 )

 if step % 100 == 0:
 print(f" Step {step}: loss={loss:.4f}, acc={accuracy:.3f}")

 # Stage 3: Rejection sampling + SFT
 print("Stage 3: Rejection sampling on diverse prompts")
 diverse_prompts = load_diverse_prompts()
 best_outputs = []
 for prompt in diverse_prompts:
 candidates = [
 model.generate(prompt) for _ in range(config.rejection_sampling_n)
 ]
 # Score and keep the best
 scored = [(c, score_response(c)) for c in candidates]
 best = max(scored, key=lambda x: x[1])
 best_outputs.append({"prompt": prompt, "response": best[0]})

 model = sft_train(model, best_outputs, epochs=1)

 # Stage 4: Final alignment RL
 print("Stage 4: General alignment with helpfulness/safety rewards")
 # Mix verifiable rewards with general preference rewards
 model = final_alignment_rl(model)

 return model

# Using open reasoning models for inference
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load an open reasoning model
model_name = "Qwen/QwQ-32B" # Open reasoning model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
 model_name, torch_dtype="bfloat16", device_map="auto"
)

# Reasoning models produce extended thinking chains
prompt = """Solve this step by step:
A train travels from City A to City B at 60 km/h. The return trip
is made at 40 km/h. What is the average speed for the round trip?"""

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(
 **inputs,
 max_new_tokens=2048,
 temperature=0.6,
 top_p=0.95,
)

response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)
# The model will produce a detailed reasoning chain:
# "Let me think about this carefully...
# Let the distance be d km.
# Time for A to B: d/60 hours
# Time for B to A: d/40 hours
# Total distance: 2d
# Total time: d/60 + d/40 = d(2+3)/120 = 5d/120 = d/24
# Average speed: 2d / (d/24) = 48 km/h
# Wait, let me verify: ..."