SuS: Strategy-aware Surprise for Intrinsic Exploration

TL;DR

The paper tackles exploration in sparse-reward reinforcement learning by shifting focus from state-level novelty to changes in behavioral strategy. It introduces Strategy-aware Surprise (SuS), which combines Strategy Stability (SS) and Strategy Surprise (SuS) computed over learned strategy embeddings to produce a balanced intrinsic reward. Empirical results on GSM8K show SuS significantly improves Pass@1 and Pass@5 over baselines, with ablations confirming that both SS and SuS are necessary for optimal performance. The approach maintains higher strategy diversity during training, suggesting robust exploration across multiple solving strategies with potential applicability beyond math reasoning domains.

Abstract

We propose Strategy-aware Surprise (SuS), a novel intrinsic motivation framework that uses pre-post prediction mismatch as a novelty signal for exploration in reinforcement learning. Unlike traditional curiosity-driven methods that rely solely on state prediction error, SuS introduces two complementary components: Strategy Stability (SS) and Strategy Surprise (SuS). SS measures consistency in behavioral strategy across temporal steps, while SuS captures unexpected outcomes relative to the agent's current strategy representation. Our combined reward formulation leverages both signals through learned weighting coefficients. We evaluate SuS on mathematical reasoning tasks using large language models, demonstrating significant improvements in both accuracy and solution diversity. Ablation studies confirm that removing either component results in at least 10% performance degradation, validating the synergistic nature of our approach. SuS achieves 17.4% improvement in Pass@1 and 26.4% improvement in Pass@5 compared to baseline methods, while maintaining higher strategy diversity throughout training.

Figures (7)

• Figure 1: Overview of the SuS architecture. The strategy encoder $E$ maps states to strategy embeddings. Strategy Stability (SS) measures embedding consistency across transitions, while Strategy Surprise (SuS) combines prediction error with strategy shift. The combined intrinsic reward $r_{int}$ guides policy learning.
• Figure 2: Learning curves comparing SuS against baselines over training epochs. (a) Pass@1 accuracy shows SuS achieving higher peak performance. (b) Pass@5 accuracy demonstrates SuS's advantage in generating diverse correct solutions.
• Figure 3: Ablation study showing the contribution of each SuS component. Removing either SS or SuS results in significant performance degradation, confirming the necessity of both components.
• Figure 4: Strategy diversity measured by cluster entropy over training. SuS maintains significantly higher diversity throughout learning compared to baselines, preventing premature convergence to narrow behavioral patterns.
• Figure 5: Evolution of intrinsic reward components over training. (a) Strategy Surprise decreases as the world model improves. (b) Success Surprise stabilizes as the policy learns consistent solution patterns.