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

1 Introduction

• The analysis characterizes mid-training as selecting an action subspace that minimizes both pruning-related approximation error and subsequent RL planning error.Pruning efficiency shapes the initial policy prior, while effects on convergence govern how much it can improve through online interactions.Source: §1 · pp. 1–2
• RA3 derives a temporal variational lower bound and iteratively discovers latent sequences with self-supervised RL before fine-tuning on bootstrapped data.The latent prior and KL penalty encourage temporal consistency and regulate the need for rational rollouts.Source: §1 · p. 2
• Experiments use Qwen and Llama base models ranging from 1B to 8B, with 3.5M code snippets totaling 1B tokens for mid-training.The reported gains span HumanEval, MBPP, and their extended variants.Source: §1 · p. 2
• RA3 improves average HumanEval and MBPP performance by 8 points over base models and 4 points over next-token prediction, while accelerating RLVR convergence and improving asymptotic performance across four benchmarks.RLVR benchmarks include HumanEval+, MBPP+, LiveCodeBench, and Codeforces.Source: §1 · p. 2

2 Background

• A language task is modeled as an MDP with states, primitive actions, rewards, and discounting; in language, states are token contexts and actions are individual tokens.Transitions append action tokens or depend on the external environment.Source: §2 · p. 2
• Next-token prediction is imitation learning on an offline expert dataset collected by rolling out the expert policy on sampled tasks.Its objective is to maximize conditional log-likelihood.Source: §2 · p. 2
• The NTP objective also applies to coarser actions such as sentence-level actions, not only token-level actions.The passage defines states, actions, sequence length, and the beginning-of-sentence token for the objective.Source: §2 · p. 2
• Mid-training is the second stage—continued pre-training to acquire reasoning foundations before RLVR—in a pipeline of pre-training, mid-training, and post-training.NTP is used across these training stages, as well as supervised fine-tuning.Source: §2 · p. 2
• RLVR uses outcome-based binary verification rewards, and the experiments adopt GRPO as their default RLVR algorithm.The reward checks whether a response matches the ground-truth answer for the prompt.Source: §2 · p. 3

3 How Mid-Training Shapes Post-Training RL

• Post-training regret separates into pruning error and RL error within the retained action space, so mid-training should minimize both.The unified action space includes primitive actions and abstractions.Source: §3 · p. 3
• Pruning efficiency determines the initial RL policy prior, while the effect on convergence determines how efficiently online interactions can further improve the policy.The analysis frames mid-training as finding useful actions while avoiding exploration of the full language space.Source: §3 · p. 3

3.1 Mid-Training Acquires Policy Priors via Action Space Pruning

• The minimal size |Zϵ| is the smallest action subset that is ϵ-optimal across tasks; in the worst case, retaining the entire finite action space requires no pruning.The theorem bounds expected pruning error with probability at least 1 − δ, and fewer demonstrations are needed as |Zϵ| and |Z| shrink.Source: §3.1 · pp. 3–4
• Action abstractions correspond to transferable skills spanning multiple tasks, yielding smaller action spaces that require less expert data to approximate a near-optimal set.This provides a stronger policy prior and reduces the burden on post-training RL compared with primitive actions.Source: §3.1 · p. 4

3.2 Mid-Training Accelerates RL Convergence

• Policy iteration reaches optimality in the pruned space in at most O((|S||Ẑ| − |S|)/(1 − γ)) iterations, with effective discount γ accounting for action duration.Longer actions reduce γ and shorten the effective planning horizon; a more compact action space also reduces the suboptimal actions online RL must eliminate.Source: §3.2 · p. 4

4 From Primitive Actions to Temporal Abstractions

• Figure 1 illustrates primitive actions and hidden temporal abstractions in expert demonstrations from web-agent and code-generation domains.The section introduces a principled method for extracting temporal abstractions from primitive actions in mid-training demonstrations.Source: §4 · p. 4, Figure 1 · p. 5

4.1 Temporal Variational Bound

• The temporal ELBO introduces latent sequences to represent hidden thoughts or intentions behind expert actions and uses a variational posterior to estimate trajectories that rationalize them.Maximizing the bound tightens the bound on the marginal likelihood of the observed actions given states.Source: §4.1 · p. 5
• An EM procedure alternates between updating the latent posterior and optimizing the action policy; its E-step treats expert-action log-likelihood as reward with a KL penalty.The M-step is imitation learning on expert trajectories bootstrapped with the inferred latents.Source: §4.1 · p. 5

4.2 Temporally-Consistent Latents as Action Abstractions

• The method recasts primitive-action likelihood maximization as optimizing an ELBO that learns from bootstrapped sequences paired with latent trajectories sampled from the variational posterior.The section motivates designing the latent prior to obtain compact sets of high-level action abstractions.Source: §4.2 · pp. 5–6
• Temporally extended abstractions require the latent to remain the same across multiple steps, and a prior with mass on the preceding latent encourages this consistency while its uniform component promotes diversity.The reused latent serves as a high-level action abstraction.Source: §4.2 · p. 6

5 RA3: A Scalable Mid-Training Algorithm

• Temporally consistent latents remain fixed for τ steps, avoiding a new rollout at every token and making it feasible to use mid-training corpora containing billions of tokens.Directly generating rationales at every timestep would require rollouts proportional to the total number of data tokens.Source: §5 · p. 6
• RA3 uses <act> to mark a temporally consistent latent and <think> to begin a new rationale; rollouts terminate on <act> and are sampled fully when <think> begins.The design reduces RL inference cost by avoiding full rollouts for reused intentions.Source: §5 · p. 6
• The KL penalty discourages new rationales when α exceeds q_act, setting a threshold for when their log-likelihood reward justifies thinking.The authors use a fixed penalty on non-<act> latents in implementation, and α can adjust additional training cost; at α = 1, RA3 degenerates to NTP.Source: §5 · p. 6
• The RL step adds a format reward that gives zero reward to incorrectly formatted outputs and optimizes the T-horizon objective with policy gradients.Advantages are calculated in a manner similar to GRPO using sampled rollouts and action-independent baselines.Source: §5 · p. 7

LLM Mid-Training

• Prior LLM mid-training uses continued pre-training on high-quality expert data, while action hierarchies have been reported to outperform primitive-action-only training [Xu et al., 2024; Chen et al., 2024; Wang et al., 2025b; Xue et al., 2025].The passage notes that synthetic reasoning distillation can introduce distributional shift and incur high costs at scale [Wang et al., 2025b].Source: §6 · p. 7

Self-Supervised RL

• RA3 shares the self-supervised RL use of expert-action log-probability as reward, but is motivated by an ELBO for temporal sequences [Zhong et al., 2025; Ruan et al., 2025; Zhou et al., 2025a; Dong et al., 2025; Wang et al., 2025a].The passage contrasts RA3 with hand-crafted reasoning-position rules and small-scale training on instructed reasoning models [Dong et al., 2025].Source: §6 · p. 7

Markov Options

• Markov options provide temporally extended action courses for learning in long-horizon tasks [Bacon et al., 2017].The cited option-transfer work studies transferability, whereas this paper focuses on LLM mid-training and its effects on post-training RL [Brunskill and Li, 2013, 2014].Source: §6 · p. 7

Experiment Setups

• Experiments focus on Python code generation across Qwen-2.5-1.5B, Llama-3.2-1B, and Llama-3.1-8B, using 3.5M snippets totaling 1B tokens.The RLVR evaluation uses GRPO and independent training runs with multiple random seeds.Source: §7 · pp. 8–9

Mid-Training Results

• In the first E step, the model quickly learns to maximize the RL reward, leaving most data and compute for bootstrapping and M-step fine-tuning.Figure 2 shows the first E-step reward curve.Source: §7 · p. 8, Figure 2 · p. 8
• After RL, the policy extracts transferable skills embedded in expert behavior, including dummy head creation and BFS.Figure 3 presents examples of the training data before and after reasoning bootstrapping.Source: §7 · p. 8, Figure 3 · p. 8
• Across four benchmarks, RA3 improves evaluation accuracy with fewer data than NTP as EM iterations proceed.The same iterations reduce cross-entropy loss; fine-tuning on reasoning-bootstrapped data accelerates learning because hidden reasoning trajectories guide primitive code-line decisions.Source: Figure 4 · p. 8, §7 · p. 9, Figure 5 · p. 9
• RA3 consistently outperforms NTP on pass@k across HumanEval, MBPP, HumanEval+, and MBPP+.Table 1 reports the mid-training comparison.Source: Table 1 · p. 9, §7 · p. 9

Post-Training RLVR Results

• RA3 learns more effectively during RLVR, with faster convergence and higher asymptotic performance than other mid-training algorithms.The evaluation reports means and standard errors across independent runs.Source: §7 · p. 9, Figure 6 · p. 10

Ablation Study

• The penalty c controls reasoning frequency: low c prompts redundant reasoning, c = 0.2 makes RA3 behave like NTP, and default c = 0.05 triggers reasoning on fewer than 40% of lines.Figure 7 tracks penalty effects on RL behavior, mean HumanEval and MBPP accuracy, latent length, and full-rollout sample ratio.Source: §7 · p. 10, Figure 7 · p. 10

• Mid-training design balances pruning the action space to form a strong policy prior with accelerating later RL convergence for online policy improvement.Source: §8 · p. 10
• RA3 uses a temporal variational bound to iteratively uncover temporally consistent reasoning trajectories via RL and fine-tune on the resulting bootstrapped data.Source: §8 · p. 10
• On code generation tasks, reasoning as action abstraction enhances initial-policy generalizability and improves post-training RLVR convergence speed and asymptotic performance.Source: §8 · p. 11

A.1 Proof of Lemma 3.2

• A policy restricted to the pruned action set is admissible in both MDPs, which share dynamics, rewards, and discount factors.The proof then says the result follows from basic algebra.Source: §A.1 · p. 13

A.2 Proof of Theorem 3.4

• With probability at least 1−δ, pruning excludes (ϵ, σ)-suboptimal action subsets, and expected pruning error is at most ϵ(1−σ)+σ/(1−γ).The proof defines suboptimality by the probability that the subset's error exceeds ϵ, then bounds the probability of retaining such a subset.Source: §A.2 · pp. 13–14

A.3 Proof of Theorem 3.5

• Policy iteration eliminates at least one suboptimal state-action pair every O(1/(1−γ)) iterations and reaches optimality within O((|S||bZ|−|S|)/(1−γ)) iterations.The proof also uses contraction of the Bellman operator and a bound on value error after n iterations.Source: §A.3 · p. 15

A.4 Proof of Theorem 4.1

• The maximum-likelihood proof introduces a sequence of latent variables and uses Jensen’s inequality.It identifies p(zt|st, zt−1) as the prior distribution of the latent zt.Source: §A.4 · pp. 15–16

A.5 Proof of Proposition 5.1

• The KL-divergence proof normalizes non-action latent probabilities by 1−qact and omits the constant contributed by the uniform prior.It expands the KL terms after separating the probability qact assigned to the action token.Source: §A.5 · pp. 16–17