cs.LGDec 1, 2025
SaveHow Does RL Post-training Induce Skill Composition? A Case Study on Countdown
Organizations: Princeton Language and Intelligence (PLI), Princeton University
Abstract
While reinforcement learning (RL) successfully enhances reasoning in large language models, its role in fostering compositional generalization (the ability to synthesize novel skills from known components) is often conflated with mere length generalization. To this end, we study what RL post-training teaches about skill composition and how the composition structure affects the skill transfer. We focus on the Countdown task (given n numbers and a target, form an expression that evaluates to the target) and analyze model solutions as expression trees, where each subtree corresponds to a reusable subtask and thus can be viewed as a "skill." Tracking tree shapes and their success rates over training, we find: (i) out-of-distribution (OOD) generalization to larger n and to unseen tree shapes, indicating compositional reuse of subtasks; (ii) the order of skill acquisition depends upon the composition structure---models master shallow balanced trees (workload is balanced between subtasks) before deep unbalanced ones, with persistent fragility on right-heavy structures for a fixed composition depth. Our diagnostic reveals what is learned, in what order, and where generalization fails, clarifying how RL-only post-training induces OOD generalization beyond what standard metrics such as pass@k reveal. To study the implications of our findings for real-world tasks, we present a preliminary study on formal theorem proving.
Figures & tables
Figure 1: From generated expression to canonical pattern . The model-generated expression is parsed into an expression tree, normalized with respect to algebraic identities, and mapped to the unique canonical pattern . Patterns are then grouped by their computational shape. Here, the root operator is “ ”, with a left sub-expression over three leaves and a right sub-expression over one leaf . The signature serves as a compact representation of the pattern’s structure at the root-level.
Figure 2: Successful length generalization, but a gap between first producing a pattern and producing it reliably. We train Qwen2.5-1.5B on problem sizes and evaluate on held-out data for . (Top) The model demonstrates strong length generalization : accuracy on the held-out problems improves over the base model (leftmost data point). (Bottom) The model quickly learns to produce most patterns at least once, but the ability to reliably produce the correct pattern (see Section 3.2 for the definition) lags significantly behind.
| Model | ||||
| 1.5B / 3B / 7B | 1.00 / 1.00 / 1.00 | 0.94 / 0.95 / 0.98 | 0.68 / 0.68 / 0.75 | 0.41 / 0.45 / 0.45 |
Table 1: Performance across model size. Pass@32 averaged across all successful seeds 6 6 6 2 seeds for 1.5B/3B and 3 for 7B. Due to compute budget, we evaluate only 1 seed per model for . ( final checkpoints for 1.5B/3B and best checkpoints for 7B) 7 7 7 We select the checkpoint right before the training reward declines. See Section B.1 for more details. on (trained) and (held out). Models learn patterns almost perfectly and generalize to larger puzzle sizes . See Section B.1 for a more complete version of the table.
Figure 3: Generalizes to entire families of held-out compositional patterns. To evaluate compositional generalization, we removed an entire family of related patterns from the training set: the base structure and all of its extensions (e.g., , ). The model successfully recovers these unseen patterns. Coverage first emerges on the held-out subpattern before generalizing to its more complex dependents, matching the typical learning hierarchy where presence precedes precision. This result demonstrates that the model can reuse learned substructures to synthesize novel operator-tree shapes it has never been explicitly trained on.
Figure 4: Composition structure (not input length) determines difficulty. Within , balanced patterns ( ; purple) are discovered (top row) and mastered more reliably (bottom row) than the unbalanced ones (red/blue). Balanced patterns have shallow trees, with solution depth of 2 equal to that of puzzles. Even when controlling for depth, right-heavy patterns ( ; blue) are harder than left-heavy patterns ( ; red). This evidence points to a “lookahead bottleneck,” the challenge of committing to a solution ahead of a complex subroutine. The entire structural hierarchy persists for . Results shown for a representative run of Qwen2.5-1.5B . See Appendix H for equivalent plots for all other runs.
| balanced | left-heavy | right-heavy | ||||
| Present in SFT dataset | 51 | 14 | 35 | |||
| Generated by RL model (ID) | 56 | +5 | 16 | +2 | 29 | -6 |
| Generated by RL model (OOD) | 68 | +17 | 19 | +5 | 13 | -22 |
Table 4: Distribution of proofs by composition structure (%). After RL post-training, the final checkpoint prefers balanced proofs and avoids right-heavy proofs. This trend is consistent with the “lookahead bottleneck” identified in Countdown : structures whose complex subroutine appears later are harder to produce reliably and are therefore less likely to be reinforced. The last row shows that the same qualitative shift persists on an out-of-distribution dataset which contains an approximately equal number of reference proofs across balanced / left-heavy / right-heavy categories.
Appendix figures & tables41 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 5: Grouping canonical patterns by tree structure ( ). Internal (blue) nodes are operators and leaves are operands. We categorize patterns by the relative size of the root’s subtrees: left-heavy, balanced, or right-heavy.
Figure 6: Number of examples per pattern in Pan et al. [2025b] . 18 out of 114 patterns do not appear in a randomly selected 10000 examples from the dataset. Few patterns (mostly from ) appear disproportionately frequently.
Figure 7: Learning dynamics at a glance (Qwen2.5-3B) trained on dataset from Pan et al. [2025a] . (Left) evaluated on our dataset; (Right) evaluated on the held-out data from Pan et al. [2025a] . Even though model accuracy smoothly increases (top), models do not reliably learn most patterns and fails to generate at least 10% of patterns (bottom). The held-out data from Pan et al. [2025a] do not include 36 possible patterns and fail to test the models on the corresponding skills.
| Pattern | Number of Examples | Pattern | Number of Examples |
| 3333 | 315 | ||
| 3905 | 791 | ||
| 2314 | 927 | ||
| 188 | 2704 | ||
| 1058 | 264 | ||
| 2546 | 647 |
Table 5: Number of training examples per pattern. Any pattern not listed has 4000 examples.
| Model | All | Avg | Pass | All | Avg | Pass |
| 1.5B | 0.00 | 0.02 | 0.42 | 0.00 | 0.00 | 0.00 |
| 1.5B (F) | 0.97 0.01 | 1.00 0.00 | 1.00 0.00 | 0.60 0.09 | 0.82 0.02 | 0.93 0.00 |
| 3B | 0.00 | 0.05 | 0.66 | 0.00 | 0.01 | 0.20 |
| 3B (F) | 0.93 0.01 | 0.99 0.01 | 1.00 0.00 | 0.56 0.06 | 0.81 0.03 | 0.94 0.01 |
| 7B | 0.00 | 0.10 | 0.88 | 0.00 | 0.02 | 0.36 |
Table 6: Performance across model size (Full version). Averaged across all successful seeds ( final (F) checkpoints for 1.5B/3B and best (B) and final (F) checkpoints for 7B). All-correct@32 (All), Average Accuracy (Avg), and Pass@32 (Pass) on (trained) and (held out). Models learn patterns almost perfectly and generalize to larger puzzle sizes , but the gap between Pass@32 and Average Accuracy increases with , signaling that the generalization is not perfect.
Table 7: Per-pattern precision averaged by pattern structure on (Full version). Averaged across all successful seeds ( final checkpoints for 1.5B/3B and best and final checkpoints for 7B). Averaged by (a) root operator; (b) pattern with root operator ; (c) pattern with root operator . Models struggle more on as the root operator. Shapes in the form of or are particularly challenging. Within each column, the highest number is boldfaced and the second highest number is underlined .
| Pattern | Explanation | Canonical Pattern | Structure |
| Itself | |||
| Direct subpattern | |||
| Direct subpattern | |||
| Direct subpattern | |||
| Direct subpattern | |||
| Direct subpattern |
Table 8: The list of patterns removed, when holding out . We additionally remove 6 patterns that contain as a subpattern and 8 more that can be collapsed into once a subtask has been solved.
| Pattern | base target | base global | RL target | RL global |
| 127 | 1,335 | 5,631 | 5,705 | |
| 13 | 677 | 1,318 | 4,871 | |
| 99 | 7,094 | 1,457 | 15,368 | |
| 4 | 140 | 147 | 2,119 | |
| 0 | 29 | 21 | 582 | |
| 27 | 1,782 | 5,078 | 12,831 |
Table 9: The number of occurrences of heldout patterns before ( base ) and after ( RL ) training. target refers to the number of occurrences of heldout patterns when they are correctly used on the puzzles that require the usage of the particular pattern, where as global refers to the number of occurrences across all puzzles in the test data. The base model rarely produces the heldout patterns to the correct target puzzles, whereas the final RL checkpoint is able to, even though it did not receive any training signal on the heldout patterns.
| Only 3 | Only 4 | No Hard | Random 30 | Random 40 | Random 50 | |
| # Successful Runs | 0 | 3 | 0 | 1 | 2 | 3 |
Table 10: Number of successful runs (out of 3) per heldout strategy: The existence of harder patterns and an overall wide range of patterns seems necessary for the 1.5B model to learn.
| All@ | Avg | Pass@ | Absent (%) | Incorrect (%) | All@ | Avg | Pass@ | |
| 1.5B | 0.47 | 0.68 | 0.82 | 4 | 4 | 0.06 | 0.25 | 0.52 |
| 3B | 0.35 | 0.69 | 0.88 | 1 | 2 | - | - | - |
Table 11: Final performance of models trained without random 50% of patterns: All-correct@32 (All@), Average Accuracy (Avg), and Pass@32 (Pass@) of models. Models can generalize to most unseen patterns. 3B model is able to generalize to more diverse patterns than 1.5B, but has a lower All-correct@32.
Figure 8: Learning dynamics at a glance (Qwen2.5-3B, prefix notation) . We see a different hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy. This provides additional evidence for our hypothesis that lookahead bottleneck causes right-heavy patterns to be more difficult for models to learn.
| Training | |||
| SFT+No CoT | 0.92 | 0.82 | 0.33 |
| SFT+Random CoT | 0.59 | 0.25 | 0.04 |
| Base Model | 0.42 | 0.00 | 0.02 |
| GRPO | 1.00 | 0.94 | 0.68 |
Table 12: Final performance of SFT-trained models. Pass@32 on (trained) and (held out). We present the performance of the base model and the GRPO-trained model as a basis for comparison.
Figure 9: Learning dynamics at a glance (Qwen2.5-1.5B with SFT) . Here, the high-precision threshold has been lowered to 50%. Even though the coverage of patterns (OOD) decrease in the following order: balanced > left-heavy > right-heavy, the model masters patterns (in-distribution) in the opposite order: right-heavy > left-heavy > balanced. This shows a fundamental difference in the learning dynamics from RL.
Figure 10: Learning dynamics at a glance (Qwen2.5-1.5B with SFT) . The model does not reliably learn any pattern.
Figure 11: Learning dynamics at a glance (Qwen2.5-1.5B with SFT+CoT) . The model does not reliably learn any pattern.
| balanced | left-heavy | right-heavy | ||||
| Present in SFT dataset | 58 | 11 | 31 | |||
| Generated by RL model (ID) | 65 | +7 | 11 | 25 | -6 | |
| Generated by RL model (OOD) | 77 | +19 | 14 | +3 | 9 | -22 |
Table 13: Distribution of proofs by compositional structure (%) under different thresholds (1/3, 2/3). The results from Table 4 still hold.
Figure 12: Qwen2.5-1.5B often learns a formatting-only shortcut: (Left) training reward; (Right) training response length of two distinct training runs. Once the model learns to collect the formatting reward (grey dotted vertical), the model attempts the puzzle without chain of thought, characterized by a sharp decrease in the response length. During some but not all runs (blue), we observe a phase transition (grey solid vertical), where the model starts generating chain of thought again.
Figure 13: When normalizing group rewards by standard deviation, negative gradients make responses longer: (Left) training reward; (Right) training response length. In the first 25 steps, training reward increases to 0.1 at the same pace across all runs. However, with standard deviation normalization (red), the response length peaks high, and the training reward stalls. Training with only positive gradients (green dotted) and negative gradients (yellow dotted) respectively mimic the behavior of training without (blue) and with (red) normalization.
| No Normalization | Normalization | |||||
| All@ | Avg | Pass@ | All@ | Avg | Pass@ | |
| 1.5B | 0.66 | 0.85 | 0.94 | 0.36 | 0.60 | 0.86 |
| 3B | 0.61 | 0.84 | 0.95 | 0.66 | 0.86 | 0.96 |
Table 14: Effect of standard deviation normalization on final performance. All-correct@32 (All@), Average Accuracy (Avg), and Pass@32 (Pass@) with or without standard deviation normalization during group rewards computation. Normalizing generally hurts final performance.
| Small | Large | Switch | |||||||
| All@ | Avg | Pass@ | All@ | Avg | Pass@ | All@ | Avg | Pass@ | |
| 1.5B | 0.66 | 0.85 | 0.94 | 0.46 | 0.78 | 0.93 | 0.46 | 0.79 | 0.95 |
| 3B | 0.61 | 0.84 | 0.95 | 0.46 | 0.78 | 0.94 | 0.48 | 0.83 | 0.98 |
Table 15: Effect of group size on final performance. All-correct@32 (All@), Average Accuracy (Avg), and Pass@32 (Pass@) on different choices of group sizes. A larger number of rollouts generally hurts final performance.
| GRPO | PPO | |||||
| All@ | Avg | Pass@ | All@ | Avg | Pass@ | |
| 3B | 0.61 | 0.84 | 0.95 | 0.65 | 0.81 | 0.90 |
| 7B | 0.62 | 0.83 | 0.94 | 0.76 | 0.85 | 0.93 |
Table 16: Effect of RL algorithm on final performance. All-correct@32 (All@), Average Accuracy (Avg), and Pass@32 (Pass@) of models. Using PPO instead of GRPO improves All-correct@32 at the cost of Pass@32.
Figure 14: Learning dynamics at a glance (Qwen3-1.7B-Base) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 15: Effect of standard deviation normalization on response length: We reproduce the results from Figure 13 .
Figure 16: Learning dynamics at a glance (Llama-3.2-3B after lightweight SFT) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 17: Learning dynamics at a glance (Qwen2.5-1.5B) with reversed canonicalization. The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 18: Learning dynamics at a glance (Qwen2.5-1.5B) without canonicalization. While coverage decreases throughout training (models converge on one of multiple equivalent expressions), the same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 19: Learning dynamics at a glance (Qwen2.5-1.5B, another seed) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 20: Learning dynamics at a glance (Qwen2.5-3B) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 21: Learning dynamics at a glance (Qwen2.5-3B, another seed) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 22: Learning dynamics at a glance (Qwen2.5-7B) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy. We take the best checkpoint at 1350 steps (grey dotted), after which training reward starts declining.
Figure 23: Learning dynamics at a glance (Qwen2.5-7B, another seed) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy. We take the best checkpoint at 950 steps (grey dotted), after which training reward starts declining.
Figure 24: Learning dynamics at a glance (Qwen2.5-7B, another seed) . The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy. We take the best checkpoint at 650 steps (grey dotted), after which training reward starts declining.
Figure 25: Generalizes to entire families of held-out compositional patterns (Qwen2.5-1.5B, another seed). The same conclusion from Figure 3 : the model can reuse learned substructures to assemble novel operator-tree shapes it has never been explicitly trained on.
Figure 26: Generalizes to entire families of held-out compositional patterns (Qwen2.5-1.5B, another seed). The same conclusion from Figure 3 : the model can reuse learned substructures to assemble novel operator-tree shapes it has never been explicitly trained on.
Figure 27: Learning dynamics at a glance (Qwen2.5-1.5B) trained on . The same conclusions from Figures 2 and 4 hold: models exhibit length generalization and master patterns progressively following a hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 28: Learning dynamics at a glance (Qwen2.5-1.5B) but high-precision defined as precision (top); (bottom). The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
Figure 29: Learning dynamics at a glance (Qwen2.5-1.5B, undersampled patterns removed) . We remove the 15 undersampled patterns ( Table 5 ) from the analysis. The same conclusion from Figure 4 holds about the hierarchy of difficulty by subtree structure: balanced left-heavy right-heavy.
| Expression | Pattern | Expression | Pattern |
Table 17: Expressions and their Canonical Forms (Patterns) (Page 1 of 4)
| Expression | Pattern | Expression | Pattern |
Table 18: Expressions and their Canonical Forms (Patterns) (Page 2 of 4)
| Expression | Pattern | Expression | Pattern |
Table 19: Expressions and their Canonical Forms (Patterns) (Page 3 of 4)
| Expression | Pattern | Expression | Pattern |
Table 20: Expressions and their Canonical Forms (Patterns) (Page 4 of 4)
More in this topic
Compositional Generalization
Compositional generalization is the ability to combine learned parts and rules in ways that transfer to unseen combinations.
See all 113 papers