Learning to Reason as Action Abstractions with Scalable Mid-Training RL

Abstract: LLMs excel with reinforcement learning (RL), but fully unlocking this potential requires a mid-training stage. An effective mid-training phase should identify a compact set of useful actions and enable fast selection among them through online RL. We formalize this intuition by presenting the first theoretical result on how mid-training shapes post-training: it characterizes an action subspace that minimizes both the value approximation error from pruning and the RL error during subsequent planning. Our analysis reveals two key determinants of mid-training effectiveness: pruning efficiency, which shapes the prior of the initial RL policy, and its impact on RL convergence, which governs the extent to which that policy can be improved via online interactions. These results suggest that mid-training is most effective when the decision space is compact and the effective horizon is short, highlighting the importance of operating in the space of action abstractions rather than primitive actions. Building on these insights, we propose Reasoning as Action Abstractions (RA3), a scalable mid-training algorithm. Specifically, we derive a sequential variational lower bound and optimize it by iteratively discovering temporally-consistent latent structures via RL, followed by fine-tuning on the bootstrapped data. Experiments on code generation tasks demonstrate the effectiveness of our approach. Across multiple base models, RA3 improves the average performance on HumanEval and MBPP by 8 and 4 points over the base model and the next-token prediction baseline. Furthermore, RA3 achieves faster convergence and higher asymptotic performance in RLVR on HumanEval+, MBPP+, LiveCodeBench, and Codeforces.

• The paper introduces RA3, which learns high-level action abstractions that compactly represent reasoning steps to improve RL performance.
• It leverages a temporal variational bound and EM-style optimization to extract and utilize these abstractions during LLM mid-training.
• Empirical results in Python code generation show RA3 outperforms baseline methods with 4–8 point improvements across multiple benchmarks.

Learning to Reason as Action Abstractions with Scalable Mid-Training RL

Overview and Motivation

This paper presents a formal analysis and practical algorithmic framework for mid-training in LLMs, specifically targeting the code generation domain. The central thesis is that mid-training should not merely imitate expert demonstrations at the primitive action (token) level, but instead discover and operate over temporally extended action abstractions—high-level, transferable skills that compactly represent reasoning steps. The authors provide the first theoretical results connecting mid-training design to post-training RL performance, and introduce Reasoning as Action Abstractions (RA3), a scalable mid-training algorithm that leverages a temporal variational lower bound and EM-style optimization to extract and utilize these abstractions.

Theoretical Foundations

The analysis formalizes the role of mid-training as identifying a compact, near-optimal action subspace that minimizes both the error from pruning the full action space and the RL error during subsequent planning. The regret decomposition (Lemma 1) splits post-training RL regret into pruning error and RL error, motivating mid-training algorithms that efficiently eliminate suboptimal actions and accelerate RL convergence.

The pruning efficiency theorem establishes that the sample complexity of pruning suboptimal actions is governed by the cardinality of the smallest near-optimal action subset. This result directly favors action abstractions over primitive actions, as abstractions correspond to transferable skills and thus yield a much smaller decision space.

The RL convergence theorem shows that temporally extended actions (options) reduce the effective planning horizon, leading to faster convergence rates in value iteration and policy gradient methods. This is formalized via the contraction factor $\overline{\gamma}$, which decreases as the duration of abstractions increases.

Figure 1: Probabilistic graphical model of the action hierarchy (left), and examples of primitive actions (blue) versus hidden high-level temporal abstractions (green) in web agent and code generation domains (middle and right).

Temporal Variational Bound and EM Optimization

To extract temporal abstractions from expert demonstrations, the authors derive a sequential ELBO for the next-token prediction objective, introducing latent variables $z_{0:T}$ that represent hidden intentions or reasoning steps. The ELBO is maximized via an EM procedure:

• E-step: RL is used to infer latent trajectories $z_{0:T}$ that rationalize expert actions, with the log-likelihood as reward and a KL penalty enforcing temporal consistency.
• M-step: Supervised fine-tuning is performed on the bootstrapped data, conditioning on the inferred latents.

Temporal consistency is enforced by designing the latent prior to favor reusing the same abstraction across multiple steps, effectively modeling options. The KL penalty is decomposed into a Bernoulli KL regularizer and an entropy term, with a hyperparameter $c$ controlling the trade-off between reasoning frequency and computational cost.

Figure 2: Examples of mid-training data before and after reasoning bootstrapping, illustrating the abstraction of transferable skills such as dummy head creation and BFS.

RA3: Scalable Mid-Training Algorithm

RA3 operationalizes these insights by alternating between RL-based reasoning extraction and supervised fine-tuning. The algorithm is designed to scale to mid-training corpora with billions of tokens by leveraging temporal consistency to minimize rollout cost. Two types of latents are used: $\langle\text{act}\rangle$ for direct actions and $\langle\text{think}\rangle$ for reasoning steps, with rollouts triggered only when new rationales are needed.

The implementation uses policy gradient for the RL step, with group-based advantage calculation, and standard NTP for the M step. Hyperparameters are chosen to balance reasoning frequency and efficiency, with ablation studies demonstrating the impact of the KL penalty $c$.

Figure 3: Bootstrapped data with reasoning reduces cross-entropy loss during M step fine-tuning.

Empirical Results

Experiments are conducted on Python code generation using Qwen and Llama base models (1B–8B parameters), with a mid-training corpus of 3.5M code snippets (1B tokens). RA3 consistently outperforms next-token prediction (NTP) and base models across HumanEval, MBPP, HumanEval+, and MBPP+ benchmarks, with average improvements of 4–8 points.

Figure 4: Evaluation results during mid-training, with accuracies averaged across four benchmarks.

Post-training RLVR experiments using GRPO show that RA3-trained models achieve faster convergence and higher asymptotic performance compared to NTP and base models, validating the theoretical predictions.

Figure 5: RLVR evaluation results (mean and standard error) for different mid-training algorithms.

Ablation studies on the KL penalty $c$ reveal that reasoning frequency and computational cost can be precisely controlled, with optimal performance achieved at intermediate values.

Figure 6: Effect of penalty $c$ on RL behavior, mean accuracy, average length of $z$, and ratio of full rollout samples.

Implementation Considerations

• Computational Requirements: RA3 is designed to scale to large mid-training datasets by minimizing the number of rollouts via temporal consistency. The RL step is implemented with asynchronous rollout engines to avoid idle time.
• Hyperparameter Tuning: The KL penalty $c$ is critical for balancing reasoning frequency and efficiency. Empirical results suggest $c=0.05$ as a robust default.
• Data Format: Reasoning abstractions are injected as comment lines in code, with format rewards ensuring syntactic correctness.
• Deployment: RA3 can be integrated into existing LLM training pipelines, replacing or augmenting standard NTP mid-training. The approach is compatible with both open-source and proprietary models.

Implications and Future Directions

The formal connection between mid-training design and post-training RL performance provides a principled foundation for future LLM training algorithms. The demonstrated benefits of action abstractions suggest that hierarchical RL and option-based reasoning should be central to scaling LLMs for complex, long-horizon tasks. Further research may explore:

• Automated discovery of abstraction hierarchies in broader domains (e.g., natural language, multi-modal tasks)
• Integration with curriculum learning and lifelong RL for continual skill acquisition
• Extension to agentic tasks with external environments and interactive feedback

Conclusion

This work establishes a rigorous theoretical and empirical framework for mid-training in LLMs, demonstrating that learning temporally extended action abstractions yields more efficient pruning, stronger policy priors, and accelerated RL convergence. The RA3 algorithm provides a scalable, practical instantiation of these principles, achieving superior performance in code generation and RLVR tasks. The results advocate for a shift from primitive action imitation to abstraction-based reasoning in future LLM training paradigms.