Robust Tool Use via Fission-GRPO
Small language models struggle when tools return errors: they often degenerate into repetitive invalid re-invocations without interpreting feedback or self-correcting. Standard RL treats errors as sparse negative signals, but vanishing gradients prevent learning recovery strategies. Fission-GRPO solves this by converting execution errors into multiplicative training signals: errors "fission" into parallel recovery attempts with diagnostic context, enabling the model to learn robust error recovery through on-policy RL.
The key insight is that errors are information: paired with correct recovery attempts, they provide supervised examples of corrective reasoning.
Core Concept
Fission-GRPO uses a three-stage framework:
- Standard Exploration: GRPO samples trajectories, computing rewards across format compliance, functional correctness, and efficiency
- Error Identification & Synthesis: Flag failures and synthesize realistic error diagnostics via learned Error Simulator
- Fission-Based Updates: Expand each error into G' parallel recovery attempts resampled on-policy, converting single failures into multiplicative training signals
This increases outcome diversity and restores meaningful within-group advantages even when most samples fail.
Architecture Overview
- GRPO Sampler: Initial trajectory exploration with policy rollout
- Error Detector: Identifies non-compliant or incorrect outputs
- Error Simulator: Learned model generating realistic error messages from failed attempts
- Recovery Resampler: Resample multiple recovery trajectories conditioned on error
- Advantage Computation: Within-group advantage based on recovery success rates
- Policy Update: GRPO gradient steps using fissioned training signals
Implementation
The method involves error detection, simulation, and fission-based resampling.
Implement error detection and categorization:
import torch
from typing import List, Dict, Tuple
class ErrorDetector:
"""Identify and categorize execution errors."""
def __init__(self, error_categories=None):
self.categories = error_categories or [
"format_error",
"runtime_error",
"logic_error",
"timeout"
]
def detect_error(self, output: str, api_response: str) -> Tuple[bool, str]:
"""Check if output is erroneous and categorize."""
# Format check
if not self.is_valid_format(output):
return True, "format_error"
# API response check
if "Error" in api_response or "error" in api_response:
return True, "runtime_error"
# Logic validation
if not self.validate_logic(output):
return True, "logic_error"
return False, "none"
def is_valid_format(self, output: str) -> bool:
"""Validate output format (JSON, function call, etc.)."""
try:
json.loads(output) # For JSON outputs
return True
except:
return False
def validate_logic(self, output: str) -> bool:
"""Check logical correctness of output."""
# Task-specific validation
return True
detector = ErrorDetector()
Implement error simulator for diagnostic generation:
class ErrorSimulator:
"""Simulate realistic error messages for recovery training."""
def __init__(self, model_name="Qwen3-32B", fine_tuned=True):
self.model = load_model(model_name)
if fine_tuned:
self.model.load_lora("error_simulator_lora")
def generate_diagnostic(self, input_query: str, failed_output: str,
error_category: str) -> str:
"""Generate realistic error message for failed attempt."""
prompt = f"""You are an API error simulator. Given a failed attempt, generate a realistic error message.
Input: {input_query}
Failed output: {failed_output}
Error type: {error_category}
Generate a realistic error message that:
- Explains what went wrong
- Suggests how to fix it (without revealing correct answer)
- Mimics actual API error format
Error message:"""
diagnostic = self.model.generate(prompt, max_tokens=100)
return diagnostic.strip()
simulator = ErrorSimulator()
Implement fission-based trajectory expansion:
def fission_grpo_training_step(policy_model, base_model, batch,
error_detector, error_simulator,
num_fission_attempts=4):
"""GRPO training with error fission for recovery."""
results = []
fissioned_trajectories = []
# Stage 1: Standard exploration
for query in batch:
# Sample trajectory
trajectory = policy_model.sample(query)
output = trajectory["output"]
api_response = trajectory["api_response"]
# Evaluate
is_error, error_type = error_detector.detect_error(output, api_response)
# Compute reward
reward = compute_reward(output, api_response, query)
results.append({
"trajectory": trajectory,
"reward": reward,
"is_error": is_error,
"error_type": error_type
})
# Stage 2: Error fission if applicable
if is_error:
# Generate diagnostic
diagnostic = error_simulator.generate_diagnostic(
query, output, error_type
)
# Augment context with error
augmented_query = f"{query}\n\nPrevious attempt failed with: {diagnostic}\n\nTry again:"
# Stage 3: Fission - sample multiple recovery attempts
for fission_idx in range(num_fission_attempts):
recovery_trajectory = policy_model.sample(augmented_query)
recovery_output = recovery_trajectory["output"]
recovery_api_response = recovery_trajectory["api_response"]
is_recovery_error, _ = error_detector.detect_error(
recovery_output, recovery_api_response
)
recovery_reward = compute_reward(
recovery_output, recovery_api_response, query
)
fissioned_trajectories.append({
"original_trajectory": trajectory,
"recovery_trajectory": recovery_trajectory,
"recovery_reward": recovery_reward,
"success": not is_recovery_error
})
# Compute advantages with fission
advantages = compute_fissioned_advantages(results, fissioned_trajectories)
# GRPO update
loss = 0
for i, result in enumerate(results):
trajectory = result["trajectory"]
advantage = advantages[i]
# Policy gradient
log_prob = compute_log_prob(policy_model, trajectory)
loss += -(log_prob * advantage).mean()
# Include fissioned trajectories from this sample
matching_fissions = [f for f in fissioned_trajectories
if f["original_trajectory"] == trajectory]
for fission in matching_fissions:
recovery_log_prob = compute_log_prob(
policy_model, fission["recovery_trajectory"]
)
fission_advantage = fission["recovery_reward"]
loss += -(recovery_log_prob * fission_advantage).mean()
# Backward pass
loss.backward()
optimizer.step()
return loss.item()
def compute_fissioned_advantages(results, fissioned_trajectories):
"""Compute advantages accounting for fission expansion."""
advantages = []
for result in results:
base_advantage = result["reward"]
# Add fission bonus if error was caught and recovered
matching_fissions = [f for f in fissioned_trajectories
if f["original_trajectory"] == result["trajectory"]]
if matching_fissions:
# Within-group advantage: how much recovery helps
recovery_success_rate = sum(
1 for f in matching_fissions if f["success"]
) / len(matching_fissions)
# Boost advantage if recovery successful
fission_bonus = recovery_success_rate * 0.5
total_advantage = base_advantage + fission_bonus
else:
total_advantage = base_advantage
advantages.append(total_advantage)
return torch.tensor(advantages)
Practical Guidance
| Aspect | Recommendation | Notes |
|---|---|---|
| Num Fission Attempts | 3-5 per error | Higher provides more training signal |
| Error Simulator Size | 32B fine-tuned | Balance realism and speed |
| Diagnostic Length | 50-150 tokens | Concise but informative |
| Fission Bonus Weight | 0.3-0.5 | Scale relative to base reward |
| Error Categories | 4-6 types | Cover main failure modes |
| Policy Model Size | 1.5B-8B (small) | Targets models struggling with errors |
When to use: For tool-using agents with APIs that return errors. For small models showing error degradation. When recovery learning is critical for success.
When NOT to use: For models with strong error handling (larger models often sufficient). When error messages are uninformative.
Common pitfalls:
- Error simulator must be realistic—validate diagnostics match actual API errors
- Too many fission attempts wastes compute—start with 3 and increase if needed
- Fission bonus weight can create strange behavior—monitor entropy and reward distributions
- Recovery context length matters—keep augmented queries reasonable
- Over-reliance on simulated errors can diverge from real errors—periodically validate on actual API
Reference
Robust Tool Use via Fission-GRPO https://arxiv.org/abs/2601.15625