Improve credit assignment in multi-objective RL by decomposing advantages into segment-specific values. Use Outcome-Conditioned Baselines to reduce cross-objective interference without expensive rollouts, enabling better training signals for multi-step completions with different reward functions per segment.
Improve credit assignment in multi-objective RL by decomposing advantages into segment-specific values. Use Outcome-Conditioned Baselines to reduce cross-objective interference without expensive rollouts, enabling better training signals for multi-step completions with different reward functions per segment.
Blockwise Advantage Estimation for Multi-Objective RL
Problem Context
Language model training often requires optimizing multiple objectives sequentially within a single completion: a math solution requires correct intermediate reasoning before a correct final answer. Standard RL treats each completion as a single unit with one advantage signal, causing objective interference where optimizing one segment undermines progress on another. The temporal structure of problems demands that later segments depend on earlier ones—a segmentation-aware credit assignment mechanism is needed.
Core Concept
Blockwise Advantage Estimation (BAE) decomposes a completion into K blocks (segments), where each block k receives an advantage A^k(i) computed from only its own objective signal r_k(i). Rather than computing advantages over the entire completion, BAE applies block-specific advantages to block-specific tokens, reducing noise and interference between different task requirements.
The key innovation is the Outcome-Conditioned Baseline (OCB), which stratifies samples by intermediate outcomes (e.g., correctness of the prefix) and computes baselines within outcome groups, avoiding expensive conditional inference.
Architecture Overview
Block segmentation: Partition completion into K segments (e.g., reasoning, answer, verification)
Intermediate outcomes: Extract discrete signals at segment boundaries (correct/incorrect, valid/invalid)
Stratification: Group completions by intermediate outcome (G_o contains all samples with outcome o)
Block-wise baseline: Compute b^k(i) = mean reward over samples in same outcome group
Advantage per block: A^k(i) = r_k(i) - b^k(i) applies only to tokens in block k
Low-variance updates: Leverage within-group statistics without additional forward passes
Implementation
Step 1: Define block structure and intermediate outcomes
Identify segments and the outcomes that determine baseline stratification.
Stratify samples and compute block-wise baselines from same-outcome groups.
classOutcomeConditionedBaseline:
"""
Compute conditional baselines by stratifying samples on intermediate outcomes.
Avoids expensive rollouts while reducing variance.
"""def__init__(self, num_blocks: int, group_by_outcome: bool = True):
self.num_blocks = num_blocks
self.group_by_outcome = group_by_outcome
defcompute_baselines(
self,
block_rewards: List[torch.Tensor], # [num_samples, num_blocks]
outcomes: List[Tuple[bool, ...]], # intermediate outcomes per sample
group_size_min: int = 2) -> List[torch.Tensor]:
"""
Compute block-wise baselines by grouping samples with same outcome.
Args:
block_rewards: Rewards for each block of each sample
outcomes: Tuple of intermediate outcomes (e.g., (reasoning_correct, answer_correct))
group_size_min: Minimum group size for baseline estimation
Returns:
List[Tensor]: Baseline values for each block [num_blocks, num_samples]
"""
num_samples = len(outcomes)
baselines = [torch.zeros(num_samples) for _ inrange(self.num_blocks)]
ifnotself.group_by_outcome:
# Simple mean baseline across all samplesfor block_idx inrange(self.num_blocks):
block_rewards_tensor = torch.stack(
[r[block_idx] ifisinstance(r, torch.Tensor) else torch.tensor(r)
for r in block_rewards]
)
baselines[block_idx] = block_rewards_tensor.mean()
return baselines
# Stratify by outcome
outcome_groups: Dict[Tuple[bool, ...], List[int]] = {}
for sample_idx, outcome inenumerate(outcomes):
if outcome notin outcome_groups:
outcome_groups[outcome] = []
outcome_groups[outcome].append(sample_idx)
# Compute baselines within outcome groupsfor outcome, group_indices in outcome_groups.items():
iflen(group_indices) < group_size_min:
continue# Skip groups too small for reliable estimationfor block_idx inrange(self.num_blocks):
group_block_rewards = torch.tensor([
float(block_rewards[idx][block_idx])
for idx in group_indices
])
baseline_value = group_block_rewards.mean()
# Assign baseline to all samples in this outcome groupfor sample_idx in group_indices:
baselines[block_idx][sample_idx] = baseline_value
return baselines
Step 3: Compute block-wise advantages
Apply block-specific advantages to block-specific tokens.
defcompute_blockwise_advantages(
block_rewards: List[torch.Tensor], # [num_blocks, num_samples]
baselines: List[torch.Tensor], # [num_blocks, num_samples]
block_masks: List[torch.Tensor], # [num_blocks, seq_length]
num_samples: int,
seq_length: int) -> torch.Tensor:
"""
Compute advantages per block and tile across tokens in each block.
Args:
block_rewards: Scalar reward per block per sample
baselines: Baseline per block per sample
block_masks: Binary masks indicating token membership
num_samples: Number of samples in batch
seq_length: Sequence length (tokens)
Returns:
Tensor of shape [num_samples, seq_length] with block-wise advantages
"""
num_blocks = len(block_rewards)
# Compute advantages per block
block_advantages = []
for block_idx inrange(num_blocks):
advantage = block_rewards[block_idx] - baselines[block_idx]
block_advantages.append(advantage)
# Expand advantages to token level
advantages_token_level = torch.zeros(num_samples, seq_length)
for sample_idx inrange(num_samples):
for block_idx inrange(num_blocks):
# Get mask for this block
block_mask = block_masks[block_idx]
# Get advantage for this sample-block pair
adv_value = block_advantages[block_idx][sample_idx]
# Assign advantage to all tokens in this block
advantages_token_level[sample_idx, block_mask] = adv_value
return advantages_token_level
Step 4: Integrate into policy gradient optimization
Full training pipeline using blockwise advantage estimation.
deftrain_with_blockwise_advantage(
model, train_loader, verifier, optimizer,
block_segmentation: BlockStructure,
num_epochs: int = 3,
group_size: int = 8,
device: str = 'cuda'):
"""
Train LLM with blockwise advantage estimation for multi-objective completion.
Args:
model: Language model to train
train_loader: Iterable of prompts
verifier: Function computing block rewards (returns list of block rewards)
optimizer: PyTorch optimizer
block_segmentation: BlockStructure defining blocks and outcome extraction
num_epochs: Number training epochs
group_size: GRPO group size
device: Training device
"""
blockwise_grpo = BlockwiseGRPO(
model, optimizer, block_segmentation.block_names.__len__(),
group_size=group_size,
block_segmentation=block_segmentation
)
segmentation = MultiObjectiveSegmentation(block_segmentation)
block_masks = segmentation.create_block_masks(seq_length=512)
for epoch inrange(num_epochs):
total_loss = 0.0
num_batches = 0for batch_idx, batch inenumerate(train_loader):
prompts = batch['prompts']
batch_size = len(prompts)
# Generate completions
completions = []
log_probs_list = []
for prompt in prompts:
completion, log_prob = model.generate_with_logprobs(
prompt, max_tokens=512
)
completions.append(completion)
log_probs_list.append(log_prob)
log_probs = torch.stack(log_probs_list).to(device)
# Extract intermediate outcomes
outcomes = segmentation.extract_outcomes(completions)
# Compute block-wise rewards
block_rewards = [[] for _ inrange(segmentation.num_blocks)]
for completion in completions:
rewards = verifier(completion) # Returns list of per-block rewardsfor block_idx, reward inenumerate(rewards):
block_rewards[block_idx].append(torch.tensor(reward, dtype=torch.float32))
block_rewards = [torch.stack(br).to(device) for br in block_rewards]
# Compute loss with blockwise advantages
loss, metrics = blockwise_grpo.compute_loss(
log_probs, block_rewards, outcomes, block_masks
)
# Backward pass
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
num_batches += 1if (batch_idx + 1) % 10 == 0:
print(f" Batch {batch_idx + 1}: loss={loss.item():.4f}, "f"baseline_vals={metrics['baseline_values']}")
avg_loss = total_loss / num_batches
print(f"Epoch {epoch + 1}: Loss={avg_loss:.4f}")
return model
Practical Guidance
When to use: Multi-step tasks with sequential objectives (math solutions, code generation with correctness at different stages, verification-augmented reasoning)
Hyperparameters:
group_size: 4-8 (GRPO grouping)
group_size_min: 2-4 (minimum outcome group for baseline estimation)
Reduces interference between sequential objectives
Lower variance in advantage estimates vs. scalar reward
No additional forward passes (outcome extraction is fast)
Particularly effective on structured multi-stage completions
Common pitfalls:
Outcome groups too small → unreliable baselines
Block boundaries misaligned with actual token semantics
Forgetting that OCB requires clear intermediate outcome signals
Not verifying outcome extraction accuracy
Scaling: Negligible overhead. Outcome extraction should be fast (regex, simple parsing).
Reference
Paper: https://arxiv.org/abs/2602.10231
Related work: GRPO, multi-task RL, credit assignment, verifiable rewards
Benchmarks: Math problem solving, code generation with intermediate verification
