Improves LLM agent decision-making by training agents to first critically evaluate actions before generating, using RL on action-pair comparisons. Develops intrinsic reasoning about action quality without requiring reflection supervision.
Agentic Critical Training: Teaching Agents to Reason About Action Quality
version
0.0.2
engine
skillxiv-v0.0.2-claude-opus-4.6
license
MIT
url
https://arxiv.org/abs/2603.08706
keywords
["Agent Training","Reinforcement Learning","Action Selection","Chain of Thought","GRPO"]
description
Improves LLM agent decision-making by training agents to first critically evaluate actions before generating, using RL on action-pair comparisons. Develops intrinsic reasoning about action quality without requiring reflection supervision.
LLM agents trained on imitation learning only learn what to do, not why certain actions are preferable. They never contrast successful actions against alternatives, leaving them unaware of action quality. Traditional reflection-based approaches attempt to address this but remain fundamentally imitative—they copy pre-generated reflections rather than autonomously reasoning about quality.
Agentic Critical Training (ACT) develops genuine action quality reasoning through a two-stage RL approach: first train the agent to critically compare action pairs via RL, then leverage this learned critical ability for direct action generation. This forces autonomous reasoning without supervised reflection supervision.
Core Concept
Stage 1 - Critical Reasoning Training:
Present expert action vs model-generated alternative at each decision point
Train via RL to select the better action
Model develops understanding of quality through comparison rewards only
No reflection text supervision—only selection outcome matters
Stage 2 - Action Generation Training:
Fine-tune on direct action generation
Model now has learned critical foundation—knows what makes actions good
Better ability to generate high-quality actions from scratch
Key insight: Selection is simpler than generation but teaches reasoning. By training on comparison first, the model internalizes quality criteria before attempting generation.
Architecture Overview
Comparison Stage: Present action pairs, train selection via RL
Reward Design: Accuracy rewards + admissibility bonus + format validation
Critical Understanding: Model learns what distinguishes good actions
Generation Fine-Tuning: Apply learned understanding to action generation
Verification Rewards: Use verifiable outcomes rather than external judges
Implementation Steps
Implement two-stage training: comparison selection followed by generation refinement.
"""Trains agents to critically evaluate actions through comparison."""
def
__init__
self, model_name="llama-2-13b", device="cuda"
self
self
self
self
def
prepare_comparison_prompt
self, state, expert_action, model_action, task
"""
Construct prompt presenting both actions for comparison.
Args:
state: current environment/problem state
expert_action: reference action known to be good
model_action: alternative action generated by model
task: task description
Returns:
prompt: formatted comparison prompt
"""
f"""Task: {task}
Current State: {state}
Two possible actions:
Action A (Reference): {expert_action}
Action B (Alternative): {model_action}
Which action is better for this situation? Respond with only "A" or "B" followed by brief reasoning."""
return
def
get_model_selection
self, prompt
"""
Get model's action selection (A or B).
Args:
prompt: comparison prompt
Returns:
selection: 'A' or 'B'
log_prob: log probability of selection
"""
"""
Compute reward for comparison selection.
Args:
selection: model's chosen action ('A' or 'B')
expert_action: known good action (Action A)
model_action: alternative action (Action B)
Returns:
reward: composite reward signal
"""
# Accuracy reward: did model select expert action?
if
'A'
'accuracy'
1.0
else
'accuracy'
1.0
# Admissibility bonus: was the non-selected action at least valid?
# This prevents excessive penalty for selecting reasonable alternatives
if
'A'
else
self
'admissibility'
0.2
if
else
0.2
# Format reward: ensure response format is correct
# (already implicitly satisfied by A/B selection)
'format'
0.0
# Total weighted reward
1.0
'accuracy'
0.3
'admissibility'
0.1
'format'
return
def
_verify_action_validity
self, action, state, task
"""Check if action is valid in this context."""
# Task-specific validation logic
# For now, simplified version
if
and
len
0
return
True
return
False
class
ComparisonRewardOptimizer
"""Optimize model for action comparison using GRPO."""
def
__init__
self, model, optimizer, trainer
self
self
self
def
training_step
self, batch_data
"""
Single GRPO training step on comparison task.
Args:
batch_data: list of {state, expert_action, model_action, task} dicts
Returns:
loss: scalar loss
"""
len
# Forward pass: get all selections
for
in
self
'state'
'expert_action'
'model_action'
'task'
self
# Compute reward for this selection
self
'state'
'expert_action'
'model_action'
'task'
self
# GRPO: group-relative advantages
1e-8
# Policy gradient: maximize expected reward
# Backward pass
self
self
1.0
self
return
Stage 2: Generation Fine-Tuning
classGenerationTrainer:
"""Fine-tune model for direct action generation after critical training."""def__init__(self, critic_model, optimizer, verifier_fn):
self.model = critic_model
self.optimizer = optimizer
self.verifier = verifier_fn # Function to verify action correctnessdefprepare_generation_prompt(self, state, task):
"""
Construct prompt for action generation.
Args:
state: problem state
task: task description
Returns:
prompt: generation prompt
"""returnf"""Task: {task}
Current State: {state}
Generate the best action for this situation:"""defgeneration_training_step(self, batch_data):
"""
GRPO training on direct action generation.
Args:
batch_data: list of {state, task, ground_truth_action} dicts
Returns:
loss: scalar loss
"""
batch_size = len(batch_data)
all_rewards = []
all_log_probs = []
for sample in batch_data:
prompt = self.prepare_generation_prompt(sample['state'], sample['task'])
input_ids = self.model.tokenizer.encode(prompt, return_tensors="pt")
with torch.enable_grad():
outputs = self.model(input_ids, output_hidden_states=True)
# Generate action greedily
generated_action = self._generate_action(outputs.logits)
all_log_probs.append(outputs.logits.sum()) # Simplified; actual version uses proper probability# Verify action correctness
is_correct = self.verifier(
generated_action,
sample['ground_truth_action'],
sample['state'],
sample['task']
)
reward = 1.0if is_correct else -1.0
all_rewards.append(reward)
# Same GRPO logic as Stage 1
all_rewards = torch.tensor(all_rewards, device=self.model.device)
advantages = all_rewards - all_rewards.mean()
advantages = advantages / (all_rewards.std() + 1e-8)
loss = -(torch.stack(all_log_probs) * advantages.detach()).mean()
self.optimizer.zero_grad()
loss.backward()
self.optimizer.step()
return loss.item()
def_generate_action(self, logits):
"""Extract action from model logits."""# Task-specific generation logicreturn"generated_action"deffull_agentic_training_pipeline(
model_name,
comparison_data,
generation_data,
num_comparison_epochs=5,
num_generation_epochs=3):
"""
Complete two-stage training pipeline.
Args:
model_name: HuggingFace model identifier
comparison_data: training data for comparison stage
generation_data: training data for generation stage
num_comparison_epochs: training iterations for Stage 1
Returns:
trained_model: agent with critical reasoning capability
"""
device = "cuda"if torch.cuda.is_available() else"cpu"# Stage 1: Critical trainingprint("[Stage 1] Critical Action Comparison Training...")
trainer = CriticalAgentTrainer(model_name, device)
model = trainer.model
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
comparison_optimizer = ComparisonRewardOptimizer(model, optimizer, trainer)
for epoch inrange(num_comparison_epochs):
epoch_loss = 0for batch in comparison_data:
loss = comparison_optimizer.training_step(batch)
epoch_loss += loss
print(f" Epoch {epoch+1}/{num_comparison_epochs} Loss: {epoch_loss:.4f}")
# Stage 2: Generation fine-tuningprint("[Stage 2] Action Generation Fine-Tuning...")
defverify_action(generated, ground_truth, state, task):
# Simplified; real version compares execution in environmentreturn generated == ground_truth
gen_trainer = GenerationTrainer(model, optimizer, verify_action)
for epoch inrange(num_generation_epochs):
epoch_loss = 0for batch in generation_data:
loss = gen_trainer.generation_training_step(batch)
epoch_loss += loss
print(f" Epoch {epoch+1}/{num_generation_epochs} Loss: {epoch_loss:.4f}")
return model
Learning rate: 1e-5 (conservative for large models)
Comparison epochs: 3-5; Generation epochs: 2-3
When to Apply:
Training LLM agents for decision-making tasks
Scenarios with clear action quality criteria
Multi-step problems where understanding action consequences matters
Tasks where reflection text is expensive or unavailable
When NOT to Apply:
Single-step tasks with no action trade-offs
Situations where expert demonstrations are scarce
Real-time systems where comparison overhead is prohibitive
Key Pitfalls:
Stage 1 too short—critical reasoning underdeveloped
Admissibility reward too high—model learns to accept poor alternatives
Using generation data for comparison training—overfits to specific solutions
Not using verifiable rewards—requires external judges
Integration Notes: Works with any causal LLM; requires pairs of (expert, alternative) actions for comparison stage; comparison data can be auto-generated or human-collected.
Evidence: Improves agent success rates 15-25% over standard imitation learning; develops genuine action quality understanding; enables agents to generalize to unseen action combinations.
