adu2021 Truncated Proximal Policy Optimization
Core Concept
Truncated PPO (T-PPO) enhances training efficiency for reasoning LLMs through truncated rollouts and extended generalized advantage estimation (EGAE). Rather than waiting for complete generation sequences before updating, T-PPO divides long sequences into fixed-window chunks and applies EGAE to compute advantages from incomplete trajectories. This approach maintains constant batch size through replacement of finished sequences while achieving 60% wall-clock time reduction and 2.5x training efficiency improvement without sacrificing convergence on mathematical reasoning tasks.
Architecture Overview
- Truncated Rollout Strategy: Divide response generation into fixed-length windows rather than waiting for completion
- Extended GAE (EGAE): Generalize advantage estimation to incomplete trajectories by assuming terminal value equals penultimate state value
- Batch Continuity: Replace finished sequences with new samples, maintaining constant batch size and GPU utilization
- Token Filtering: Selectively filter training tokens to enable independent policy/value optimization
- Successive Batching: Continuously feed new prompts while earlier ones still generate
Implementation
Step 1: Implement Truncated Rollout Strategy
Divide long sequences into manageable chunks:
import torch
import torch.nn as nn
from typing import List, Dict
class TruncatedRolloutBuffer:
"""
Manages truncated rollouts: divides generation into windows.
"""
def __init__(self, window_size=128, max_total_length=512):
self.window_size = window_size
self.max_total_length = max_total_length
def create_windows(self, prompt_ids, max_new_tokens):
"""
Divide expected generation into fixed-length windows.
Args:
prompt_ids: [batch, prompt_len] token IDs
max_new_tokens: maximum tokens to generate per prompt
Returns:
window_config: dict with window boundaries and settings
"""
batch_size = prompt_ids.shape[0]
num_windows = (max_new_tokens + self.window_size - 1) // self.window_size
window_config = {
'window_size': self.window_size,
'num_windows': num_windows,
'batch_size': batch_size,
'max_new_tokens': max_new_tokens
}
return window_config
def process_truncated_generation(self, prompt_ids, generation_outputs,
window_idx):
"""
Extract tokens for a specific generation window.
Args:
prompt_ids: [batch, prompt_len]
generation_outputs: full generated sequences (may not be complete)
window_idx: which window to extract
Returns:
window_tokens: [batch, window_size] tokens in this window
is_complete: [batch] boolean indicating if sequence completed
"""
start_idx = window_idx * self.window_size
end_idx = (window_idx + 1) * self.window_size
# Extract window
window_tokens = generation_outputs[:, start_idx:end_idx]
# Pad if incomplete
if window_tokens.shape[1] < self.window_size:
padding = torch.full(
(window_tokens.shape[0], self.window_size - window_tokens.shape[1]),
fill_value=0 # or pad token ID
)
window_tokens = torch.cat([window_tokens, padding], dim=1)
# Check completion
is_complete = (window_tokens != 0).sum(dim=1) < self.window_size
return window_tokens, is_complete
Step 2: Implement Extended GAE
Generalize GAE to incomplete trajectories:
class ExtendedGeneralizedAdvantageEstimation:
"""
Extended GAE: handles incomplete (truncated) trajectories.
Key assumption: V(s_l) = V(s_{l-1}) for ungenerated states.
"""
def __init__(self, gamma=0.99, gae_lambda=0.95):
self.gamma = gamma
self.gae_lambda = gae_lambda
def compute_advantages_truncated(self, rewards, values, dones,
trajectory_length):
"""
Compute advantages assuming incomplete trajectories.
Args:
rewards: [batch, seq_len] rewards at each position
values: [batch, seq_len] state value estimates
dones: [batch, seq_len] episode termination flags
trajectory_length: list of actual lengths per batch
Returns:
advantages: [batch, seq_len] advantage estimates
returns: [batch, seq_len] return estimates
"""
batch_size, seq_len = rewards.shape
advantages = torch.zeros_like(rewards)
gae = 0
# Iterate backwards through time
for t in reversed(range(seq_len)):
if t == seq_len - 1:
# Last position: assume next value equals current for ungenerated
next_value = values[:, t].clone()
# Check if trajectory continues beyond window
next_non_terminal = 1 - dones[:, t]
else:
next_value = values[:, t + 1]
next_non_terminal = 1 - dones[:, t]
delta = (
rewards[:, t] +
self.gamma * next_value * next_non_terminal -
values[:, t]
)
# GAE computation with temporal decay
gae = (
delta +
self.gamma * self.gae_lambda * next_non_terminal * gae
)
advantages[:, t] = gae
# Compute returns
returns = advantages + values
return advantages, returns
def compute_advantages_complete(self, rewards, values, dones):
"""
Standard GAE for completed trajectories (reference).
"""
batch_size, seq_len = rewards.shape
advantages = torch.zeros_like(rewards)
gae = 0
# Bootstrap from final value
next_value = values[:, -1].clone()
for t in reversed(range(seq_len)):
if t == seq_len - 1:
next_non_terminal = 1 - dones[:, t]
next_value = values[:, t]
else:
next_non_terminal = 1 - dones[:, t]
next_value = values[:, t + 1]
delta = (
rewards[:, t] +
self.gamma * next_value * next_non_terminal -
values[:, t]
)
gae = (
delta +
self.gamma * self.gae_lambda * next_non_terminal * gae
)
advantages[:, t] = gae
returns = advantages + values
return advantages, returns
Step 3: Implement Token Filtering
Selectively filter tokens for optimization:
class TokenFilter:
"""
Filters training tokens to improve efficiency.
"""
@staticmethod
def filter_high_advantage(tokens, advantages, log_probs,
percentile_threshold=50):
"""
Focus on tokens with high absolute advantage.
Args:
tokens: [batch, seq_len] token IDs
advantages: [batch, seq_len] advantage estimates
log_probs: [batch, seq_len] log probabilities
percentile_threshold: keep tokens above this percentile
Returns:
filtered_mask: [batch, seq_len] boolean mask
"""
# Compute percentile threshold
advantage_threshold = torch.kthvalue(
advantages.reshape(-1),
int(advantages.numel() * (100 - percentile_threshold) / 100)
).values
# Create mask
filtered_mask = torch.abs(advantages) > advantage_threshold
return filtered_mask
@staticmethod
def filter_redundant_tokens(log_probs, entropy_threshold=0.1):
"""
Filter low-entropy (confident) tokens to reduce redundancy.
Args:
log_probs: [batch, seq_len, vocab_size]
entropy_threshold: keep if entropy > threshold
Returns:
filtered_mask: [batch, seq_len] boolean mask
"""
# Compute entropy
probs = torch.softmax(log_probs, dim=-1)
entropy = -(probs * torch.log(probs + 1e-8)).sum(dim=-1)
# Filter
filtered_mask = entropy > entropy_threshold
return filtered_mask
@staticmethod
def filter_composite(advantages, log_probs, weight_advantage=0.7):
"""
Combine multiple filtering criteria.
"""
mask_advantage = TokenFilter.filter_high_advantage(
None, advantages, None, percentile_threshold=40
)
mask_entropy = TokenFilter.filter_redundant_tokens(
log_probs, entropy_threshold=0.15
)
# Weighted combination
combined_mask = (
weight_advantage * mask_advantage.float() +
(1 - weight_advantage) * mask_entropy.float()
) > 0.5
return combined_mask
Step 4: Implement Successive Batching
Maintain constant batch size with continuous sampling:
class SuccessiveBatchSampler:
"""
Samples batches continuously, replacing finished sequences.
Maintains constant batch size for GPU efficiency.
"""
def __init__(self, prompt_dataset, batch_size=32, window_size=128):
self.prompt_dataset = prompt_dataset
self.batch_size = batch_size
self.window_size = window_size
self.current_idx = 0
self.active_sequences = []
def get_next_batch(self):
"""
Return batch with new samples replacing finished ones.
Returns:
batch: dict with prompts, generation_states
"""
batch = {
'new_prompts': [],
'in_progress': self.active_sequences.copy()
}
# Add new sequences to replace finished ones
num_new_needed = self.batch_size - len(self.active_sequences)
for _ in range(num_new_needed):
if self.current_idx >= len(self.prompt_dataset):
self.current_idx = 0 # Cycle through dataset
prompt = self.prompt_dataset[self.current_idx]
batch['new_prompts'].append(prompt)
self.active_sequences.append({
'prompt': prompt,
'generated_tokens': [],
'generation_step': 0
})
self.current_idx += 1
return batch
def update_generation_state(self, window_outputs, finished_mask):
"""
Update tracking of in-progress generations.
Args:
window_outputs: [batch, window_size] generated tokens
finished_mask: [batch] boolean indicating finished sequences
"""
# Remove finished sequences
self.active_sequences = [
seq for i, seq in enumerate(self.active_sequences)
if not finished_mask[i]
]
Step 5: Implement T-PPO Training Loop
Complete training with truncation and successive batching:
class TruncatedPPOTrainer:
"""
Full T-PPO trainer with truncated rollouts and batch continuity.
"""
def __init__(self, model, value_model, window_size=128, batch_size=32):
self.model = model # policy model
self.value_model = value_model
self.window_size = window_size
self.batch_size = batch_size
self.egae = ExtendedGeneralizedAdvantageEstimation()
self.rollout_buffer = TruncatedRolloutBuffer(window_size)
self.batch_sampler = SuccessiveBatchSampler(
[], batch_size=batch_size, window_size=window_size
)
def training_step(self, batch_prompts, batch_responses, window_idx,
verification_fn):
"""
Single training step on a generation window.
Args:
batch_prompts: [batch, prompt_len]
batch_responses: [batch, total_len] responses generated so far
window_idx: which window to train on
verification_fn: function to verify answer correctness
"""
# Extract window
window_tokens, is_complete = self.rollout_buffer.process_truncated_generation(
batch_prompts, batch_responses, window_idx
)
# Compute rewards based on verification
rewards = torch.zeros(self.batch_size, self.window_size)
for i in range(self.batch_size):
response_text = self.model.tokenizer.decode(batch_responses[i])
is_correct = verification_fn(response_text)
# Sparse reward: only at end of window or completion
if is_complete[i] or (window_idx == self.rollout_buffer.window_size - 1):
rewards[i, -1] = 1.0 if is_correct else 0.0
# Compute values using value model
with torch.no_grad():
value_outputs = self.value_model(window_tokens)
values = value_outputs.squeeze(-1) # [batch, window_size]
# Compute advantages
advantages, returns = self.egae.compute_advantages_truncated(
rewards,
values,
~is_complete.unsqueeze(1), # dones flags
self.window_size
)
# Get policy log probs
policy_outputs = self.model(window_tokens)
log_probs = policy_outputs.log_probs # [batch, window_size]
# Filter tokens if desired
token_filter = TokenFilter()
filter_mask = token_filter.filter_composite(
advantages, policy_outputs.logits
)
# PPO loss with clipping
ratio = torch.exp(log_probs - policy_outputs.old_log_probs)
clipped_ratio = torch.clamp(ratio, 1 - 0.2, 1 + 0.2)
loss_actor = -torch.min(
ratio * advantages,
clipped_ratio * advantages
)[filter_mask].mean()
# Value loss
loss_value = torch.nn.functional.mse_loss(values, returns)
# Total loss
loss = loss_actor + 0.5 * loss_value
return loss, {
'actor_loss': loss_actor.item(),
'value_loss': loss_value.item(),
'filtered_fraction': filter_mask.float().mean().item()
}
def train_epoch(self, dataloader, num_epochs=1):
"""
Train for one epoch with successive batching.
"""
optimizer = torch.optim.Adam(
list(self.model.parameters()) +
list(self.value_model.parameters()),
lr=1e-4
)
for epoch in range(num_epochs):
epoch_loss = 0
epoch_metrics = {}
for batch_idx, (prompts, responses) in enumerate(dataloader):
# Determine number of windows
max_len = responses.shape[1]
num_windows = (max_len + self.window_size - 1) // self.window_size
# Train on each window
for window_idx in range(num_windows):
loss, metrics = self.training_step(
prompts, responses, window_idx,
verification_fn=self._verify_correctness
)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(
list(self.model.parameters()) +
list(self.value_model.parameters()),
1.0
)
optimizer.step()
epoch_loss += loss.item()
for k, v in metrics.items():
if k not in epoch_metrics:
epoch_metrics[k] = []
epoch_metrics[k].append(v)
# Log results
avg_loss = epoch_loss / (len(dataloader) * num_windows)
print(f"Epoch {epoch + 1}, Loss: {avg_loss:.4f}")
for k, v in epoch_metrics.items():
print(f" {k}: {sum(v) / len(v):.4f}")
def _verify_correctness(self, response_text):
"""Verify if response is correct (domain-specific)"""
return True # Placeholder
Practical Guidance
- Window Size: Larger windows (256) reduce overhead but increase memory; start with 128
- EGAE Assumption: V(s_l) = V(s_{l-1}) works well empirically; validate on your tasks
- Token Filtering: Apply selectively; too aggressive filtering can bias updates
- Batch Continuity: Requires prompt dataset large enough to avoid repetition
- Speedup Target: Expect 2-3x speedup; measure actual wall-clock time, not just theoretical
- Hyperparameters: Use standard PPO settings (clip=0.2, GAE lambda=0.95); adjust window size first
- Evaluation: Benchmark on long-sequence reasoning tasks (AIME, MATH500, code generation)
- Comparison: Always compare to standard PPO on same tasks to validate efficiency gains
Reference
Paper: arXiv:2506.15050 Key metrics: 60% wall-clock reduction, 2.5x training efficiency on math reasoning EGAE improvement: Handles incomplete trajectories enabling continuous batching Related work: PPO, GAE, policy optimization, training efficiency