ExTra: Exploratory Trajectory Optimization for Language Model Reinforcement Learning

2026-06-23

AI Summary

This work addresses the challenge that language models often fail to explore effectively in reinforcement learning for reasoning due to task difficulty polarization: on easy tasks, they generate low-diversity correct trajectories that yield insufficient gradient signals, while on hard tasks, they produce entirely incorrect outputs with no positive rewards. To overcome this, the authors propose the ExTra framework, which extends GRPO with a trajectory-based exploration mechanism that uniquely combines embedding-space diversity rewards and entropy-guided regeneration of intermediate trajectory prefixes to dynamically enhance exploration efficiency. Evaluated on six mathematical reasoning benchmarks, ExTra significantly outperforms GRPO, achieving approximately 5-point gains in pass@1 and 7-point gains in pass@16 with the Qwen3-1.7B model.

Abstract

Reinforcement Learning with Verifiable Rewards (RLVR) for language-model reasoning can fail at both extremes of task difficulty: easy prompts often produce all-correct, low-diversity rollout groups with little gradient signal, while hard prompts can produce all-incorrect groups with no positive reward. We introduce ExTra (Exploratory Trajectory Optimization), a GRPO-compatible framework that extracts exploration signals from the model's own rollouts. ExTra combines two mechanisms: (i) a novelty reward that adds embedding-based diversity bonuses after GRPO normalization, rewarding diverse correct solutions; and (ii) entropy-guided prefix regeneration, which scores partial trajectories using entropy signals and continues exploration from promising intermediate steps. Across six mathematical reasoning benchmarks, ExTra improves Qwen3-1.7B over GRPO by about +5 points on pass@1 and +7 points on pass@16, showing that trajectory-level exploration signals can improve both single-sample accuracy and inference-time coverage.

Problem

• Reinforcement Learning
• Language Model
• Exploration
• Trajectory Optimization
• Reasoning

Innovation

• Exploratory Trajectory Optimization
• GRPO
• novelty reward
• entropy-guided prefix regeneration
• Reinforcement Learning with Verifiable Rewards