Difficulty: ⭐⭐☆☆☆ Intermediate
Source file:apex/alignment/combined_reward.py
You will learn: How APEX-1 combines outcome, process, and constitutional rewards into one signal for GRPO.
Each reward component captures a different aspect of response quality:
| Signal | What It Measures |
|---|---|
| Outcome reward | Did the response correctly answer the question? |
| Process reward | Did the reasoning steps make logical sense? |
| Constitutional reward | Is the response safe and ethical? |
Using only one signal:
- Outcome only: Model learns to get right answers but may use flawed reasoning or be unsafe.
- Process only: Model reasons carefully but may still give wrong final answers.
- Constitutional only: Model is safe but may be unhelpful.
Combining all three gives a holistic quality signal.
Default weights:
-
$\lambda_{outcome} = 1.0 - \lambda_{process} - \lambda_{cai}$ (remainder after PRM and CAI) $\lambda_{process} = 0.3$ $\lambda_{cai} = 0.3$ - Implied
$\lambda_{outcome} = 0.4$
All rewards are normalised to
Each raw reward has a different scale. Before combining:
Wait — actually for APEX-1's binary/bounded signals, we normalise differently:
-
Outcome reward (from RewardModel): already in
$[0, 1]$ via sigmoid -
Process reward (from PRM): product of sigmoid scores → in
$[0, 1]$ -
Constitutional reward: fraction of principles passed → in
$[0, 1]$
All three are naturally in
"""
Combined Reward Signal for GRPO Alignment.
Combines three reward signals:
1. Outcome reward (from RewardModel)
2. Process reward (from PRM)
3. Constitutional reward (from ConstitutionalAI)
Used in GRPO training to provide a holistic quality signal.
"""
import logging
from dataclasses import dataclass
from typing import Optional
import torch
logger = logging.getLogger(__name__)
@dataclass
class CombinedRewardOutput:
"""Result from the combined reward computation."""
combined_reward: float # Final scalar reward in [0, 1]
outcome_reward: float # Component from outcome reward model
process_reward: float # Component from PRM
cai_reward: float # Component from Constitutional AI
outcome_weight: float
process_weight: float
cai_weight: float
class CombinedRewardModel:
"""Tri-signal reward aggregator for GRPO.
Combines:
- Outcome reward from RewardModel (did the answer correct?)
- Process reward from PRM (was the reasoning sound?)
- Constitutional reward from CAI (is it safe?)
Args:
reward_model: Trained RewardModel instance.
process_rm: Trained ProcessRewardModel instance (or None).
constitutional_ai: ConstitutionalAI instance (or None).
lambda_process: Weight for process reward (default: 0.3).
lambda_cai: Weight for constitutional reward (default: 0.3).
device: Computation device.
"""
def __init__(
self,
reward_model,
process_rm=None,
constitutional_ai=None,
lambda_process: float = 0.3,
lambda_cai: float = 0.3,
device: str = "cpu",
):
self.reward_model = reward_model
self.process_rm = process_rm
self.constitutional_ai = constitutional_ai
self.device = device
# Compute outcome weight as the remainder
self.lambda_process = lambda_process
self.lambda_cai = lambda_cai
self.lambda_outcome = max(0.0, 1.0 - lambda_process - lambda_cai)
logger.info(
"Combined reward weights — outcome: %.2f, process: %.2f, cai: %.2f",
self.lambda_outcome, self.lambda_process, self.lambda_cai,
)
def compute(
self,
input_ids: torch.Tensor, # Full sequence token IDs [1, S]
prompt_text: str, # Prompt text (for CAI)
response_text: str, # Response text (for CAI, PRM)
prefix_len: int = 0, # Prompt length (for PRM forward)
) -> CombinedRewardOutput:
"""Compute the combined reward for a response.
Args:
input_ids: Encoded (prompt + response) token IDs.
prompt_text: Raw prompt text for CAI evaluation.
response_text: Raw response text for CAI evaluation.
prefix_len: Where the response starts in input_ids.
Returns:
CombinedRewardOutput with all components and the final combined score.
"""
input_ids = input_ids.to(self.device)
# ── 1. Outcome Reward ─────────────────────────────────────────────
# RewardModel scores the full (prompt + response) sequence
with torch.no_grad():
outcome_score = self.reward_model(input_ids) # [1]
# Clamp to [0, 1] (reward model uses sigmoid head but could overflow)
outcome_reward = float(outcome_score.squeeze().clamp(0.0, 1.0).item())
# ── 2. Process Reward ─────────────────────────────────────────────
process_reward = 0.5 # Neutral default (0.5 = no opinion)
if self.process_rm is not None and self.lambda_process > 0.0:
try:
with torch.no_grad():
prm_output = self.process_rm(
input_ids, prefix_len=prefix_len
)
# process_reward: [B] → take batch 0
process_reward = float(
prm_output["process_reward"][0].clamp(0.0, 1.0).item()
)
logger.debug(
"PRM: process_reward=%.3f, n_steps=%d",
process_reward,
prm_output["n_steps"][0] if prm_output["n_steps"] else 0,
)
except Exception as e:
# Do not crash GRPO if PRM fails — use neutral score
logger.warning("PRM scoring failed: %s — using 0.5", str(e))
process_reward = 0.5
# ── 3. Constitutional Reward ──────────────────────────────────────
cai_reward = 1.0 # Default: fully compliant (if CAI not configured)
if self.constitutional_ai is not None and self.lambda_cai > 0.0:
try:
cai_result = self.constitutional_ai.evaluate(response_text)
cai_reward = float(cai_result["constitutional_score"])
logger.debug(
"CAI: score=%.3f, violations=%d",
cai_reward,
cai_result["n_violations"],
)
except Exception as e:
logger.warning("CAI evaluation failed: %s — using 1.0", str(e))
cai_reward = 1.0
# ── 4. Combine Signals ────────────────────────────────────────────
combined = (
self.lambda_outcome * outcome_reward +
self.lambda_process * process_reward +
self.lambda_cai * cai_reward
)
# Clamp combined to [0, 1] for numerical safety
combined = max(0.0, min(1.0, combined))
return CombinedRewardOutput(
combined_reward=combined,
outcome_reward=outcome_reward,
process_reward=process_reward,
cai_reward=cai_reward,
outcome_weight=self.lambda_outcome,
process_weight=self.lambda_process,
cai_weight=self.lambda_cai,
)# In GRPO training:
combined_rm = CombinedRewardModel(reward_model, prm, constitutional_ai)
for prompt_ids in prompts:
prompt_text = tokenizer.decode(prompt_ids[0].tolist())
response_ids_list = []
rewards_list = []
# Generate G responses
for _ in range(G):
output = generator.generate(prompt_ids)
response_text = tokenizer.decode(output.token_ids)
full_ids = torch.cat([prompt_ids, torch.tensor([output.token_ids])], dim=1)
# Compute combined reward for this response
reward_output = combined_rm.compute(
full_ids, prompt_text=prompt_text, response_text=response_text,
prefix_len=prompt_ids.shape[1],
)
rewards_list.append(reward_output.combined_reward)
response_ids_list.append(torch.tensor([output.token_ids]))
# Compute group-relative advantages and update policy
rewards = torch.tensor(rewards_list)
grpo_training_step(model, ref_model, optimizer, prompt_ids,
response_ids_list, rewards, prompt_len=...)| Training Config | Reasoning Acc | Safety Pass Rate |
|---|---|---|
| Outcome only | 70% | 65% |
| Outcome + CAI | 69% | 92% |
| Outcome + PRM | 81% | 66% |
| All three (APEX-1) | 83% | 93% |
The tri-signal system provides both better reasoning AND better safety simultaneously.
Next: 30 — Utilities →