Learning to Reason as Action Abstractions with Scalable Mid-Training RL
Proposes RA3, a mid-training algorithm that learns action abstractions via temporal ELBO, improving code generation by 8 points on average.
Key Findings
Methodology
This work develops a theoretical framework analyzing how mid-training shapes the action subspace for subsequent RL. It introduces a sequential variational lower bound (ELBO) to learn latent high-level actions, enforcing temporal consistency. The approach combines RL for discovering structured abstractions with supervised fine-tuning, forming an iterative EM process. The core components include action space pruning, latent structure learning, and policy optimization, with specific algorithms like Group Relative Policy Optimization (GRPO) employed for RL updates.
Key Results
- RA3 achieves an average performance increase of 8 points on HumanEval and 4 points on MBPP across multiple models, surpassing next-token prediction baselines. It accelerates RL convergence by over 30%, reaching higher asymptotic performance on extended benchmarks such as HumanEval+ and MBPP+. The experiments validate that learning high-level action abstractions during mid-training significantly reduces sample complexity and improves generalization.
- The method effectively learns structured, temporally consistent latent actions, which serve as transferable skills across tasks. Empirical data show that models trained with RA3 require fewer samples to reach optimal policies, demonstrating improved sample efficiency and robustness. The approach also reduces inference costs by leveraging temporal abstractions, enabling scalable training on billions of tokens.
- Theoretical analysis confirms that action abstraction quality directly influences RL convergence rate and policy optimality. The experiments highlight the importance of action space pruning, latent temporal consistency, and iterative EM updates, establishing a new paradigm for efficient, scalable mid-training of large language models.
Significance
This research bridges the gap between hierarchical RL theory and large language model training, providing a principled method for learning high-level action abstractions during mid-training. It addresses the longstanding challenge of managing large action spaces, offering a scalable solution that enhances sample efficiency, accelerates convergence, and improves task generalization. The framework’s theoretical guarantees and empirical success demonstrate its potential to transform how models acquire reasoning and decision-making skills, with broad implications for AI applications such as code synthesis, reasoning, and autonomous agents. By enabling models to autonomously learn transferable skills, RA3 paves the way for more efficient, adaptable, and intelligent systems.
Technical Contribution
The paper introduces a novel variational framework based on a sequential ELBO to learn temporally consistent latent actions, formalizing the relationship between action abstraction quality and RL convergence. It provides theoretical bounds on pruning efficiency and convergence speed, demonstrating how high-level abstractions reduce the decision space and shorten the effective planning horizon. The iterative EM algorithm combines RL-based latent discovery with supervised fine-tuning, offering a scalable, self-supervised approach for large models. These contributions extend hierarchical RL principles into the realm of large-scale language modeling, with rigorous theoretical analysis and practical algorithms.
Novelty
This work is the first to systematically integrate action abstraction learning into mid-training of large language models via a temporal variational lower bound. Unlike prior approaches relying on handcrafted skills or synthetic distillation, RA3 autonomously learns structured, transferable high-level actions directly from expert demonstrations. Its theoretical analysis linking pruning efficiency and RL convergence, combined with a scalable EM-based algorithm, sets it apart from existing methods, marking a significant advancement in hierarchical RL applied to language modeling.
AI Executive Summary
This paper introduces RA3, a scalable mid-training algorithm that learns high-level action abstractions through a temporal variational lower bound (ELBO). Traditional large language models rely heavily on primitive token-level actions, leading to vast decision spaces and slow RL convergence. RA3 addresses this by extracting structured, temporally consistent latent actions that serve as transferable skills, significantly reducing the complexity of subsequent reinforcement learning. The core innovation lies in formulating a variational inference framework that jointly optimizes latent structures and policy parameters via an iterative EM process, combining RL-based latent discovery with supervised fine-tuning.
Theoretical analysis demonstrates that the quality of action abstractions directly influences RL convergence speed and policy optimality. Specifically, pruning efficiency improves as the size of the action subspace shrinks, and longer temporally extended actions accelerate RL convergence by effectively shortening the planning horizon. Empirical results on code generation benchmarks, including HumanEval and MBPP, show that RA3 improves average performance by 8 and 4 points respectively, across models ranging from 1B to 8B parameters. The method not only enhances accuracy but also speeds up training and inference, making it highly scalable.
Broader implications include enabling models to autonomously learn transferable skills, reducing reliance on extensive expert data, and improving sample efficiency. This approach paves the way for more intelligent, adaptable systems capable of complex reasoning and decision-making. Despite current limitations such as computational overhead and dependence on demonstration quality, RA3 offers a promising direction for future research in hierarchical RL and large-scale language modeling. Overall, this work significantly advances the theoretical understanding and practical implementation of action abstraction learning during mid-training, with wide-ranging applications in AI development.
Deep Analysis
Background
近年来，随着大规模预训练模型的崛起，强化学习（RL）在提升模型推理和生成能力方面逐渐成为研究热点。早期工作如GPT系列、BERT等主要依赖自监督学习，但在复杂推理任务中，单纯的下一词预测存在局限。引入RL后，模型能通过奖励信号优化策略，增强推理能力。中期训练阶段，利用专家示范进行策略微调，逐步引入动作层级结构，提升模型的抽象能力。相关研究如动作层次（Options）和动作抽象（Hierarchical RL）为本工作提供理论基础。尽管如此，如何在大规模模型中高效学习高层次动作抽象，仍是当前挑战。现有方法多依赖人工设计或有限示例，难以充分利用海量数据。本文突破在于提出基于时间变分的潜在结构学习框架，为中期训练引入高效的动作抽象机制。
Core Problem
核心问题在于大模型中动作空间庞大，直接强化学习难以高效收敛，且样本需求巨大。传统方法多在原子动作层面操作，导致探索空间过大，训练成本高昂。中期训练阶段的目标是提取高层次的动作抽象，减少决策复杂度，但缺乏理论指导，导致效果不稳定。如何设计既能有效剪枝动作空间，又能保证后续RL收敛速度的中期训练策略，是当前亟待解决的问题。本文试图通过理论分析，明确动作剪枝的效率与RL收敛的关系，为设计高效的中期训练算法提供依据。
Innovation
主要创新包括：1）提出基于时间变分下界（ELBO）的潜在结构学习框架，有效捕获高层次动作抽象；2）引入时间一致性机制，确保潜在变量的连续性和合理性；3）设计迭代EM流程，将RL与微调结合，实现自监督学习与策略优化的协同提升；4）理论分析动作剪枝效率与RL收敛速度的关系，为动作空间设计提供指导。这些创新使得模型能在中期训练阶段自主学习高层次策略，显著提升样本效率和泛化能力。
Methodology
- �� 设计时间变分ELBO，最大化专家示范数据的潜在结构似然；
- �� 引入潜在变量z，表示高层次动作抽象，定义其先验分布以保证时间一致性；
- �� 采用EM算法，E步通过RL优化潜在变量的后验分布，M步通过微调策略实现动作生成；
- �� 利用RL发现潜在动作的时间一致性，减少采样成本；
- �� 结合奖励设计，强化潜在抽象的表达能力，提升模型推理能力；
- �� 通过潜在空间剪枝，有效缩小动作决策空间，加快RL收敛速度。
Experiments
在Python代码生成任务中，选用Qwen-2.5-1.5B、Llama-3.2-1B和Llama-3.1-8B模型，使用3.5M代码片段作为中期训练数据。对比基线包括纯下一词预测和传统微调，评估指标为准确率和损失。采用多轮EM迭代，每轮400次梯度更新，前40次为RL策略梯度。通过AB测试验证RA3在HumanEval、MBPP及扩展任务中的性能提升，观察到平均提升8点和4点，收敛速度明显加快。还分析了潜在结构学习对样本效率的影响，验证了动作抽象的有效性。
Results
RA3在多个模型和任务中实现了显著性能提升，平均8点在HumanEval，4点在MBPP，超越传统微调和下一词预测基线。模型在RLVR任务中的收敛速度提升30%以上，极限性能也优于对比方法。潜在结构学习显著减少了样本需求，提高了模型的推理能力。实验还验证了动作抽象在长序列任务中的优势，模型表现出更强的泛化能力和鲁棒性。
Applications
该方法适用于大规模模型的代码生成、推理和决策任务。通过中期训练学习高层次动作抽象，模型能更高效地进行复杂任务规划，减少样本和计算成本。未来可扩展到多模态任务、长序列推理，推动AI自主策略学习，降低工业应用门槛。
Limitations & Outlook
当前方法依赖大量专家示范，数据偏差可能影响潜在结构的学习效果。潜在变量的推断复杂，计算成本较高，特别是在大模型中。对于极端复杂或多模态任务，潜在抽象设计仍需优化。未来需解决潜在空间的泛化能力和推断效率问题。
Plain Language Accessible to non-experts
想象你在学习做菜，平时只会做简单的菜，比如炒蛋或煮面。这些都是基本动作，但如果你想做一道复杂的菜，比如火锅或烤鸡，就需要把做菜的步骤拆分成更高层次的“技能”——比如准备调料、腌制、火候控制等。这些高层次技能可以反复用在不同菜肴中，帮你更快做出好菜。本文的算法就像教你学会这些“厨艺技能”，让模型在学习过程中掌握一些“高层次的动作”，而不是每次都从头开始。这样，模型就能更快理解任务、减少试错，最终做出更好的“菜”。通过学习这些“技能”，模型在处理复杂任务时变得更聪明、更高效，就像你学会了厨房里的“秘密武器”。
ELI14 Explained like you're 14
想象你在学校学新技能，比如打篮球。刚开始，你只会投篮或运球这些基本动作，但如果你学会了“传球”、“防守”这些更高层次的技能，你就能更快地玩好比赛。每次比赛都像一场大任务，里面有很多步骤。现在，假设你能把这些步骤变成一些“高层次的策略”，比如“组织进攻”或“防守反击”。这样，你就不用每次都从零开始思考，而是用这些“策略”来指导行动。论文里的算法就像教模型学会这些“策略”，让它在面对复杂问题时，能用更聪明的方法解决。这样，模型不仅能做得更快，还能做得更好，就像你变成了篮球高手！
Abstract
Large language models 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.